|
|
|
@ -45,10 +45,6 @@ class Dim : public Array<int64_t, D> {
|
|
|
|
|
HOSTDEVICE explicit Dim(int64_t head, Args... args)
|
|
|
|
|
: BaseClass(head, args...) {}
|
|
|
|
|
|
|
|
|
|
/** Construct a Dim from a linear index and size. Uses Fortran order
|
|
|
|
|
* indexing. */
|
|
|
|
|
HOSTDEVICE Dim(int64_t idx, const Dim<D>& size);
|
|
|
|
|
|
|
|
|
|
/** Construct a Dim with each dimension set to the given index */
|
|
|
|
|
HOSTDEVICE explicit Dim(int64_t idx) { this->Fill(idx); }
|
|
|
|
|
|
|
|
|
@ -57,181 +53,12 @@ class Dim : public Array<int64_t, D> {
|
|
|
|
|
HOST std::string to_string() const;
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
namespace detail {
|
|
|
|
|
template <int kStart, int kEnd, bool kStop>
|
|
|
|
|
struct FortranOrderIndexingConstructorFunctor {
|
|
|
|
|
HOSTDEVICE inline static void Run(const int64_t* in, int64_t* idx,
|
|
|
|
|
int64_t* out) {
|
|
|
|
|
out[kStart] = (*idx) % in[kStart];
|
|
|
|
|
(*idx) /= in[kStart];
|
|
|
|
|
FortranOrderIndexingConstructorFunctor<kStart + 1, kEnd,
|
|
|
|
|
kStart + 1 == kEnd>::Run(in, idx,
|
|
|
|
|
out);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <int kStart, int kEnd>
|
|
|
|
|
struct FortranOrderIndexingConstructorFunctor<kStart, kEnd, true> {
|
|
|
|
|
HOSTDEVICE inline static void Run(const int64_t* in, int64_t* idx,
|
|
|
|
|
int64_t* out) {}
|
|
|
|
|
};
|
|
|
|
|
} // namespace detail
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE Dim<D>::Dim(int64_t idx, const Dim<D>& size) {
|
|
|
|
|
detail::FortranOrderIndexingConstructorFunctor<0, D, D == 0>::Run(
|
|
|
|
|
size.Get(), &idx, this->GetMutable());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int idx, int D>
|
|
|
|
|
HOSTDEVICE inline int64_t get(const Dim<D>& dim) {
|
|
|
|
|
return dim[idx];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int idx, int D>
|
|
|
|
|
HOSTDEVICE inline int64_t& get(Dim<D>& dim) { // NOLINT
|
|
|
|
|
return dim[idx];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline int64_t get(const Dim<D>& dim, int idx) {
|
|
|
|
|
return dim[idx];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline int64_t& get(Dim<D>& dim, int idx) { // NOLINT
|
|
|
|
|
return dim[idx];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Dot product of two dims
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline int64_t linearize(const Dim<D>& a, const Dim<D>& b) {
|
|
|
|
|
return UnrollProduct<D>::Run(a.Get(), b.Get());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Product of a Dim
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline int64_t product(const Dim<D>& a) {
|
|
|
|
|
return UnrollProduct<D>::Run(a.Get());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Is 0 <= idx_i < size_i for all i?
|
|
|
|
|
namespace detail {
|
|
|
|
|
template <int kStart, int kEnd, bool kStop>
|
|
|
|
|
struct ContainedFunctor {
|
|
|
|
|
HOSTDEVICE static inline bool Run(const int64_t* idx, const int64_t* size) {
|
|
|
|
|
return (idx[kStart] >= 0 && idx[kStart] < size[kStart]) &&
|
|
|
|
|
ContainedFunctor<kStart + 1, kEnd, kStart + 1 == kEnd>::Run(idx,
|
|
|
|
|
size);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <int kStart, int kEnd>
|
|
|
|
|
struct ContainedFunctor<kStart, kEnd, true> {
|
|
|
|
|
HOSTDEVICE static constexpr inline bool Run(const int64_t* idx,
|
|
|
|
|
const int64_t* size) {
|
|
|
|
|
return true;
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
} // namespace detail
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline bool contained(const Dim<D>& idx, const Dim<D>& size) {
|
|
|
|
|
return detail::ContainedFunctor<0, D, D == 0>::Run(idx.Get(), size.Get());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* \brief Compute exclusive prefix-multiply of a Dim.
|
|
|
|
|
*/
|
|
|
|
|
namespace detail {
|
|
|
|
|
template <int kStart, int kEnd, bool kStop>
|
|
|
|
|
struct ExPrefixMulFunctor {
|
|
|
|
|
HOSTDEVICE static inline void Run(const int64_t* in, int64_t* out) {
|
|
|
|
|
kStart == 0 ? out[kStart] = 1 : out[kStart] =
|
|
|
|
|
out[kStart - 1] * in[kStart - 1];
|
|
|
|
|
detail::ExPrefixMulFunctor<kStart + 1, kEnd, kStart + 1 == kEnd>::Run(in,
|
|
|
|
|
out);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <int kStart, int kEnd>
|
|
|
|
|
struct ExPrefixMulFunctor<kStart, kEnd, true> {
|
|
|
|
|
HOSTDEVICE static inline void Run(const int64_t* in, int64_t* out) {}
|
|
|
|
|
};
|
|
|
|
|
} // namespace detail
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline Dim<D> ex_prefix_mul(const Dim<D>& src) {
|
|
|
|
|
Dim<D> ret;
|
|
|
|
|
detail::ExPrefixMulFunctor<0, D, D == 0>::Run(src.Get(), ret.GetMutable());
|
|
|
|
|
return ret;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Add two dimensions together
|
|
|
|
|
*/
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline Dim<D> dim_plus(const Dim<D>& a, const Dim<D>& b) {
|
|
|
|
|
Dim<D> ret;
|
|
|
|
|
UnrollAdd<D>::Run(a.Get(), b.Get(), ret.GetMutable());
|
|
|
|
|
return ret;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline Dim<D> operator+(const Dim<D>& lhs, const Dim<D>& rhs) {
|
|
|
|
|
return dim_plus(lhs, rhs);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Multiply two dimensions together
|
|
|
|
|
*/
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE inline Dim<D> dim_mult(const Dim<D>& a, const Dim<D>& b) {
|
|
|
|
|
Dim<D> ret;
|
|
|
|
|
UnrollMul<D>::Run(a.Get(), b.Get(), ret.GetMutable());
|
|
|
|
|
return ret;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE Dim<D> operator*(const Dim<D>& lhs, const Dim<D>& rhs) {
|
|
|
|
|
return dim_mult(lhs, rhs);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* \brief Normalize strides to ensure any dimension with extent 1
|
|
|
|
|
* has stride 0.
|
|
|
|
|
*
|
|
|
|
|
* \param size Dim object containing the size of an array
|
|
|
|
|
* \param stride Dim object containing stride of an array
|
|
|
|
|
* \return Dim object the same size as \p size with normalized strides
|
|
|
|
|
*
|
|
|
|
|
*/
|
|
|
|
|
namespace detail {
|
|
|
|
|
template <int kStart, int kEnd, bool kStop>
|
|
|
|
|
struct NormalizeStridesFunctor {
|
|
|
|
|
HOSTDEVICE static void Run(const int64_t* size, const int64_t* stride,
|
|
|
|
|
int64_t* ret) {
|
|
|
|
|
ret[kStart] = (size[kStart] == 1 ? 0 : stride[kStart]);
|
|
|
|
|
NormalizeStridesFunctor<kStart + 1, kEnd, kStart + 1 == kEnd>::Run(
|
|
|
|
|
size, stride, ret);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <int kStart, int kEnd>
|
|
|
|
|
struct NormalizeStridesFunctor<kStart, kEnd, true> {
|
|
|
|
|
HOSTDEVICE static void Run(const int64_t* size, const int64_t* stride,
|
|
|
|
|
int64_t* ret) {}
|
|
|
|
|
};
|
|
|
|
|
} // namespace detail
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE Dim<D> normalize_strides(const Dim<D>& size, const Dim<D>& stride) {
|
|
|
|
|
Dim<D> ret;
|
|
|
|
|
detail::NormalizeStridesFunctor<0, D, D == 0>::Run(size.Get(), stride.Get(),
|
|
|
|
|
ret.GetMutable());
|
|
|
|
|
return ret;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Helper function to create a Dim
|
|
|
|
|
*
|
|
|
|
@ -265,20 +92,6 @@ HOST std::string Dim<D>::to_string() const {
|
|
|
|
|
return stream.str();
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int D>
|
|
|
|
|
HOSTDEVICE Dim<D> linear_to_dimension(int linear_index, const Dim<D>& extents) {
|
|
|
|
|
Dim<D> result;
|
|
|
|
|
|
|
|
|
|
for (int i = 0; i < D - 1; ++i) {
|
|
|
|
|
result[i] = linear_index % extents[i];
|
|
|
|
|
linear_index /= extents[i];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
result[D - 1] = linear_index;
|
|
|
|
|
|
|
|
|
|
return result;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <int D, typename T1, typename T2>
|
|
|
|
|
inline void static_dim_assign(const T1* in, T2* out) {
|
|
|
|
|
UnrollAssign<D>::Run(in, out);
|
|
|
|
|