File size: 2,433 Bytes
c1af2fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
#pragma once
namespace at::native { inline namespace CPU_CAPABILITY {
// n: number of function arguments (arity)
// traits: function_traits (see FunctionTraits.h)
// s: index of scalar argument or -1
template <int n, int stride_index, typename traits, int s=-1>
struct IsContiguous {
static bool eval(const int64_t* strides) {
using type = typename traits::template arg<n - 1>::type;
return strides[stride_index] == (s == n ? 0 : sizeof(type)) &&
IsContiguous<n - 1, stride_index - 1, traits, s>::eval(strides);
}
};
// will be called when there is an output exists
template <typename traits, int s>
struct IsContiguous<0, 0, traits, s> {
static bool eval(const int64_t* strides) {
return strides[0] == sizeof(typename traits::result_type);
}
};
// will be called when there is no output
template <typename traits, int s>
struct IsContiguous<0, -1, traits, s> {
static bool eval(const int64_t* /*strides*/) {
return true;
}
};
// output and all inputs are contiguous
template <
typename traits,
std::enable_if_t<std::is_void_v<typename traits::result_type>>* =
nullptr>
static inline bool is_contiguous(const int64_t* strides) {
return IsContiguous<traits::arity, traits::arity - 1, traits>::eval(strides);
}
template <typename traits,
std::enable_if_t<!std::is_void_v<typename traits::result_type>>* = nullptr>
static inline bool is_contiguous(const int64_t* strides) {
return IsContiguous<traits::arity, traits::arity, traits>::eval(strides);
}
// input at `s` is scalar (stride 0); output and other inputs are contiguous
// NB: output is typically at strides[0] so first input corresponds to s=1
template <typename traits, int s,
std::enable_if_t<std::is_void_v<typename traits::result_type>>* = nullptr>
static inline bool is_contiguous_scalar(const int64_t* strides) {
static_assert(s > 0 && s <= traits::arity, "scalar argument index out of bounds");
return IsContiguous<traits::arity, traits::arity - 1, traits, s>::eval(strides);
}
template <typename traits, int s,
std::enable_if_t<!std::is_void_v<typename traits::result_type>>* = nullptr>
static inline bool is_contiguous_scalar(const int64_t* strides) {
static_assert(s > 0 && s <= traits::arity, "scalar argument index out of bounds");
return IsContiguous<traits::arity, traits::arity, traits, s>::eval(strides);
}
}}
|