leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

CUTLASS 3.3.0 (#1167)

Browse Source 
* Release 3.3.0

Adds support for mixed precision GEMMs On Hopper and Ampere
Adds support for < 16B aligned GEMMs on Hopper
Enhancements to EVT
Enhancements to Python interface
Enhancements to Sub-byte type handling in CuTe
Several other bug-fixes and performance improvements.

* minor doc update
This commit is contained in:
Pradeep Ramani
2023-11-02 11:09:05 -04:00
committed by GitHub
parent 922fb5108b
commit c008b4aea8
263 changed files with 16214 additions and 5008 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/copy.hpp
+25 -6
View File
@@ -130,6 +130,17 @@ copy_if(PrdTensor const& pred,
// copy_if -- Predicated CopyAtom
//
namespace detail {
// Trait that detects if atom's traits has a member function with(bool)
template<typename, typename Enable = void>
constexpr bool has_with_bool = false;
template<typename T>
constexpr bool has_with_bool<T, cute::void_t<decltype(declval<typename T::Traits>().with(declval<bool>()))>> = true;
} // end namespace detail
template <class... CopyArgs,
class PredTensor,
class SrcEngine, class SrcLayout,
@@ -150,8 +161,14 @@ copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
auto dst_v = group_modes<1,R>(dst);
CUTE_UNROLL
for (int i = 0; i < size<1>(src_v); ++i) {
if (pred(i)) {
copy_atom.call(src_v(_,i), dst_v(_,i));
// If copy traits can be transformed with a predicate value, do it, otherwise branch here
if constexpr (detail::has_with_bool<Copy_Atom<CopyArgs...>>) {
copy_atom.with(pred(i)).call(src_v(_,i), dst_v(_,i));
}
else {
if (pred(i)) {
copy_atom.call(src_v(_,i), dst_v(_,i));
}
}
}
}
@@ -169,15 +186,17 @@ void
copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
using SrcType = typename SrcEngine::element_type;
using DstType = typename DstEngine::element_type;
if constexpr (sizeof(SrcType) == sizeof(DstType) && sizeof(VecType) > sizeof(DstType))
{
/* @pre is_aligned<N>(src.data()) &&
* is_aligned<N>(dst.data())
*/
auto src_v = recast<VecType const>(src);
auto dst_v = recast<VecType >(dst);
using SrcVecType = conditional_t<is_volatile_v<SrcType>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<DstType>, VecType volatile, VecType >;
auto src_v = recast<SrcVecType>(src);
auto dst_v = recast<DstVecType>(dst);
#if 0
if (thread0()) {
include/cute/algorithm/functional.hpp
+70
View File
@@ -170,6 +170,76 @@ CUTE_NAMED_BINARY_OP(min_fn, cute::min);
#undef CUTE_BINARY_OP
#undef CUTE_NAMED_BINARY_OP
/**********/
/** Fold **/
/**********/
#define CUTE_FOLD_OP(NAME,OP) \
struct NAME##_unary_rfold { \
template <class... T> \
CUTE_HOST_DEVICE constexpr \
auto operator()(T&&... t) const { \
return (t OP ...); \
} \
}; \
struct NAME##_unary_lfold { \
template <class... T> \
CUTE_HOST_DEVICE constexpr \
auto operator()(T&&... t) const { \
return (... OP t); \
} \
}; \
struct NAME##_binary_rfold { \
template <class U, class... T> \
CUTE_HOST_DEVICE constexpr \
auto operator()(U&& u, T&&... t) const { \
return (t OP ... OP u); \
} \
}; \
struct NAME##_binary_lfold { \
template <class U, class... T> \
CUTE_HOST_DEVICE constexpr \
auto operator()(U&& u, T&&... t) const { \
return (u OP ... OP t); \
} \
}
CUTE_FOLD_OP(plus, +);
CUTE_FOLD_OP(minus, -);
CUTE_FOLD_OP(multiplies, *);
CUTE_FOLD_OP(divides, /);
CUTE_FOLD_OP(modulus, %);
CUTE_FOLD_OP(plus_assign, +=);
CUTE_FOLD_OP(minus_assign, -=);
CUTE_FOLD_OP(multiplies_assign, *=);
CUTE_FOLD_OP(divides_assign, /=);
CUTE_FOLD_OP(modulus_assign, %=);
CUTE_FOLD_OP(bit_and, &);
CUTE_FOLD_OP(bit_or, |);
CUTE_FOLD_OP(bit_xor, ^);
CUTE_FOLD_OP(left_shift, <<);
CUTE_FOLD_OP(right_shift, >>);
CUTE_FOLD_OP(bit_and_assign, &=);
CUTE_FOLD_OP(bit_or_assign, |=);
CUTE_FOLD_OP(bit_xor_assign, ^=);
CUTE_FOLD_OP(left_shift_assign, <<=);
CUTE_FOLD_OP(right_shift_assign, >>=);
CUTE_FOLD_OP(logical_and, &&);
CUTE_FOLD_OP(logical_or, ||);
CUTE_FOLD_OP(equal_to, ==);
CUTE_FOLD_OP(not_equal_to, !=);
CUTE_FOLD_OP(greater, >);
CUTE_FOLD_OP(less, <);
CUTE_FOLD_OP(greater_equal, >=);
CUTE_FOLD_OP(less_equal, <=);
#undef CUTE_FOLD_OP
/**********/
/** Meta **/
/**********/
include/cute/arch/copy.hpp
+13 -3
View File
@@ -48,11 +48,21 @@ struct UniversalCopy
using SRegisters = S[1];
using DRegisters = D[1];
template<class S_, class D_>
CUTE_HOST_DEVICE static constexpr void
copy(S const& src,
D & dst)
copy(S_ const& src,
D_ & dst)
{
dst = src;
dst = static_cast<D>(static_cast<S>(src));
}
// Accept mutable temporaries
template<class S_, class D_>
CUTE_HOST_DEVICE static constexpr void
copy(S_ const& src,
D_ && dst)
{
copy(src, dst);
}
};
include/cute/arch/copy_sm80.hpp
+60
View File
@@ -96,6 +96,66 @@ struct SM80_CP_ASYNC_CACHEGLOBAL
}
};
/// Copy via cp.async with caching at all levels
template <class TS, class TD = TS>
struct SM80_CP_ASYNC_CACHEALWAYS_ZFILL
{
using SRegisters = TS[1];
using DRegisters = TD[1];
static_assert(sizeof(TS) == sizeof(TD), "cp.async requires sizeof(src_value_type) == sizeof(dst_value_type)");
static_assert(sizeof(TS) == 4 || sizeof(TS) == 8 || sizeof(TS) == 16, "cp.async sizeof(TS) is not supported");
CUTE_HOST_DEVICE static void
copy(TS const& gmem_src,
TD & smem_dst,
bool pred)
{
#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
TS const* gmem_ptr = &gmem_src;
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst);
int src_size = pred ? sizeof(TS) : 0;
asm volatile("cp.async.ca.shared.global [%0], [%1], %2, %3;\n"
:: "r"(smem_int_ptr),
"l"(gmem_ptr),
"n"(sizeof(TS)),
"r"(src_size));
#else
CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled");
#endif
}
};
/// Copy via cp.async with caching at global level
template <class TS, class TD = TS>
struct SM80_CP_ASYNC_CACHEGLOBAL_ZFILL
{
using SRegisters = TS[1];
using DRegisters = TD[1];
static_assert(sizeof(TS) == sizeof(TD), "cp.async requires sizeof(src_value_type) == sizeof(dst_value_type)");
static_assert(sizeof(TS) == 4 || sizeof(TS) == 8 || sizeof(TS) == 16, "cp.async sizeof(TS) is not supported");
CUTE_HOST_DEVICE static void
copy(TS const& gmem_src,
TD & smem_dst,
bool pred)
{
#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
TS const* gmem_ptr = &gmem_src;
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst);
int src_size = pred ? sizeof(TS) : 0;
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n"
:: "r"(smem_int_ptr),
"l"(gmem_ptr),
"n"(sizeof(TS)),
"r"(src_size));
#else
CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Establishes an ordering w.r.t previously issued cp.async instructions. Does not block.
include/cute/arch/copy_sm90_tma.hpp
+1 -1
View File
@@ -785,7 +785,7 @@ tma_store_arrive() {
#endif
}
// Wait on prior N (Count) TMA_STORE instructions to complete
// Wait until at most Count committed TMA_STOREs are pending and all prior commits are complete
template <int Count>
CUTE_HOST_DEVICE static void
tma_store_wait() {
include/cute/atom/copy_traits.hpp
+17 -3
View File
@@ -92,6 +92,21 @@ struct Copy_Traits<DefaultCopy>
using RefLayout = SrcLayout;
};
namespace detail {
template <class Operation,
class PtrS, int... Is,
class PtrD, int... Id>
CUTE_HOST_DEVICE constexpr
void
copy_explode(PtrS&& s, int_sequence<Is...>,
PtrD&& d, int_sequence<Id...>)
{
return Operation::copy(s[Is]..., d[Id]...);
}
} // end namespace detail
//
// Generic copy_unpack for any Copy_Traits
//
@@ -123,9 +138,8 @@ copy_unpack(Copy_Traits<Operation, Args...> const&,
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumDst>{},
"In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
detail::explode(Operation::copy,
rS, make_int_sequence<RegNumSrc>{},
rD, make_int_sequence<RegNumDst>{});
detail::copy_explode<Operation>(rS, make_int_sequence<RegNumSrc>{},
rD, make_int_sequence<RegNumDst>{});
}
//
include/cute/atom/copy_traits_sm80.hpp
+96
View File
@@ -51,6 +51,13 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<S,D>>
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Construct a zfill variant with a given predicate value
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
with(bool pred) const {
return {pred};
}
};
template <class S, class D>
@@ -66,6 +73,95 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<S,D>>
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Construct a zfill variant with a given predicate value
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
with(bool pred) const {
return {pred};
}
};
template <class S, class D>
struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,Int<sizeof_bits<S>::value>>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,Int<sizeof_bits<D>::value>>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Predicate value that determines whether to load or zfill
bool pred = false;
// Overload copy_unpack for zfill variant to pass the predicate into the op
template <class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr
void
copy_unpack(Copy_Traits const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_gmem<TS>::value, "Expected gmem source for cp.async.");
static_assert(is_smem<TD>::value, "Expected smem destination for cp.async.");
Tensor rS = recast<S>(src);
Tensor rD = recast<D>(dst);
CUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
"In CopyAtom, src layout doesn't vectorize into registers. This src layout is incompatible with this tiled copy.");
CUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
"In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>::copy(rS[0], rD[0], traits.pred);
}
};
template <class S, class D>
struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,Int<sizeof_bits<S>::value>>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,Int<sizeof_bits<D>::value>>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Predicate value that determines whether to load or zfill
bool pred = false;
// Overload copy_unpack for zfill variant to pass the predicate into the op
template <class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr
void
copy_unpack(Copy_Traits const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_gmem<TS>::value, "Expected gmem source for cp.async.");
static_assert(is_smem<TD>::value, "Expected smem destination for cp.async.");
Tensor rS = recast<S>(src);
Tensor rD = recast<D>(dst);
CUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
"In CopyAtom, src layout doesn't vectorize into registers. This src layout is incompatible with this tiled copy.");
CUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
"In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>::copy(rS[0], rD[0], traits.pred);
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cute/atom/copy_traits_sm90_tma.hpp
+52 -23
View File
@@ -79,6 +79,14 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
copy_unpack_(void const* const dst_ptr,
Coord const& src_coord, seq<Is...>) const
{
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr);
#endif
SM90_TMA_LOAD::copy(&tma_desc_, tma_load_mbar_,
dst_ptr, get<Is>(src_coord)...);
}
@@ -185,6 +193,14 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
copy_unpack_(void const* const dst_ptr,
Coord const& src_coord, seq<Is...>) const
{
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr);
#endif
SM90_TMA_LOAD_MULTICAST::copy(&tma_desc_, tma_load_mbar_, multicast_mask_,
dst_ptr, get<Is>(src_coord)...);
}
@@ -298,6 +314,14 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
copy_unpack_(void const* const src_ptr,
Coord const& dst_coord, seq<Is...>) const
{
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr);
#endif
SM90_TMA_STORE::copy(&tma_desc_,
src_ptr, get<Is>(dst_coord)...);
}
@@ -354,8 +378,8 @@ struct Copy_Traits<SM90_BULK_COPY_G2S, NumBitsPerTMA, OpArgs...>
"Extra arguments not set. Set .with() before use.");
static_assert(is_gmem<TS>::value, "Expected gmem src for SM90_BULK_COPY_G2S");
static_assert(is_smem<TD>::value, "Expected smem dst for SM90_BULK_COPY_G2S");
SM90_BULK_COPY_G2S::copy(src.data().get(), *get<0>(traits.bulk_load_mbar_),
dst.data().get(), int32_t(NumBitsPerTMA::value / 8));
SM90_BULK_COPY_G2S::copy(raw_pointer_cast(src.data()), *get<0>(traits.bulk_load_mbar_),
raw_pointer_cast(dst.data()), int32_t(NumBitsPerTMA::value / 8));
}
// Record the memory barrier for the instruction
@@ -390,7 +414,7 @@ struct Copy_Traits<SM90_BULK_COPY_S2G, NumBitsPerTMA>
{
static_assert(is_smem<TS>::value, "Expected smem src for SM90_BULK_COPY_S2G");
static_assert(is_gmem<TD>::value, "Expected gmem dst for SM90_BULK_COPY_S2G");
SM90_BULK_COPY_S2G::copy(src.data().get(), dst.data().get(), int32_t(NumBitsPerTMA::value / 8));
SM90_BULK_COPY_S2G::copy(raw_pointer_cast(src.data()), raw_pointer_cast(dst.data()), int32_t(NumBitsPerTMA::value / 8));
}
};
@@ -497,7 +521,7 @@ coalesce_256(Tensor<Engine,Layout> const& tensor)
// and construct a TMA Descriptor for the resulting instruction
// At the same time, construct the Tma Tensor's Stride to generate
// the TMA coordinates that the instruction consumes.
//
//
template <class TmaInternalType,
class GEngine, class GLayout,
class SShape, class SStride,
@@ -518,7 +542,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
// Perform the tiling to the gmem vector again, but with indirections to the gtensor modes
auto gbasis = make_identity_layout(shape(gtensor));
auto tile_gbasis_tmp = gbasis.compose(smem_inv_h);
// Instead of the recast (gbasis doesn't have type info), replace the shape with the already-recasted shape
// tma_box_shape:gmem_mode
auto tile_gbasis = make_layout(shape(tile_gstride), stride(tile_gbasis_tmp));
@@ -530,8 +554,8 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
// NOTE This is essentially ArithmeticTuple complement...
// NOTE in pursuit of implementing an ArithmeticTuple logical_divide for smem_inv_h
auto tile_gbasis_remaining_stride = filter_tuple(flatten(shape (gtensor_T)), flatten(stride(gtensor_T)),
flatten(stride(gbasis)),
[&](auto s, auto d, auto e)
flatten(stride(gbasis)),
[&](auto s, auto d, auto e)
{
if constexpr (is_constant<1, decltype(s)>::value || is_constant<0, decltype(d)>::value) {
return cute::tuple<>{}; // If size-1 or stride-0, then don't append
@@ -551,7 +575,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
auto tma_gbasis_tile = tile_gbasis.compose(make_layout(wrap(shape(tma_gstride))));
// Append the remaining basis modes that contribute to the TMA with size-1
auto tma_gbasis_full = make_layout(tuple_cat(wrap( shape(tma_gbasis_tile)), wrap(repeat<tile_gbasis_remaining_rank>(Int<1>{}))),
auto tma_gbasis_full = make_layout(tuple_cat(wrap( shape(tma_gbasis_tile)), wrap(repeat<tile_gbasis_remaining_rank>(Int<1>{}))),
tuple_cat(wrap(stride(tma_gbasis_tile)), wrap(tile_gbasis_remaining_stride)));
// Group the trailing modes to make this max rank-5 -- TMA rank limitation
@@ -570,7 +594,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
//
// TMA desc creation
//
//
constexpr int tma_dim = decltype(rank(tma_gbasis))::value;
@@ -579,7 +603,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
//
void* gmem_address = (void*) raw_pointer_cast(gtensor_T.data());
auto gmem_layout = gtensor_T.layout();
auto gmem_layout = gtensor_T.layout();
cute::array<uint64_t, 5> gmem_prob_shape = {1,1,1,1,1};
cute::array<uint64_t, 5> gmem_prob_stride = {0,0,0,0,0};
@@ -665,20 +689,20 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
//
// Construct the descriptor
//
TmaDescriptor tma_desc = {0};
//
// TMA general info
//
#if (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
CUtensorMapDataType tma_format = TMA::to_CUtensorMapDataType<TmaInternalType>();
CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE;
CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_L2_128B;
CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
// TMA smem swizzle type
CUtensorMapSwizzle smem_swizzle = TMA::to_CUtensorMapSwizzle(get_tma_swizzle_bits(swizzle));
CUresult result = cuTensorMapEncodeTiled(
@@ -694,7 +718,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
smem_swizzle,
tma_l2Promotion,
tma_oobFill);
if (result != CUDA_SUCCESS) {
std::cerr << "TMA Desc Addr: " << &tma_desc
<< "\nformat " << tma_format
@@ -711,8 +735,11 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
std::cerr << "Error: Failed to initialize the TMA descriptor " << result << std::endl;
assert(false);
}
#endif // (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
auto recast_ratio = cute::ratio(Int<sizeof_bits<typename GEngine::value_type>::value>{},
Int<sizeof_bits< TmaInternalType>::value>{});
// Finally, get the inverse permutation of the E<i> bases for the mocked gmem stride
// NOTE This is essentially ArithmeticTuple inverse...
auto gmem_stride_bases = transform_leaf(stride(gbasis), [&](auto ei) {
@@ -727,9 +754,9 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
[[maybe_unused]] auto j = find_if(tma_gbasis_stride, [&](auto tma_stride_j) { return any_of(tma_stride_j, [&](auto dj) { return dj == EI{}; }); });
if constexpr (decltype(j == rank(tma_gbasis_stride))::value) {
return Int<0>{}; // If not-found, return arithmetic identity -- no contribution to the TMA
} else
} else
if constexpr (decltype(j == Int<0>{})::value) {
auto scale = ratio(size(tma_gstride), size(smem_inv_h)) * basis_get(ei, stride(gtensor));
auto scale = recast_ratio * basis_get(ei, stride(gtensor));
return E<j>{} * scale; // Return TMA Coord basis -- with a recast scale factor
} else
if constexpr (decltype(rank<j>(tma_gbasis_stride) == Int<1>{})::value) {
@@ -959,21 +986,23 @@ template <class TmaInternalType,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tile,
class CTA_Tiler,
class Cluster_Size>
CUTE_HOST_RTC
auto
make_tma_copy(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tile const& cta_tile,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler);
auto cta_t_tile = make_layout(cluster_size);
return detail::make_tma_copy_tiled<TmaInternalType>(copy_op,
gtensor,
slayout,
make_layout(cluster_size),
make_identity_layout(cta_tile));
cta_t_tile,
cta_v_tile);
}
// Explicit defaulting
include/cute/atom/copy_traits_sm90_tma_swizzle.hpp
+3 -2
View File
@@ -37,12 +37,13 @@
#include <cuda.h>
#endif
#include "cute/arch/copy_sm90_desc.hpp"
#include "cute/swizzle_layout.hpp"
#include <cute/arch/copy_sm90_desc.hpp>
#include <cute/swizzle_layout.hpp>
namespace cute::detail {
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
TMA::SmemSwizzleBits
get_tma_swizzle_bits(Swizzle<B,M,S>)
{
include/cute/atom/mma_atom.hpp
+13 -6
View File
@@ -155,7 +155,8 @@ struct MMA_Atom<MMA_Traits<Args...>>
if constexpr (has_dereference<FrgTypeA>::value) {
// If the intended FrgTypeA is a view (of the current tensor), forward the whole
static_assert(is_same<get_raw_type_t<ValTypeA>, typename remove_cvref_t<ATensor>::value_type>::value, "Expecting ValTypeA type");
static_assert(is_same<ValTypeA, typename remove_cvref_t<ATensor>::value_type>::value
, "Expecting ValTypeA type");
return make_tensor<FrgTypeA>(std::forward<ATensor>(atensor));
} else {
// Else, the intended FrgTypeA is a value type, construct a new tensor with a fragment layout
@@ -176,7 +177,8 @@ struct MMA_Atom<MMA_Traits<Args...>>
if constexpr (has_dereference<FrgTypeB>::value) {
// If the intended FrgTypeB is a view (of the current tensor), forward the whole
static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value, "Expecting ValTypeB type");
static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value
, "Expecting ValTypeB type");
return make_tensor<FrgTypeB>(std::forward<BTensor>(btensor));
} else {
// Else, the intended FrgTypeB is a value type, construct a new tensor with a fragment layout
@@ -224,6 +226,11 @@ struct TiledMMA : MMA_Atom
// thr_idx -> (ThrV,ThrM,ThrN,ThrK)
using TidLayout = decltype(right_inverse(ThrLayoutVMNK{}));
CUTE_HOST_DEVICE constexpr auto
get_thr_layout_vmnk() const {
return ThrLayoutVMNK{};
}
// Tile a tensor or a layout from shape
// (M,N,...)
// to shape
@@ -295,8 +302,8 @@ struct TiledMMA : MMA_Atom
thrfrg_A(ATensor&& atensor)
{
CUTE_STATIC_ASSERT_V(rank(atensor) >= Int<2>{});
CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
CUTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
//UTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(left_inverse(get<0>(PermutationsMNK{})),
@@ -353,8 +360,8 @@ struct TiledMMA : MMA_Atom
thrfrg_B(BTensor&& btensor)
{
CUTE_STATIC_ASSERT_V(rank(btensor) >= Int<2>{});
CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(left_inverse(get<1>(PermutationsMNK{})),
include/cute/atom/mma_traits_sm90_gmma.hpp
+40 -34
View File
@@ -117,16 +117,22 @@ using Layout_SW128_Atom = typename conditional<tnsp == GMMA::Major::MN,
Layout_K_SW128_Atom<Type>>::type;
//
// Tensor to LayoutType utility
// Tensor (position-dependent swizzle) to LayoutType utility
//
// smem_ptr_swizzle LayoutType
template <int B, int M, int S, class Shape, class Stride>
template <class Engine, class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
LayoutType
layout_type(Tensor<ViewEngine<smem_ptr_swizzle<const uint128_t, Swizzle<B,M,S>>>,
Layout<Shape,Stride>> const&)
layout_type(Tensor<Engine, Layout<Shape,Stride>> const&)
{
static_assert(is_same<uint128_t, typename Engine::value_type>::value,
"Expected uint128_t type in LayoutType conversion.");
using Swizzle = get_swizzle_t<Engine>;
constexpr int B = Swizzle::num_bits;
constexpr int M = Swizzle::num_base;
constexpr int S = Swizzle::num_shft;
static_assert(M == 4, "Unsupported layout swizzle");
static_assert(0 <= B && B <= 3, "Unsupported layout swizzle");
static_assert(S == 3, "Unsupported layout swizzle");
@@ -140,16 +146,6 @@ layout_type(Tensor<ViewEngine<smem_ptr_swizzle<const uint128_t, Swizzle<B,M,S>>>
return LayoutType::INTERLEAVE; // ERROR
}
// smem_ptr non-swizzled LayoutType
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
LayoutType
layout_type(Tensor<ViewEngine<smem_ptr<const uint128_t>>,
Layout<Shape,Stride>> const&)
{
return LayoutType::INTERLEAVE;
}
///////////////////////////////////////////////////////////////////////////////
// Construction method for GMMA Descriptors
///////////////////////////////////////////////////////////////////////////////
@@ -211,7 +207,7 @@ make_gmma_desc(Tensor<TEngine,TLayout> const& tensor)
desc.bitfield.layout_type_ = uint8_t(LAYOUT_TYPE);
// Start address (4LSB not included)
uint32_t start_address = cast_smem_ptr_to_uint(u128_tensor.data().get());
uint32_t start_address = cast_smem_ptr_to_uint(raw_pointer_cast(u128_tensor.data()));
desc.bitfield.start_address_ = start_address >> 4;
constexpr uint8_t base_offset = 0;
@@ -314,57 +310,67 @@ make_gmma_desc(Tensor<TEngine,TLayout> const& tensor)
struct DescriptorIterator
{
using reference = GmmaDescriptor;
using element_type = GmmaDescriptor;
using value_type = GmmaDescriptor;
GmmaDescriptor desc_;
// Dereference returns the GmmaDescriptor
CUTE_HOST_DEVICE constexpr
GmmaDescriptor const& operator*() const { return desc_; }
reference operator*() const { return desc_; }
// Advance and return a new GmmaDescriptor
template <class Index>
CUTE_HOST_DEVICE constexpr
GmmaDescriptor operator[](Index const& i) const { return *(*this + i); }
reference operator[](Index const& i) const { return *(*this + i); }
// Return an advanced iterator
template <class Index>
CUTE_HOST_DEVICE constexpr
DescriptorIterator operator+(Index const& offset) const
{
return { GmmaDescriptor {desc_ + uint64_t(offset)} };
return { GmmaDescriptor{desc_ + uint64_t(offset)} };
}
CUTE_HOST_DEVICE friend void
print(DescriptorIterator const&) { printf("GMMA::DescriptorIterator"); }
print(DescriptorIterator) { printf("GMMA::DescriptorIterator"); }
};
template <class T>
CUTE_HOST_DEVICE constexpr
GmmaDescriptor
raw_pointer_cast(DescriptorIterator const& ptr) {
return ptr.desc_;
}
// Recast a DescriptorIterator Tensor to uint64_t, it's RegType in mma_unpack
template <class NewT>
CUTE_HOST_DEVICE constexpr
DescriptorIterator
recast_ptr(DescriptorIterator const& iter) {
static_assert(is_same<NewT, uint64_t>::value, "Can only cast GmmaDescriptorIterator to uint64_t.");
return iter; // Do nothing, it will still dereference to GmmaDescriptor and decay to uint64_t
}
// The GMMA Traits below have custom fragment type flags for their smem desc tensors.
// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired.
template <GMMA::Major>
struct smem_desc : DescriptorIterator {};
// Recast a DescriptorIterator Tensor to uint64_t, it's RegType
template <class TLayout, class NewT>
CUTE_HOST_DEVICE constexpr
auto
recast(Tensor<ViewEngine<DescriptorIterator>,TLayout> const& tensor, type_list<NewT>)
{
static_assert(is_same<NewT, uint64_t>::value, "Can only cast descriptors to uint64_t.");
return make_tensor(tensor.data(), Layout<_1,_0>{});
}
} // end namespace GMMA
// Customization point for creating a GMMA::smem_desc Tensor
template <GMMA::Major MajorMode>
struct MakeTensor<GMMA::smem_desc<MajorMode>>
{
template <class Engine, class Layout>
template <class TEngine, class TLayout>
CUTE_HOST_DEVICE constexpr auto
operator()(Tensor<Engine,Layout> const& smem_tensor)
operator()(Tensor<TEngine,TLayout> const& smem_tensor)
{
static_assert(is_smem<Engine>::value, "Expected SMEM Tensor to construct a GMMA Desc Tensor");
static_assert(is_smem<TEngine>::value, "Expected SMEM Tensor to construct a GMMA Desc Tensor");
return make_tensor(GMMA::DescriptorIterator{GMMA::make_gmma_desc<MajorMode>(tensor<0>(smem_tensor))},
recast<uint128_t const>(smem_tensor).layout());
replace<0>(recast<uint128_t const>(smem_tensor).layout(), Layout<_1,_0>{}));
}
};
include/cute/container/array.hpp
+13 -11
View File
@@ -41,13 +41,14 @@ namespace cute
template <class T, size_t N>
struct array
{
using value_type = T;
using element_type = T;
using value_type = remove_cv_t<T>;
using size_type = size_t;
using difference_type = ptrdiff_t;
using reference = value_type&;
using const_reference = const value_type&;
using pointer = value_type*;
using const_pointer = const value_type*;
using reference = element_type&;
using const_reference = const element_type&;
using pointer = element_type*;
using const_pointer = const element_type*;
using iterator = pointer;
using const_iterator = const_pointer;
@@ -190,20 +191,21 @@ struct array
}
}
value_type __elems_[N > 0 ? N : 1];
element_type __elems_[N];
};
template <class T>
struct array<T, 0>
{
using value_type = T;
using element_type = T;
using value_type = remove_cv_t<T>;
using size_type = size_t;
using difference_type = ptrdiff_t;
using reference = value_type&;
using const_reference = const value_type&;
using pointer = value_type*;
using const_pointer = const value_type*;
using reference = element_type&;
using const_reference = const element_type&;
using pointer = element_type*;
using const_pointer = const element_type*;
using const_iterator = const_pointer;
using iterator = const_iterator;
include/cute/container/array_subbyte.hpp
+169 -102
View File
@@ -39,11 +39,18 @@
#include <cute/numeric/int.hpp> // sizeof_bits
#include <cute/numeric/integral_constant.hpp>
#include <cute/container/bit_field.hpp> // dummy_type
namespace cute
{
template <class T>
struct is_subbyte {
static constexpr bool value = sizeof_bits_v<T> < 8;
};
template <class T>
constexpr bool is_subbyte_v = is_subbyte<T>::value;
//
// Underlying subbyte storage type
//
@@ -53,43 +60,44 @@ using subbyte_storage_type_t = conditional_t<(sizeof_bits_v<T> <= 8), uint8_t,
conditional_t<(sizeof_bits_v<T> <= 32), uint32_t,
conditional_t<(sizeof_bits_v<T> <= 64), uint64_t,
conditional_t<(sizeof_bits_v<T> <= 128), uint128_t,
dummy_type>>>>>;
T>>>>>;
template <class T>
struct subbyte_iterator;
template <class T> struct subbyte_iterator;
template <class, class> struct swizzle_ptr;
//
// subbyte_reference
// Proxy object for sub-byte element references
//
template <class T>
struct subbyte_reference
struct subbyte_reference
{
// Iterator Element type (const or non-const)
using element_type = T;
// Iterator Value type without type qulifier.
// Iterator Value type without type qualifier.
using value_type = remove_cv_t<T>;
// Storage type (const or non-const)
using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
static_assert(!is_same_v<storage_type, dummy_type>, "Storage type is not supported");
static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
static_assert(sizeof_bits_v<element_type> <= sizeof_bits_v<storage_type>,
"Size of Element must not be greater than Storage.");
// Number of logical elements per stored object
static constexpr uint8_t ElementsPerStoredItem = sizeof_bits_v<storage_type> / sizeof_bits_v<element_type>;
// Bitmask for covering one item
static constexpr storage_type BitMask = storage_type((storage_type(1) << sizeof_bits_v<element_type>) - 1);
private:
// Bitmask for covering one item
static constexpr storage_type BitMask = storage_type(storage_type(-1) >> (sizeof_bits_v<storage_type> - sizeof_bits_v<element_type>));
// Flag for fast branching on straddled elements
static constexpr bool is_storage_unaligned = ((sizeof_bits_v<storage_type> % sizeof_bits_v<element_type>) != 0);
friend class subbyte_iterator<T>;
// Pointer to storage element
storage_type* ptr_ = nullptr;
// Index into elements packed into storage_type element. RI: 0 <= idx_ < ElementsPerStoredItem
// Bit index of value_type starting position within storage_type element.
// RI: 0 <= idx_ < sizeof_bit<storage_type>
uint8_t idx_ = 0;
// Ctor
@@ -100,38 +108,73 @@ private:
public:
// Copy Ctor
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
subbyte_reference(subbyte_reference const& other) {
*this = element_type(other);
}
// Copy Assignment
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
subbyte_reference& operator=(subbyte_reference const& other) {
return *this = element_type(other);
}
// Dtor
~subbyte_reference() = default;
// Assignment
template<class T_=element_type>
template <class T_ = element_type>
CUTE_HOST_DEVICE constexpr
enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(element_type x) {
enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(element_type x)
{
static_assert(is_same_v<T_, element_type>, "Do not specify template arguments!");
storage_type item = (reinterpret_cast<storage_type const &>(x) & BitMask);
storage_type kUpdateMask = storage_type(~(BitMask << (idx_ * sizeof_bits_v<element_type>)));
*ptr_ = storage_type((*ptr_ & kUpdateMask) | (item << (idx_ * sizeof_bits_v<element_type>)));
storage_type item = (reinterpret_cast<storage_type const&>(x) & BitMask);
// Update the current storage element
storage_type bit_mask_0 = storage_type(BitMask << idx_);
ptr_[0] = storage_type((ptr_[0] & ~bit_mask_0) | (item << idx_));
// If value_type is unaligned with storage_type (static) and this is a straddled value (dynamic)
if (is_storage_unaligned && idx_ + sizeof_bits_v<value_type> > sizeof_bits_v<storage_type>) {
uint8_t straddle_bits = uint8_t(sizeof_bits_v<storage_type> - idx_);
storage_type bit_mask_1 = storage_type(BitMask >> straddle_bits);
// Update the next storage element
ptr_[1] = storage_type((ptr_[1] & ~bit_mask_1) | (item >> straddle_bits));
}
return *this;
}
// Comparison of referenced values
CUTE_HOST_DEVICE constexpr friend
bool operator==(subbyte_reference const& x, subbyte_reference const& y) { return x.get() == y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() != y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator< (subbyte_reference const& x, subbyte_reference const& y) { return x.get() < y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator> (subbyte_reference const& x, subbyte_reference const& y) { return x.get() > y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator<=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() <= y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator>=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() >= y.get(); }
// Value
CUTE_HOST_DEVICE
element_type get() const {
element_type get() const
{
if constexpr (is_same_v<bool, value_type>) { // Extract to bool -- potentially faster impl
return bool((*ptr_) & (BitMask << (idx_ * sizeof_bits_v<element_type>)));
return bool((*ptr_) & (BitMask << idx_));
} else { // Extract to element_type
storage_type item = storage_type((*ptr_ >> (idx_ * sizeof_bits_v<element_type>)) & BitMask);
return reinterpret_cast<element_type &>(item);
// Extract from the current storage element
auto item = storage_type((ptr_[0] >> idx_) & BitMask);
// If value_type is unaligned with storage_type (static) and this is a straddled value (dynamic)
if (is_storage_unaligned && idx_ + sizeof_bits_v<value_type> > sizeof_bits_v<storage_type>) {
uint8_t straddle_bits = uint8_t(sizeof_bits_v<storage_type> - idx_);
storage_type bit_mask_1 = storage_type(BitMask >> straddle_bits);
// Extract from the next storage element
item |= storage_type((ptr_[1] & bit_mask_1) << straddle_bits);
}
return reinterpret_cast<element_type&>(item);
}
}
@@ -142,63 +185,77 @@ public:
}
};
//
// subbyte_iterator
// Random-access iterator over subbyte references
//
template <class T>
struct subbyte_iterator
struct subbyte_iterator
{
// Iterator Element type (const or non-const)
using element_type = T;
// Iterator Value type without type qulifier.
// Iterator Value type without type qualifier.
using value_type = remove_cv_t<T>;
// Storage type (const or non-const)
using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
// Reference proxy type
using reference = subbyte_reference<element_type>;
static_assert(!is_same_v<storage_type, dummy_type>, "Storage type is not supported");
static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
static_assert(sizeof_bits_v<element_type> <= sizeof_bits_v<storage_type>,
"Size of Element must not be greater than Storage.");
// Number of logical elements per stored object
static constexpr uint8_t ElementsPerStoredItem = sizeof_bits_v<storage_type> / sizeof_bits_v<element_type>;
private:
template <class, class> friend class swizzle_ptr;
// Pointer to storage element
storage_type* ptr_ = nullptr;
// Index into elements packed into storage_type element. RI: 0 <= idx_ < ElementsPerStoredItem
// Bit index of value_type starting position within storage_type element.
// RI: 0 <= idx_ < sizeof_bit<storage_type>
uint8_t idx_ = 0;
public:
// Ctor
subbyte_iterator() = default;
// Ctor
template <class PointerType>
CUTE_HOST_DEVICE constexpr
subbyte_iterator(PointerType* ptr, uint8_t idx = 0): ptr_(reinterpret_cast<storage_type*>(ptr)), idx_(idx) { }
subbyte_iterator(PointerType* ptr, uint8_t idx = 0) : ptr_(reinterpret_cast<storage_type*>(ptr)), idx_(idx) { }
subbyte_iterator() = default;
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator++() {
++idx_;
if (idx_ == ElementsPerStoredItem) {
++ptr_;
idx_ = 0;
}
reference operator*() const {
return reference(ptr_, idx_);
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator+=(uint64_t k) {
k = sizeof_bits_v<value_type> * k + idx_;
ptr_ += k / sizeof_bits_v<storage_type>;
idx_ = k % sizeof_bits_v<storage_type>;
return *this;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator--() {
if (idx_) {
--idx_;
} else {
--ptr_;
idx_ = ElementsPerStoredItem - 1;
subbyte_iterator operator+(uint64_t k) const {
return subbyte_iterator(ptr_, idx_) += k;
}
CUTE_HOST_DEVICE constexpr
reference operator[](uint64_t k) const {
return *(*this + k);
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator++() {
idx_ += sizeof_bits_v<value_type>;
if (idx_ >= sizeof_bits_v<storage_type>) {
++ptr_;
idx_ -= sizeof_bits_v<storage_type>;
}
return *this;
}
@@ -210,6 +267,17 @@ public:
return ret;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator--() {
if (idx_ >= sizeof_bits_v<value_type>) {
idx_ -= sizeof_bits_v<value_type>;
} else {
--ptr_;
idx_ += sizeof_bits_v<storage_type> - sizeof_bits_v<value_type>;
}
return *this;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator operator--(int) {
subbyte_iterator ret(*this);
@@ -217,37 +285,45 @@ public:
return ret;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator+=(uint64_t k) {
k += idx_;
ptr_ += k / ElementsPerStoredItem;
idx_ = k % ElementsPerStoredItem;
return *this;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator operator+(uint64_t k) const {
return subbyte_iterator(ptr_,idx_) += k;
}
CUTE_HOST_DEVICE constexpr
reference operator*() const {
return reference(ptr_, idx_);
}
CUTE_HOST_DEVICE constexpr
reference operator[](uint64_t k) const {
return *(*this + k);
}
CUTE_HOST_DEVICE constexpr
friend bool operator==(subbyte_iterator const& x, subbyte_iterator const& y) {
CUTE_HOST_DEVICE constexpr friend
bool operator==(subbyte_iterator const& x, subbyte_iterator const& y) {
return x.ptr_ == y.ptr_ && x.idx_ == y.idx_;
}
CUTE_HOST_DEVICE constexpr friend
bool operator< (subbyte_iterator const& x, subbyte_iterator const& y) {
return x.ptr_ < y.ptr_ || (x.ptr_ == y.ptr_ && x.idx_ < y.idx_);
}
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
CUTE_HOST_DEVICE constexpr friend
bool operator<=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator> (subbyte_iterator const& x, subbyte_iterator const& y) { return (y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator>=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x < y); }
CUTE_HOST_DEVICE constexpr
friend bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) {
return !(x == y);
// Conversion to raw pointer with loss of subbyte index
CUTE_HOST_DEVICE constexpr friend
T* raw_pointer_cast(subbyte_iterator const& x) {
assert(x.idx_ == 0);
return reinterpret_cast<T*>(x.ptr_);
}
// Conversion to NewT_ with possible loss of subbyte index
template <class NewT_>
CUTE_HOST_DEVICE constexpr friend
auto recast_ptr(subbyte_iterator const& x) {
using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
if constexpr (is_subbyte<NewT>::value) { // Making subbyte_iter, preserve the subbyte idx
return subbyte_iterator<NewT>(x.ptr_, x.idx_);
} else { // Not subbyte, assume/assert subbyte idx 0
return reinterpret_cast<NewT*>(raw_pointer_cast(x));
}
CUTE_GCC_UNREACHABLE;
}
CUTE_HOST_DEVICE friend void print(subbyte_iterator x) {
printf("subptr[%db](%p.%u)", int(sizeof_bits<T>::value), x.ptr_, x.idx_);
}
};
@@ -281,26 +357,20 @@ struct array_subbyte
// Storage type (const or non-const)
using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
static_assert(!is_same_v<storage_type, dummy_type>, "Storage type is not supported");
// Number of logical elements per stored object
static constexpr uint8_t ElementsPerStoredItem = sizeof_bits_v<storage_type> / sizeof_bits_v<T>;
// Bitmask for covering one item
static constexpr storage_type BitMask = ((storage_type(1) << sizeof_bits<T>::value) - 1);
// Number of storage elements
static constexpr size_type StorageElements = (N + ElementsPerStoredItem - 1) / ElementsPerStoredItem;
static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
private:
// Number of storage elements, ceil_div
static constexpr size_type StorageElements = (N * sizeof_bits_v<value_type> + sizeof_bits_v<storage_type> - 1) / sizeof_bits_v<storage_type>;
// Internal storage
storage_type storage[StorageElements];
public:
CUTE_HOST_DEVICE constexpr
array_subbyte() { }
array_subbyte() {}
CUTE_HOST_DEVICE constexpr
array_subbyte(array_subbyte const& x) {
@@ -334,20 +404,11 @@ public:
}
}
// Efficient fill method
CUTE_HOST_DEVICE constexpr
void fill(T const& value) {
storage_type item = (reinterpret_cast<storage_type const&>(value) & BitMask);
// Reproduce the value over the bits of the storage item
CUTE_UNROLL
for (size_type s = sizeof_bits_v<T>; s < sizeof_bits_v<storage_type>; s *= 2) {
item |= item << s;
}
CUTE_UNROLL
for (size_type i = 0; i < StorageElements; ++i) {
storage[i] = item;
for (size_type i = 0; i < N; ++i) {
at(i) = value;
}
}
@@ -428,12 +489,12 @@ public:
CUTE_HOST_DEVICE constexpr
iterator end() {
return iterator(storage + N / ElementsPerStoredItem, N % ElementsPerStoredItem);
return iterator(storage) + N;
}
CUTE_HOST_DEVICE constexpr
const_iterator end() const {
return const_iterator(storage + N / ElementsPerStoredItem, N % ElementsPerStoredItem);
return const_iterator(storage) + N;
}
CUTE_HOST_DEVICE constexpr
@@ -509,6 +570,12 @@ T&& get(array_subbyte<T,N>&& a)
namespace CUTE_STL_NAMESPACE
{
template <class T>
struct is_reference<cute::subbyte_reference<T>>
: CUTE_STL_NAMESPACE::true_type
{};
template <class T, size_t N>
struct tuple_size<cute::array_subbyte<T,N>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, N>
include/cute/container/bit_field.hpp
+3 -9
View File
@@ -72,16 +72,10 @@ struct bit_field
// Number of bits in data_[idx] used for NumBits if straddling, else 0
static constexpr uint32_t bit_hi = (idx + 1 < N) ? (storage_type_bits - bit_lo) : 0;
private:
// MSVC issues warning C4293 ("shift count negative or too big, undefined behavior")
// if we use NumBits directly in the shift expression, even if the shift occurs
// in the branch of a ternary expression where NumBits is known to be less than
// the number of bits of the value being shifted.
static constexpr uint32_t MollifiedNumBits = NumBits > 63u ? 63u : NumBits;
public:
// NumBits mask
static constexpr value_type mask = (NumBits < 64u) ? ((uint64_t(1) << MollifiedNumBits) - 1) : uint64_t(-1);
static constexpr value_type mask = value_type(uint64_t(-1) >> (64u - NumBits));
// NumBits mask for BitStart
static constexpr storage_type mask_lo = storage_type(mask) << bit_lo;
// NumBits mask for leftover bits in data_[idx+1] if straddling, else 0
@@ -93,7 +87,7 @@ public:
CUTE_HOST_DEVICE constexpr
value_type get() const {
storage_type result = (data_[idx] & mask_lo) >> bit_lo;
if constexpr (bit_hi) {
if constexpr (bit_hi != 0) {
result |= (data_[idx+1] & mask_hi) << bit_hi;
}
return static_cast<value_type>(result);
@@ -104,7 +98,7 @@ public:
void set(value_type x) {
storage_type item = static_cast<storage_type>(x & mask);
data_[idx] = static_cast<storage_type>((data_[idx] & ~mask_lo) | (item << bit_lo));
if constexpr (bit_hi) {
if constexpr (bit_hi != 0) {
data_[idx+1] = static_cast<storage_type>((data_[idx+1] & ~mask_hi) | (item >> bit_hi));
}
}
include/cute/int_tuple.hpp
+26 -16
View File
@@ -40,7 +40,7 @@
/** IntTuple is an integer or a tuple of IntTuples.
* This file holds utilities for working with IntTuples,
* but does not hold a concrete concept or class of IntTuple.
*/
*/
namespace cute
{
@@ -49,7 +49,7 @@ namespace cute
// Even though is_tuple<Integral> is false and tuple_size<Integral> doesn't compile,
// CuTe defines rank(Integral) as 1, so it's useful for get<0>(Integral) to return its input
template <size_t I, class T, __CUTE_REQUIRES(cute::is_integral<cute::remove_cvref_t<T>>::value)>
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(T&& t) noexcept
{
@@ -59,7 +59,7 @@ get(T&& t) noexcept
// Custom recursive get for anything that implements get<I>(.) (for a single integer I).
template <size_t I0, size_t I1, size_t... Is, class T>
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(T&& t) noexcept
{
@@ -218,19 +218,29 @@ static constexpr int depth_v = depth_t<Tuple>::value;
// product
//
template <class IntTuple>
CUTE_HOST_DEVICE constexpr
auto
product(IntTuple const& a)
// Implementation of product (see below) as a function object
struct Product
{
if constexpr (is_tuple<IntTuple>::value) {
return cute::apply(a, [](auto const&... v){ return (Int<1>{} * ... * product(v)); });
} else {
return a;
}
template <class IntTuple>
CUTE_HOST_DEVICE constexpr
auto
operator()(IntTuple const& a) const
{
if constexpr (is_tuple<IntTuple>::value) {
if constexpr (tuple_size<IntTuple>::value == 0) {
return Int<1>{};
} else {
return cute::transform_apply(a, Product{}, multiplies_unary_lfold{});
}
} else {
return a;
}
CUTE_GCC_UNREACHABLE;
}
CUTE_GCC_UNREACHABLE;
}
};
// Callable product function object
CUTE_INLINE_CONSTANT Product product;
// Return a rank(t) tuple @a result such that get<i>(@a result) = product(get<i>(@a t))
template <class Tuple>
@@ -259,7 +269,7 @@ size(IntTuple const& a)
if constexpr (sizeof...(Is) == 0) {
return product(a);
} else {
return product(get<Is...>(a));
return size(get<Is...>(a));
}
CUTE_GCC_UNREACHABLE;
@@ -361,7 +371,7 @@ shape_div(IntTupleA const& a, IntTupleB const& b)
if constexpr (is_static<IntTupleA>::value && is_static<IntTupleB>::value) {
static_assert(IntTupleA::value % IntTupleB::value == 0 || IntTupleB::value % IntTupleA::value == 0, "Static shape_div failure");
return C<shape_div(IntTupleA::value, IntTupleB::value)>{};
} else { // int int
} else { // int int
//assert(a % b == 0 || b % a == 0); // Wave dynamic assertion
return a / b != 0 ? a / b : signum(a) * signum(b); // Division with rounding away from zero
}
include/cute/layout.hpp
+33 -30
View File
@@ -1034,7 +1034,7 @@ complement(Shape const& shape, Stride const& stride, CoSizeHi const& cosize_hi)
// Should just be a sort and a fold...
// Then we could even handle dynamic strides (but they would destroy all static strides)
auto [shape_, stride_, result_shape_, result_stride] =
auto [shape_, stride_, result_shape_, result_stride] =
fold(make_seq<R-1>{},
cute::make_tuple(shape, stride, cute::make_tuple(), cute::make_tuple(Int<1>{})),
[](auto const& init, auto i)
@@ -1094,7 +1094,7 @@ CUTE_HOST_DEVICE constexpr
auto
inverse_seq(Shape const& shape, Stride const& stride, seq<Is...>)
{
auto next_I = find_if(stride, [](auto a) { return is_constant<NextStride, decltype(a)>{}; });
auto next_I = cute::find_if(stride, [](auto a) { return is_constant<NextStride, decltype(a)>{}; });
if constexpr (next_I == decltype(rank(stride))::value) {
return seq<Is...>{};
@@ -1197,22 +1197,16 @@ auto
max_common_layout(Layout<ShapeA,StrideA> const& a,
Layout<ShapeB,StrideB> const& b)
{
if constexpr (is_static<ShapeA>::value && is_static<StrideA>::value &&
is_static<ShapeB>::value && is_static<StrideB>::value)
{
Layout inv_b = right_inverse(b);
Layout common = coalesce(composition(a, inv_b));
Layout inv_b = right_inverse(b);
Layout common = coalesce(composition(a, inv_b));
if constexpr (is_constant<1, decltype(stride<0>(common))>::value) {
// Truncate to the size of the contiguous vector (static stride-1 mode)
return composition(inv_b, layout<0>(common));
} else {
return Layout<_1,_0>{};
}
// NOTE: If one of the layouts is dynamic, we can't prove alignment+vectorization is valid
// We assume dynamic shapes/strides obey alignment requirements (i.e. are large and multiples of the vector)
if constexpr (is_static<decltype(shape<0>(common))>::value &&
is_constant<1, decltype(stride<0>(common))>::value) {
// Truncate to the size of the contiguous vector (static stride-1 mode)
return composition(inv_b, layout<0>(common));
} else {
// CASE: One of the layouts is dynamic, can't prove alignment+vectorization is valid
// NOTE: Could weaken if we assume dynamic shapes/strides obey alignment requirements
// (i.e. are large and multiples of the vector)
return Layout<_1,_0>{};
}
}
@@ -1231,21 +1225,15 @@ auto
max_common_vector(Layout<ShapeA,StrideA> const& a,
Layout<ShapeB,StrideB> const& b)
{
if constexpr (is_static<ShapeA>::value && is_static<StrideA>::value &&
is_static<ShapeB>::value && is_static<StrideB>::value)
{
Layout common = coalesce(composition(a, right_inverse(b)));
Layout common = coalesce(composition(a, right_inverse(b)));
if constexpr (is_constant<1, decltype(stride<0>(common))>::value) {
// Truncate to the size of the contiguous vector (static stride-1 mode)
return shape<0>(common);
} else {
return Int<1>{};
}
// NOTE: If one of the layouts is dynamic, we can't prove alignment+vectorization is valid
// We assume dynamic shapes/strides obey alignment requirements (i.e. are large and multiples of the vector)
if constexpr (is_static<decltype(shape<0>(common))>::value &&
is_constant<1, decltype(stride<0>(common))>::value) {
// Truncate to the size of the contiguous vector (static stride-1 mode)
return shape<0>(common);
} else {
// CASE: One of the layouts is dynamic, can't prove alignment+vectorization is valid
// NOTE: Could weaken if we assume dynamic shapes/strides obey alignment requirements
// (i.e. are large and multiples of the vector)
return Int<1>{};
}
@@ -1412,6 +1400,21 @@ tiled_divide(Layout<LShape,LStride> const& layout,
return div(_, repeat<R>(_));
}
// Same as zipped_divide, but unpacks both modes: (BLK_A,BLK_B,...,a,b,...,x,y)
template <class LShape, class LStride,
class Tile>
CUTE_HOST_DEVICE constexpr
auto
flat_divide(Layout<LShape,LStride> const& layout,
Tile const& tile)
{
auto div = zipped_divide(layout, tile);
auto R0 = rank<0>(div);
auto R1 = rank<1>(div);
return div(repeat<R0>(_), repeat<R1>(_));
}
//
// Logical product
//
@@ -1606,7 +1609,7 @@ template <class OldType, class NewType,
class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
recast(Layout<Shape,Stride> const& layout)
recast_layout(Layout<Shape,Stride> const& layout)
{
if constexpr (sizeof_bits<NewType>::value == sizeof_bits<OldType>::value) {
return layout;
include/cute/layout_composed.hpp
+1 -1
View File
@@ -573,7 +573,7 @@ template <class OldType, class NewType,
class A, class O, class B>
CUTE_HOST_DEVICE constexpr
auto
recast(ComposedLayout<A,O,B> const& layout)
recast_layout(ComposedLayout<A,O,B> const& layout)
{
if constexpr (sizeof(NewType) == sizeof(OldType)) {
return layout;
include/cute/numeric/arithmetic_tuple.hpp
+49 -36
View File
@@ -126,24 +126,18 @@ operator+(tuple<T...> const& t, ArithmeticTuple<U...> const& u) {
template <auto t, class... U>
CUTE_HOST_DEVICE constexpr
auto
ArithmeticTuple<U...> const&
operator+(C<t>, ArithmeticTuple<U...> const& u) {
if constexpr (t == 0) {
return u;
} else {
static_assert(t == 0, "Artihmetic tuple op+ error!");
}
static_assert(t == 0, "Artihmetic tuple op+ error!");
return u;
}
template <class... T, auto u>
CUTE_HOST_DEVICE constexpr
auto
ArithmeticTuple<T...> const&
operator+(ArithmeticTuple<T...> const& t, C<u>) {
if constexpr (u == 0) {
return t;
} else {
static_assert(u == 0, "Artihmetic tuple op+ error!");
}
static_assert(u == 0, "Artihmetic tuple op+ error!");
return t;
}
//
@@ -153,30 +147,41 @@ operator+(ArithmeticTuple<T...> const& t, C<u>) {
template <class ArithTuple>
struct ArithmeticTupleIterator
{
using value_type = ArithTuple;
using element_type = ArithTuple;
using reference = ArithTuple;
ArithTuple coord_;
CUTE_HOST_DEVICE constexpr
ArithmeticTupleIterator() : coord_() {}
CUTE_HOST_DEVICE constexpr
ArithmeticTupleIterator(ArithTuple const& coord) : coord_(coord) {}
ArithmeticTupleIterator(ArithTuple const& coord = {}) : coord_(coord) {}
CUTE_HOST_DEVICE constexpr
ArithTuple const& operator*() const { return coord_; }
template <class Coord>
CUTE_HOST_DEVICE constexpr
auto operator+(Coord const& c) const {
return ArithmeticTupleIterator<decltype(coord_ + c)>(coord_ + c);
}
auto operator[](Coord const& c) const { return *(*this + c); }
template <class Coord>
CUTE_HOST_DEVICE constexpr
auto operator[](Coord const& c) const { return *(*this + c); }
auto operator+(Coord const& c) const {
return ArithmeticTupleIterator<decltype(coord_ + c)>(coord_ + c);
}
};
template <class ArithTuple>
CUTE_HOST_DEVICE void print(ArithmeticTupleIterator<ArithTuple> const& iter) {
printf("ArithTuple"); print(iter.coord_);
template <class Tuple>
CUTE_HOST_DEVICE constexpr
auto
make_inttuple_iter(Tuple const& t) {
return ArithmeticTupleIterator(as_arithmetic_tuple(t));
}
template <class T0, class T1, class... Ts>
CUTE_HOST_DEVICE constexpr
auto
make_inttuple_iter(T0 const& t0, T1 const& t1, Ts const&... ts) {
return make_tuple_iter(cute::make_tuple(t0, t1, ts...));
}
//
@@ -211,7 +216,7 @@ struct is_integral<ScaledBasis<T,N>> : true_type {};
// Get the scalar T out of a ScaledBasis
template <class SB>
CUTE_HOST_DEVICE constexpr auto
basis_value(SB const& e)
basis_value(SB const& e)
{
if constexpr (is_scaled_basis<SB>::value) {
return basis_value(e.value());
@@ -224,7 +229,7 @@ basis_value(SB const& e)
// Apply the N... pack to another Tuple
template <class SB, class Tuple>
CUTE_HOST_DEVICE constexpr auto
basis_get(SB const& e, Tuple const& t)
basis_get(SB const& e, Tuple const& t)
{
if constexpr (is_scaled_basis<SB>::value) {
return basis_get(e.value(), get<SB::mode()>(t));
@@ -448,36 +453,44 @@ template <auto t, class U, int M>
CUTE_HOST_DEVICE constexpr
auto
operator+(C<t>, ScaledBasis<U,M> const& u) {
if constexpr (t == 0) {
return u;
} else {
static_assert(t == 0, "ScaledBasis op+ error!");
}
static_assert(t == 0, "ScaledBasis op+ error!");
return u;
}
template <class T, int N, auto u>
CUTE_HOST_DEVICE constexpr
auto
operator+(ScaledBasis<T,N> const& t, C<u>) {
if constexpr (u == 0) {
return t;
} else {
static_assert(u == 0, "ScaledBasis op+ error!");
}
static_assert(u == 0, "ScaledBasis op+ error!");
return t;
}
//
// Display utilities
//
template <class ArithTuple>
CUTE_HOST_DEVICE void print(ArithmeticTupleIterator<ArithTuple> const& iter)
{
printf("ArithTuple"); print(iter.coord_);
}
template <class T, int N>
CUTE_HOST_DEVICE void print(ScaledBasis<T,N> const& e) {
CUTE_HOST_DEVICE void print(ScaledBasis<T,N> const& e)
{
print(e.value()); printf("@%d", N);
}
#if !defined(__CUDACC_RTC__)
template <class ArithTuple>
CUTE_HOST std::ostream& operator<<(std::ostream& os, ArithmeticTupleIterator<ArithTuple> const& iter)
{
return os << "ArithTuple" << iter.coord_;
}
template <class T, int N>
CUTE_HOST std::ostream& operator<<(std::ostream& os, ScaledBasis<T,N> const& e) {
CUTE_HOST std::ostream& operator<<(std::ostream& os, ScaledBasis<T,N> const& e)
{
return os << e.value() << "@" << N;
}
#endif
include/cute/numeric/int.hpp
+5
View File
@@ -140,6 +140,11 @@ struct sizeof_bits<integer_subbyte<Bits,Signed>> {
static constexpr size_t value = Bits;
};
template <int Bits, bool Signed>
struct sizeof_bits<cutlass::integer_subbyte<Bits,Signed>> {
static constexpr size_t value = Bits;
};
template <class T>
static constexpr int sizeof_bits_v = sizeof_bits<T>::value;
include/cute/numeric/integer_subbyte.hpp
+2
View File
@@ -36,6 +36,8 @@
#include <cstdint>
#endif
#include <cutlass/integer_subbyte.h>
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
include/cute/numeric/integral_constant.hpp
+1 -1
View File
@@ -85,7 +85,7 @@ struct is_integral<integral_constant<T,v>> : true_type {};
// is_static detects if an (abstract) value is defined completely by it's type (no members)
template <class T>
struct is_static : bool_constant<is_empty<T>::value> {};
struct is_static : bool_constant<is_empty<remove_cvref_t<T>>::value> {};
template <class T>
constexpr bool is_static_v = is_static<T>::value;
include/cute/numeric/integral_ratio.hpp
+19 -15
View File
@@ -40,12 +40,16 @@ namespace cute
{
/** Compile-time rational arithmetic type.
* Like cute::C for std::integral_constant, cute::R for std::ratio has a short name
* Like cute::C for std::integral_constant, cute::R for std::ratio has a short name
* for error messages and compile times.
* The static data members @a num and @a den represent the reduced numerator and denominator
* of the rational value. Thus, two cute::R types with different @a n or @a d are distinct types
* even if they represent the same rational value. A cute::R exposes the reduced canonical type
* via its type member. That is, cute::R<3,6>::type is cute::R<1,2> and cute::R<6,3>::type is cute::C<2>
* of the rational value. Thus, two cute::R types with different @a n or @a d are distinct types
* even if they represent the same rational value.
* A cute::R exposes the reduced canonical type via its ::type member.
* That is, cute::R<3,6>::type is cute::R<1,2> and cute::R<6,3>::type is cute::C<2>.
* A cute::R<n,d>::value can be used much like any other trait::value. It can be involved in
* arithmetic expressions (according to the operator-overloads for cute::C and cute::R,
* though these may be incomplete) but with a potential rational value rather than an integral value.
*/
template <auto n, auto d>
class R {
@@ -53,7 +57,7 @@ class R {
static constexpr auto an = abs(n);
static constexpr auto ad = abs(d);
static constexpr auto g = gcd(an, ad);
public:
static constexpr auto num = signum(n) * signum(d) * an / g;
static constexpr auto den = ad / g;
@@ -63,28 +67,28 @@ class R {
template <auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<a,b>::type
typename R<a,b>::type
ratio(C<a>, C<b>) {
return {};
}
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
typename R<a*x,b*y>::type
typename R<a*x,b*y>::type
operator*(R<a,b>, R<x,y>) {
return {};
}
template <auto a, auto b, auto c>
CUTE_HOST_DEVICE constexpr
typename R<a*c,b>::type
typename R<a*c,b>::type
operator*(R<a,b>, C<c>) {
return {};
}
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<a*c,b>::type
typename R<a*c,b>::type
operator*(C<c>, R<a,b>) {
return {};
}
@@ -109,28 +113,28 @@ operator*(R<a,b>, C const& c) {
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
typename R<a*y+b*x, b*y>::type
typename R<a*y+b*x, b*y>::type
operator+(R<a,b>, R<x,y>) {
return {};
}
template <auto a, auto b, auto c>
CUTE_HOST_DEVICE constexpr
typename R<a+c*b,b>::type
typename R<a+c*b,b>::type
operator+(R<a,b>, C<c>) {
return {};
}
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<a+c*b,b>::type
typename R<a+c*b,b>::type
operator+(C<c>, R<a,b>) {
return {};
}
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
bool_constant<R<a,b>::num == R<x,y>::num && R<a,b>::den == R<x,y>::den>
bool_constant<R<a,b>::num == R<x,y>::num && R<a,b>::den == R<x,y>::den>
operator==(R<a,b>, R<x,y>) {
return {};
}
@@ -144,14 +148,14 @@ operator==(R<a,b>, C<c>) {
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
bool_constant<R<a,b>::num == c && R<a,b>::den == 1>
bool_constant<R<a,b>::num == c && R<a,b>::den == 1>
operator==(C<c>, R<a,b>) {
return {};
}
template <auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<abs(a),abs(b)>::type
typename R<abs(a),abs(b)>::type
abs(R<a,b>) {
return {};
}
include/cute/numeric/math.hpp
+4 -2
View File
@@ -130,6 +130,8 @@ has_single_bit(T x) {
}
// Smallest number of bits needed to represent the given value
// For x == 0, this is 0
// For x != 0, this is 1 + floor(log2(x))
// bit_width( 0b0000 ) = 0
// bit_width( 0b0001 ) = 1
// bit_width( 0b0010 ) = 2
@@ -203,7 +205,7 @@ CUTE_HOST_DEVICE constexpr
T
rotl(T x, int s) {
constexpr int N = numeric_limits<T>::digits;
return s == 0 ? x : s > 0 ? (x << s) | (x >> (N - s)) : rotr(x, -s);
return static_cast<T>(s == 0 ? x : s > 0 ? (x << s) | (x >> (N - s)) : rotr(x, -s));
}
// Computes the result of circular bitwise right-rotation
@@ -212,7 +214,7 @@ CUTE_HOST_DEVICE constexpr
T
rotr(T x, int s) {
constexpr int N = numeric_limits<T>::digits;
return s == 0 ? x : s > 0 ? (x >> s) | (x << (N - s)) : rotl(x, -s);
return static_cast<T>(s == 0 ? x : s > 0 ? (x >> s) | (x << (N - s)) : rotl(x, -s));
}
// Counts the number of consecutive 0 bits, starting from the most significant bit
include/cute/pointer.hpp
+158 -240
View File
@@ -33,308 +33,232 @@
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/numeric/int.hpp> // sizeof_bits
#include <cute/numeric/math.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/container/array_subbyte.hpp>
#include <cute/pointer_base.hpp>
#include <cute/pointer_swizzle.hpp>
namespace cute
{
//
// has_dereference to determine if a type is a pointer concept
// recast_ptr<T> -- Create an iterator over values of type T.
// For most types this will simply be T*, but certain types require more care.
// Subbyte Types: uint2_t, uint4_t, etc
// Requires construction of a subbyte_iterator<T> in order to properly
// resolve each element in byte-addressed memory.
//
template <class T, class = void>
struct has_dereference : false_type {
};
template <class T>
struct has_dereference<T, void_t<decltype(*declval<T>())>> : true_type {
};
template <class T>
template <class NewT>
CUTE_HOST_DEVICE constexpr
T*
raw_pointer_cast(T* ptr) {
return ptr;
auto
recast_ptr(void* ptr)
{
if constexpr (is_subbyte<NewT>::value) {
return subbyte_iterator<NewT>(ptr);
} else {
return reinterpret_cast<NewT*>(ptr);
}
CUTE_GCC_UNREACHABLE;
}
//
// Extract the physical type from a logical elem type.
//
template <class T>
struct get_raw_type
{
using type = T;
};
template <class T>
using get_raw_type_t = typename get_raw_type<T>::type;
//
// Pointer categories
//
template <class T>
struct is_gmem : false_type {};
template <class T>
struct is_smem : false_type {};
// Anything that is not gmem or smem is rmem
template <class T>
struct is_rmem : bool_constant< not (is_gmem<T>::value || is_smem<T>::value)> {};
//
// A very simplified wrapper for pointers -- use for constructing tagged pointers
//
template <class T, class DerivedType>
struct device_ptr
{
using value_type = T;
static const uint32_t ElementsPerStoredItem = sizeof(T) * 8 / sizeof_bits_v<T>;
CUTE_HOST_DEVICE constexpr
device_ptr(T* ptr) : ptr_(ptr) {}
CUTE_HOST_DEVICE constexpr
T* get() const { return ptr_; }
CUTE_HOST_DEVICE constexpr
T& operator*() const { return *ptr_; }
template <class Index>
CUTE_HOST_DEVICE constexpr
T& operator[](Index const& i) const {
static_assert(sizeof_bits_v<T> >= 8, "Use subbyte_iterator to access the element");
return ptr_[i];
}
template <class Index>
CUTE_HOST_DEVICE constexpr
DerivedType operator+(Index const& i) const { return {ptr_ + i / ElementsPerStoredItem}; }
CUTE_HOST_DEVICE constexpr friend
ptrdiff_t operator-(device_ptr<T,DerivedType> const& a,
device_ptr<T,DerivedType> const& b) {
return a.ptr_ - b.ptr_;
}
T* ptr_;
};
template <class T, class D>
template <class NewT>
CUTE_HOST_DEVICE constexpr
T*
raw_pointer_cast(device_ptr<T,D> ptr) {
return ptr.get();
auto
recast_ptr(void const* ptr)
{
if constexpr (is_subbyte<NewT>::value) {
return subbyte_iterator<NewT const>(ptr);
} else {
return reinterpret_cast<NewT const*>(ptr);
}
CUTE_GCC_UNREACHABLE;
}
// Disambiguate nullptr
template <class NewT>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(decltype(nullptr)) { // nullptr_t
return recast_ptr<NewT>(static_cast<NewT*>(nullptr));
}
//
// gmem_ptr
//
template <class T>
struct gmem_ptr : device_ptr<T, gmem_ptr<T>> {
using device_ptr<T, gmem_ptr<T>>::device_ptr;
template <class P>
struct gmem_ptr : iter_adaptor<P, gmem_ptr<P>> {
using iter_adaptor<P, gmem_ptr<P>>::iter_adaptor;
};
template <class T>
template <class T, class = void>
struct is_gmem : false_type {};
template <class P> // Found the gmem
struct is_gmem<gmem_ptr<P>> : true_type {};
template <class P> // Recurse on ::iterator, if possible
struct is_gmem<P, void_t<typename P::iterator>> : is_gmem<typename P::iterator> {};
// Idempotent gmem tag on an iterator
template <class Iterator>
CUTE_HOST_DEVICE constexpr
gmem_ptr<T>
make_gmem_ptr(T* ptr) {
return {ptr};
auto
make_gmem_ptr(Iterator iter) {
if constexpr (is_gmem<Iterator>::value) {
return iter;
} else {
return gmem_ptr<Iterator>{iter};
}
CUTE_GCC_UNREACHABLE;
}
// Explicitly typed construction from a raw pointer
template <class T>
CUTE_HOST_DEVICE constexpr
gmem_ptr<T>
auto
make_gmem_ptr(void* ptr) {
return {reinterpret_cast<T*>(ptr)};
return make_gmem_ptr(recast_ptr<T>(ptr));
}
// Explicitly typed construction from a raw pointer
template <class T>
CUTE_HOST_DEVICE constexpr
gmem_ptr<T const>
auto
make_gmem_ptr(void const* ptr) {
return {reinterpret_cast<T const*>(ptr)};
return make_gmem_ptr(recast_ptr<T const>(ptr));
}
// nullptr_t overloads are needed because otherwise,
// make_gmem_ptr<float>(nullptr) will be ambiguous,
// as std::nullptr_t can be converted to any pointer
// or pointer to member type.
// nullptr_t overload for make_gmem_ptr<float>(nullptr) disambiguation
template <class T>
CUTE_HOST_DEVICE constexpr
gmem_ptr<T>
auto
make_gmem_ptr(decltype(nullptr)) { // nullptr_t
return {static_cast<T*>(nullptr)};
return make_gmem_ptr(recast_ptr<T>(nullptr));
}
template <class T>
struct is_gmem<gmem_ptr<T>> : true_type {};
// The gmem tag is invariant over type-recast
template <class NewT, class P>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(gmem_ptr<P> const& ptr) {
return make_gmem_ptr(recast_ptr<NewT>(ptr.get()));
}
//
// smem_ptr
//
template <class T>
struct smem_ptr : device_ptr<T, smem_ptr<T>> {
using device_ptr<T, smem_ptr<T>>::device_ptr;
template <class P>
struct smem_ptr : iter_adaptor<P, smem_ptr<P>> {
using iter_adaptor<P, smem_ptr<P>>::iter_adaptor;
};
template <class T>
template <class T, class = void>
struct is_smem : false_type {};
template <class P> // Found the smem
struct is_smem<smem_ptr<P>> : true_type {};
template <class P> // Recurse on ::iterator, if possible
struct is_smem<P, void_t<typename P::iterator>> : is_smem<typename P::iterator> {};
// Idempotent smem tag on an iterator
template <class Iterator>
CUTE_HOST_DEVICE constexpr
smem_ptr<T>
make_smem_ptr(T* ptr) {
return {ptr};
auto
make_smem_ptr(Iterator iter) {
if constexpr (is_smem<Iterator>::value) {
return iter;
} else {
return smem_ptr<Iterator>{iter};
}
CUTE_GCC_UNREACHABLE;
}
// Make a smem swizzle pointer, common operation
template <class Iterator, class Swizzle>
CUTE_HOST_DEVICE constexpr
auto
make_smem_ptr(Iterator ptr, Swizzle sw)
{
return make_swizzle_ptr(make_smem_ptr(ptr), sw);
}
// Explicitly typed construction from a raw pointer
template <class T>
CUTE_HOST_DEVICE constexpr
smem_ptr<T>
auto
make_smem_ptr(void* ptr) {
return {reinterpret_cast<T*>(ptr)};
return make_smem_ptr(recast_ptr<T>(ptr));
}
// Explicitly typed construction from a raw pointer
template <class T>
CUTE_HOST_DEVICE constexpr
smem_ptr<T const>
auto
make_smem_ptr(void const* ptr) {
return {reinterpret_cast<T const*>(ptr)};
return make_smem_ptr(recast_ptr<T const>(ptr));
}
template <class T>
struct is_smem<smem_ptr<T>> : true_type {};
// The smem tag is invariant over type-recast
template <class NewT, class P>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(smem_ptr<P> const& ptr) {
return make_smem_ptr(recast_ptr<NewT>(ptr.get()));
}
//
// rmem_ptr
//
template <class T>
struct rmem_ptr : device_ptr<T, rmem_ptr<T>> {
using device_ptr<T, rmem_ptr<T>>::device_ptr;
template <class P>
struct rmem_ptr : iter_adaptor<P, rmem_ptr<P>> {
using iter_adaptor<P, rmem_ptr<P>>::iter_adaptor;
};
template <class T>
// Anything that is not gmem or smem is rmem
template <class T, class = void>
struct is_rmem : bool_constant<not (is_gmem<T>::value || is_smem<T>::value)> {};
template <class P>
struct is_rmem<rmem_ptr<P>> : true_type {};
// Idempotent rmem tag on an iterator
template <class Iterator>
CUTE_HOST_DEVICE constexpr
rmem_ptr<T>
make_rmem_ptr(T* ptr) {
return {ptr};
auto
make_rmem_ptr(Iterator iter) {
if constexpr (is_rmem<Iterator>::value) {
return iter;
} else {
return rmem_ptr<Iterator>{iter};
}
CUTE_GCC_UNREACHABLE;
}
// Explicitly typed construction from a raw pointer
template <class T>
CUTE_HOST_DEVICE constexpr
rmem_ptr<T>
auto
make_rmem_ptr(void* ptr) {
return {reinterpret_cast<T*>(ptr)};
return make_rmem_ptr(recast_ptr<T>(ptr));
}
// Explicitly typed construction from a raw pointer
template <class T>
CUTE_HOST_DEVICE constexpr
rmem_ptr<T const>
auto
make_rmem_ptr(void const* ptr) {
return {reinterpret_cast<T const*>(ptr)};
return make_rmem_ptr(recast_ptr<T const>(ptr));
}
template <class T>
struct is_rmem<rmem_ptr<T>> : true_type {};
//
// counting iterator -- quick and dirty
//
struct counting
{
using index_type = int;
using value_type = index_type;
CUTE_HOST_DEVICE constexpr
counting() : n_(0) {}
CUTE_HOST_DEVICE constexpr
counting(index_type const& n) : n_(n) {}
CUTE_HOST_DEVICE constexpr
index_type operator[](index_type const& i) const { return n_ + i; }
CUTE_HOST_DEVICE constexpr
index_type const& operator*() const { return n_; }
CUTE_HOST_DEVICE constexpr
counting operator+(index_type const& i) const { return {n_ + i}; }
CUTE_HOST_DEVICE constexpr
counting& operator++() { ++n_; return *this; }
CUTE_HOST_DEVICE constexpr
bool operator==(counting const& other) const { return n_ == other.n_; }
CUTE_HOST_DEVICE constexpr
bool operator!=(counting const& other) const { return n_ != other.n_; }
CUTE_HOST_DEVICE constexpr
bool operator< (counting const& other) const { return n_ < other.n_; }
index_type n_;
};
//
// recast
//
template <class NewT, class T>
// The rmem tag is invariant over type-recast
template <class NewT, class P>
CUTE_HOST_DEVICE constexpr
auto
recast(T* ptr) {
return reinterpret_cast<NewT*>(ptr);
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(T const* ptr) {
return reinterpret_cast<NewT const*>(ptr);
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(gmem_ptr<T> const& ptr) {
return make_gmem_ptr(recast<NewT>(ptr.ptr_));
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(gmem_ptr<T const> const& ptr) {
return make_gmem_ptr(recast<NewT const>(ptr.ptr_));
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(smem_ptr<T> const& ptr) {
return make_smem_ptr(recast<NewT>(ptr.ptr_));
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(smem_ptr<T const> const& ptr) {
return make_smem_ptr(recast<NewT const>(ptr.ptr_));
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(rmem_ptr<T> const& ptr) {
return make_rmem_ptr(recast<NewT>(ptr.ptr_));
}
template <class NewT, class T>
CUTE_HOST_DEVICE constexpr
auto
recast(rmem_ptr<T const> const& ptr) {
return make_rmem_ptr(recast<NewT const>(ptr.ptr_));
recast_ptr(rmem_ptr<P> const& ptr) {
return make_rmem_ptr(recast_ptr<NewT>(ptr.get()));
}
//
@@ -342,46 +266,40 @@ recast(rmem_ptr<T const> const& ptr) {
//
template <class T>
CUTE_HOST_DEVICE void print(T const* const ptr)
CUTE_HOST_DEVICE void print(gmem_ptr<T> ptr)
{
printf("raw_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr);
printf("gmem_"); print(ptr.get());
}
template <class T>
CUTE_HOST_DEVICE void print(gmem_ptr<T> const& ptr)
CUTE_HOST_DEVICE void print(smem_ptr<T> ptr)
{
printf("gmem_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr.get());
printf("smem_"); print(ptr.get());
}
template <class T>
CUTE_HOST_DEVICE void print(smem_ptr<T> const& ptr)
CUTE_HOST_DEVICE void print(rmem_ptr<T> ptr)
{
printf("smem_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr.get());
}
template <class T>
CUTE_HOST_DEVICE void print(rmem_ptr<T> const& ptr)
{
printf("rmem_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr.get());
printf("rmem_"); print(ptr.get());
}
#if !defined(__CUDACC_RTC__)
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, gmem_ptr<T> const& ptr)
CUTE_HOST std::ostream& operator<<(std::ostream& os, gmem_ptr<T> ptr)
{
return os << "gmem_ptr_" << int(sizeof_bits<T>::value) << "b";
return os << "gmem_[" << int(sizeof_bits<iter_value_t<T>>::value) << "b]";
}
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr<T> const& ptr)
CUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr<T> ptr)
{
return os << "smem_ptr_" << int(sizeof_bits<T>::value) << "b";
return os << "smem_[" << int(sizeof_bits<iter_value_t<T>>::value) << "b]";
}
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, rmem_ptr<T> const& ptr)
CUTE_HOST std::ostream& operator<<(std::ostream& os, rmem_ptr<T> ptr)
{
return os << "rmem_ptr_" << int(sizeof_bits<T>::value) << "b";
return os << "rmem_[" << int(sizeof_bits<iter_value_t<T>>::value) << "b]";
}
#endif // !defined(__CUDACC_RTC__)
include/cute/pointer_base.hpp
+247
View File
@@ -0,0 +1,247 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/int.hpp> // sizeof_bits
namespace cute
{
//
// C++20 <iterator> iterator_traits
//
namespace detail {
// Default reference type of an iterator
template <class T, class = void>
struct iter_ref { using type = decltype(*declval<T&>()); };
// Prefer to propagate ::reference
template <class T>
struct iter_ref<T,void_t<typename T::reference>> { using type = typename T::reference; };
} // end namespace detail
template <class T>
using iter_reference = detail::iter_ref<T>;
template <class T>
using iter_reference_t = typename iter_reference<T>::type;
namespace detail {
// Default element_type of an iterator
template <class T, class = void>
struct iter_e { using type = remove_reference_t<typename iter_ref<T>::type>; };
// Prefer to propagate ::element_type
template <class T>
struct iter_e<T,void_t<typename T::element_type>> { using type = typename T::element_type; };
} // end namespace detail
template <class T>
using iter_element = detail::iter_e<T>;
template <class T>
using iter_element_t = typename iter_element<T>::type;
namespace detail {
// Default value_type of an iterator
template <class T, class = void>
struct iter_v { using type = remove_cv_t<typename iter_e<T>::type>; };
// Prefer to propagate ::value_type
template <class T>
struct iter_v<T,void_t<typename T::value_type>> { using type = typename T::value_type; };
} // end namespace detail
template <class T>
using iter_value = detail::iter_v<T>;
template <class T>
using iter_value_t = typename iter_value<T>::type;
template <class Iterator>
struct iterator_traits {
using reference = iter_reference_t<Iterator>;
using element_type = iter_element_t<Iterator>;
using value_type = iter_value_t<Iterator>;
};
//
// has_dereference to determine if a type is an iterator concept
//
namespace detail {
template <class T, class = void>
struct has_dereference : CUTE_STL_NAMESPACE::false_type {};
template <class T>
struct has_dereference<T, void_t<decltype(*declval<T&>())>> : CUTE_STL_NAMESPACE::true_type {};
} // end namespace detail
template <class T>
using has_dereference = detail::has_dereference<T>;
//
// raw_pointer_cast
//
template <class T>
CUTE_HOST_DEVICE constexpr
T*
raw_pointer_cast(T* ptr) {
return ptr;
}
//
// A very simplified iterator adaptor.
// Derived classed may override methods, but be careful to reproduce interfaces exactly.
// Clients should never have an instance of this class. Do not write methods that take this as a param.
//
template <class Iterator, class DerivedType>
struct iter_adaptor
{
using iterator = Iterator;
using reference = typename iterator_traits<iterator>::reference;
using element_type = typename iterator_traits<iterator>::element_type;
using value_type = typename iterator_traits<iterator>::value_type;
iterator ptr_;
CUTE_HOST_DEVICE constexpr
iter_adaptor(iterator ptr = {}) : ptr_(ptr) {}
CUTE_HOST_DEVICE constexpr
reference operator*() const { return *ptr_; }
template <class Index>
CUTE_HOST_DEVICE constexpr
reference operator[](Index const& i) const { return ptr_[i]; }
template <class Index>
CUTE_HOST_DEVICE constexpr
DerivedType operator+(Index const& i) const { return {ptr_ + i}; }
CUTE_HOST_DEVICE constexpr
iterator get() const { return ptr_; }
CUTE_HOST_DEVICE constexpr
friend bool operator==(DerivedType const& x, DerivedType const& y) { return x.ptr_ == y.ptr_; }
CUTE_HOST_DEVICE constexpr
friend bool operator!=(DerivedType const& x, DerivedType const& y) { return x.ptr_ != y.ptr_; }
CUTE_HOST_DEVICE constexpr
friend bool operator< (DerivedType const& x, DerivedType const& y) { return x.ptr_ < y.ptr_; }
CUTE_HOST_DEVICE constexpr
friend bool operator<=(DerivedType const& x, DerivedType const& y) { return x.ptr_ <= y.ptr_; }
CUTE_HOST_DEVICE constexpr
friend bool operator> (DerivedType const& x, DerivedType const& y) { return x.ptr_ > y.ptr_; }
CUTE_HOST_DEVICE constexpr
friend bool operator>=(DerivedType const& x, DerivedType const& y) { return x.ptr_ >= y.ptr_; }
};
template <class I, class D>
CUTE_HOST_DEVICE constexpr
auto
raw_pointer_cast(iter_adaptor<I,D> const& x) {
return raw_pointer_cast(x.ptr_);
}
//
// counting iterator -- quick and dirty
//
template <class T = int>
struct counting_iterator
{
using index_type = T;
using value_type = T;
using reference = T;
index_type n_;
CUTE_HOST_DEVICE constexpr
counting_iterator(index_type n = 0) : n_(n) {}
CUTE_HOST_DEVICE constexpr
index_type operator*() const { return n_; }
CUTE_HOST_DEVICE constexpr
index_type operator[](index_type i) const { return n_ + i; }
CUTE_HOST_DEVICE constexpr
counting_iterator operator+(index_type i) const { return {n_ + i}; }
CUTE_HOST_DEVICE constexpr
counting_iterator& operator++() { ++n_; return *this; }
CUTE_HOST_DEVICE constexpr
counting_iterator operator++(int) { counting_iterator ret = *this; ++n_; return ret; }
CUTE_HOST_DEVICE constexpr
friend bool operator==(counting_iterator const& x, counting_iterator const& y) { return x.n_ == y.n_; }
CUTE_HOST_DEVICE constexpr
friend bool operator!=(counting_iterator const& x, counting_iterator const& y) { return x.n_ != y.n_; }
CUTE_HOST_DEVICE constexpr
friend bool operator< (counting_iterator const& x, counting_iterator const& y) { return x.n_ < y.n_; }
CUTE_HOST_DEVICE constexpr
friend bool operator<=(counting_iterator const& x, counting_iterator const& y) { return x.n_ <= y.n_; }
CUTE_HOST_DEVICE constexpr
friend bool operator> (counting_iterator const& x, counting_iterator const& y) { return x.n_ > y.n_; }
CUTE_HOST_DEVICE constexpr
friend bool operator>=(counting_iterator const& x, counting_iterator const& y) { return x.n_ >= y.n_; }
};
template <class T>
CUTE_HOST_DEVICE constexpr
T
raw_pointer_cast(counting_iterator<T> const& x) {
return x.n_;
}
//
// Display utilities
//
template <class T>
CUTE_HOST_DEVICE void print(T const* const ptr)
{
printf("ptr[%db](%p)", int(sizeof_bits<T>::value), ptr);
}
template <class T>
CUTE_HOST_DEVICE void print(counting_iterator<T> ptr)
{
printf("counting_iter_"); print(ptr.n_);
}
#if !defined(__CUDACC_RTC__)
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, counting_iterator<T> ptr)
{
return os << "counting_iter_" << ptr.n_;
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace cute
include/cute/pointer_flagged.hpp
+146
View File
@@ -0,0 +1,146 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/arch/util.hpp> // cast_smem_ptr_to_uint
#include <cute/pointer.hpp>
#include <cute/pointer_swizzle.hpp>
#include <cute/swizzle_layout.hpp>
#include <cute/tensor.hpp>
namespace cute
{
//
// Stand-in Swizzle Layout
// A model of a nullptr smem_ptr<T> with B == sizeof_bits<T>::value
// That represents an unset pointer. This is a placeholder type that is waiting for an smem_ptr
//
template <int Bits>
struct smem_ptr_flag_bits : Int<0> {};
using smem_ptr_flag = smem_ptr_flag_bits<1>;
// A flagged construction method to transform ComposedLayout
// Make a swizzle pointer tensor and check that the intended type size matches
template <class Iterator, class SwizzleFn, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
make_tensor(Iterator const& ptr,
ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
{
static_assert(is_smem<Iterator>::value, "Expected smem.");
static_assert(B == sizeof_bits<iter_value_t<Iterator>>::value, "Expected a B-bit pointer type.");
return make_tensor(make_smem_ptr(ptr.get(), layout.layout_a()),
layout.layout_b());
}
// NOTE: To preserve smem_ptr_flag_bits under recast ops
template <int N, class SwizzleFn, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
upcast(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
{
return composition(layout.layout_a(), smem_ptr_flag_bits<B*N>{}, upcast<N>(layout.layout_b()));
}
template <int N, class SwizzleFn, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
downcast(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
{
return composition(layout.layout_a(), smem_ptr_flag_bits<B/N>{}, downcast<N>(layout.layout_b()));
}
//
// Conversion with swizzle_layout
//
template <class T, class SwizzleFn, int B, class Layout>
CUTE_HOST_DEVICE
auto
as_position_independent_swizzle_layout(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
{
return composition(recast_layout<uint8_t,uint_bit_t<B>>(layout.layout_a()), Int<0>{}, layout.layout_b());
}
template <class Tensor>
CUTE_HOST_DEVICE
auto
as_position_independent_swizzle_tensor(Tensor&& tensor)
{
static_assert(is_smem<remove_cvref_t<Tensor>>::value, "Expected smem tensor.");
using SwizzleFn = get_swizzle_t<remove_cvref_t<Tensor>>;
if constexpr (SwizzleFn::num_bits == 0) {
return tensor;
} else {
#if !defined(NDEBUG)
{
uint32_t address = cast_smem_ptr_to_uint(raw_pointer_cast(std::forward<Tensor>(tensor).data()));
uint32_t mask = ((uint32_t(1) << SwizzleFn::num_base) - 1) | SwizzleFn::swizzle_code;
assert((address & mask) == 0); // Alignment to the Base, Z, and Y of Swizzle
}
#endif
using T = typename remove_cvref_t<Tensor>::value_type;
// Recast swizzle from acting on byte-addressed pointers to elements of type-T
auto new_swizzle = recast_layout<uint8_t, T>(SwizzleFn{});
// Strip off everything and create a new smem_ptr for type-T
auto new_ptr = make_smem_ptr<T>(raw_pointer_cast(std::forward<Tensor>(tensor).data()));
return make_tensor(new_ptr, composition(new_swizzle, Int<0>{}, tensor.layout()));
}
CUTE_GCC_UNREACHABLE;
}
//
// Display utilities
//
// Capture and cast smem_ptr_flag Layouts to offset-0 layouts
template <class SwizzleFn, int B, class Layout>
CUTE_HOST_DEVICE
void
print_latex(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
{
print_latex(as_position_independent_swizzle_layout(layout));
}
template <int B>
CUTE_HOST_DEVICE void print(smem_ptr_flag_bits<B> ptr)
{
printf("smem_ptr[%db](unset)", B);
}
} // end namespace cute
include/cute/pointer_swizzle.hpp
+172
View File
@@ -0,0 +1,172 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp> // iterator_traits
#include <cute/container/array_subbyte.hpp>
#include <cute/pointer_base.hpp>
#include <cute/swizzle.hpp>
/* This implements a swizzle pointer of the form
* InvolutionFn o PtrAdd
* where the InvolutionFn need not be linear.
*
* This differs subtly from swizzle_layout because the smem pointer is used
* as the offset. That means that swizzle_layout will implement position-independent
* swizzle layouts, while swizzle_ptr implements position-dependent swizzle tensors.
* Arch chose to design hardware with position-dependent swizzles.
*
* For clarity:
* NormalLayout : DeRef <- PtrAdd <- [Layout]
* ComposedLayout: DeRef <- PtrAdd <- [Swizzle <- OffsetAdd <- Layout]
* SwizzlePtr : [DeRef <- Swizzle <- PtrAdd] <- Layout
*
* Furthermore, for known swizzles, this pointer attempts to decay itself
* to a normal-pointer with a new layout containing dynamic or static strides.
* This is possible by determining the subdomain of the InvolutionFn
* that is identity and testing if the Layout's codomain is contained
* within it.
*/
namespace cute
{
// concept SwizzleFn {
// CUTE_HOST_DEVICE constexpr static uint apply(uint);
// }
// See Swizzle<B,M,S> in swizzle.hpp for common swizzle-functions.
template <class SwizzleFn, class Iterator>
struct swizzle_ptr : iter_adaptor<Iterator,swizzle_ptr<SwizzleFn,Iterator>>
{
using iterator = Iterator;
using reference = typename iterator_traits<iterator>::reference;
using element_type = typename iterator_traits<iterator>::element_type;
using value_type = typename iterator_traits<iterator>::value_type;
using iter_adaptor<Iterator,swizzle_ptr<SwizzleFn,Iterator>>::iter_adaptor;
template <class Iter>
CUTE_HOST_DEVICE constexpr static
Iter apply_swizzle(Iter ptr) {
return {apply_swizzle(ptr.get())};
}
template <class T>
CUTE_HOST_DEVICE constexpr static
T* apply_swizzle(T* ptr) {
return reinterpret_cast<T*>(SwizzleFn::apply(reinterpret_cast<uintptr_t>(ptr)));
}
template <class T>
CUTE_HOST_DEVICE constexpr static
subbyte_iterator<T> apply_swizzle(subbyte_iterator<T> ptr) {
return {apply_swizzle(ptr.ptr_), ptr.idx_};
}
CUTE_HOST_DEVICE constexpr
reference operator*() const {
return *apply_swizzle(this->get());
}
template <class Int>
CUTE_HOST_DEVICE constexpr
reference operator[](Int const& i) const {
return *apply_swizzle(this->get() + i);
}
};
template <class T, class = void> // Default No-Swizzle
struct get_swizzle { using type = Swizzle<0,4,3>; };
template <class SwizzleFn, class P> // Found the SwizzleFn
struct get_swizzle<swizzle_ptr<SwizzleFn,P>> { using type = SwizzleFn; };
template <class T> // Recurse into anything with a ::iterator
struct get_swizzle<T, void_t<typename T::iterator>> : get_swizzle<typename T::iterator> {};
template <class Iter>
using get_swizzle_t = typename get_swizzle<Iter>::type;
template <class Iterator, class SwizzleFn>
CUTE_HOST_DEVICE constexpr
swizzle_ptr<SwizzleFn,Iterator>
make_swizzle_ptr(Iterator ptr, SwizzleFn) {
return {ptr};
}
// Swizzle-0 specialization for immediate decay
template <class Iterator, int M, int S>
CUTE_HOST_DEVICE constexpr
Iterator
make_swizzle_ptr(Iterator ptr, Swizzle<0,M,S>) {
return ptr;
}
//
// Recast
//
template <class SwizzleFn, class P>
CUTE_HOST_DEVICE constexpr
auto
raw_pointer_cast(swizzle_ptr<SwizzleFn,P> const& ptr) {
return raw_pointer_cast(ptr.get());
}
// SwizzleFn operates on the pointer address, so it doesn't care about the type
template <class NewT, class SwizzleFn, class P>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(swizzle_ptr<SwizzleFn,P> const& ptr) {
return make_swizzle_ptr(recast_ptr<NewT>(ptr.get()), SwizzleFn{});
}
//
// Display utilities
//
template <class SwizzleFn, class P>
CUTE_HOST_DEVICE void print(swizzle_ptr<SwizzleFn,P> ptr)
{
print(SwizzleFn{}); printf("_"); print(ptr.get());
}
#if !defined(__CUDACC_RTC__)
template <class SwizzleFn, class P>
CUTE_HOST std::ostream& operator<<(std::ostream& os, swizzle_ptr<SwizzleFn,P> ptr)
{
return os << SwizzleFn{} << "_" << ptr.get();
}
#endif
} // end namespace cute
include/cute/swizzle.hpp
+14 -100
View File
@@ -124,92 +124,6 @@ composition(Swizzle<B0,M0,S0>, Swizzle<B1,M1,S1>)
//return ComposedFn<Swizzle<B0,M0,S0>, Swizzle<B1,M1,S1>>{};
}
//
// Inverse
//
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
Swizzle<B,M,S>
right_inverse(Swizzle<B,M,S> const& sw)
{
return sw;
}
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
Swizzle<B,M,S>
left_inverse(Swizzle<B,M,S> const& sw)
{
return sw;
}
// Kludge -- Probably want an OffsetFn<T> here instead
template <class T, __CUTE_REQUIRES(is_integral<T>::value)>
CUTE_HOST_DEVICE constexpr
auto
right_inverse(T const& t)
{
return -t;
}
// Kludge -- Probably want an OffsetFn<T> here instead
template <class T, __CUTE_REQUIRES(is_integral<T>::value)>
CUTE_HOST_DEVICE constexpr
auto
left_inverse(T const& t)
{
return -t;
}
//
// Upcast and Downcast
//
template <int N, int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
upcast(Swizzle<B,M,S> const& swizzle)
{
static_assert(has_single_bit(N), "N must be a power of two");
constexpr int log2_n = bit_width(uint32_t(N)) - 1;
constexpr int NewM = M - log2_n;
if constexpr (NewM >= 0) {
return Swizzle<B,NewM,S>{};
} else {
return Swizzle<cute::max(B+NewM,0), 0, S>{};
}
CUTE_GCC_UNREACHABLE;
}
template <int N, int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
downcast(Swizzle<B,M,S> const& swizzle)
{
static_assert(has_single_bit(N), "N must be a power of two");
constexpr int log2_n = bit_width(uint32_t(N)) - 1;
return Swizzle<B,(M + log2_n),S>{};
}
template <class OldType, class NewType,
int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
recast(Swizzle<B,M,S> const& swizzle)
{
if constexpr (sizeof_bits<NewType>::value == sizeof_bits<OldType>::value) {
return swizzle;
} else if constexpr (sizeof_bits<NewType>::value > sizeof_bits<OldType>::value) {
static_assert(sizeof_bits<NewType>::value % sizeof_bits<OldType>::value == 0, "NewType must be a multiple of OldType");
return upcast<sizeof_bits<NewType>::value/sizeof_bits<OldType>::value>(swizzle);
} else if constexpr (sizeof_bits<NewType>::value < sizeof_bits<OldType>::value) {
static_assert(sizeof_bits<OldType>::value % sizeof_bits<NewType>::value == 0, "NewType must be a divisor of OldType");
return downcast<sizeof_bits<OldType>::value/sizeof_bits<NewType>::value>(swizzle);
}
}
//
// Utility for slicing and swizzle "offsets"
//
@@ -218,8 +132,8 @@ recast(Swizzle<B,M,S> const& swizzle)
// consumed and which bits are free. Furthermore, it is useful to know whether
// each of these bits is known statically or dynamically.
// MixedBits is an 32-bit unsigned integer class where some bits are known statically
// and some bits are known dynamically. These sets of bits are disjoint and it is
// MixedBits is an 32-bit unsigned integer class where some bits are known statically
// and some bits are known dynamically. These sets of bits are disjoint and it is
// known statically which bits are known dynamically.
// MixedBits can only be manipulated through bitwise operations
@@ -524,6 +438,12 @@ to_mixed_bits(Layout const& layout, Coord const& coord)
// Display utilities
//
template <int B, int M, int S>
CUTE_HOST_DEVICE void print(Swizzle<B,M,S> const&)
{
printf("Sw<%d,%d,%d>", B, M, S);
}
template <uint32_t S, uint32_t F>
CUTE_HOST_DEVICE void print(MixedBits<S,F> const& m)
{
@@ -531,23 +451,17 @@ CUTE_HOST_DEVICE void print(MixedBits<S,F> const& m)
}
#if !defined(__CUDACC_RTC__)
template <int B, int M, int S>
CUTE_HOST std::ostream& operator<<(std::ostream& os, Swizzle<B,M,S> const&)
{
return os << "Sw<" << B << "," << M << "," << S << ">";
}
template <uint32_t S, class D, uint32_t F>
CUTE_HOST std::ostream& operator<<(std::ostream& os, MixedBits<S,F> const& m)
{
return os << "M_" << S << "|(" << m.dynamic_int_ << "&" << F << ")=" << uint32_t(m);
}
template <int B, int M, int S>
CUTE_HOST_DEVICE void print(Swizzle<B,M,S> const&)
{
print("S<%d,%d,%d>", B, M, S);
}
template <int B, int M, int S>
CUTE_HOST std::ostream& operator<<(std::ostream& os, Swizzle<B,M,S> const&)
{
return os << "S<" << B << "," << M << "," << S << ">";
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace cute
include/cute/swizzle_layout.hpp
+84
View File
@@ -147,6 +147,7 @@ get_swizzle_portion(Layout<Shape,Stride>)
// Get the "non-swizzle" part of a composed layout,
// which is the underlying (non-composed) Layout.
template <int B, int M, int S, class Offset, class LayoutB>
CUTE_HOST_DEVICE constexpr
auto
get_nonswizzle_portion(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& slayout)
{
@@ -155,6 +156,7 @@ get_nonswizzle_portion(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& slay
// The non-swizzle part of a non-swizzled layout is just the Layout.
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
get_nonswizzle_portion(Layout<Shape,Stride> const& slayout)
{
@@ -361,6 +363,88 @@ left_inverse(ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
}
}
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
Swizzle<B,M,S>
right_inverse(Swizzle<B,M,S> const& sw)
{
return sw;
}
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
Swizzle<B,M,S>
left_inverse(Swizzle<B,M,S> const& sw)
{
return sw;
}
// Kludge -- Probably want an OffsetFn<T> here instead
template <class T, __CUTE_REQUIRES(is_integral<T>::value)>
CUTE_HOST_DEVICE constexpr
auto
right_inverse(T const& t)
{
return -t;
}
// Kludge -- Probably want an OffsetFn<T> here instead
template <class T, __CUTE_REQUIRES(is_integral<T>::value)>
CUTE_HOST_DEVICE constexpr
auto
left_inverse(T const& t)
{
return -t;
}
//
// Upcast and Downcast
//
template <int N, int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
upcast(Swizzle<B,M,S> const& swizzle)
{
static_assert(has_single_bit(N), "N must be a power of two");
constexpr int log2_n = bit_width(uint32_t(N)) - 1;
constexpr int NewM = M - log2_n;
if constexpr (NewM >= 0) {
return Swizzle<B,NewM,S>{};
} else {
return Swizzle<cute::max(B+NewM,0), 0, S>{};
}
CUTE_GCC_UNREACHABLE;
}
template <int N, int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
downcast(Swizzle<B,M,S> const& swizzle)
{
static_assert(has_single_bit(N), "N must be a power of two");
constexpr int log2_n = bit_width(uint32_t(N)) - 1;
return Swizzle<B,(M + log2_n),S>{};
}
template <class OldType, class NewType,
int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
recast_layout(Swizzle<B,M,S> const& swizzle)
{
if constexpr (sizeof_bits<NewType>::value == sizeof_bits<OldType>::value) {
return swizzle;
} else if constexpr (sizeof_bits<NewType>::value > sizeof_bits<OldType>::value) {
static_assert(sizeof_bits<NewType>::value % sizeof_bits<OldType>::value == 0, "NewType must be a multiple of OldType");
return upcast<sizeof_bits<NewType>::value/sizeof_bits<OldType>::value>(swizzle);
} else if constexpr (sizeof_bits<NewType>::value < sizeof_bits<OldType>::value) {
static_assert(sizeof_bits<OldType>::value % sizeof_bits<NewType>::value == 0, "NewType must be a divisor of OldType");
return downcast<sizeof_bits<OldType>::value/sizeof_bits<NewType>::value>(swizzle);
}
}
//
// Other operations
//
include/cute/swizzle_ptr.hpp
-303
View File
@@ -1,303 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/arch/util.hpp>
#include <cute/swizzle_layout.hpp>
#include <cute/tensor.hpp>
#include <cute/pointer.hpp>
#include <cute/container/array.hpp>
#include <cute/numeric/int.hpp>
/* This implements a swizzle pointer of the form
* InvolutionFn o PtrAdd
* where the InvolutionFn need not be linear.
*
* This differs subtly from swizzle_layout because the smem pointer is used
* as the offset. That means that swizzle_layout will implement position-independent
* swizzle layouts, while swizzle_ptr implements position-dependent swizzle tensors.
* Arch chose to design hardware with position-dependent swizzles.
*
* For clarity:
* NormalLayout : DeRef <- PtrAdd <- [Layout]
* ComposedLayout: DeRef <- PtrAdd <- [Swizzle <- OffsetAdd <- Layout]
* SwizzlePtr : [DeRef <- Swizzle <- PtrAdd] <- Layout
*
* Furthermore, for known swizzles, this pointer attempts to decay itself
* to a normal-pointer with a new layout containing dynamic or static strides.
* This is possible by determining the subdomain of the InvolutionFn
* that is identity and testing if the Layout's codomain is contained
* within it.
*/
namespace cute
{
template <class T, class Swizzle>
struct smem_ptr_swizzle
{
static_assert(is_empty<Swizzle>::value, "Swizzle can't have state.");
static const uint32_t ElementsPerStoredItem = sizeof(T) * 8 / sizeof_bits_v<T>;
CUTE_HOST_DEVICE constexpr
T* get() const
{
return ptr_;
}
CUTE_HOST_DEVICE constexpr static
Swizzle get_swizzle()
{
return {};
}
CUTE_HOST_DEVICE constexpr static
T* apply_swizzle(T* ptr)
{
return reinterpret_cast<T*>(Swizzle::apply(reinterpret_cast<uintptr_t>(ptr)));
}
CUTE_HOST_DEVICE constexpr
T& operator*() const
{
return *apply_swizzle(get());
}
template <class Int>
CUTE_HOST_DEVICE constexpr
T& operator[](Int const& i) const
{
static_assert(sizeof_bits_v<T> >= 8, "Use subbyte_iterator to access the element");
return *apply_swizzle(get() + i);
}
template <class Int>
CUTE_HOST_DEVICE constexpr
smem_ptr_swizzle operator+(Int const& i) const
{
return {ptr_ + i / ElementsPerStoredItem};
}
T* ptr_;
};
template <class T, class S>
struct is_smem<smem_ptr_swizzle<T,S>> : true_type {};
// Make a swizzle pointer
template <class T, class Swizzle>
CUTE_HOST_DEVICE constexpr
auto
make_smem_ptr(T* ptr, Swizzle const&)
{
return smem_ptr_swizzle<T,Swizzle>{ptr};
}
// Specialization for immediate decay
template <class T, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
make_smem_ptr(T* ptr, Swizzle<0,M,S> const&)
{
return make_smem_ptr(ptr);
}
// A model of a nullptr smem_ptr<T> with B == sizeof_bits<T>::value
// That represents an unset pointer. This is a placeholder type that is waiting for an smem_ptr
template <int Bits>
struct smem_ptr_flag_bits : Int<0> {};
using smem_ptr_flag = smem_ptr_flag_bits<1>;
// A flagged construction method to transform ComposedLayout
// Make a swizzle pointer tensor and check that the intended type size matches
template <class T, class Swizzle, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
make_tensor(smem_ptr<T> const& ptr,
ComposedLayout<Swizzle,smem_ptr_flag_bits<B>,Layout> const& layout)
{
static_assert(B == sizeof_bits<T>::value, "Expected a B-bit pointer type.");
return make_tensor(make_smem_ptr(ptr.get(), layout.layout_a()),
layout.layout_b());
}
// NOTE: To preserve smem_ptr_flag_bits under recast ops
template <int N, class Swizzle, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
upcast(ComposedLayout<Swizzle,smem_ptr_flag_bits<B>,Layout> const& layout)
{
return composition(layout.layout_a(), smem_ptr_flag_bits<B*N>{}, upcast<N>(layout.layout_b()));
}
template <int N, class Swizzle, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
downcast(ComposedLayout<Swizzle,smem_ptr_flag_bits<B>,Layout> const& layout)
{
return composition(layout.layout_a(), smem_ptr_flag_bits<B/N>{}, downcast<N>(layout.layout_b()));
}
//
// Recast
// Swizzle operates on the pointer address, so it doesn't care about the type
//
template <class NewT, class T, class Swizzle>
CUTE_HOST_DEVICE constexpr
auto
recast(smem_ptr_swizzle<T,Swizzle> const& ptr)
{
return smem_ptr_swizzle<NewT,Swizzle>{recast<NewT>(ptr.ptr_)};
}
template <class NewT, class T, class Swizzle>
CUTE_HOST_DEVICE constexpr
auto
recast(smem_ptr_swizzle<T const,Swizzle> const& ptr)
{
return smem_ptr_swizzle<NewT const,Swizzle>{recast<NewT const>(ptr.ptr_)};
}
template <class T, class Swizzle>
CUTE_HOST_DEVICE constexpr
T*
raw_pointer_cast(smem_ptr_swizzle<T,Swizzle> ptr) {
return ptr.get();
}
//
// Conversion with swizzle_layout
//
template <class T, class Swizzle, int B, class Layout>
CUTE_HOST_DEVICE
auto
as_position_independent_swizzle_layout(ComposedLayout<Swizzle,smem_ptr_flag_bits<B>,Layout> const& layout)
{
return composition(recast<uint_bit_t<8>,uint_bit_t<B>>(layout.layout_a()), Int<0>{}, layout.layout_b());
}
template <class T, class Swizzle, class Layout>
CUTE_HOST_DEVICE
auto
as_position_independent_swizzle_tensor(Tensor<ViewEngine<smem_ptr_swizzle<T,Swizzle>>, Layout> const& tensor)
{
{
uint32_t address = cast_smem_ptr_to_uint(tensor.data().get());
uint32_t mask = ((uint32_t(1) << Swizzle::num_base) - 1) & (Swizzle::swizzle_code);
assert((address & mask) == 0); // Alignment to the Base, Z, and Y of Swizzle
}
auto new_swizzle = recast<uint_bit_t<8>,uint_bit_t<sizeof_bits_v<T>>>(tensor.data().get_swizzle());
return make_tensor(make_smem_ptr(tensor.data().get()), composition(new_swizzle, Int<0>{}, tensor.layout()));
}
template <class T, class Swizzle, class Layout>
CUTE_HOST_DEVICE
auto
as_position_independent_swizzle_tensor(Tensor<ViewEngine<smem_ptr_swizzle<T,Swizzle>>, Layout>& tensor)
{
{
[[maybe_unused]] uint32_t address = cast_smem_ptr_to_uint(tensor.data().get());
[[maybe_unused]] uint32_t mask = ((uint32_t(1) << Swizzle::num_base) - 1) & (Swizzle::swizzle_code);
assert((address & mask) == 0); // Alignment to the Base, Z, and Y of Swizzle
}
auto new_swizzle = recast<uint_bit_t<8>,uint_bit_t<sizeof_bits_v<T>>>(tensor.data().get_swizzle());
return make_tensor(make_smem_ptr(tensor.data().get()), composition(new_swizzle, Int<0>{}, tensor.layout()));
}
template <class T, class Swizzle, class Layout>
CUTE_HOST_DEVICE
auto
as_position_independent_swizzle_tensor(Tensor<ViewEngine<smem_ptr_swizzle<T,Swizzle>>, Layout>&& tensor)
{
return as_position_independent_swizzle_tensor(tensor);
}
// Pass through everything else
// Used if the tensor doesn't have a swizzled layout, e.g. Layout_MN_INTER_Atom, Layout_K_INTER_Atom
template <class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
as_position_independent_swizzle_tensor(Tensor<Engine, Layout> const& tensor)
{
return tensor;
}
template <class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
as_position_independent_swizzle_tensor(Tensor<Engine, Layout>&& tensor)
{
return tensor;
}
//
// Print
//
// Capture and cast smem_ptr_flag Layouts to offset-0 layouts
template <class Swizzle, int B, class Layout>
CUTE_HOST_DEVICE
void
print_latex(ComposedLayout<Swizzle,smem_ptr_flag_bits<B>,Layout> const& layout)
{
auto new_swizzle = recast<uint_bit_t<8>,uint_bit_t<B>>(layout.layout_a());
print_latex(composition(new_swizzle, Int<0>{}, layout.layout_b()));
}
template <int B>
CUTE_HOST_DEVICE void print(smem_ptr_flag_bits<B> const& ptr)
{
printf("smem_ptr_%db(unset)", B);
}
template <class T, int B, int M, int S>
CUTE_HOST_DEVICE void print(smem_ptr_swizzle<T,Swizzle<B,M,S>> const& ptr)
{
printf("smem_ptr_S<%d,%d,%d>_%db(%p)", B, M, S, int(sizeof_bits<T>::value), ptr.get());
}
#if !defined(__CUDACC_RTC__)
template <class T, int B, int M, int S>
CUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr_swizzle<T,Swizzle<B,M,S>> const&)
{
return os << "smem_ptr_S<" << B << "," << M << "," << S << ">_" << int(sizeof_bits<T>::value) << "b";
}
#endif
} // end namespace cute
include/cute/tensor.hpp
+216 -183
View File
@@ -33,16 +33,16 @@
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/container/tuple.hpp>
#include <cute/container/array_aligned.hpp>
#include <cute/container/array_subbyte.hpp>
#include <cute/container/type_list.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/numeric/integer_sequence.hpp>
#include <cute/container/tuple.hpp>
#include <cute/container/array_aligned.hpp>
#include <cute/container/array_subbyte.hpp>
#include <cute/pointer.hpp>
#include <cute/layout.hpp>
#include <cute/tile.hpp>
#include <cute/pointer.hpp>
namespace cute
{
@@ -52,63 +52,54 @@ namespace cute
//
// concept Engine {
// using value_type = ;
// using iterator = ;
// using value_type = ;
// using element_type = ;
// using reference = ;
// iterator begin();
// };
template <class T, int N>
using ArrayEngine = typename conditional<(sizeof_bits<T>::value % 8 == 0),
array_aligned<T,N>,
array_subbyte<T,N>>::type;
struct ArrayEngine
{
using Storage = typename conditional<(sizeof_bits<T>::value % 8 == 0),
array_aligned<T,N>,
array_subbyte<T,N>>::type;
using iterator = typename Storage::iterator;
using reference = typename iterator_traits<iterator>::reference;
using element_type = typename iterator_traits<iterator>::element_type;
using value_type = typename iterator_traits<iterator>::value_type;
Storage storage_;
CUTE_HOST_DEVICE constexpr auto begin() const { return storage_.begin(); }
CUTE_HOST_DEVICE constexpr auto begin() { return storage_.begin(); }
};
template <class Iterator>
struct ViewEngine
{
using value_type = typename cute::remove_cvref<decltype(*declval<Iterator>())>::type;
using iterator = Iterator;
using iterator = Iterator;
using reference = typename iterator_traits<iterator>::reference;
using element_type = typename iterator_traits<iterator>::element_type;
using value_type = typename iterator_traits<iterator>::value_type;
iterator storage_;
CUTE_HOST_DEVICE constexpr
iterator const&
begin() const {
return storage_;
}
CUTE_HOST_DEVICE constexpr
iterator&
begin() {
return storage_;
}
CUTE_HOST_DEVICE constexpr iterator const& begin() const { return storage_; }
CUTE_HOST_DEVICE constexpr iterator & begin() { return storage_; }
};
template <class Iter>
struct is_rmem<ViewEngine<Iter>> : is_rmem<Iter> {};
template <class Iter>
struct is_smem<ViewEngine<Iter>> : is_smem<Iter> {};
template <class Iter>
struct is_gmem<ViewEngine<Iter>> : is_gmem<Iter> {};
template <class Iterator>
struct ConstViewEngine
{
using value_type = typename cute::remove_cvref<decltype(*declval<Iterator>())>::type;
using iterator = Iterator;
using iterator = Iterator;
using reference = typename iterator_traits<iterator>::reference;
using element_type = typename iterator_traits<iterator>::element_type;
using value_type = typename iterator_traits<iterator>::value_type;
iterator storage_;
CUTE_HOST_DEVICE constexpr
iterator const&
begin() const {
return storage_;
}
CUTE_HOST_DEVICE constexpr iterator const& begin() const { return storage_; }
};
template <class Iter>
struct is_rmem<ConstViewEngine<Iter>> : is_rmem<Iter> {};
template <class Iter>
struct is_smem<ConstViewEngine<Iter>> : is_smem<Iter> {};
template <class Iter>
struct is_gmem<ConstViewEngine<Iter>> : is_gmem<Iter> {};
//
// Tensor
//
@@ -116,14 +107,13 @@ struct is_gmem<ConstViewEngine<Iter>> : is_gmem<Iter> {};
template <class Engine, class Layout>
struct Tensor
{
using value_type = typename Engine::value_type;
//using pointer = typename engine_traits<Engine>::pointer;
//using const_pointer = typename engine_traits<Engine>::const_pointer;
//using reference = typename engine_traits<Engine>::reference;
//using const_reference = typename engine_traits<Engine>::const_reference;
using iterator = typename Engine::iterator;
using value_type = typename Engine::value_type;
using element_type = typename Engine::element_type;
using reference = typename Engine::reference;
using engine_type = Engine;
using layout_type = Layout;
using engine_type = Engine;
using layout_type = Layout;
CUTE_HOST_DEVICE constexpr
Tensor() {}
@@ -323,18 +313,11 @@ struct Tensor
cute::tuple<layout_type, engine_type> rep_;
};
template <class Layout>
template <class T>
struct is_tensor : false_type {};
template <class Engine, class Layout>
struct is_tensor<Tensor<Engine,Layout>> : true_type {};
template <class Engine, class Layout>
struct is_rmem<Tensor<Engine,Layout>> : is_rmem<Engine> {};
template <class Engine, class Layout>
struct is_smem<Tensor<Engine,Layout>> : is_smem<Engine> {};
template <class Engine, class Layout>
struct is_gmem<Tensor<Engine,Layout>> : is_gmem<Engine> {};
// Customization point for creation of owning and non-owning Tensors
template <class T>
struct MakeTensor
@@ -471,8 +454,7 @@ CUTE_HOST_DEVICE constexpr
auto
make_counting_tensor(Layout const& layout)
{
return make_tensor(ArithmeticTupleIterator(as_arithmetic_tuple(repeat_like(coshape(layout), Int<0>{}))),
layout);
return make_tensor(make_inttuple_iter(repeat_like(coshape(layout), Int<0>{})), layout);
}
//
@@ -665,16 +647,14 @@ group_modes(Tensor&& tensor)
// -- doesn't check dynamic integer divisibility
// -- doesn't check alignment
// A tagged version for dispatching
template <class NewType, class Tensor,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
template <class NewType, class Tensor>
CUTE_HOST_DEVICE constexpr
auto
recast(Tensor&& tensor, type_list<NewType>)
recast(Tensor&& tensor)
{
using OldType = typename remove_cvref_t<Tensor>::value_type;
auto old_layout = tensor.layout();
auto new_layout = recast<OldType,NewType>(old_layout);
auto new_layout = recast_layout<OldType,NewType>(old_layout);
// If this is an upcast of a normal Layout with static negative strides, then offset as well
if constexpr (sizeof(OldType) < sizeof(NewType) && not is_composed_layout<decltype(old_layout)>::value) {
@@ -682,38 +662,14 @@ recast(Tensor&& tensor, type_list<NewType>)
auto extent_diff = transform(shape_diff, flatten(old_layout.stride()), multiplies{});
auto offset = fold(extent_diff, Int<0>{}, [](auto const& i, auto const& a) { return i + cute::min(a,Int<0>{}); });
return make_tensor(recast<NewType>(std::forward<Tensor>(tensor).data() + offset), new_layout);
return make_tensor(recast_ptr<NewType>(std::forward<Tensor>(tensor).data() + offset), new_layout);
} else {
return make_tensor(recast<NewType>(std::forward<Tensor>(tensor).data() ), new_layout);
return make_tensor(recast_ptr<NewType>(std::forward<Tensor>(tensor).data() ), new_layout);
}
CUTE_GCC_UNREACHABLE;
}
template <class NewType, class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
recast(Tensor<Engine,Layout> const& tensor)
{
return recast(tensor, type_list<NewType>{});
}
template <class NewType, class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
recast(Tensor<Engine,Layout>& tensor)
{
return recast(tensor, type_list<NewType>{});
}
template <class NewType, class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
recast(Tensor<Engine,Layout>&& tensor)
{
return recast(std::forward<Tensor<Engine,Layout>>(tensor), type_list<NewType>{});
}
//
// max_common_vector
//
@@ -736,13 +692,12 @@ max_common_vector(Tensor<SrcEngine,SrcLayout> const& a,
{
using SrcType = typename Tensor<SrcEngine,SrcLayout>::value_type;
using DstType = typename Tensor<DstEngine,DstLayout>::value_type;
using SrcRef = decltype(*(a.data()));
using DstRef = decltype(*(b.data()));
using SrcRef = typename Tensor<SrcEngine,SrcLayout>::reference;
using DstRef = typename Tensor<SrcEngine,SrcLayout>::reference;
// Determine if vectorization candidates at all
if constexpr (// Should be the same value_types, else the copy is also performing a cast
sizeof(SrcType) == sizeof(DstType) &&
sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
// The types should be trivially copyable so that vectorization is valid
is_trivially_copyable<SrcType>::value &&
is_trivially_copyable<DstType>::value &&
@@ -759,144 +714,222 @@ max_common_vector(Tensor<SrcEngine,SrcLayout> const& a,
}
//
// Key algebraic operations
// Key algebraic operations -- Divide and Product
//
template <class Tensor, class Tile,
// Apply a Tiler to the Tensor.
//
// Consider a Tensor with shape (A,B,x,y)
// And a Tiler that is:
//
// * A Layout with shape (BLK_A,BLK_B)
// ** Result Tensor shape ((BLK_A,BLK_B),Rest).
// ** That is, the Tensor and Tile are treated as 1D for the tiling.
// ** See logical_divide(Layout,Layout)
//
// * A Tile<Layout...> with shape <BLK_A,BLK_B>
// ** Result Tensor shape ((BLK_A,a),(BLK_B,b),x,y).
// ** Each mode of the Tile<Layout...> is applied to the corresponding mode of the Tensor.
// ** See logical_divide(Layout,Tuple)
//
// * A Shape (BLK_A,BLK_B)
// ** Result Tensor shape ((BLK_A,a),(BLK_B,b),x,y).
// ** Equivalent to applying Tile<BLK_A:_1,BLK_B:_1>.
// ** See logical_divide(Layout,Tuple) and logical_divide(Layout,Int)
//
// Note that the Tile<Layout...>/Shape Tilers must be weakly_congruent to the Tensor
template <class Tensor, class Tiler,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
logical_divide(Tensor && tensor,
Tile const& tile)
Tiler const& tiler) // Layout or Tile<Layout...> or Shape
{
return make_tensor(std::forward<Tensor>(tensor).data(),
logical_divide(tensor.layout(), tile));
logical_divide(tensor.layout(), tiler));
}
// zipped_divide is logical_divide with modes gathered into standard form ((BLK_A,BLK_B),(a,b))
template <class Tensor, class Tile,
// zipped_divide is logical_divide with Tiler modes and Rest modes gathered together: (Tiler,Rest)
// When Tiler is Layout, this has no effect as logical_divide results in the same.
// When Tiler is Tile<Layout...> or Shape, this zips modes into standard form ((BLK_A,BLK_B),(a,b,x,y))
template <class Tensor, class Tiler,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
zipped_divide(Tensor && tensor,
Tile const& tile) // Layout or Tile<Layout...>
zipped_divide(Tensor && tensor,
Tiler const& tiler) // Layout or Tile<Layout...> or Shape
{
return make_tensor(std::forward<Tensor>(tensor).data(),
zipped_divide(tensor.layout(), tile));
zipped_divide(tensor.layout(), tiler));
}
// tiled_divide is logical_divide with the second output mode flattened ((BLK_A,BLK_B),a,b)
template <class Tensor, class Tile,
// tiled_divide is zipped_divide with the second output mode flattened ((BLK_A,BLK_B),a,b,x,y)
template <class Tensor, class Tiler,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
tiled_divide(Tensor && tensor,
Tile const& tile) // Layout or Tile<Layout...>
tiled_divide(Tensor && tensor,
Tiler const& tiler) // Layout or Tile<Layout...> or Shape
{
return make_tensor(std::forward<Tensor>(tensor).data(),
tiled_divide(tensor.layout(), tile));
tiled_divide(tensor.layout(), tiler));
}
// flat_divide is zipped_divide with the both modes flattened (BLK_A,BLK_B,a,b,x,y)
template <class Tensor, class Tiler,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
flat_divide(Tensor && tensor,
Tiler const& tiler) // Layout or Tile<Layout...> or Shape
{
return make_tensor(std::forward<Tensor>(tensor).data(),
flat_divide(tensor.layout(), tiler));
}
// logical_product on a Tensor doesn't make sense since it often increases cosize
// though this might make sense for creating Tensors with broadcasted (stride-0) modes
//
// Logical Divide utilities: local_partition and local_tile
// Tensor partitioning utilities
//
template <class Tensor, class Tile, class Coord,
// Apply a Tiler to the Tensor, then slice out one of those tiles by slicing into the "Rest" modes.
// With an inner_partition, you get everything that's inside the Tiler. Everything that the Tiler is pointing to.
// Split the modes of tensor according to the Tiler
// zipped_divide returns something like ((BLK_A,BLK_B,...),(a,b,...,x,y))
// Then slice into the second mode (the "Rest" mode) with Coord
template <class Tensor, class Tiler, class Coord,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
local_partition(Tensor && tensor,
Tile const& tile,
Coord const& coord)
inner_partition(Tensor && tensor,
Tiler const& tiler,
Coord const& coord)
{
constexpr int R1 = decltype(rank(tensor))::value;
auto tensor_tiled = zipped_divide(std::forward<Tensor>(tensor), tiler);
constexpr int R0 = decltype(rank<0>(tensor_tiled))::value;
// Split the modes of tensor according to the modes of tile
// zipped_divide returns something like ((VEC_A,VEC_B,...),(a,b,...))
// The_coord is the coord into the first mode, flatten the rest
return zipped_divide(std::forward<Tensor>(tensor), tile)(coord, repeat<R1>(_));
// The coord slices into the second mode (the "rest" mode), flatten the first
if constexpr (is_tuple<Coord>::value) {
// Append trailing modes if coord is tuple
constexpr int R1 = decltype(rank<1>(tensor_tiled))::value;;
return tensor_tiled(repeat<R0>(_), append<R1>(coord,_));
} else {
// Flat indexing if coord is not tuple
return tensor_tiled(repeat<R0>(_), coord);
}
}
template <class Tensor, class Tile, class Coord, class Projection,
// Apply a Tiler to the Tensor, then slice out the remainder by slicing into the "Tile" modes.
// With an outer_partition, you get everything that's outside the Tiler. The layout of the Tile in the Tensor.
// Split the modes of tensor according to the Tiler
// zipped_divide returns something like ((BLK_A,BLK_B,...),(a,b,...,x,y))
// Then slice into the first mode (the "Tile" mode) with Coord
template <class Tensor, class Tiler, class Coord,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
local_partition(Tensor && tensor,
Tile const& tile,
Coord const& coord,
Projection const& proj)
outer_partition(Tensor && tensor,
Tiler const& tiler,
Coord const& coord)
{
return local_partition(std::forward<Tensor>(tensor),
dice(proj, tile),
dice(proj, coord));
auto tensor_tiled = zipped_divide(std::forward<Tensor>(tensor), tiler);
constexpr int R1 = decltype(rank<1>(tensor_tiled))::value;
// The coord slices into the first mode (the "tile" mode), flatten the second
if constexpr (is_tuple<Coord>::value) {
// Append trailing modes if coord is tuple
constexpr int R0 = decltype(rank<0>(tensor_tiled))::value;
return tensor_tiled(append<R0>(coord,_), repeat<R1>(_));
} else {
// Flat indexing if coord is not tuple
return tensor_tiled(coord, repeat<R1>(_));
}
}
// Special case with Layout and Integral that extracts the coord first
// e.g. local_partition(tensor, ThrLayout, threadIdx.x)
template <class Tensor, class LShape, class LStride, class Index,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value &&
is_integral<Index>::value)>
// Tile a tensor according to @a tiler and use @a coord to index into the remainder, keeping the tile.
// This is typical at the CTA level where tiles of data are extracted:
// Tensor data = ... // ( M, N)
// Tensor cta_data = local_tile(data, Shape<_32,_64>{}, make_coord(blockIdx.x,blockIdx.y)); // (_32,_64)
template <class Tensor, class Tiler, class Coord,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
local_tile(Tensor && tensor,
Tiler const& tiler, // tiler to apply
Coord const& coord) // coord to slice into "remainder"
{
return inner_partition(std::forward<Tensor>(tensor),
tiler,
coord);
}
// Same as above, but with a projection parameter to strip out unwanted tiling modes for convenience
// when using projections of the same tiler.
// This is typical at the CTA level where tiles of data are extracted as projections:
// Tensor dataA = ... // (M,K)
// Tensor dataB = ... // (N,K)
// Tensor dataC = ... // (M,N)
// auto cta_tiler = Shape<_32, _64, _4>{};
// auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _);
// Tensor ctaA = local_tile(dataA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (_32,_4,k)
// Tensor ctaB = local_tile(dataA, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (_64,_4,k)
// Tensor ctaC = local_tile(dataA, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (_32,_64)
template <class Tensor, class Tiler, class Coord, class Proj,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE
auto
local_partition(Tensor && tensor,
Layout<LShape,LStride> const& tile,
Index const& index)
local_tile(Tensor && tensor,
Tiler const& tiler, // tiler to apply
Coord const& coord, // coord to slice into "remainder"
Proj const& proj) // projection to apply to tiler and coord
{
return local_partition(std::forward<Tensor>(tensor),
return local_tile(std::forward<Tensor>(tensor),
dice(proj, tiler),
dice(proj, coord));
}
// Tile a tensor according to the flat shape of a layout that provides the coordinate of the target index.
// This is typical at the Thread level where data is partitioned across repeated patterns of threads:
// Tensor data = ... // (_16,_64)
// Tensor thr_data = local_partition(data, Layout<Shape<_2,_16>>{}, thr_idx); // ( _8, _4)
template <class Tensor, class LShape, class LStride, class Index,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE
auto
local_partition(Tensor && tensor,
Layout<LShape,LStride> const& tile, // coord -> index
Index const& index) // index to slice for
{
static_assert(is_integral<Index>::value);
return outer_partition(std::forward<Tensor>(tensor),
product_each(shape(tile)),
tile.get_flat_coord(index));
}
// Special case with Layout and Integral that extracts the coord first
// e.g. local_partition(tensor, ThrLayout, threadIdx.x, Step<_1,X,_1>{})
// Same as above, but with a projection parameter to strip out unwanted tiling modes for convenience
// when using projections of the same tiler.
// This is typical at the Thread level where data is partitioned across projected layouts of threads:
// Tensor dataA = ... // (M,K)
// Tensor dataB = ... // (N,K)
// Tensor dataC = ... // (M,N)
// auto thr_layout = Layout<Shape<_2,_16,_1>, Stride<_16,_1,_0>>{};
// Tensor thrA = local_partition(dataA, thr_layout, thr_idx, Step<_1, X,_1>{}); // (M/2,K/1)
// Tensor thrB = local_partition(dataB, thr_layout, thr_idx, Step< X,_1,_1>{}); // (N/16,K/1)
// Tensor thrC = local_partition(dataC, thr_layout, thr_idx, Step<_1,_1, X>{}); // (M/2,N/16)
template <class Tensor, class LShape, class LStride, class Index, class Projection,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value &&
is_integral<Index>::value)>
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE
auto
local_partition(Tensor && tensor,
Layout<LShape,LStride> const& tile,
Index const& index,
Projection const& proj)
local_partition(Tensor && tensor,
Layout<LShape,LStride> const& tile, // coord -> index
Index const& index, // index to slice for
Projection const& proj)
{
return local_partition(std::forward<Tensor>(tensor),
dice(proj, product_each(shape(tile))),
dice(proj, tile).get_flat_coord(index));
}
template <class Tensor, class Tile, class Coord,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE constexpr
auto
local_tile(Tensor && tensor,
Tile const& tile,
Coord const& coord)
{
constexpr int R0 = decltype(rank(tile))::value;
constexpr int R1 = decltype(rank(tensor))::value;
// Split the modes of tensor according to the modes of tile
// zipped_divide returns something like ((VEC_A,VEC_B,...),(a,b,...))
// The padded_coord is the coord into the second mode, flatten the rest
return zipped_divide(std::forward<Tensor>(tensor), tile)(repeat<R0>(_), append<R1>(coord,_));
}
template <class Tensor, class Tile, class Coord, class Proj,
__CUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
CUTE_HOST_DEVICE
auto
local_tile(Tensor && tensor,
Tile const& tile,
Coord const& coord,
Proj const& proj)
{
return local_tile(std::forward<Tensor>(tensor),
dice(proj, tile),
dice(proj, coord));
dice(proj, tile),
index);
}
//
@@ -906,7 +939,7 @@ local_tile(Tensor && tensor,
template <class Engine, class Layout>
CUTE_HOST_DEVICE void print(Tensor<Engine,Layout> const& tensor)
{
print(tensor.data()); print(" o "); print(tensor.layout());
print(tensor.data()); print(" o "); print(tensor.layout());
}
template <class Engine, class Layout>
@@ -951,8 +984,6 @@ CUTE_HOST_DEVICE void print_tensor(Tensor<Engine,Layout> const& tensor)
}
}
#if !defined(__CUDACC_RTC__)
template <class Engine, class Layout>
CUTE_HOST std::ostream& print_tensor_os(std::ostream& os, Tensor<Engine,Layout> const& tensor)
@@ -1008,7 +1039,9 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, Tensor<Engine,Layout> const
// Extended Engines
//
#include <cute/swizzle_ptr.hpp>
#include <cute/pointer_swizzle.hpp>
#include <cute/pointer_flagged.hpp>
//
// Tensor Algorithms
//
include/cute/tensor_predicate.hpp
+16
View File
@@ -60,4 +60,20 @@ struct TrivialPredTensor
}
};
template <class Fn>
struct FunctionPredTensor
{
CUTE_HOST_DEVICE constexpr
FunctionPredTensor(Fn const& fn) : fn_(fn) {}
template <class... Coords>
CUTE_HOST_DEVICE constexpr
auto
operator()(Coords const&... coords) const {
return fn_(coords...);
}
Fn const& fn_;
};
} // end namespace cute
include/cute/underscore.hpp
+48 -4
View File
@@ -95,6 +95,28 @@ using has_int0 = has_elem<Tuple, Int<0>>;
// Slice keeps only the elements of Tuple B that are paired with an Underscore
//
namespace detail {
template <class A, class B>
CUTE_HOST_DEVICE constexpr
auto
lift_slice(A const& a, B const& b)
{
if constexpr (is_tuple<A>::value) {
static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
return filter_tuple(a, b, [](auto const& x, auto const& y) { return lift_slice(x,y); });
} else if constexpr (is_underscore<A>::value) {
return cute::tuple<B>{b};
} else {
return cute::tuple<>{};
}
CUTE_GCC_UNREACHABLE;
}
} // end namespace detail
// Entry point overrides the lifting so that slice(_,b) == b
template <class A, class B>
CUTE_HOST_DEVICE constexpr
auto
@@ -102,9 +124,9 @@ slice(A const& a, B const& b)
{
if constexpr (is_tuple<A>::value) {
static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
return filter_tuple(a, b, [](auto const& x, auto const& y) { return slice(x,y); });
return filter_tuple(a, b, [](auto const& x, auto const& y) { return detail::lift_slice(x,y); });
} else if constexpr (is_underscore<A>::value) {
return cute::tuple<B>{b};
return b;
} else {
return cute::tuple<>{};
}
@@ -116,6 +138,28 @@ slice(A const& a, B const& b)
// Dice keeps only the elements of Tuple B that are paired with an Int
//
namespace detail {
template <class A, class B>
CUTE_HOST_DEVICE constexpr
auto
lift_dice(A const& a, B const& b)
{
if constexpr (is_tuple<A>::value) {
static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
return filter_tuple(a, b, [](auto const& x, auto const& y) { return lift_dice(x,y); });
} else if constexpr (is_underscore<A>::value) {
return cute::tuple<>{};
} else {
return cute::tuple<B>{b};
}
CUTE_GCC_UNREACHABLE;
}
} // end namespace detail
// Entry point overrides the lifting so that dice(1,b) == b
template <class A, class B>
CUTE_HOST_DEVICE constexpr
auto
@@ -123,11 +167,11 @@ dice(A const& a, B const& b)
{
if constexpr (is_tuple<A>::value) {
static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
return filter_tuple(a, b, [](auto const& x, auto const& y) { return dice(x,y); });
return filter_tuple(a, b, [](auto const& x, auto const& y) { return detail::lift_dice(x,y); });
} else if constexpr (is_underscore<A>::value) {
return cute::tuple<>{};
} else {
return cute::tuple<B>{b};
return b;
}
CUTE_GCC_UNREACHABLE;
include/cute/util/type_traits.hpp
+4
View File
@@ -74,7 +74,10 @@ using CUTE_STL_NAMESPACE::is_void_v;
using CUTE_STL_NAMESPACE::is_base_of;
using CUTE_STL_NAMESPACE::is_base_of_v;
using CUTE_STL_NAMESPACE::is_const;
using CUTE_STL_NAMESPACE::is_const_v;
using CUTE_STL_NAMESPACE::is_volatile;
using CUTE_STL_NAMESPACE::is_volatile_v;
// using CUTE_STL_NAMESPACE::true_type;
// using CUTE_STL_NAMESPACE::false_type;
@@ -115,6 +118,7 @@ template <class T>
using is_std_integral = CUTE_STL_NAMESPACE::is_integral<T>;
using CUTE_STL_NAMESPACE::is_empty;
using CUTE_STL_NAMESPACE::is_empty_v;
using CUTE_STL_NAMESPACE::invoke_result_t;