Skip to content

Commit

Permalink
Loop: add "offset" argument for 4D indexing
Browse files Browse the repository at this point in the history 
NOT to be used for productiong, since it is a horrible hack
  • Loading branch information
@rhaas80
rhaas80 committed Jan 12, 2025
1 parent 61ee7d1 commit a6c3a9c
Showing 1 changed file with 134 additions and 0 deletions.
134 changes: 134 additions & 0 deletions Loop/src/loop.hxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1073,6 +1073,18 @@ template <typename T> struct GF3D2 {
CCTK_DEVICE CCTK_HOST T &restrict operator()(int i, int j, int k) const {
return ptr[linear(i, j, k)];
}
CCTK_DEVICE CCTK_HOST T &restrict operator()(const GF3D2index &index, const int var) const {
#ifdef CCTK_DEBUG
assert(index.layout == this->layout);
#endif
return ptr[var + index.linear()];
}
CCTK_DEVICE CCTK_HOST T &restrict operator()(const vect<int, dim> &I, const int var) const {
return ptr[var + linear(I)];
}
CCTK_DEVICE CCTK_HOST T &restrict operator()(int i, int j, int k, const int var) const {
return ptr[var + linear(i, j, k)];
}
CCTK_DEVICE CCTK_HOST void store(const GF3D2index &index,
const T &value) const {
ptr[index.linear()] = value;
Expand All @@ -1081,6 +1093,14 @@ template <typename T> struct GF3D2 {
const T &value) const {
ptr[linear(I)] = value;
}
CCTK_DEVICE CCTK_HOST void store(const GF3D2index &index, const int var,
const T &value) const {
ptr[var + index.linear()] = value;
}
CCTK_DEVICE CCTK_HOST void store(const vect<int, dim> &I, const int var,
const T &value) const {
ptr[var + linear(I)] = value;
}
#if 0
struct simd_reference {
using element_type = std::remove_cv_t<T>;
Expand Down Expand Up @@ -1191,6 +1211,42 @@ template <typename T> struct GF3D2 {
const vect<int, dim> &I, const U &other) const {
return Arith::masko_loadu(mask, &(*this)(I), other);
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D2index &index, const int var) const {
return Arith::maskz_loadu(mask, &(*this)(index, var));
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D2index &index, const int var,
const Arith::simd<std::remove_cv_t<T> > &other) const {
return Arith::masko_loadu(mask, &(*this)(index, var), other);
}
template <typename U,
std::enable_if_t<std::is_convertible_v<T, U> > * = nullptr>
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D2index &index, const int var, const U &other) const {
return Arith::masko_loadu(mask, &(*this)(index, var), other);
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const vect<int, dim> &I, const int var) const {
return Arith::maskz_loadu(mask, &(*this)(I, var));
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const vect<int, dim> &I, const int var,
const Arith::simd<std::remove_cv_t<T> > &other) const {
return Arith::masko_loadu(mask, &(*this)(I, var), other);
}
template <typename U,
std::enable_if_t<std::is_convertible_v<T, U> > * = nullptr>
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const vect<int, dim> &I, const int var, const U &other) const {
return Arith::masko_loadu(mask, &(*this)(I, var), other);
}
CCTK_DEVICE CCTK_HOST void store(const Arith::simdl<T> &mask,
const vect<int, dim> &I,
const Arith::simd<T> &value) const {
Expand All @@ -1201,6 +1257,16 @@ template <typename T> struct GF3D2 {
const Arith::simd<T> &value) const {
mask_storeu(mask, &ptr[index.linear()], value);
}
CCTK_DEVICE CCTK_HOST void store(const Arith::simdl<T> &mask,
const vect<int, dim> &I, const int var,
const Arith::simd<T> &value) const {
mask_storeu(mask, &ptr[linear(I)], value);
}
CCTK_DEVICE CCTK_HOST void store(const Arith::simdl<T> &mask,
const GF3D2index &index, const int var,
const Arith::simd<T> &value) const {
mask_storeu(mask, &ptr[var + index.linear()], value);
}
};

////////////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -1429,6 +1495,13 @@ template <typename T> struct GF3D5 {
return ptr[index.linear()];
}
CCTK_DEVICE CCTK_HOST constexpr T &restrict
operator()(const GF3D5index &index, const int var) const {
#ifdef CCTK_DEBUG
assert(index.layout == this->layout);
#endif
return ptr[var + index.linear()];
}
CCTK_DEVICE CCTK_HOST constexpr T &restrict
operator()(const GF3D5layout &layout, const vect<int, dim> &I) const {
return (*this)(GF3D5index(layout, I));
}
Expand All @@ -1445,6 +1518,15 @@ template <typename T> struct GF3D5 {
const T &value) const {
operator()(layout, I) = value;
}
CCTK_DEVICE CCTK_HOST void store(const GF3D5index &index, const int var,
const T &value) const {
operator()(index, var) = value;
}
CCTK_DEVICE CCTK_HOST void store(const GF3D5layout &layout,
const vect<int, dim> &I, const int var,
const T &value) const {
operator()(layout, I, var) = value;
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5index &index) const {
Expand Down Expand Up @@ -1527,6 +1609,46 @@ template <typename T> struct GF3D5 {
const U &other) const {
return (*this)(mask, GF3D5index(layout, I), other);
}

// 4D array
CCTK_DEVICE CCTK_HOST simd_reference
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5index &index, const int var) const {
return simd_reference(&(*this)(index, var), mask);
}
CCTK_DEVICE CCTK_HOST simd_reference
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5layout &layout, const vect<int, dim> &I, const int var) const {
return (*this)(mask, GF3D5index(layout, I), var);
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5index &index, const int var,
const Arith::simd<std::remove_cv_t<T> > &other) const {
return Arith::masko_loadu(mask, &(*this)(index, var), other);
}
template <typename U,
std::enable_if_t<std::is_convertible_v<T, U> > * = nullptr>
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5index &index, const int var, const U &other) const {
return Arith::masko_loadu(mask, &(*this)(index, var), other);
}
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5layout &layout, const vect<int, dim> &I, const int var,
const Arith::simd<std::remove_cv_t<T> > &other) const {
return (*this)(mask, GF3D5index(layout, I), var, other);
}
template <typename U,
std::enable_if_t<std::is_convertible_v<T, U> > * = nullptr>
CCTK_DEVICE CCTK_HOST Arith::simd<std::remove_cv_t<T> >
operator()(const Arith::simdl<std::remove_cv_t<T> > &mask,
const GF3D5layout &layout, const vect<int, dim> &I, const int var,
const U &other) const {
return (*this)(mask, GF3D5index(layout, I), var, other);
}

CCTK_DEVICE CCTK_HOST void store(const Arith::simdl<T> &mask,
const GF3D5index &index,
const Arith::simd<T> &value) const {
Expand All @@ -1538,6 +1660,18 @@ template <typename T> struct GF3D5 {
const Arith::simd<T> &value) const {
store(mask, GF3D5index(layout, I), value);
}

CCTK_DEVICE CCTK_HOST void store(const Arith::simdl<T> &mask,
const GF3D5index &index, const int var,
const Arith::simd<T> &value) const {
mask_storeu(mask, &(*this)(index, var), value);
}
CCTK_DEVICE CCTK_HOST void store(const Arith::simdl<T> &mask,
const GF3D5layout &layout,
const vect<int, dim> &I, const int var,
const Arith::simd<T> &value) const {
store(mask, GF3D5index(layout, I), var, value);
}
};

template <typename T> struct is_GF3D5 : std::false_type {};
Expand Down

0 comments on commit a6c3a9c

Please sign in to comment.