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

CUTLASS 3.2.1 (#1113)

Browse Source 
* Updates for 3.2.1 release.

* Minor fix in gemm op profiler for raster order.

* Add scheduler mapping for raster order in the kernels.
This commit is contained in:
ANIKET SHIVAM
2023-09-26 17:24:26 -04:00
committed by GitHub
parent e0aaa3c3b3
commit 90d3b0fb18
428 changed files with 22252 additions and 21761 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/axpby.hpp
+8 -1
View File
@@ -68,7 +68,14 @@ axpby(Alpha const& alpha,
Beta const& beta,
Tensor<YEngine, YLayout> & y)
{
auto isBetaZero = (beta == Int<0>{});
auto isBetaZero = [&] () {
if constexpr (is_complex<Beta>::value) {
return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
}
else {
return beta == Int<0>{};
}
} ();
CUTE_UNROLL
for (int i = 0; i < size(x); ++i) {
include/cute/algorithm/gemm.hpp
+1 -5
View File
@@ -218,7 +218,6 @@ gemm(MMA_Atom<MMA> const& mma,
CUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C)); // AM == CM
CUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C)); // BN == CN
CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D));
gemm(mma,
D, // (M,N)
make_tensor(A.data(), append<2>(A.layout())), // (M,1)
@@ -253,7 +252,7 @@ gemm(MMA_Atom<MMA> const& mma,
CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutC_TV{}) == Int<1>{});
CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutA_TV{}) == Int<1>{});
CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutB_TV{}) == Int<1>{});
gemm(mma,
make_tensor(D.data(), prepend<3>(D.layout())), // (1,M,N)
make_tensor(A.data(), prepend<3>(A.layout())), // (1,M,K)
@@ -282,7 +281,6 @@ gemm(MMA_Atom<MMA> const& mma,
CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C)); // AM == CM
CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN
CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
auto M = size<1>(A);
auto N = size<1>(B);
// REGISTER .reuse OPTIMIZATIONS
@@ -409,7 +407,6 @@ gemm(MMA_Atom<MMA> const& mma,
CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN
CUTE_STATIC_ASSERT_V(size<2>(A) == size<2>(B)); // AK == BK
CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
auto K = size<2>(A);
CUTE_UNROLL
@@ -454,7 +451,6 @@ gemm(MMA_Atom<MMA> const& mma,
CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutC_TV{}) == Int<1>{});
CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutA_TV{}) == Int<1>{});
CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutB_TV{}) == Int<1>{});
gemm(mma,
make_tensor(D.data(), prepend<3>(D.layout())), // (1,M,N)
make_tensor(A.data(), prepend<3>(A.layout())), // (1,M,K)
include/cute/algorithm/tuple_algorithms.hpp
+53 -22
View File
@@ -140,7 +140,11 @@ CUTE_HOST_DEVICE constexpr
auto
transform_apply(T&& t, F&& f, G&& g)
{
return detail::tapply(static_cast<T&&>(t), f, g, tuple_seq<T>{});
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
return detail::tapply(static_cast<T&&>(t), f, g, tuple_seq<T>{});
} else {
return g(f(static_cast<T&&>(t)));
}
}
template <class T0, class T1, class F, class G>
@@ -148,7 +152,11 @@ CUTE_HOST_DEVICE constexpr
auto
transform_apply(T0&& t0, T1&& t1, F&& f, G&& g)
{
return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), f, g, tuple_seq<T0>{});
if constexpr (is_tuple<remove_cvref_t<T0>>::value) {
return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), f, g, tuple_seq<T0>{});
} else {
return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1)));
}
}
template <class T0, class T1, class T2, class F, class G>
@@ -156,7 +164,11 @@ CUTE_HOST_DEVICE constexpr
auto
transform_apply(T0&& t0, T1&& t1, T2&& t2, F&& f, G&& g)
{
return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), static_cast<T2&&>(t2), f, g, tuple_seq<T0>{});
if constexpr (is_tuple<remove_cvref_t<T0>>::value) {
return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), static_cast<T2&&>(t2), f, g, tuple_seq<T0>{});
} else {
return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1), static_cast<T2&&>(t2)));
}
}
//
@@ -306,21 +318,16 @@ transform_leaf(T0 const& t0, T1 const& t1, F&& f)
namespace detail {
template <class T, class F>
CUTE_HOST_DEVICE constexpr
auto
find_if(T const& t, F&& f, seq<>)
{
return cute::integral_constant<int, tuple_size<T>::value>{};
}
template <class T, class F, int I, int... Is>
CUTE_HOST_DEVICE constexpr
auto
find_if(T const& t, F&& f, seq<I,Is...>)
{
if constexpr (decltype(f(get<I>(t)))::value) {
return cute::integral_constant<int, I>{};
return cute::C<I>{};
} else
if constexpr (sizeof...(Is) == 0) {
return cute::C<I+1>{};
} else {
return find_if(t, f, seq<Is...>{});
}
@@ -338,7 +345,7 @@ find_if(T const& t, F&& f)
if constexpr (is_tuple<T>::value) {
return detail::find_if(t, f, tuple_seq<T>{});
} else {
return cute::integral_constant<int, decltype(f(t))::value ? 0 : 1>{};
return cute::C<decltype(f(t))::value ? 0 : 1>{};
}
CUTE_GCC_UNREACHABLE;
@@ -355,12 +362,12 @@ find(T const& t, X const& x)
template <class T, class F>
CUTE_HOST_DEVICE constexpr
auto
none_of(T const& t, F&& f)
any_of(T const& t, F&& f)
{
if constexpr (is_tuple<T>::value) {
return cute::integral_constant<bool, decltype(find_if(t, f))::value == tuple_size<T>::value>{};
return detail::apply(cute::transform(t, f), [&] (auto const&... a) { return (false_type{} || ... || a); }, tuple_seq<T>{});
} else {
return not f(t);
return f(t);
}
CUTE_GCC_UNREACHABLE;
@@ -372,8 +379,7 @@ auto
all_of(T const& t, F&& f)
{
if constexpr (is_tuple<T>::value) {
auto not_f = [&](auto const& a) { return not f(a); };
return cute::integral_constant<bool, decltype(find_if(t, not_f))::value == tuple_size<T>::value>{};
return detail::apply(t, [&] (auto const&... a) { return (true_type{} && ... && f(a)); }, tuple_seq<T>{});
} else {
return f(t);
}
@@ -384,9 +390,9 @@ all_of(T const& t, F&& f)
template <class T, class F>
CUTE_HOST_DEVICE constexpr
auto
any_of(T const& t, F&& f)
none_of(T const& t, F&& f)
{
return not none_of(t, f);
return not any_of(t, f);
}
//
@@ -410,6 +416,14 @@ filter_tuple(T0 const& t0, T1 const& t1, F&& f)
return transform_apply(t0, t1, f, [](auto const&... a) { return cute::tuple_cat(a...); });
}
template <class T0, class T1, class T2, class F>
CUTE_HOST_DEVICE constexpr
auto
filter_tuple(T0 const& t0, T1 const& t1, T2 const& t2, F&& f)
{
return transform_apply(t0, t1, t2, f, [](auto const&... a) { return cute::tuple_cat(a...); });
}
//
// Fold (Reduce, Accumulate)
// (t, v, f) => f(...f(f(v,t_0),t_1),...,t_n)
@@ -595,6 +609,13 @@ unwrap(T const& t)
//
// Flatten a hierarchical tuple to a tuple of depth one.
//
//
template <class T>
struct is_flat : true_type {};
template <class... Ts>
struct is_flat<tuple<Ts...>> : bool_constant<(true && ... && (not is_tuple<Ts>::value))> {};
template <class T>
CUTE_HOST_DEVICE constexpr
@@ -602,7 +623,12 @@ auto
flatten_to_tuple(T const& t)
{
if constexpr (is_tuple<T>::value) {
return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
if constexpr (is_flat<T>::value) {
return t;
} else
{
return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
}
} else {
return cute::make_tuple(t);
}
@@ -616,7 +642,12 @@ auto
flatten(T const& t)
{
if constexpr (is_tuple<T>::value) {
return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
if constexpr (is_flat<T>::value) {
return t;
} else
{
return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
}
} else {
return t;
}
include/cute/arch/copy_sm90_desc.hpp
+1 -1
View File
@@ -177,7 +177,7 @@ to_CUtensorMapSwizzle(SmemSwizzleBits const& t) {
#if (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
using TmaDescriptor = CUtensorMap;
#else
using TmaDescriptor = struct { char bytes[128]; };
using TmaDescriptor = struct alignas(64) { char bytes[128]; };
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Initiates a TensorMap Prefetch
include/cute/arch/mma_sm80.hpp
+14 -3
View File
@@ -37,8 +37,19 @@
// Config
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
# define CUTE_ARCH_MMA_SM80_ENABLED
#if (__CUDA_ARCH__ <= 900)
#define CUTE_ARCH_MMA_B1_AND_SM80_ENABLED
#endif
#if (__CUDA_ARCH__ <= 890)
#define CUTE_ARCH_MMA_B1_XOR_SM80_ENABLED
#endif
#endif
namespace cute {
////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -2044,7 +2055,7 @@ struct SM80_8x8x128_S32U1U1S32_TN_XORPOPC
uint32_t const& b0,
uint32_t const& c0, uint32_t const& c1)
{
#if defined(CUTE_ARCH_MMA_SM80_ENABLED)
#if defined(CUTE_ARCH_MMA_B1_XOR_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.xor.popc "
"{%0, %1},"
@@ -2077,7 +2088,7 @@ struct SM80_16x8x128_S32U1U1S32_TN_XORPOPC
uint32_t const& b0,
uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3)
{
#if defined(CUTE_ARCH_MMA_SM80_ENABLED)
#if defined(CUTE_ARCH_MMA_B1_XOR_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.xor.popc "
"{%0, %1, %2, %3},"
@@ -2110,7 +2121,7 @@ struct SM80_16x8x256_S32U1U1S32_TN_XORPOPC
uint32_t const& b0, uint32_t const& b1,
uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3)
{
#if defined(CUTE_ARCH_MMA_SM80_ENABLED)
#if defined(CUTE_ARCH_MMA_B1_XOR_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc "
"{%0, %1, %2, %3},"
include/cute/arch/mma_sm90.hpp
+4 -3
View File
@@ -38,6 +38,7 @@
// Config
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
# define CUTE_ARCH_MMA_SM90_ENABLED
# define CUTE_ARCH_MMA_F64_SM90_ENABLED
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -60,7 +61,7 @@ struct SM90_16x8x4_F64F64F64F64_TN
double const& b0,
double const& c0, double const& c1, double const& c2, double const& c3)
{
#if defined(CUTE_ARCH_MMA_SM90_ENABLED)
#if defined(CUTE_ARCH_MMA_F64_SM90_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64"
"{%0, %1, %2, %3},"
@@ -93,7 +94,7 @@ struct SM90_16x8x8_F64F64F64F64_TN
double const& b0, double const& b1,
double const& c0, double const& c1, double const& c2, double const& c3)
{
#if defined(CUTE_ARCH_MMA_SM90_ENABLED)
#if defined(CUTE_ARCH_MMA_F64_SM90_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k8.row.col.f64.f64.f64.f64"
"{%0, %1, %2, %3},"
@@ -127,7 +128,7 @@ struct SM90_16x8x16_F64F64F64F64_TN
double const& b0, double const& b1, double const& b2, double const& b3,
double const& c0, double const& c1, double const& c2, double const& c3)
{
#if defined(CUTE_ARCH_MMA_SM90_ENABLED)
#if defined(CUTE_ARCH_MMA_F64_SM90_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f64.f64.f64.f64"
"{%0, %1, %2, %3},"
include/cute/arch/mma_sm90_desc.hpp
+7 -7
View File
@@ -86,22 +86,22 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) {
union GmmaDescriptor
{
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
GmmaDescriptor() noexcept : desc_(0) {}
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
GmmaDescriptor(uint64_t desc) noexcept : desc_(desc) {}
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
GmmaDescriptor(GmmaDescriptor const& t) noexcept : desc_(t.desc_) {}
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
GmmaDescriptor(GmmaDescriptor && t) noexcept : desc_(t.desc_) {}
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
GmmaDescriptor& operator=(GmmaDescriptor const& t) noexcept {
desc_ = t.desc_;
return *this;
}
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
GmmaDescriptor& operator=(GmmaDescriptor && t) noexcept {
desc_ = t.desc_;
return *this;
include/cute/atom/copy_traits_sm90_tma.hpp
+303 -171
View File
@@ -38,9 +38,17 @@
#include <cute/atom/copy_traits.hpp>
#include <cute/atom/copy_atom.hpp>
#include <cute/numeric/integral_ratio.hpp>
namespace cute
{
template <class GmemStrides_, class TmaGBasis_, class TmaSwizzle_>
struct AuxTmaParams {
using GmemStrides = GmemStrides_;
GmemStrides g_stride_;
};
//////////////////////////////////////////////////////////////////////////////
///////////////////////////// TMA_LOAD ///////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
@@ -88,14 +96,14 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
{
static_assert(is_smem<TD>::value, "Expected smem dst for SM90_TMA_LOAD");
traits.copy_unpack_(raw_pointer_cast(dst.data()), src.data().coord_, tuple_seq<decltype(src.data().coord_)>{});
traits.copy_unpack_(cute::raw_pointer_cast(dst.data()), src.data().coord_, tuple_seq<decltype(src.data().coord_)>{});
}
};
// The non-executable SM90_TMA_LOAD with tma_desc and no tma_mbar
// Use .with(tma_mbar) to construct an executable version
template <class NumBitsPerTMA, class GmemStrides>
struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, GmemStrides>
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_1>;
@@ -109,7 +117,8 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, GmemStrides>
// SM90_TMA_LOAD arguments
TmaDescriptor tma_desc_;
GmemStrides g_stride_;
using AuxParams = AuxParams_;
AuxParams aux_params_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
@@ -133,8 +142,8 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, GmemStrides>
CUTE_HOST_DEVICE constexpr
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, g_stride_));
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM90_TMA_LOAD before calling .with()
@@ -190,12 +199,12 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
{
static_assert(is_smem<TD>::value, "Expected smem dst for SM90_TMA_LOAD_MULTICAST");
traits.copy_unpack_(raw_pointer_cast(dst.data()), src.data().coord_, tuple_seq<decltype(src.data().coord_)>{});
traits.copy_unpack_(cute::raw_pointer_cast(dst.data()), src.data().coord_, tuple_seq<decltype(src.data().coord_)>{});
}
};
template <class NumBitsPerTMA, class GmemStrides>
struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, GmemStrides>
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_1>;
@@ -209,7 +218,8 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, GmemStrides>
// SM90_TMA_LOAD_MULTICAST arguments
TmaDescriptor tma_desc_;
GmemStrides g_stride_;
using AuxParams = AuxParams_;
AuxParams aux_params_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
@@ -230,8 +240,8 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, GmemStrides>
CUTE_HOST_DEVICE constexpr
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, g_stride_));
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM90_TMA_LOAD_MULTICAST before calling .with()
@@ -248,8 +258,8 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, GmemStrides>
//////////////////////////////////////////////////////////////////////////////
// The executable SM90_TMA_STORE with tma_desc
template <class NumBitsPerTMA, class GmemStrides>
struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, GmemStrides>
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_1>;
@@ -263,7 +273,8 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, GmemStrides>
// SM90_TMA_STORE arguments
TmaDescriptor tma_desc_;
GmemStrides g_stride_;
using AuxParams = AuxParams_;
AuxParams aux_params_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
@@ -277,8 +288,8 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, GmemStrides>
CUTE_HOST_DEVICE constexpr
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, g_stride_));
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
template <class Coord, int... Is>
@@ -305,7 +316,7 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, GmemStrides>
static_assert(is_smem<TS>::value, "Expected smem src for SM90_TMA_STORE");
//static_assert(is_gmem<TD>::value, "Expected gmem dst for SM90_TMA_STORE"); // TMA spoofed src tensor
traits.copy_unpack_(raw_pointer_cast(src.data()), dst.data().coord_, tuple_seq<decltype(dst.data().coord_)>{});
traits.copy_unpack_(cute::raw_pointer_cast(src.data()), dst.data().coord_, tuple_seq<decltype(dst.data().coord_)>{});
}
};
@@ -417,9 +428,78 @@ struct Copy_Traits<SM90_BULK_COPY_AUTO, OpArgs...>
namespace detail {
// Use a smem2gmode map to read through the GMEM tensor
// and construct a TMA Descriptor for the resulting instruction
template <class GEngine, class GLayout,
// Custom version of coalesce that greedily combines modes only up to size-256
// Look at each element and the back of the stack (in order of priority)
// back(NewLayout) get<I>(OldLayout)
// s0:d0 _1:d1 => continue
// _1:d0 s1:d1 => replace_back s1:d1
// s0:d0 s1:s0*d0 => replace_back s0*s1:d0 if s0*s1 <= 256
// s0:d0 s1:d1 => append s1:d1
//
// @pre OldShape and OldStride are flat
template <int I, class OldShape, class OldStride, class NewShape, class NewStride>
CUTE_HOST_DEVICE constexpr
auto
coalesce_256_impl(OldShape const& old_shape, OldStride const& old_stride,
NewShape const& new_shape, NewStride const& new_stride)
{
if constexpr (I == rank_v<OldShape>) {
// Base case, we're done
if constexpr (is_constant<1, NewShape>::value) {
return Layout<_1,_0>{};
} else {
return Layout<NewShape,NewStride>{new_shape,new_stride};
}
} else if constexpr (is_constant<1, decltype(get<I>(old_shape))>::value) {
// shape<I>(layout) == _1, skip it and continue
return coalesce_256_impl<I+1>(old_shape, old_stride, new_shape, new_stride);
} else if constexpr (is_constant<1, NewShape>::value) {
// Replace our shape-1 with anything (Can only happen on input new_shape/new_stride)
return coalesce_256_impl<I+1>(old_shape, old_stride, get<I>(old_shape), get<I>(old_stride));
} else if constexpr (is_constant<true, decltype(back(new_shape) * back(new_stride) == get<I>(old_stride) &&
get<I>(old_shape) * back(new_shape) <= Int<256>{})>::value) {
// Merge modes because the shapes and strides match and the merge is 256 or less
return coalesce_256_impl<I+1>(old_shape, old_stride,
replace_back(new_shape, get<I>(old_shape) * back(new_shape)),
new_stride);
} else {
// Can't replace or merge, so append a new mode
return coalesce_256_impl<I+1>(old_shape, old_stride,
append(new_shape, get<I>(old_shape)),
append(new_stride, get<I>(old_stride)));
}
CUTE_GCC_UNREACHABLE;
}
// Combine all the modes that are possible to combine
// Does not respect the profile of the layout, but does preserve total size
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
coalesce_256(Layout<Shape,Stride> const& layout)
{
auto flat_shape = flatten(layout.shape());
auto flat_stride = flatten(layout.stride());
return coalesce_256_impl<1>(flat_shape, flat_stride, get<0>(flat_shape), get<0>(flat_stride));
}
template <class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
coalesce_256(Tensor<Engine,Layout> const& tensor)
{
return make_tensor(tensor.data(), coalesce_256(tensor.layout()));
}
// Use a smem_inv_h to read through the GMEM 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,
int B, int M, int S>
CUTE_HOST_RTC
@@ -428,63 +508,78 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
Layout<SShape,SStride> const& smem_inv_h, // smem_idx to hier gmode
Swizzle<B,M,S> const& swizzle) // Swizzle fn on smem_idx
{
using T = typename GEngine::value_type;
// This is the gmem "vector" that corresponds to the smem vector in memory (smem_box_shape):(gmem_prob_stride)
Tensor tma_gstride = recast<T>(gtensor.compose(smem_inv_h));
// If the sizes of smem_inv_h and tma_gstride don't match, then a non-trivial recast was performed.
// In that case, require that the recasted modes all have size-1 so TMA can identity them and skip them.
for_each(zip(flatten(shape(smem_inv_h)), flatten(shape(tma_gstride))), [] (auto s_and_g) {
auto [s,g] = s_and_g;
CUTE_STATIC_ASSERT_V(s == g or g == Int<1>{},
"A non-trivial recast was performed, but TMA cannot identify which modes to leave out.");
});
// The smem vector is the same units as gtensor, so compose first and then recast
// tma_val_idx:gmem_strides
Tensor tile_gstride = recast<TmaInternalType>(gtensor.compose(smem_inv_h));
// Coalesce modes up to size-256 (the maximum TMA box extent in units of TmaInternalType)
// tma_box_shape:gmem_strides
Tensor tma_gstride = coalesce_256(tile_gstride);
// Perform the tiling to the gmem vector again, but with indirections to the gtensor modes
auto gbasis = make_identity_layout(shape(gtensor));
auto tma_gbasis_tile_tmp = gbasis.compose(smem_inv_h);
// Instead of the recast (gbasis doesn't have type info), replace the shape with the already-recasted shape and coalesce out any size-1 modes
auto tma_gbasis_tile = coalesce(make_layout(shape(tma_gstride), stride(tma_gbasis_tile_tmp)));
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));
// Recast the original tensor for shape inspections
auto glayout_T = recast<T>(gtensor).layout();
auto gtensor_T = recast<TmaInternalType>(gtensor);
// Find missing bases that don't belong to a size-1 mode of the recast input
// Find missing bases that don't appear in tile_gbasis
// NOTE This is essentially ArithmeticTuple complement...
// NOTE in persuit of implementing an ArithmeticTuple logical_divide for smem_inv_h
auto tma_gbasis_full = fold(zip(flatten(shape(glayout_T)), flatten(stride(gbasis))), tma_gbasis_tile,
[](auto tma_g, auto s_and_d) {
auto [s,d] = s_and_d;
auto k = find(stride(tma_g), d); // Find the basis in tma_gstride
if constexpr (decltype(k != rank(tma_g) || is_constant<1, decltype(s)>{})::value) {
// If d was found or s is static-1, then don't append
return tma_g;
// 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)
{
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
} else {
// Else, append the missing basis
return append(tma_g, make_layout(Int<1>{}, d));
using E = decltype(e);
auto has_e = any_of(stride(tile_gbasis), [] (auto tb) { return tb == E{}; });
if constexpr (decltype(has_e)::value) {
return cute::tuple<>{}; // If d was found, then don't append
} else {
return cute::tuple<E>(e); // Else, this is missing so append
}
}
});
auto tile_gbasis_remaining_rank = rank(tile_gbasis_remaining_stride);
// Group the trailing modes to make this max rank-5
// "Coalesce" the tile basis into a compatible shape with the tma
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>{}))),
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
// tma_box_shape:gmem_mode
auto tma_gbasis = group<cute::min(rank(tma_gbasis_full),4),-1>(tma_gbasis_full);
#if 0
print("gtensor : "); print(gtensor); print("\n");
print("smem_inv_h : "); print(smem_inv_h); print("\n");
print("gtensor : "); print(gtensor); print("\n");
print("tile_gstride : "); print(tile_gstride); print("\n");
print("tma_gstride : "); print(tma_gstride); print("\n");
print("gbasis : "); print(gbasis); print("\n");
print("tma_gb_tile : "); print(tma_gbasis_tile ); print("\n");
print("tile_gbasis : "); print(tile_gbasis); print("\n");
print("tma_gbasis : "); print(tma_gbasis); print("\n");
#endif
//
// TMA desc creation
//
constexpr int tma_dim = decltype(rank(tma_gbasis))::value;
//
// TMA gmem desc info
//
void* gmem_address = (void*) raw_pointer_cast(gtensor.data());
void* gmem_address = (void*) raw_pointer_cast(gtensor_T.data());
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};
@@ -492,12 +587,12 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
for_each(make_seq<tma_dim>{}, [&](auto i) {
for_each(stride<i>(tma_gbasis), [&](auto ej) {
// Problem stride
uint64_t stride_j = basis_get(ej, stride(glayout_T)) * sizeof(T);
uint64_t stride_j = ceil_div(basis_get(ej, stride(gmem_layout)) * sizeof_bits_v<TmaInternalType>, 8);
uint64_t old_stride = gmem_prob_stride[i];
gmem_prob_stride[i] = gcd(gmem_prob_stride[i], stride_j);
// Problem shape
uint64_t shape_j = basis_get(ej, shape(glayout_T));
uint64_t shape_j = basis_get(ej, shape(gmem_layout));
if (gmem_prob_stride[i] != 0) {
// Recurrence: g_shape = (s_i - 1) * (d_i / gcd_j d_j) + 1
gmem_prob_shape[i] = (gmem_prob_shape[i]-1) * (old_stride / gmem_prob_stride[i])
@@ -522,8 +617,8 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
assert(gmem_prob_shape[4] >= (uint64_t(1))); // Size must be min 1
assert(gmem_prob_shape[4] <= (uint64_t(1) << 32)); // Size must be max 2^32
// TMA descriptor does not store the zeroth stride and assumes it is sizeof(T) == one element.
assert(gmem_prob_stride[0] == sizeof(T) && "Majorness of smem doesn't match majorness of gmem");
// TMA descriptor does not store the zeroth stride and assumes it is 1 (TmaInternalType element).
assert(gmem_prob_stride[0] == sizeof(TmaInternalType) && "Majorness of smem doesn't match majorness of gmem");
assert((gmem_prob_stride[1]) < (uint64_t(1) << 40)); // Stride must be max 2^40
assert((gmem_prob_stride[1] & 0b1111) == 0); // Stride must be multiple of 16B (128b)
@@ -545,14 +640,16 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
smem_box_shape[i] *= size<i>(tma_gbasis);
});
assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1
assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1
assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1
assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1
assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[0] >= (uint32_t(1))); // Size must be min 1
assert(smem_box_shape[0] <= (uint32_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[1] >= (uint32_t(1))); // Size must be min 1
assert(smem_box_shape[1] <= (uint32_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[2] >= (uint32_t(1))); // Size must be min 1
assert(smem_box_shape[2] <= (uint32_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[3] >= (uint32_t(1))); // Size must be min 1
assert(smem_box_shape[3] <= (uint32_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_shape[4] >= (uint32_t(1))); // Size must be min 1
assert(smem_box_shape[4] <= (uint32_t(1) << 8)); // Size must be max 2^8 = 256
assert(smem_box_stride[0] >= (uint32_t(1))); // Stride must be min 1
assert(smem_box_stride[0] <= (uint32_t(8))); // Stride must be max 2^3 = 8
@@ -565,88 +662,101 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The original GM
assert(smem_box_stride[4] >= (uint32_t(1))); // Stride must be min 1
assert(smem_box_stride[4] <= (uint32_t(8))); // Stride must be max 2^3 = 8
//
// Construct the descriptor
//
TmaDescriptor tma_desc = {0};
//
// TMA general info
//
#if (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
CUtensorMapDataType tma_format = TMA::to_CUtensorMapDataType<T>();
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(
&tma_desc,
tma_format,
tma_dim,
gmem_address,
gmem_prob_shape.data(),
gmem_prob_stride.data() + 1, // gmem_prob_stride[0] implicitly 1
smem_box_shape.data(),
smem_box_stride.data(),
tma_interleave,
smem_swizzle,
tma_l2Promotion,
tma_oobFill);
if (result != CUDA_SUCCESS) {
std::cerr << "TMA Desc Addr: " << &tma_desc
<< "\nformat " << tma_format
<< "\ndim " << tma_dim
<< "\ngmem_address " << gmem_address
<< "\nglobalDim " << gmem_prob_shape
<< "\nglobalStrides " << gmem_prob_stride
<< "\nboxDim " << smem_box_shape
<< "\nelementStrides " << smem_box_stride
<< "\ninterleave " << tma_interleave
<< "\nswizzle " << smem_swizzle
<< "\nl2Promotion " << tma_l2Promotion
<< "\noobFill " << tma_oobFill << std::endl;
std::cerr << "Error: Failed to initialize the TMA descriptor " << result << std::endl;
assert(false);
}
#endif // (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
//
// 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(
&tma_desc,
tma_format,
tma_dim,
gmem_address,
gmem_prob_shape.data(),
gmem_prob_stride.data() + 1, // gmem_prob_stride[0] implicitly 1
smem_box_shape.data(),
smem_box_stride.data(),
tma_interleave,
smem_swizzle,
tma_l2Promotion,
tma_oobFill);
if (result != CUDA_SUCCESS) {
std::cerr << "TMA Desc Addr: " << &tma_desc
<< "\nformat " << tma_format
<< "\ndim " << tma_dim
<< "\ngmem_address " << gmem_address
<< "\nglobalDim " << gmem_prob_shape
<< "\nglobalStrides " << gmem_prob_stride
<< "\nboxDim " << smem_box_shape
<< "\nelementStrides " << smem_box_stride
<< "\ninterleave " << tma_interleave
<< "\nswizzle " << smem_swizzle
<< "\nl2Promotion " << tma_l2Promotion
<< "\noobFill " << tma_oobFill << std::endl;
std::cerr << "Error: Failed to initialize the TMA descriptor " << result << std::endl;
assert(false);
}
#endif // (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
// 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) {
auto si = basis_get(ei, shape(glayout_T));
auto di = basis_get(ei, stride(glayout_T));
auto tma_gbasis_stride = stride(tma_gbasis);
// Find j such that E<i> is in stride<j>(tma_gbasis)
[[maybe_unused]] auto j = find_if(tma_gbasis_stride, [&](auto tma_stride_j) { return any_of(tma_stride_j, [&](auto dj) { return dj == ei; }); });
// Return the TMA basis this gmode contributes to
if constexpr (is_constant<1, decltype(si)>::value || decltype(j == rank(tma_gbasis_stride))::value) {
return Int<0>{}; // Return arithmetic identity -- no contribution to the TMA
} else
if constexpr (decltype(rank<j>(tma_gbasis_stride) == Int<1>{})::value) {
return E<j>{}; // We know that the scale factor is Int<1>{}
auto si = basis_get(ei, shape(gmem_layout));
auto di = basis_get(ei, stride(gmem_layout));
if constexpr (is_constant<1, decltype(si)>::value || is_constant<0, decltype(di)>::value) {
return Int<0>{}; // If size-1 or stride-0, return arithmetic identity -- no contribution to the TMA
} else {
return E<j>{} * int32_t(di * sizeof(T) / cute::max(gmem_prob_stride[j], 16));
auto tma_gbasis_stride = stride(tma_gbasis);
// Find j such that E<i> is in stride<j>(tma_gbasis)
using EI = decltype(ei);
[[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
if constexpr (decltype(j == Int<0>{})::value) {
auto scale = ratio(size(tma_gstride), size(smem_inv_h)) * 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) {
return E<j>{}; // Return TMA Coord basis -- known scale of Int<1>{}
} else {
int32_t scale = ceil_div(int32_t(di * sizeof_bits_v<TmaInternalType> / cute::max(gmem_prob_stride[j], 16)), 8);
return E<j>{} * scale; // Return TMA Coord basis -- with a dynamic scale factor
}
}
});
#if 0
print("gmem_stride_bases : "); print(gmem_stride_bases); print("\n");
#endif
#if 0
print("tma_gbasis : "); print(gmem_stride_bases); print("\n");
#endif
return cute::make_tuple(tma_desc, gmem_stride_bases);
using AuxParams = AuxTmaParams<decltype(gmem_stride_bases),
decltype(tma_gbasis),
decltype(swizzle)>;
return cute::make_tuple(tma_desc, AuxParams{gmem_stride_bases});
}
// The "logical TMA tid" is a map from the CTA rank to its logical id
// within the instruction. It works like a mask or ordering on the
// CTAs. For non-multicast TMA, all CTAs should map to 0. For
// multicast TMA of size 4, CTAs will be mapped to {0,1,2,3}.
template <class CopyOp,
template <class TmaInternalType,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class TShape, class TStride,
@@ -657,7 +767,7 @@ make_tma_copy_tiled(CopyOp,
Tensor<GEngine,GLayout> const& gtensor, // Full GMEM Tensor
SLayout const& slayout, // CTA Tile of SMEM
Layout<TShape,TStride> const& cta_t_map, // T: CTA thr idx -> logical TMA tid
Layout<VShape,VStride> const& cta_v_map) // V: CTA val idx -> gmem coord
Layout<VShape,VStride> const& cta_v_map) // V: CTA val idx -> gmem mode
{
//
// TMA parameter checking
@@ -673,18 +783,19 @@ make_tma_copy_tiled(CopyOp,
//
// Invert the smem to get the largest contiguous vector in the smem layout
// smem idx -> smem coord
auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout));
// trunc_smem_idx -> trunc_smem_coord
// Map from smem idx to a gmem mode
// Compose with the V-Map to convert smem coord (CTA val idx) to gmem mode
// smem idx -> gmem mode
auto sidx_to_gmode = coalesce(composition(cta_v_map, inv_smem_layout));
#if 0
print("g_layout : "); print(gtensor.layout()); print("\n");
print("g_tensor : "); print(gtensor); print("\n");
print("s_layout : "); print(slayout); print("\n");
print("cta_t_map : "); print(cta_t_map); print("\n");
print("cta_v_map : "); print(cta_v_map); print("\n");
print("inv_smem : "); print(inv_smem_layout); print("\n");
print("inv_s_layout : "); print(inv_smem_layout); print("\n");
print("sidx_to_gmode : "); print(sidx_to_gmode); print("\n");
#endif
@@ -693,9 +804,11 @@ make_tma_copy_tiled(CopyOp,
//
// Generate a TupleBasis for the gtensor
// gmem coord -> gmem coord
auto glayout_basis = make_identity_layout(shape(gtensor));
// Tile the modes of gtensor with the truncated cta_v_map o inv_smem_layout_trunc
// smem idx -> gmem coord
auto tma_layout_full = flatten(composition(glayout_basis, sidx_to_gmode));
// Truncate any incompatibilities -- no starting in the middle of gmodes
@@ -704,61 +817,60 @@ make_tma_copy_tiled(CopyOp,
return not is_constant<1,decltype(v)>{};
});
static_assert(smem_rank > 0, "Could not find a common tile-gmem vectorization. Does the Tile select out major GMEM modes?");
// TMA uses a maximum of 5 modes
// If the gtensor has more than 5 modes, we need to reserve the last TMA-mode as a "multimode"
constexpr int smem_tma_rank = cute::min(int(smem_rank), (rank(tma_layout_full) > Int<5>{} ? 4 : 5));
// Keep only the static-1 basis modes into gmem
auto tma_layout_trunc = take<0,smem_tma_rank>(tma_layout_full);
auto tma_layout_trunc = take<0,smem_rank>(tma_layout_full);
// Split according to the portion each multicast CTA will be responsible for
auto tma_layout_vt = logical_divide(tma_layout_trunc, shape_div(size(tma_layout_trunc), cosize(cta_t_map)));
// Keep only the portion each multicast CTA will be responsible for
auto tma_layout_v = composition(tma_layout_trunc, shape_div(size(tma_layout_trunc), cosize(cta_t_map)));
#if 0
print("glayout_basis : "); print(glayout_basis); print("\n");
print("tma_layout_full : "); print(tma_layout_full); print("\n");
print("tma_layout_trunc: "); print(tma_layout_trunc); print("\n");
print("tma_layout_vt : "); print(tma_layout_vt); print("\n");
print("tma_layout_v : "); print(tma_layout_v); print("\n");
#endif
//
// Construct the TMA Desc and GMEM mode ordering
// Construct the TMA Desc and the strides of the TMA Tensor
//
auto [tma_desc, gmem_stride_bases] = detail::make_tma_copy_desc(gtensor, layout<0>(tma_layout_vt), get_swizzle_portion(slayout));
auto [tma_desc, aux_params] = detail::make_tma_copy_desc<TmaInternalType>(gtensor,
tma_layout_v,
get_swizzle_portion(slayout));
//
// Construct the Copy_Traits
//
using T = typename GEngine::value_type;
constexpr int num_bits_per_tma = decltype(size<0>(tma_layout_vt))::value * sizeof(T) * 8;
using Traits = Copy_Traits<CopyOp, cute::C<num_bits_per_tma>, decltype(gmem_stride_bases)>;
constexpr int num_bits_per_tma = decltype(size(tma_layout_trunc))::value * sizeof_bits_v<T>;
using Traits = Copy_Traits<CopyOp, cute::C<num_bits_per_tma>, decltype(aux_params)>;
using Atom = Copy_Atom<Traits, T>;
Traits tma_traits{tma_desc, aux_params};
#if 0
print("num_bits : "); print(NumBitsPerTMA{}); print("\n");
print("g_stride_bases: "); print(gmem_stride_bases); print("\n");
print("num_bits_per_tma : "); print(num_bits_per_tma); print("\n");
print("g_stride_bases : "); print(tma_traits.aux_params_.g_stride_); print("\n");
#endif
Traits tma_traits{tma_desc, gmem_stride_bases};
//
// Construct the TiledCopy
//
auto cta_tiler = product_each(shape(cta_v_map));
// (CTA V, CTA T) -> smem_coord
auto layout_vt = composition(inv_smem_layout, make_layout(shape(tma_layout_vt)));
// Scale that up to cover all of the smem_coords
//
// The smem vector might not cover all of the tile,
// so multiply it up to cover the entire tile.
// "T" here (the parallel index) is a CTA index.
auto layout_VT = tile_to_shape(layout_vt, make_shape(size(cta_v_map)/size<1>(layout_vt), size<1>(layout_vt)));
// Flip it and change the domain of the T from logical thr to thr_idx
auto layout_TV = make_layout(composition(layout<1>(layout_VT), cta_t_map), layout<0>(layout_VT));
// CTA V -> smem_coord
auto layout_v = composition(inv_smem_layout, size(tma_layout_trunc));
auto layout_V = tile_to_shape(make_layout(layout_v), size(cta_v_map));
// CTA T -> smem idx
auto layout_t = make_layout(cosize(cta_t_map), shape_div(size(tma_layout_trunc), cosize(cta_t_map)));
// CTA TID -> smem coord
auto layout_T = composition(inv_smem_layout, composition(layout_t, cta_t_map));
// Combine with the T mapping
auto layout_TV = make_layout(layout_T, layout_V);
#if 0
print("cta_tiler : "); print(cta_tiler); print("\n");
@@ -766,8 +878,7 @@ make_tma_copy_tiled(CopyOp,
print("layout_TV : "); print(layout_TV); print("\n");
#endif
using T = typename GEngine::value_type;
return TiledCopy<Copy_Atom<Traits,T>, decltype(layout_TV), decltype(cta_tiler)>{tma_traits};
return TiledCopy<Atom, decltype(layout_TV), decltype(cta_tiler)>{tma_traits};
}
} // end namespace detail
@@ -844,6 +955,28 @@ make_tma_copy_tiled(CopyOp,
copy(tma.with(barrier, mcast_mask), tAgA, tAsA); // copy with supporting TMA params
*/
template <class TmaInternalType,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tile,
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,
Cluster_Size const& cluster_size)
{
return detail::make_tma_copy_tiled<TmaInternalType>(copy_op,
gtensor,
slayout,
make_layout(cluster_size),
make_identity_layout(cta_tile));
}
// Explicit defaulting
template <class CopyOp,
class GEngine, class GLayout,
class SLayout,
@@ -857,15 +990,14 @@ make_tma_copy(CopyOp const& copy_op,
CTA_Tile const& cta_tile,
Cluster_Size const& cluster_size)
{
return detail::make_tma_copy_tiled(copy_op,
gtensor,
slayout,
make_layout(cluster_size),
make_identity_layout(cta_tile));
using TmaInternalType = typename GEngine::value_type;
return make_tma_copy<TmaInternalType>(copy_op,
gtensor,
slayout,
cta_tile,
cluster_size);
}
// Explicit defaulting
template <class CopyOp,
class GEngine, class GLayout,
class SLayout>
include/cute/atom/mma_atom.hpp
+1 -1
View File
@@ -155,7 +155,7 @@ 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<ValTypeA, typename remove_cvref_t<ATensor>::value_type>::value, "Expecting ValTypeA type");
static_assert(is_same<get_raw_type_t<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
include/cute/atom/mma_traits_sm75.hpp
+2 -2
View File
@@ -49,11 +49,11 @@ struct MMA_Traits<SM75_16x8x8_F32F16F16F32_TN>
using Shape_MNK = Shape<_16,_8,_8>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _2,_2>>,
Stride<Stride<_32,_2>,Stride<_16,_1>>>;
Stride<Stride<_32,_1>,Stride<_16,_8>>>;
using BLayout = Layout<Shape <Shape < _4,_8>,_2>,
Stride<Stride<_16,_1>,_8>>;
using CLayout = Layout<Shape <Shape < _4,_8>,Shape < _2,_2>>,
Stride<Stride<_32,_2>,Stride<_16,_1>>>;
Stride<Stride<_32,_1>,Stride<_16,_8>>>;
};
///////////////////////////////////////////////////////////////////////////////
include/cute/config.hpp
+1 -1
View File
@@ -30,7 +30,7 @@
**************************************************************************************************/
#pragma once
#if defined(__CUDA_ARCH__) || defined(_NVHPC_CUDA) || defined(__clang__)
#if defined(__CUDA_ARCH__) || defined(_NVHPC_CUDA)
# define CUTE_HOST_DEVICE __forceinline__ __host__ __device__
# define CUTE_DEVICE __forceinline__ __device__
# define CUTE_HOST __forceinline__ __host__
include/cute/container/bit_field.hpp
+9 -1
View File
@@ -72,8 +72,16 @@ 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 < 64) ? ((uint64_t(1) << NumBits) - 1) : uint64_t(-1);
static constexpr value_type mask = (NumBits < 64u) ? ((uint64_t(1) << MollifiedNumBits) - 1) : uint64_t(-1);
// 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
include/cute/container/tuple.hpp
+35 -3
View File
@@ -76,6 +76,10 @@ namespace detail
template <size_t N, class T, bool IsEmpty = is_empty<T>::value>
struct EBO;
template <class T, size_t N, bool B>
CUTE_HOST_DEVICE constexpr C<N> findt(EBO<N, T, B> const&)
{ return {}; }
// Specialization for types T that have no data;
// the "static tuple leaf." Valid T here include
// integral_constant<U, Value>, Int<Value>,
@@ -218,6 +222,20 @@ get(tuple<T...>&& t) noexcept
return detail::getv<I>(static_cast<tuple<T...>&&>(t));
}
//
// find a type X within a cute::tuple
// Requires X to be unique in tuple
// Returns a static integer
//
template <class X, class... T>
CUTE_HOST_DEVICE constexpr
auto
find(tuple<T...> const& t) noexcept
{
return detail::findt<X>(t);
}
//
// Custom is_tuple trait simply checks the existence of tuple_size
// and assumes std::get<I>(.), std::tuple_element<I,.>
@@ -225,7 +243,7 @@ get(tuple<T...>&& t) noexcept
namespace detail {
template <class T>
auto has_tuple_size( T*) -> integral_constant<bool, 0 <= tuple_size<T>::value>;
auto has_tuple_size( T*) -> bool_constant<(0 <= tuple_size<T>::value)>;
auto has_tuple_size(...) -> false_type;
} // end namespace detail
@@ -347,6 +365,14 @@ tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4,
return cute::make_tuple(get<I0>(t0)..., get<I1>(t1)..., get<I2>(t2)..., get<I3>(t3)..., get<I4>(t4)...);
}
template<class T0, class T1>
struct tuple_cat_static;
template<class... T0s, class... T1s>
struct tuple_cat_static<tuple<T0s...>, tuple<T1s...>> {
using type = tuple<T0s..., T1s...>;
};
} // end namespace detail
CUTE_HOST_DEVICE constexpr
@@ -370,9 +396,15 @@ CUTE_HOST_DEVICE constexpr
auto
tuple_cat(T0 const& t0, T1 const& t1)
{
return detail::tuple_cat(t0, t1,
if constexpr (is_static<T0>::value && is_static<T1>::value &&
is_tuple<T0>::value && is_tuple<T1>::value) {
return typename detail::tuple_cat_static<T0, T1>::type{};
} else
{
return detail::tuple_cat(t0, t1,
make_index_sequence<tuple_size<T0>::value>{},
make_index_sequence<tuple_size<T1>::value>{});
}
}
template <class T0, class T1, class T2>
@@ -416,7 +448,7 @@ CUTE_HOST_DEVICE constexpr
auto
tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4, T5 const& t5, Ts const&... ts)
{
return cute::tuple_cat(cute::tuple_cat(t0,t1,t2,t3,t4), t5, ts...);
return cute::tuple_cat(cute::tuple_cat(t0,t1,t2,t3,t4), cute::tuple_cat(t5, ts...));
}
#endif
include/cute/int_tuple.hpp
+50 -61
View File
@@ -37,29 +37,20 @@
#include <cute/algorithm/tuple_algorithms.hpp>
#include <cute/numeric/integral_constant.hpp>
/** 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
{
template <class... Ts>
using IntTuple = cute::tuple<Ts...>;
// Construct an IntTuple with all value-elements
template <class... Ts>
CUTE_HOST_DEVICE constexpr
IntTuple<Ts...>
make_int_tuple(Ts const&... t)
{
return {t...};
}
// CuTe does not treat integers as tuples.
// For example, is_tuple is false, and tuple_size doesn't compile.
// Nevertheless, CuTe defines rank(Integral) as 1
// (where "Integral" is a shorthand for either run-time integers
// or CuTe's compile-time integer constants),
// so therefore get<0>(Integral) just returns its input.
// Implementation of get<0>(Integral).
// 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 decltype(auto)
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(T&& t) noexcept
{
static_assert(I == 0, "Index out of range");
@@ -67,23 +58,12 @@ 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 Tuple>
CUTE_HOST_DEVICE constexpr decltype(auto)
get(Tuple&& t) noexcept
template <size_t I0, size_t I1, size_t... Is, class T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(T&& t) noexcept
{
using get_I0_result_t = cute::remove_cvref_t<decltype(cute::get<I0>(static_cast<Tuple&&>(t)))>;
if constexpr (cute::is_integral<get_I0_result_t>::value) {
// Help MSVC deduce that the inner get<I0>(...) call is not a "local variable or temporary."
// The above if constexpr test repeats the constraint on the above get(T&&) overload.
// get<0, 0, ..., 0>(t) for cute::integral (either one of the built-in integer types like int,
// or one of CuTe's compile-time constant types) t, and for one or more zeros, just returns t.
static_assert(I1 == 0, "Index I1 is out of range");
static_assert(((Is == 0) && ...), "At least one index in Is is out of range");
return get<I0>(static_cast<Tuple&&>(t));
}
else {
return get<I1, Is...>(get<I0>(static_cast<Tuple&&>(t)));
}
return get<I1, Is...>(get<I0>(static_cast<T&&>(t)));
}
//
@@ -347,6 +327,16 @@ ceil_div(IntTupleA const& a, IntTupleB const& b)
}
/** Division for Shapes
* Case Tuple Tuple:
* Perform shape_div element-wise
* Case Tuple Int:
* Fold the division of b across each element of a
* Example: shape_div((4,5,6),40) -> shape_div((1,5,6),10) -> shape_div((1,1,6),2) -> (1,1,3)
* Case Int Tuple:
* Return shape_div(a, product(b))
* Case Int Int:
* Enforce the divisibility condition a % b == 0 || b % a == 0 when possible
* Return a / b with rounding away from 0 (that is, 1 or -1 when a < b)
*/
template <class IntTupleA, class IntTupleB>
CUTE_HOST_DEVICE constexpr
@@ -357,39 +347,28 @@ shape_div(IntTupleA const& a, IntTupleB const& b)
if constexpr (is_tuple<IntTupleB>::value) { // tuple tuple
static_assert(tuple_size<IntTupleA>::value == tuple_size<IntTupleB>::value, "Mismatched ranks");
return transform(a, b, [](auto const& x, auto const& y) { return shape_div(x,y); });
} else { // tuple int
} else { // tuple int
auto const [result, rest] = fold(a, cute::make_tuple(cute::make_tuple(), b),
[] (auto const& init, auto const& ai) {
return cute::make_tuple(append(get<0>(init), shape_div(ai, get<1>(init))), shape_div(get<1>(init), ai));
});
return result;
}
} else {
if constexpr (is_tuple<IntTupleB>::value) { // int tuple
return shape_div(a, product(b));
} else { // int int
//assert(a % b == 0 || b % a == 0);
return a / b != 0 ? a / b : signum(a) * signum(b); // divide with rounding away from zero
}
} else
if constexpr (is_tuple<IntTupleB>::value) { // int tuple
return shape_div(a, product(b));
} else
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
//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
}
CUTE_GCC_UNREACHABLE;
}
/** Division for Shapes that are static constants
* @pre t % u == 0 || u % t == 0
* @result if t % u == 0, then t / u
* if u % t == 0, then signum(t) * signum(u)
*/
template <class T, T t, class U, U u>
CUTE_HOST_DEVICE constexpr
constant<decltype(shape_div(t,u)), shape_div(t,u)>
shape_div(constant<T, t> const&, constant<U, u> const&)
{
static_assert(t % u == 0 || u % t == 0, "Static shape_div failure");
return {};
}
/** Minimum for Shapes
*/
template <class IntTupleA, class IntTupleB>
@@ -581,7 +560,7 @@ make_int_tuple(Indexable const& t, int n, T const& init)
/** Fill the dynamic values of a Tuple with values from another Tuple
* \code
* auto params = make_int_tuple(6,3,4);
* auto params = make_tuple(6,3,4);
* cute::tuple<Int<1>, cute::tuple<int, int, Int<3>>, int, Int<2>> result;
* fill_int_tuple_from(result, params); // (_1,(6,3,_3),4,_2)
* \endcode
@@ -893,7 +872,8 @@ increment(Coord& coord, Shape const& shape)
struct ForwardCoordIteratorSentinal
{};
// A forward iterator for a coordinate that starts from zero and goes to shape
// A forward iterator for a starting coordinate in a shape's domain, and a shape.
// The starting coordinate may be zero but need not necessarily be.
template <class Coord, class Shape>
struct ForwardCoordIterator
{
@@ -905,7 +885,7 @@ struct ForwardCoordIterator
CUTE_HOST_DEVICE constexpr
ForwardCoordIterator& operator++() { increment(coord, shape); return *this; }
// Sentinal for the end of the implied range
// Sentinel for the end of the implied range
CUTE_HOST_DEVICE constexpr
bool operator< (ForwardCoordIteratorSentinal const&) const { return back(coord) < back(shape); }
CUTE_HOST_DEVICE constexpr
@@ -924,6 +904,15 @@ struct ForwardCoordIterator
Shape const& shape;
};
// A forward iterator for a coordinate that starts from a provided coordinate
template <class Shape, class Coord>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Coord const& coord, Shape const& shape)
{
return ForwardCoordIterator<Coord,Shape>{coord,shape};
}
// A forward iterator for a coordinate that starts from zero
template <class Shape>
CUTE_HOST_DEVICE constexpr
@@ -931,7 +920,7 @@ auto
make_coord_iterator(Shape const& shape)
{
auto coord = repeat_like(shape, int(0));
return ForwardCoordIterator<decltype(coord),Shape>{coord,shape};
return make_coord_iterator(coord, shape);
}
} // end namespace cute
include/cute/layout.hpp
+30 -31
View File
@@ -43,16 +43,16 @@ namespace cute
// Aliases
template <class... Shapes>
using Shape = IntTuple<Shapes...>;
using Shape = cute::tuple<Shapes...>;
template <class... Strides>
using Stride = IntTuple<Strides...>;
using Stride = cute::tuple<Strides...>;
template <class... Strides>
using Step = IntTuple<Strides...>;
using Step = cute::tuple<Strides...>;
template <class... Coords>
using Coord = IntTuple<Coords...>;
using Coord = cute::tuple<Coords...>;
template <class... Ts>
CUTE_HOST_DEVICE constexpr
@@ -1034,29 +1034,29 @@ 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 result = fold(make_seq<R-1>{},
cute::make_tuple(shape, stride, cute::make_tuple(), cute::make_tuple(Int<1>{})),
[](auto const& init, auto i)
{
auto curr_stride = cute::min(get<1>(init));
auto curr_idx = find(get<1>(init), curr_stride);
auto curr_shape = get<curr_idx>(get<0>(init));
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)
{
auto [shape, stride, result_shape, result_stride] = init;
auto min_stride = cute::min(stride);
auto min_idx = find(stride, min_stride);
return cute::make_tuple(remove<curr_idx>(get<0>(init)), // Remove the curr shape
remove<curr_idx>(get<1>(init)), // Remove the curr stride
append(get<2>(init), curr_stride / get<3,i>(init)), // new shape = curr_stride / last_stride
append(get<3>(init), curr_shape * curr_stride)); // new stride = curr_shape * curr_stride
});
return cute::make_tuple(remove<min_idx>(shape), // Remove the min_idx from shape
remove<min_idx>(stride), // Remove the min_idx from stride
append(result_shape , min_stride / get<i>(result_stride)), // new shape = min_stride / last_stride
append(result_stride, get<min_idx>(shape) * min_stride)); // new stride = curr_shape * min_stride
});
// Append the last shape mode
auto result_stride = get<3>(result);
auto result_shape = append(get<2>(result), get<1,0>(result) / back(result_stride)); // new shape = curr_stride / last_stride
auto result_shape = append(result_shape_, get<0>(stride_) / get<R-1>(result_stride)); // new shape = min_stride / last_stride
// Compute the rest_stride
auto rest_stride = get<0,0>(result) * get<1,0>(result);
//return make_layout(append(result_shape, ceil_div(cosize_hi, rest_stride)), append(result_stride, rest_stride));
// Jump into coalesce and append (ceil_div(cosize_hi, rest_stride), rest_stride)
return detail::bw_coalesce<R-1>(result_shape, result_stride, ceil_div(cosize_hi, rest_stride), rest_stride);
// Compute the rest_shape and rest_stride
auto rest_stride = get<0>(shape_) * get<0>(stride_);
auto rest_shape = ceil_div(cosize_hi, rest_stride);
// Jump into coalesce and append (rest_shape, rest_stride)
return detail::bw_coalesce<R-1>(result_shape, result_stride, rest_shape, rest_stride);
}
CUTE_GCC_UNREACHABLE;
@@ -1608,16 +1608,15 @@ CUTE_HOST_DEVICE constexpr
auto
recast(Layout<Shape,Stride> const& layout)
{
if constexpr (sizeof(NewType) == sizeof(OldType)) {
if constexpr (sizeof_bits<NewType>::value == sizeof_bits<OldType>::value) {
return layout;
} else if constexpr (sizeof(NewType) > sizeof(OldType)) {
static_assert(sizeof(NewType) % sizeof(OldType) == 0, "NewType must be a multiple of OldType");
return upcast<sizeof(NewType)/sizeof(OldType)>(layout);
} else if constexpr (sizeof(NewType) < sizeof(OldType)) {
static_assert(sizeof(OldType) % sizeof(NewType) == 0, "NewType must be a divisor of OldType");
return downcast<sizeof(OldType)/sizeof(NewType)>(layout);
} 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>(layout);
} 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>(layout);
}
CUTE_GCC_UNREACHABLE;
}
include/cute/numeric/arithmetic_tuple.hpp
+2 -4
View File
@@ -387,8 +387,7 @@ abs(ScaledBasis<T,N> const& e) {
}
// Multiplication
template <class A, int N, class T,
__CUTE_REQUIRES(cute::is_integral<A>::value)>
template <class A, int N, class T>
CUTE_HOST_DEVICE constexpr
auto
operator*(A const& a, ScaledBasis<T,N> const& e) {
@@ -396,8 +395,7 @@ operator*(A const& a, ScaledBasis<T,N> const& e) {
return ScaledBasis<decltype(r),N>{r};
}
template <int N, class T, class B,
__CUTE_REQUIRES(cute::is_integral<B>::value)>
template <int N, class T, class B>
CUTE_HOST_DEVICE constexpr
auto
operator*(ScaledBasis<T,N> const& e, B const& b) {
include/cute/numeric/complex.hpp
+12 -132
View File
@@ -30,97 +30,19 @@
**************************************************************************************************/
#pragma once
#include <cute/util/type_traits.hpp>
//#if defined(__CUDA_ARCH__)
//# include <cuda/std/complex>
//#else
//# include <complex>
//#endif
// Suppress warnings for code in Thrust headers.
#if defined(_MSC_VER)
// We check for MSVC first, because MSVC also defines __GNUC__.
// It's common for non-GCC compilers that emulate GCC's behavior
// to define __GNUC__.
//
// thrust/complex.h triggers MSVC's warning on conversion
// from double to float (or const float) ("possible loss of data").
// MSVC treats this as an error by default (at least with
// CUTLASS's default CMake configuration).
#pragma warning( push )
#pragma warning( disable : 4244 )
#elif defined(__GNUC__)
// With GCC + CUDA 11.4, builds show spurious "-Wconversion"
// warnings on line 656 of thrust/detail/type_traits.h.
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wconversion"
#endif
#if defined(__CUDACC_RTC__)
#include <cuda/std/complex>
#else
#include <thrust/complex.h>
#endif
#if defined(_MSC_VER)
#pragma warning( pop )
#elif defined(__GNUC__)
#pragma GCC diagnostic pop
#endif
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cutlass/complex.h>
namespace cute
{
//#if defined(__CUDA_ARCH__)
//template <class T>
//using complex = cuda::std::complex<T>;
//#else
//template <class T>
//using complex = std::complex<T>;
//#endif
//template <class T>
//using complex = thrust::complex<T>;
#if defined(__CUDACC_RTC__)
using cuda::std::complex;
#else
using thrust::complex;
#endif
template <class T>
CUTE_HOST_DEVICE
T real(complex<T> const& z) {
return z.real();
}
template <class T>
CUTE_HOST_DEVICE
T imag(complex<T> const& z) {
return z.imag();
}
template <class T>
CUTE_HOST_DEVICE
complex<T> conj(complex<T> const& z) {
return complex<T>(real(z), -imag(z));
}
// cute::conj forwards scalars
template <class T>
CUTE_HOST_DEVICE
T conj(T z) {
return z;
}
//CUTE_HOST_DEVICE constexpr
//float conj(float z) { return z; }
//CUTE_HOST_DEVICE constexpr
//double conj(double z) { return z; }
using cutlass::complex;
using cutlass::is_complex;
using cutlass::RealType;
using cutlass::real;
using cutlass::imag;
using cutlass::conj;
/// Fused multiply-add for complex numbers
template <class T>
@@ -131,10 +53,10 @@ fma(complex<T> & d,
complex<T> const& b,
complex<T> const& c)
{
d.real(c.real() + a.real() * b.real());
d.imag(c.imag() + a.real() * b.imag());
d.real(d.real() - a.imag() * b.imag());
d.imag(d.imag() + a.imag() * b.real());
d.real(fma( a.real(), b.real(), c.real()));
d.imag(fma( a.real(), b.imag(), c.imag()));
d.real(fma(-a.imag(), b.imag(), d.real()));
d.imag(fma( a.imag(), b.real(), d.imag()));
}
/// Fused multiply-add for triplets
@@ -148,46 +70,4 @@ fma(complex<T> const& a,
return fma(c, a, b, c);
}
/// Used to determine the real-valued underlying type of a numeric type T
template <class T>
struct RealType {
using Type = T;
};
/// Partial specialization for complex-valued type
template <class T>
struct RealType<complex<T>> {
using Type = T;
};
//////////////////////////////////////////////////////////////////////////////////////////////////
template <class T>
struct is_complex {
static bool const value = false;
};
template <class T>
struct is_complex<complex<T>> {
static bool const value = true;
};
//////////////////////////////////////////////////////////////////////////////////////////////////
// Display utilities
#if !defined(__CUDACC_RTC__)
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, complex<T> const& z)
{
T _r = z.real();
T _i = z.imag();
if (bool(_i)) {
return os << _r << "+i" << _i;
} else {
return os << _r;
}
}
#endif
} // end namespace cute
include/cute/numeric/integral_constant.hpp
+36 -23
View File
@@ -30,15 +30,14 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/math.hpp>
#include "cute/util/print.hpp"
#include "cute/util/type_traits.hpp"
#include "cute/numeric/math.hpp"
namespace cute
{
// Short name for fast compilation
// A constant value: short name and type-deduction for fast compilation
template <auto v>
struct C {
using type = C<v>;
@@ -48,29 +47,40 @@ struct C {
CUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
};
// Deprecate
template <class T, T v>
using constant = C<v>;
template <class T, T v>
using integral_constant = C<v>;
template <bool b>
using bool_constant = C<b>;
using true_type = bool_constant<true>;
using false_type = bool_constant<false>;
// A more std:: conforming integral_constant that enforces type but interops with C<v>
template <class T, T v>
struct integral_constant : C<v> {
using type = integral_constant<T,v>;
static constexpr T value = v;
using value_type = T;
// Disambiguate C<v>::operator value_type()
//CUTE_HOST_DEVICE constexpr operator value_type() const noexcept { return value; }
CUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
};
//
// Traits
//
// Use cute::is_std_integral<T> to match built-in integral types (int, int64_t, unsigned, etc)
// Use cute::is_integral<T> to match both built-in integral types AND constant<T,t>
// Use cute::is_integral<T> to match both built-in integral types AND static integral types.
template <class T>
struct is_integral : bool_constant<is_std_integral<T>::value> {};
template <auto v>
struct is_integral<C<v>> : true_type {};
struct is_integral<C<v> > : true_type {};
template <class T, T v>
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)
@@ -80,20 +90,22 @@ struct is_static : bool_constant<is_empty<T>::value> {};
template <class T>
constexpr bool is_static_v = is_static<T>::value;
// is_constant detects if a type is a constant<T,v> and if v is equal to a value
// is_constant detects if a type is a static integral type and if v is equal to a value
template <auto n, class T>
struct is_constant : false_type {};
template <auto n, class T>
struct is_constant<n, T const > : is_constant<n,T> {};
template <auto n, class T>
struct is_constant<n, T const&> : is_constant<n,T> {};
template <auto n, class T>
struct is_constant<n, T &> : is_constant<n,T> {};
template <auto n, class T>
struct is_constant<n, T &&> : is_constant<n,T> {};
template <auto n, auto v>
struct is_constant<n, C<v> > : bool_constant<v == n> {};
template <auto n, auto v>
struct is_constant<n, C<v> const > : bool_constant<v == n> {};
template <auto n, auto v>
struct is_constant<n, C<v> const&> : bool_constant<v == n> {};
template <auto n, auto v>
struct is_constant<n, C<v> &> : bool_constant<v == n> {};
template <auto n, auto v>
struct is_constant<n, C<v> &&> : bool_constant<v == n> {};
struct is_constant<n, C<v> > : bool_constant<v == n> {};
template <auto n, class T, T v>
struct is_constant<n, integral_constant<T,v>> : bool_constant<v == n> {};
//
// Specializations
@@ -403,9 +415,10 @@ conditional_return(TrueType const& t, FalseType const& f) {
// Display utilities
//
template <auto t>
CUTE_HOST_DEVICE void print(C<t> const&) {
printf("_%d", int(t));
template <auto Value>
CUTE_HOST_DEVICE void print(C<Value>) {
printf("_");
::cute::print(Value);
}
#if !defined(__CUDACC_RTC__)
include/cute/numeric/integral_ratio.hpp
+175
View File
@@ -0,0 +1,175 @@
/***************************************************************************************************
* 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/math.hpp>
#include <cute/numeric/integral_constant.hpp>
namespace cute
{
/** Compile-time rational arithmetic type.
* 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>
*/
template <auto n, auto d>
class R {
static_assert(d != 0);
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;
// RI: den >= 1 && gcd(abs(num),den) == 1
using type = typename conditional<num == 0 || den == 1, C<num>, R<num,den>>::type;
};
template <auto a, auto b>
CUTE_HOST_DEVICE constexpr
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
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
operator*(R<a,b>, C<c>) {
return {};
}
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<a*c,b>::type
operator*(C<c>, R<a,b>) {
return {};
}
// Product with dynamic type needs to produce an integer...
template <class C, auto a, auto b,
__CUTE_REQUIRES(cute::is_std_integral<C>::value)>
CUTE_HOST_DEVICE constexpr
auto
operator*(C const& c, R<a,b>) {
return c * R<a,b>::num / R<a,b>::den;
}
// Product with dynamic type needs to produce an integer...
template <auto a, auto b, class C,
__CUTE_REQUIRES(cute::is_std_integral<C>::value)>
CUTE_HOST_DEVICE constexpr
auto
operator*(R<a,b>, C const& c) {
return c * R<a,b>::num / R<a,b>::den;
}
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
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
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
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>
operator==(R<a,b>, R<x,y>) {
return {};
}
template <auto a, auto b, auto c>
CUTE_HOST_DEVICE constexpr
bool_constant<R<a,b>::num == c && R<a,b>::den == 1>
operator==(R<a,b>, C<c>) {
return {};
}
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
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
abs(R<a,b>) {
return {};
}
//
// Display utilities
//
template <auto a, auto b>
CUTE_HOST_DEVICE void print(R<a,b>) {
print(C<a>{}); print("/"); print(C<b>{});
}
#if !defined(__CUDACC_RTC__)
template <auto a, auto b>
CUTE_HOST std::ostream& operator<<(std::ostream& os, R<a,b>) {
return os << "_" << C<a>{} << "/" << C<b>{};
}
#endif
} // end namespace cute
include/cute/numeric/math.hpp
+17 -19
View File
@@ -73,11 +73,26 @@ abs(T const& t) {
CUTE_GCC_UNREACHABLE;
}
// Returns 1 if x > 0, -1 if x < 0, and 0 if x is zero.
template <class T,
__CUTE_REQUIRES(is_arithmetic<T>::value)>
CUTE_HOST_DEVICE constexpr
int
signum(T const& x) {
if constexpr (is_signed<T>::value) {
return (T(0) < x) - (x < T(0));
} else {
return T(0) < x;
}
CUTE_GCC_UNREACHABLE;
}
//
// C++17 <numeric> operations
//
// Greatest common divisor of two integers
// Greatest common divisor of two positive integers
template <class T, class U,
__CUTE_REQUIRES(is_std_integral<T>::value &&
is_std_integral<U>::value)>
@@ -92,7 +107,7 @@ gcd(T t, U u) {
}
}
// Least common multiple of two integers
// Least common multiple of two positive integers
template <class T, class U,
__CUTE_REQUIRES(is_std_integral<T>::value &&
is_std_integral<U>::value)>
@@ -280,23 +295,6 @@ shiftr(T x, int s) {
return s >= 0 ? (x >> s) : (x << -s);
}
// Returns 1 if x > 0, -1 if x < 0, and 0 if x is zero.
template <class T,
__CUTE_REQUIRES(is_unsigned<T>::value)>
CUTE_HOST_DEVICE constexpr
int
signum(T const& x) {
return T(0) < x;
}
template <class T,
__CUTE_REQUIRES(not is_unsigned<T>::value)>
CUTE_HOST_DEVICE constexpr
int
signum(T const& x) {
return (T(0) < x) - (x < T(0));
}
// Safe divide
// @pre t % u == 0
// @result t / u
include/cute/pointer.hpp
+27 -9
View File
@@ -58,6 +58,19 @@ raw_pointer_cast(T* ptr) {
return ptr;
}
//
// 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
//
@@ -79,6 +92,8 @@ 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) {}
@@ -91,11 +106,14 @@ struct device_ptr
template <class Index>
CUTE_HOST_DEVICE constexpr
T& operator[](Index const& i) const { return ptr_[i]; }
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}; }
DerivedType operator+(Index const& i) const { return {ptr_ + i / ElementsPerStoredItem}; }
CUTE_HOST_DEVICE constexpr friend
ptrdiff_t operator-(device_ptr<T,DerivedType> const& a,
@@ -326,44 +344,44 @@ recast(rmem_ptr<T const> const& ptr) {
template <class T>
CUTE_HOST_DEVICE void print(T const* const ptr)
{
printf("raw_ptr_%db(%p)", int(8*sizeof(T)), ptr);
printf("raw_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr);
}
template <class T>
CUTE_HOST_DEVICE void print(gmem_ptr<T> const& ptr)
{
printf("gmem_ptr_%db(%p)", int(8*sizeof(T)), ptr.get());
printf("gmem_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr.get());
}
template <class T>
CUTE_HOST_DEVICE void print(smem_ptr<T> const& ptr)
{
printf("smem_ptr_%db(%p)", int(8*sizeof(T)), ptr.get());
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(8*sizeof(T)), ptr.get());
printf("rmem_ptr_%db(%p)", int(sizeof_bits<T>::value), ptr.get());
}
#if !defined(__CUDACC_RTC__)
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, gmem_ptr<T> const& ptr)
{
return os << "gmem_ptr_" << int(8*sizeof(T)) << "b";
return os << "gmem_ptr_" << int(sizeof_bits<T>::value) << "b";
}
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr<T> const& ptr)
{
return os << "smem_ptr_" << int(8*sizeof(T)) << "b";
return os << "smem_ptr_" << int(sizeof_bits<T>::value) << "b";
}
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, rmem_ptr<T> const& ptr)
{
return os << "rmem_ptr_" << int(8*sizeof(T)) << "b";
return os << "rmem_ptr_" << int(sizeof_bits<T>::value) << "b";
}
#endif // !defined(__CUDACC_RTC__)
include/cute/stride.hpp
+3
View File
@@ -75,6 +75,9 @@ crd2idx_itt(CInt const& coord,
{
if constexpr (sizeof...(Is) == 0) { // Avoid recursion and mod on single/last iter
return crd2idx(coord, get<I0>(shape), get<I0>(stride));
} else if constexpr (is_constant<0, CInt>::value) {
return crd2idx(_0{}, get<I0>(shape), get<I0>(stride))
+ (_0{} + ... + crd2idx(_0{}, get<Is>(shape), get<Is>(stride)));
} else { // General case
return crd2idx(coord % product(get<I0>(shape)), get<I0>(shape), get<I0>(stride))
+ crd2idx_itt(coord / product(get<I0>(shape)), shape, stride, seq<Is...>{});
include/cute/swizzle.hpp
+85 -86
View File
@@ -218,41 +218,40 @@ 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 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 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
// Abstract value: StaticInt | (dynamic_int_ & StaticFlags)
template <uint32_t StaticInt = 0,
class DynamicType = uint32_t,
uint32_t StaticFlags = 0> // 0: static, 1: dynamic
template <uint32_t StaticInt,
uint32_t StaticFlags> // 0: static, 1: dynamic
struct MixedBits
{
// Representation invariants
static_assert(StaticFlags != 0, "Should be at least one dynamic bit in MixedBits.");
static_assert((StaticInt & StaticFlags) == 0, "No static/dynamic overlap allowed in MixedBits.");
// assert((dynamic_int_ & ~F) == 0);
DynamicType dynamic_int_;
uint32_t dynamic_int_;
// assert((dynamic_int_ & ~StaticFlags) == 0);
CUTE_HOST_DEVICE constexpr operator uint32_t() const noexcept { return StaticInt | dynamic_int_; }
};
template <class S, S s, class DynamicType, class F, F f>
// Return a value representing (C<s>{} | (d & C<f>)) potentially using MixedBits to track s and f.
// This maker does allow ((s & f) != 0) and enforces the MixedBits invariant before creation.
template <auto s, class DynamicType, auto f>
CUTE_HOST_DEVICE constexpr
auto
make_mixed_bits(constant<S,s> const&, DynamicType const& d, constant<F,f> const&)
make_mixed_bits(C<s>, DynamicType const& d, C<f>)
{
static_assert(is_integral<DynamicType>::value);
if constexpr (is_static<DynamicType>::value) {
static_assert((s & DynamicType::value & f) == 0, "No static/dynamic overlap allowed.");
return constant<S,s>{} | (d & constant<F,f>{}); // Just return a static int
} else if constexpr (f == 0) {
return constant<S,s>{}; // Just return a static int
constexpr uint32_t new_f = uint32_t(f) & ~uint32_t(s); // StaticBits take precedence, M<0,f>{d} | C<s>{}
if constexpr (new_f == 0 || is_static<DynamicType>::value) {
return C<s>{} | (d & C<new_f>{}); // Just return a static int
} else {
return MixedBits<s, DynamicType, f>{d & f}; // MixedBits
return MixedBits<s, new_f>{uint32_t(d) & new_f}; // MixedBits
}
CUTE_GCC_UNREACHABLE;
@@ -263,28 +262,28 @@ make_mixed_bits(constant<S,s> const&, DynamicType const& d, constant<F,f> const&
//
// Equality
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator==(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const&)
operator==(MixedBits<S0,F0> const& m, C<S1>)
{
return (S0 == (S1 & ~F0)) && (m.dynamic_int_ == (S1 & F0));
return (S0 == (uint32_t(S1) & ~F0)) && (m.dynamic_int_ == (uint32_t(S1) & F0));
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator==(constant<TS1,S1> const& s, MixedBits<S0,D0,F0> const& m)
operator==(C<S1> s, MixedBits<S0,F0> const& m)
{
return m == s;
}
// Bitwise AND
template <uint32_t S0, class D0, uint32_t F0,
uint32_t S1, class D1, uint32_t F1>
template <uint32_t S0, uint32_t F0,
uint32_t S1, uint32_t F1>
CUTE_HOST_DEVICE constexpr
auto
operator&(MixedBits<S0,D0,F0> const& m0, MixedBits<S1,D1,F1> const& m1)
operator&(MixedBits<S0,F0> const& m0, MixedBits<S1,F1> const& m1)
{
// Truth table for (S0,D0,F0) & (S1,D1,F1) -> (S,D,F)
// S0D0F0 | 0X0 | 001 | 011 | 1X0 |
@@ -294,36 +293,36 @@ operator&(MixedBits<S0,D0,F0> const& m0, MixedBits<S1,D1,F1> const& m1)
// 011 | 0X0 | 001 | 011 | 011 |
// 1X0 | 0X0 | 001 | 011 | 1X0 |
return make_mixed_bits(constant<uint32_t,S0 & S1>{},
return make_mixed_bits(C<S0 & S1>{},
//(S0 | m0.dynamic_int_) & (S1 | m1.dynamic_int_),
((S1 & F0) & m0.dynamic_int_) | ((S0 & F1) & m1.dynamic_int_) | (m0.dynamic_int_ & m1.dynamic_int_),
constant<uint32_t,(S1 & F0) | (S0 & F1) | (F0 & F1)>{});
C<(S1 & F0) | (S0 & F1) | (F0 & F1)>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator&(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const&)
operator&(MixedBits<S0,F0> const& m, C<S1>)
{
return make_mixed_bits(constant<uint32_t,S0 & S1>{},
return make_mixed_bits(C<S0 & uint32_t(S1)>{},
m.dynamic_int_,
constant<uint32_t,S1 & F0>{});
C<F0 & uint32_t(S1)>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator&(constant<TS1,S1> const& s, MixedBits<S0,D0,F0> const& m)
operator&(C<S1> s, MixedBits<S0,F0> const& m)
{
return m & s;
}
// Bitwise OR
template <uint32_t S0, class D0, uint32_t F0,
uint32_t S1, class D1, uint32_t F1>
template <uint32_t S0, uint32_t F0,
uint32_t S1, uint32_t F1>
CUTE_HOST_DEVICE constexpr
auto
operator|(MixedBits<S0,D0,F0> const& m0, MixedBits<S1,D1,F1> const& m1)
operator|(MixedBits<S0,F0> const& m0, MixedBits<S1,F1> const& m1)
{
// Truth table for (S0,D0,F0) | (S1,D1,F1) -> (S,D,F)
// S0D0F0 | 0X0 | 001 | 011 | 1X0 |
@@ -333,35 +332,35 @@ operator|(MixedBits<S0,D0,F0> const& m0, MixedBits<S1,D1,F1> const& m1)
// 011 | 011 | 011 | 011 | 1X0 |
// 1X0 | 1X0 | 1X0 | 1X0 | 1X0 |
return make_mixed_bits(constant<uint32_t,S0 | S1>{},
return make_mixed_bits(C<S0 | S1>{},
((~S1 & F0) & m0.dynamic_int_) | ((~S0 & F1) & m1.dynamic_int_),
constant<uint32_t,(~S0 & F1) | (~S1 & F0)>{});
C<(~S0 & F1) | (~S1 & F0)>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator|(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const&)
operator|(MixedBits<S0,F0> const& m, C<S1>)
{
return make_mixed_bits(constant<uint32_t,S0 | S1>{},
return make_mixed_bits(C<S0 | uint32_t(S1)>{},
m.dynamic_int_,
constant<uint32_t,~S1 & F0>{});
C<F0 & ~uint32_t(S1)>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator|(constant<TS1,S1> const& s, MixedBits<S0,D0,F0> const& m)
operator|(C<S1> s, MixedBits<S0,F0> const& m)
{
return m | s;
}
// Bitwise XOR
template <uint32_t S0, class D0, uint32_t F0,
uint32_t S1, class D1, uint32_t F1>
template <uint32_t S0, uint32_t F0,
uint32_t S1, uint32_t F1>
CUTE_HOST_DEVICE constexpr
auto
operator^(MixedBits<S0,D0,F0> const& m0, MixedBits<S1,D1,F1> const& m1)
operator^(MixedBits<S0,F0> const& m0, MixedBits<S1,F1> const& m1)
{
// Truth table for (S0,D0,F0) ^ (S1,D1,F1) -> (S,D,F)
// S0D0F0 | 0X0 | 001 | 011 | 1X0 |
@@ -371,53 +370,53 @@ operator^(MixedBits<S0,D0,F0> const& m0, MixedBits<S1,D1,F1> const& m1)
// 011 | 011 | 011 | 001 | 001 |
// 1X0 | 1X0 | 011 | 001 | 0X0 |
return make_mixed_bits(constant<uint32_t,(~S0 & S1 & ~F0) | (S0 & ~S1 & ~F1)>{},
return make_mixed_bits(C<(~S0 & S1 & ~F0) | (S0 & ~S1 & ~F1)>{},
(S0 | m0.dynamic_int_) ^ (S1 | m1.dynamic_int_),
constant<uint32_t,F0 | F1>{});
C<F0 | F1>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator^(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const&)
operator^(MixedBits<S0,F0> const& m, C<S1>)
{
return make_mixed_bits(constant<uint32_t,(~S0 & S1 & ~F0) | (S0 & ~S1)>{},
(S0 | m.dynamic_int_) ^ S1,
constant<uint32_t,F0>{});
return make_mixed_bits(C<(~S0 & uint32_t(S1) & ~F0) | (S0 & ~uint32_t(S1))>{},
(S0 | m.dynamic_int_) ^ uint32_t(S1),
C<F0>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator^(constant<TS1,S1> const& s, MixedBits<S0,D0,F0> const& m)
operator^(C<S1> s, MixedBits<S0,F0> const& m)
{
return m ^ s;
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator<<(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const&)
operator<<(MixedBits<S0,F0> const& m, C<S1>)
{
return make_mixed_bits(constant<uint32_t,(S0 << S1)>{},
return make_mixed_bits(C<(S0 << S1)>{},
m.dynamic_int_ << S1,
constant<uint32_t,(F0 << S1)>{});
C<(F0 << S1)>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
operator>>(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const&)
operator>>(MixedBits<S0,F0> const& m, C<S1>)
{
return make_mixed_bits(constant<uint32_t,(S0 >> S1)>{},
return make_mixed_bits(C<(S0 >> S1)>{},
m.dynamic_int_ >> S1,
constant<uint32_t,(F0 >> S1)>{});
C<(F0 >> S1)>{});
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
shiftl(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const& s)
shiftl(MixedBits<S0,F0> const& m, C<S1> s)
{
if constexpr (S1 >= 0) {
return m << s;
@@ -426,10 +425,10 @@ shiftl(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const& s)
}
}
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
shiftr(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const& s)
shiftr(MixedBits<S0,F0> const& m, C<S1> s)
{
if constexpr (S1 >= 0) {
return m >> s;
@@ -442,24 +441,24 @@ shiftr(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const& s)
// upcast and downcast
//
template <uint32_t S0, class D0, uint32_t F0, class TS1, TS1 S1>
template <uint32_t S0, uint32_t F0, auto S1>
CUTE_HOST_DEVICE constexpr
auto
safe_div(MixedBits<S0,D0,F0> const& m, constant<TS1,S1> const& s)
safe_div(MixedBits<S0,F0> const& m, C<S1> s)
{
static_assert(has_single_bit(S1), "Only divide MixedBits by powers of two.");
return make_mixed_bits(safe_div(constant<uint32_t,S0>{}, s),
static_assert(has_single_bit(uint32_t(S1)), "Only divide MixedBits by powers of two.");
return make_mixed_bits(safe_div(C<S0>{}, s),
safe_div(m.dynamic_int_, s),
safe_div(constant<uint32_t,F0>{}, s));
safe_div(C<F0>{}, s));
}
template <uint32_t N, uint32_t S0, class D0, uint32_t F0>
template <uint32_t N, uint32_t S0, uint32_t F0>
CUTE_HOST_DEVICE constexpr
auto
upcast(MixedBits<S0,D0,F0> const& m)
upcast(MixedBits<S0,F0> const& m)
{
static_assert(has_single_bit(N), "Only divide MixedBits by powers of two.");
return safe_div(m, constant<uint32_t,N>{});
return safe_div(m, C<N>{});
}
template <uint32_t N, class T, __CUTE_REQUIRES(cute::is_integral<T>::value)>
@@ -467,18 +466,18 @@ CUTE_HOST_DEVICE constexpr
auto
upcast(T const& m)
{
return safe_div(m, constant<uint32_t,N>{});
return safe_div(m, C<N>{});
}
template <uint32_t N, uint32_t S0, class D0, uint32_t F0>
template <uint32_t N, uint32_t S0, uint32_t F0>
CUTE_HOST_DEVICE constexpr
auto
downcast(MixedBits<S0,D0,F0> const& m)
downcast(MixedBits<S0,F0> const& m)
{
static_assert(has_single_bit(N), "Only scale MixedBits by powers of two.");
return make_mixed_bits(constant<uint32_t,S0 * N>{},
return make_mixed_bits(C<S0 * N>{},
m.dynamic_int_ * N,
constant<uint32_t,F0 * N>{});
C<F0 * N>{});
}
template <uint32_t N, class T, __CUTE_REQUIRES(cute::is_integral<T>::value)>
@@ -486,7 +485,7 @@ CUTE_HOST_DEVICE constexpr
auto
downcast(T const& m)
{
return m * constant<uint32_t, N>{};
return m * C<N>{};
}
//
@@ -525,17 +524,17 @@ to_mixed_bits(Layout const& layout, Coord const& coord)
// Display utilities
//
template <uint32_t S, class D, uint32_t F>
CUTE_HOST_DEVICE void print(MixedBits<S,D,F> const& m)
template <uint32_t S, uint32_t F>
CUTE_HOST_DEVICE void print(MixedBits<S,F> const& m)
{
printf("M_%u|(%u&%u)=%u", S, uint32_t(m.dynamic_int_), F, uint32_t(m));
printf("M_%u|(%u&%u)=%u", S, m.dynamic_int_, F, uint32_t(m));
}
#if !defined(__CUDACC_RTC__)
template <uint32_t S, class D, uint32_t F>
CUTE_HOST std::ostream& operator<<(std::ostream& os, MixedBits<S,D,F> const& m)
CUTE_HOST std::ostream& operator<<(std::ostream& os, MixedBits<S,F> const& m)
{
return os << "M_" << S << "|(" << uint32_t(m.dynamic_int_) << "&" << F << ")=" << uint32_t(m);
return os << "M_" << S << "|(" << m.dynamic_int_ << "&" << F << ")=" << uint32_t(m);
}
template <int B, int M, int S>
include/cute/swizzle_layout.hpp
+2
View File
@@ -128,6 +128,7 @@ namespace detail {
// Get just the Swizzle part of a composed layout.
template <int B, int M, int S, class Offset, class LayoutB>
CUTE_HOST_DEVICE constexpr
auto
get_swizzle_portion(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB>)
{
@@ -136,6 +137,7 @@ get_swizzle_portion(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB>)
// A non-swizzled layout's "Swizzle part" is the identity swizzle.
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
get_swizzle_portion(Layout<Shape,Stride>)
{
include/cute/swizzle_ptr.hpp
+6 -3
View File
@@ -70,6 +70,8 @@ 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
{
@@ -98,6 +100,7 @@ struct smem_ptr_swizzle
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);
}
@@ -105,7 +108,7 @@ struct smem_ptr_swizzle
CUTE_HOST_DEVICE constexpr
smem_ptr_swizzle operator+(Int const& i) const
{
return {ptr_ + i};
return {ptr_ + i / ElementsPerStoredItem};
}
T* ptr_;
@@ -286,14 +289,14 @@ CUTE_HOST_DEVICE void print(smem_ptr_flag_bits<B> const& ptr)
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(8*sizeof(T)), ptr.get());
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(8*sizeof(T)) << "b";
return os << "smem_ptr_S<" << B << "," << M << "," << S << ">_" << int(sizeof_bits<T>::value) << "b";
}
#endif
include/cute/util/print.hpp
+57 -13
View File
@@ -33,7 +33,6 @@
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/integral_constant.hpp>
//
// CUDA compatible print and printf
@@ -119,16 +118,68 @@ get_format(double) {
CUTE_HOST_DEVICE
void
print(char const& c) {
print(char c) {
printf("%c", c);
}
template <class T,
__CUTE_REQUIRES(is_std_integral<T>::value)>
CUTE_HOST_DEVICE
void
print(T const& a) {
printf("%d", int(a));
print(signed char a) {
printf("%hhd", a);
}
CUTE_HOST_DEVICE
void
print(unsigned char a) {
printf("%hhu", a);
}
CUTE_HOST_DEVICE
void
print(short a) {
printf("%hd", a);
}
CUTE_HOST_DEVICE
void
print(unsigned short a) {
printf("%hu", a);
}
CUTE_HOST_DEVICE
void
print(int a) {
printf("%d", a);
}
CUTE_HOST_DEVICE
void
print(unsigned int a) {
printf("%u", a);
}
CUTE_HOST_DEVICE
void
print(long a) {
printf("%ld", a);
}
CUTE_HOST_DEVICE
void
print(unsigned long a) {
printf("%lu", a);
}
CUTE_HOST_DEVICE
void
print(long long a) {
printf("%lld", a);
}
CUTE_HOST_DEVICE
void
print(unsigned long long a) {
printf("%llu", a);
}
template <class... T>
@@ -138,13 +189,6 @@ print(char const* format, T const&... t) {
printf(format, t...);
}
template <class... T>
CUTE_HOST_DEVICE
void
print(T const&... t) {
(print(t), ...);
}
CUTE_HOST_DEVICE
void
print(char const* format) {