leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 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) {
include/cutlass/arch/mma_sm75.h
+44 -60
View File
@@ -130,7 +130,7 @@ struct Mma<
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -196,7 +196,7 @@ struct Mma<
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -257,13 +257,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -318,13 +317,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.u8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -379,14 +377,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.s8.u8 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -441,13 +437,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -461,7 +456,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = S8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,16>,
gemm::GemmShape<8, 8, 16>,
32,
int8_t,
layout::RowMajor,
@@ -471,7 +466,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,16>;
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
@@ -508,13 +503,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -522,7 +516,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = U8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,16>,
gemm::GemmShape<8, 8, 16>,
32,
uint8_t,
layout::RowMajor,
@@ -532,7 +526,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,16>;
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
@@ -569,13 +563,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.u8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -583,7 +576,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = S8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,16>,
gemm::GemmShape<8, 8, 16>,
32,
int8_t,
layout::RowMajor,
@@ -593,7 +586,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,16>;
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
@@ -630,13 +623,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -644,7 +636,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = U8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,16>,
gemm::GemmShape<8, 8, 16>,
32,
uint8_t,
layout::RowMajor,
@@ -654,7 +646,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,16>;
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
@@ -691,13 +683,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.u8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -711,7 +702,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = S4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
int4b_t,
layout::RowMajor,
@@ -721,7 +712,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = int4b_t;
using LayoutA = layout::RowMajor;
@@ -751,19 +742,19 @@ struct Mma<
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.s4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -771,7 +762,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = U4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
uint4b_t,
layout::RowMajor,
@@ -781,7 +772,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = uint4b_t;
using LayoutA = layout::RowMajor;
@@ -818,13 +809,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.u4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -832,7 +822,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = S4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
int4b_t,
layout::RowMajor,
@@ -842,7 +832,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = int4b_t;
using LayoutA = layout::RowMajor;
@@ -879,13 +869,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -893,7 +882,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = U4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
uint4b_t,
layout::RowMajor,
@@ -903,7 +892,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = uint4b_t;
using LayoutA = layout::RowMajor;
@@ -940,13 +929,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.u4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -960,7 +948,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = S4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
int4b_t,
layout::RowMajor,
@@ -970,7 +958,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = int4b_t;
using LayoutA = layout::RowMajor;
@@ -1007,13 +995,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.s4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -1021,7 +1008,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = U4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
uint4b_t,
layout::RowMajor,
@@ -1031,7 +1018,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = uint4b_t;
using LayoutA = layout::RowMajor;
@@ -1068,13 +1055,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.u4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -1082,7 +1068,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = S4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
int4b_t,
layout::RowMajor,
@@ -1092,7 +1078,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = int4b_t;
using LayoutA = layout::RowMajor;
@@ -1129,13 +1115,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.s4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -1143,7 +1128,7 @@ struct Mma<
/// Matrix multiply-add operation: S32 = U4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<8,8,32>,
gemm::GemmShape<8, 8, 32>,
32,
uint4b_t,
layout::RowMajor,
@@ -1153,7 +1138,7 @@ struct Mma<
layout::RowMajor,
OpMultiplyAddSaturate> {
using Shape = gemm::GemmShape<8,8,32>;
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = uint4b_t;
using LayoutA = layout::RowMajor;
@@ -1190,13 +1175,12 @@ struct Mma<
asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.u4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -1287,7 +1271,7 @@ struct Mma<
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0); // WMMA must be supported to issue binary matrix multiply-accumulate instructions.
CUTLASS_NOT_IMPLEMENTED(); // WMMA must be supported to issue binary matrix multiply-accumulate instructions.
#endif // defined(CUTLASS_ARCH_WMMA_ENABLED)
include/cutlass/arch/mma_sm80.h
+13 -4
View File
@@ -53,7 +53,16 @@
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
#define CUTLASS_ARCH_MMA_SM80_ENABLED
#if (__CUDA_ARCH__ <= 900)
#define CUTLASS_ARCH_MMA_B1_AND_SM80_ENABLED
#endif
#if (__CUDA_ARCH__ <= 890)
#define CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED
#endif
#endif
#endif
////////////////////////////////////////////////////////////////////////////////
@@ -2084,7 +2093,7 @@ struct Mma<
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
#if defined(CUTLASS_ARCH_MMA_B1_AND_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
@@ -2149,7 +2158,7 @@ struct Mma<
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
#if defined(CUTLASS_ARCH_MMA_B1_AND_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
@@ -2220,7 +2229,7 @@ struct Mma<
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
#if defined(CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
@@ -2244,7 +2253,7 @@ struct Mma<
CUTLASS_UNUSED(d);
assert(0);
#endif // defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
#endif // defined(CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED)
}
};
include/cutlass/array.h
+15 -1
View File
@@ -33,10 +33,24 @@
and is safe to use in a union.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/functional.h"
#include "cutlass/numeric_types.h"
#include "cutlass/numeric_size.h"
#include "cutlass/half.h"
#include "cutlass/integer_subbyte.h"
#include "cutlass/tfloat32.h"
#include "cutlass/bfloat16.h"
#include "cutlass/half.h"
namespace cutlass {
include/cutlass/array_subbyte.h
+9
View File
@@ -32,6 +32,15 @@
\brief Statically sized array of elements that accommodates all CUTLASS-supported numeric types
and is safe to use in a union.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
include/cutlass/barrier.h
+1 -1
View File
@@ -342,7 +342,7 @@ struct SyncManager {
CUTLASS_DEVICE
static
void wait_lt(uint32_t, void *lock_ptr, int thread_idx, int flag_idx, int count) {
BarrierSync::wait_lt_helper(lock_ptr, thread_idx, flag_idx, count);
BarrierSync::wait_lt(lock_ptr, thread_idx, flag_idx, count);
}
CUTLASS_DEVICE
include/cutlass/bfloat16.h
+11
View File
@@ -33,6 +33,17 @@
\brief Defines a proxy class for storing non-standard 16-bit floating point values with
8 bits of exponent and 7 bit of mantissa.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#if defined(__CUDACC_RTC__)
include/cutlass/complex.h
+42 -14
View File
@@ -300,6 +300,13 @@ class complex
CUTLASS_HOST_DEVICE
T &imag() { return _imag; }
/// Set the real part of the complex number
CUTLASS_HOST_DEVICE
void real(T real) { _real = real; }
/// Set the imaginary part of the complex number
CUTLASS_HOST_DEVICE
void imag(T imag) { _imag = imag; }
#if !defined(__CUDACC_RTC__)
/// Converts to cuFloatComplex
@@ -431,34 +438,55 @@ CUTLASS_HOST_DEVICE R norm_accumulate(T const &x, R const & accumulator) {
/// Norm accumulate specialized for complex types
template <typename T, typename R>
CUTLASS_HOST_DEVICE R norm_accumulate(complex<T> const &z, R const &accumulator) {
return accumulator + static_cast<R>(real(z)) * static_cast<R>(real(z)) +
return accumulator + static_cast<R>(real(z)) * static_cast<R>(real(z)) +
static_cast<R>(imag(z)) * static_cast<R>(imag(z));
}
/// Returns the complex conjugate
CUTLASS_HOST_DEVICE float conj(float const &z) {
return z;
}
/// Returns the complex conjugate
CUTLASS_HOST_DEVICE double conj(double const &z) {
return z;
}
CUTLASS_HOST_DEVICE half_t conj(half_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE int32_t conj(int32_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint32_t conj(uint32_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE int4b_t conj(int4b_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint4b_t conj(uint4b_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE bfloat16_t conj(bfloat16_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint1b_t conj(uint1b_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE tfloat32_t conj(tfloat32_t const& z) {
return z;
}
/// Returns the complex conjugate
template <typename T>
CUTLASS_HOST_DEVICE complex<T> conj(complex<T> const &z) {
return complex<T>(real(z), -imag(z));
}
/// Indentity transform for non-complex types
template <typename T>
CUTLASS_HOST_DEVICE T conj(T const &z) {
static_assert( !platform::is_same<T, cuComplex>::value &&
!platform::is_same<T, cuDoubleComplex>::value &&
!platform::is_same<T, cutlass::complex<double>>::value &&
!platform::is_same<T, cutlass::complex<float>>::value, "May not be a complex data type");
return z;
}
/// Projects the complex number z onto the Riemann sphere
template <typename T>
@@ -511,10 +539,10 @@ CUTLASS_HOST_DEVICE complex<T> sin(complex<T> const &z) {
return (exp(-z) - exp(z)) * complex<T>(T(0), T(1) / T(2));
}
/// Comparison
/// Comparison
template <typename T>
CUTLASS_HOST_DEVICE bool operator<(complex<T> const &lhs, complex<T> const &rhs) {
return true;
return true;
}
//////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/conv/conv2d_problem_size.h
+14 -6
View File
@@ -44,13 +44,22 @@
Map tensor sizes (Conv2d -> ImplicitGemm) : implicit_gemm_tensor_[a|b|c]_size(ConvolutionOperator)
Map tensor problem sizes (Conv2d -> ImplicitGemm): implicit_gemm_problem_size(ConvolutionOperator)
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/tensor_coord.h"
#include "cutlass/fast_math.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/gemm_enumerated_types.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/conv/convolution.h"
#include "cutlass/functional.h"
@@ -80,7 +89,7 @@ struct Conv2dProblemSize {
public:
CUTLASS_HOST_DEVICE
Conv2dProblemSize():
Conv2dProblemSize():
N(0), H(0), W(0), C(0), P(0), Q(0), K(0), R(0), S(0),
pad_h(0), pad_w(0), stride_h(1), stride_w(1), dilation_h(1), dilation_w(1),
mode(Mode::kConvolution), split_k_slices(1), groups(1) { }
@@ -125,7 +134,7 @@ public:
int split_k_slices = 1,
int groups = 1
):
N(N), H(H), W(W), C(C), K(K), R(R), S(S), P(P), Q(Q),
N(N), H(H), W(W), C(C), P(P), Q(Q), K(K), R(R), S(S),
pad_h(pad_h), pad_w(pad_w), stride_h(stride_h), stride_w(stride_w),
dilation_h(dilation_h), dilation_w(dilation_w),
mode(mode), split_k_slices(split_k_slices), groups (groups) { }
@@ -145,11 +154,11 @@ public:
int groups = 1
):
N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()),
P(output_size.h()), Q(output_size.w()),
K(filter_size.n()), R(filter_size.h()), S(filter_size.w()),
pad_h(padding[0]), pad_w(padding[2]),
stride_h(stride.row()), stride_w(stride.column()),
dilation_h(dilation.row()), dilation_w(dilation.column()),
P(output_size.h()), Q(output_size.w()),
mode(mode), split_k_slices(split_k_slices), groups(groups) {}
/// Constructs convolution problem size from cutlass Tensor4DCoord and MatrixCoord
@@ -188,8 +197,8 @@ public:
int groups = 1
):
N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()),
P(output_size.h()), Q(output_size.w()),
K(filter_size.n()), R(filter_size.h()), S(filter_size.w()),
P(output_size.h()), Q(output_size.w()),
pad_h(R / 2), pad_w(S / 2), stride_h(1), stride_w(1),
dilation_h(1), dilation_w(1),
mode(mode), split_k_slices(split_k_slices), groups(groups) {}
@@ -486,7 +495,6 @@ int depthwise_gemm_k_iterations(
CUTLASS_HOST_DEVICE
int implicit_gemm_k_iterations_per_channel(
Operator conv_operator,
int threadblock_K,
Conv2dProblemSize const &problem_size,
IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic) {
include/cutlass/conv/conv3d_problem_size.h
+30 -22
View File
@@ -44,6 +44,15 @@
Map tensor sizes (Conv3d -> ImplicitGemm) : implicit_gemm_tensor_[a|b|c]_size(ConvolutionOperator)
Map tensor problem sizes (Conv3d -> ImplicitGemm): implicit_gemm_problem_size(ConvolutionOperator)
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
@@ -80,11 +89,11 @@ struct Conv3dProblemSize : public Conv2dProblemSize {
public:
CUTLASS_HOST_DEVICE
Conv3dProblemSize():
Conv2dProblemSize(),
D(0), T(0), Z(0),
pad_d(0),
stride_d(1),
dilation_d(1),
Conv2dProblemSize() { }
dilation_d(1) { }
/// Constructor for default padding, stride, dilation, and split-K
CUTLASS_HOST_DEVICE
@@ -102,10 +111,10 @@ public:
int R,
int S,
Mode mode
):
):
Conv2dProblemSize(N, H, W, C, P, Q, K, R, S, mode),
D(D), T(T), Z(Z),
pad_d(T / 2), stride_d(1), dilation_d(1),
Conv2dProblemSize(N, H, W, C, P, Q, K, R, S, mode) { }
pad_d(T / 2), stride_d(1), dilation_d(1) { }
/// Constructor
CUTLASS_HOST_DEVICE
@@ -134,15 +143,15 @@ public:
Mode mode,
int split_k_slices = 1,
int groups = 1
):
D(D), T(T), Z(Z),
pad_d(pad_d), stride_d(stride_d), dilation_d(dilation_d),
):
Conv2dProblemSize(
N, H, W, C, K, R, S, P, Q,
pad_h, pad_w,
stride_h, stride_w,
dilation_h, dilation_w,
mode, split_k_slices, groups) { }
N, H, W, C, K, R, S, P, Q,
pad_h, pad_w,
stride_h, stride_w,
dilation_h, dilation_w,
mode, split_k_slices, groups),
D(D), T(T), Z(Z),
pad_d(pad_d), stride_d(stride_d), dilation_d(dilation_d) { }
/// Constructs convolution problem size from cutlass Tensor5DCoord and Coord3D
// set *user-defined* output size and sets Z, P, and Q (include all data members in ctor)
@@ -158,8 +167,6 @@ public:
int split_k_slices = 1,
int groups = 1
):
D(input_size.d()), T(filter_size.d()), Z(output_size.d()),
pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0]),
Conv2dProblemSize(
{input_size.n(), input_size.h(), input_size.w(), input_size.c()},
{filter_size.n(), filter_size.h(), filter_size.w(), filter_size.c()},
@@ -167,8 +174,9 @@ public:
{stride[1], stride[2]},
{dilation[1], dilation[2]},
{output_size.n(), output_size.h(), output_size.w(), output_size.c()},
mode, split_k_slices, groups
) { }
mode, split_k_slices, groups),
D(input_size.d()), T(filter_size.d()), Z(output_size.d()),
pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0]) { }
/// Constructs convolution problem size from cutlass Tensor5DCoord and Coord3D
// *computes* output size and sets Z, P and Q (include all data members in ctor)
@@ -183,18 +191,18 @@ public:
int split_k_slices = 1,
int groups = 1
):
D(input_size.d()), T(filter_size.d()),
pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0]),
Conv2dProblemSize(
{input_size.n(), input_size.h(), input_size.w(), input_size.c()},
{filter_size.n(), filter_size.h(), filter_size.w(), filter_size.c()},
{padding[1], padding[1], padding[2], padding[2]},
{stride[1], stride[2]},
{dilation[1], dilation[2]},
mode, split_k_slices, groups
) {
mode, split_k_slices, groups),
D(input_size.d()), T(filter_size.d()),
pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0])
{
// set output Z
Z = ((D + pad_d * 2 - T * dilation_d) / stride_d) + 1;
Z = ((D + pad_d * 2 - T * dilation_d) / stride_d) + 1;
}
/// Equality operator (ignores mode and split_k_slice)
include/cutlass/conv/convolution.h
+11 -1
View File
@@ -70,13 +70,23 @@ Map elements' data types (ImplicitGemm -> Conv): GemmToConvElementMap
Map elements' data types (Conv -> ImplicitGemm): ConvToGemmElementMap
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/layout/tensor.h"
#include "cutlass/tensor_coord.h"
#include "cutlass/fast_math.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/gemm_enumerated_types.h"
#include "cutlass/matrix_coord.h"
namespace cutlass {
include/cutlass/conv/kernel/direct_convolution.h
+1 -1
View File
@@ -142,7 +142,7 @@ struct DirectConvolutionParams {
ThreadblockShape::kN);
gemm_k_iterations_per_channel = implicit_gemm_k_iterations_per_channel(
kConvolutionalOperator, ThreadblockShape::kK, args.problem_size, kIteratorAlgorithm);
kConvolutionalOperator, args.problem_size, kIteratorAlgorithm);
ThreadblockSwizzle threadblock_swizzle;
include/cutlass/conv/kernel/implicit_gemm_convolution.h
+1 -1
View File
@@ -250,7 +250,7 @@ struct ImplicitGemmConvolution {
ThreadblockShape::kN);
gemm_k_iterations_per_channel = implicit_gemm_k_iterations_per_channel(
kConvolutionalOperator, ThreadblockShape::kK, args.problem_size, kIteratorAlgorithm);
kConvolutionalOperator, args.problem_size, kIteratorAlgorithm);
ThreadblockSwizzle threadblock_swizzle;
include/cutlass/conv/threadblock/threadblock_swizzle.h
+6 -6
View File
@@ -95,11 +95,11 @@ struct StridedDgradHorizontalThreadblockSwizzle :
/// Returns the shape of the problem in units of logical tiles
/// For ImplicitGemmConvolution Conv2d problem size: conv_operator(NPQK, NHWC, KRSC)
CUTLASS_HOST_DEVICE
gemm::GemmCoord get_tiled_shape(
static gemm::GemmCoord get_tiled_shape(
cutlass::conv::Operator conv_operator,
cutlass::conv::Conv2dProblemSize const &problem_size,
gemm::GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
gemm::GemmCoord implicit_gemm_problem_size =
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
@@ -136,11 +136,11 @@ struct StridedDgradIdentityThreadblockSwizzle :
/// Returns the shape of the problem in units of logical tiles
/// For ImplicitGemmConvolution Conv2d problem size: conv_operator(NPQK, NHWC, KRSC)
CUTLASS_HOST_DEVICE
gemm::GemmCoord get_tiled_shape(
static gemm::GemmCoord get_tiled_shape(
cutlass::conv::Operator conv_operator,
cutlass::conv::Conv2dProblemSize const &problem_size,
gemm::GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
gemm::GemmCoord implicit_gemm_problem_size =
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
@@ -174,10 +174,10 @@ struct DepthwiseDirect2dConvIdentityThreadblockSwizzle
/// Returns the shape of the problem in units of logical tiles
CUTLASS_HOST_DEVICE
gemm::GemmCoord get_tiled_shape(cutlass::conv::Operator conv_operator,
static gemm::GemmCoord get_tiled_shape(cutlass::conv::Operator conv_operator,
cutlass::conv::Conv2dProblemSize const &problem_size,
gemm::GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
gemm::GemmCoord implicit_gemm_problem_size =
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
include/cutlass/coord.h
+10
View File
@@ -32,6 +32,16 @@
\brief A Coord is a coordinate of arbitrary rank into a tensor or matrix
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#if defined(__CUDACC_RTC__)
include/cutlass/core_io.h
+9 -2
View File
@@ -31,7 +31,15 @@
/*! \file
\brief Helpers for printing cutlass/core objects
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include <iostream>
@@ -45,11 +53,10 @@
#include "cutlass/matrix_shape.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/tensor_view.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/gemm_enumerated_types.h"
#include "cutlass/conv/convolution.h"
#include "cutlass/conv/conv2d_problem_size.h"
#include "cutlass/conv/conv3d_problem_size.h"
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Output operator for CUDA built-in dim3 type
include/cutlass/cutlass.h
+10
View File
@@ -33,6 +33,16 @@
\brief Basic include for CUTLASS.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/detail/helper_macros.hpp"
include/cutlass/detail/helper_macros.hpp
+14
View File
@@ -141,4 +141,18 @@ namespace cutlass {
#define CUTLASS_THREAD_LOCAL
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
#if (201700L <= __cplusplus)
#define CUTLASS_CONSTEXPR_IF_CXX17 constexpr
#define CUTLASS_CXX17_OR_LATER 1
#else
#define CUTLASS_CONSTEXPR_IF_CXX17
#define CUTLASS_CXX17_OR_LATER 0
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
}; // namespace cutlass
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/detail/layout.hpp
+17
View File
@@ -239,6 +239,23 @@ check_alignment(Shape const & shape, Stride const & stride) {
: get_contiguous_shape(cute::get<1>(shape), cute::get<1>(stride)) % Alignment == 0;
}
// Check if tensor shape satisfies a given major alignment
template<int B, int M, int S>
CUTLASS_HOST_DEVICE constexpr
size_t
alignment_for_swizzle(cute::Swizzle<B, M, S>) {
static_assert(B >= 0 and M >= 0);
return size_t(1) << size_t(B + M + cute::abs(S));
}
template<class Layout>
CUTLASS_HOST_DEVICE constexpr
size_t
alignment_for_swizzle(Layout layout) {
return alignment_for_swizzle(cute::detail::get_swizzle_portion(layout));
}
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::detail
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+91 -1
View File
@@ -259,6 +259,95 @@ struct Sm90TmaBuilderImpl {
>;
};
///////////////////////////////////////////////////////////////////////////////
// Descriptor classes for defining EVT nodes
// Some of the epilogue visitor nodes require non-intuitive template arguments
// such as CopyOpS2R for AuxLoad node. Traditionaly, these are resolved by the
// builder classes. Here we provide a set of descriptor classes that resolve
// these template arguments from more intuitive types such as Stride, Layout
// Get TileShape, EpilogueTile, Dispatch Policy, StagesC, and STagesD
template<
typename TileShape_MNK,
typename EpilogueTileType,
typename ElementC,
typename ElementD,
typename Schedule
>
struct EpilogueDescriptor {
using TileShape = TileShape_MNK;
using EpilogueTile =
decltype(
detail::sm90_compute_tile_shape_or_override<
ElementD, EpilogueTileType, Schedule
>()
);
using DispatchPolicy =
decltype(
detail::sm90_get_tma_dispatch_policy<
TileShape_MNK, EpilogueTile,
ElementC, ElementD, Schedule
>()
);
constexpr static int StagesC = DispatchPolicy::StagesC;
constexpr static int StagesD = DispatchPolicy::StagesD;
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxLoad node
template<
typename EpilogueDescriptor,
typename StrideOrLayoutTag,
typename ElementAux
>
struct AuxLoadDescriptor {
constexpr static int Stages = EpilogueDescriptor::StagesC;
using EpilogueTile = typename EpilogueDescriptor::EpilogueTile;
using Element = ElementAux;
using Stride = cutlass::detail::TagToStrideC_t<StrideOrLayoutTag>;
using SmemLayoutAtom =
decltype(
detail::sm90_get_epilogue_smem_swizzle_layout_atom<
Stride, ElementAux, typename EpilogueDescriptor::EpilogueTile
>()
);
using CopyOpS2R =
decltype(detail::sm90_get_smem_load_op_for_source<Stride, ElementAux>());
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxStore node
template<
typename EpilogueDescriptor,
typename StrideOrLayoutTag,
typename ElementAux
>
struct AuxStoreDescriptor {
constexpr static int Stages = EpilogueDescriptor::StagesD;
using EpilogueTile = typename EpilogueDescriptor::EpilogueTile;
using Element = ElementAux;
using Stride = cutlass::detail::TagToStrideC_t<StrideOrLayoutTag>;
using SmemLayoutAtom =
decltype(
detail::sm90_get_epilogue_smem_swizzle_layout_atom<
Stride, ElementAux, typename EpilogueDescriptor::EpilogueTile
>()
);
using CopyOpR2S =
decltype(detail::sm90_get_smem_store_op_for_accumulator<Stride, ElementAux>());
};
template<
typename EpilogueDescriptor,
typename ElementVector
>
struct RowBroadcastDescriptor {
constexpr static int Stages = ceil_div(
EpilogueDescriptor::StagesC,
size(shape_div(take<0, 2>(typename EpilogueDescriptor::TileShape{}), typename EpilogueDescriptor::EpilogueTile{}))
) + 1;
using Element = ElementVector;
};
} // namespace detail
///////////////////////////////////////////////////////////////////////////////
@@ -426,7 +515,8 @@ private:
ElementD,
GmemLayoutTagD,
AlignmentD,
EpilogueSchedule
EpilogueSchedule,
FusionOperation
>;
public:
include/cutlass/epilogue/collective/collective_builder.hpp
+1
View File
@@ -45,6 +45,7 @@ struct EpilogueTileAuto {};
// Used to let the builder pick the epilogue schedule automatically.
// Can be overridden with kernel schedule tags in cutlass/gemm/dispatch_policy.hpp
struct EpilogueScheduleAuto {};
struct EpilogueIm2ColScheduleAuto {};
template <
class ArchTag,
include/cutlass/epilogue/collective/detail.hpp
+2 -2
View File
@@ -126,14 +126,14 @@ public:
CUTLASS_HOST_DEVICE
static constexpr int
get_load_pipe_increment([[maybe_unused]] TileShapeMNK) {
return 0;
return 1;
}
template<class TileShapeMNK>
CUTLASS_HOST_DEVICE
static constexpr int
get_store_pipe_increment([[maybe_unused]] TileShapeMNK) {
return 0;
return 1;
}
CUTLASS_DEVICE
include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp
+73 -56
View File
@@ -40,6 +40,7 @@
#include "cutlass/epilogue/collective/detail.hpp"
#include "cutlass/epilogue/thread/scale_type.h"
#include "cutlass/epilogue/fusion/callbacks.hpp"
#include "cutlass/detail/layout.hpp"
#include "cutlass/trace.h"
#include "cute/tensor.hpp"
@@ -119,40 +120,52 @@ public:
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]");
private:
using InternalElementC = cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
using SmemElementC = cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
constexpr static int StagesC = StagesC_;
constexpr static int StagesD = StagesD_;
constexpr static bool ReuseSmemC = ReuseSmemC_;
constexpr static bool is_source_supported = not cute::is_void_v<ElementC>;
// internal optimization to reuse C shared memory for storing D
using SmemLayoutAtomBitsC = decltype(downcast<sizeof_bits<InternalElementC>::value>(SmemLayoutAtomC{}));
using SmemLayoutAtomBitsD = decltype(downcast<sizeof_bits<ElementD>::value>(SmemLayoutAtomD{}));
constexpr static bool support_smem_reuse = is_source_supported &&
sizeof(InternalElementC) == sizeof(ElementD) &&
StrideC{} == StrideD{} &&
StagesD <= StagesC &&
cute::is_same_v<SmemLayoutAtomBitsC,SmemLayoutAtomBitsD>;
constexpr static bool ReuseSmemC = DispatchPolicy::ReuseSmemC;
static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met");
constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
public:
using SmemLayoutC = decltype(tile_to_shape(
SmemLayoutAtomC{},
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesC>{}),
cute::conditional_t<is_m_major_C, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
using SmemLayoutD = decltype(tile_to_shape(
SmemLayoutAtomD{},
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesD>{}),
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<ReuseSmemC ? StagesC : StagesD>{}),
cute::conditional_t<is_m_major_D, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
constexpr static bool support_smem_reuse = is_source_supported && StagesD <= StagesC
&& cosize(take<0,2>(SmemLayoutC{})) == cosize(take<0,2>(SmemLayoutD{}));
static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met");
constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{});
constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{});
struct TensorStorageWithC {
alignas(SmemAlignmentC) array_aligned<SmemElementC, size(SmemLayoutC{})> smem_C;
alignas(SmemAlignmentD) array_aligned<ElementD, size(SmemLayoutD{})> smem_D;
using FusionStorage = typename FusionCallbacks::SharedStorage;
FusionStorage thread;
};
struct TensorStorageWithoutC {
alignas(SmemAlignmentD) array_aligned<ElementD, size(SmemLayoutD{})> smem_D;
using FusionStorage = typename FusionCallbacks::SharedStorage;
FusionStorage thread;
};
public:
// TMA pipeline for loading C
using LoadPipeline = cutlass::PipelineTransactionAsync<StagesC>;
using LoadPipelineState = cutlass::PipelineState<StagesC>;
constexpr static uint32_t TmaTransactionBytes =
size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof(InternalElementC));
size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof(SmemElementC));
// TMA pipeline for storing D
using StorePipeline = cute::conditional_t<ReuseSmemC,
@@ -161,17 +174,9 @@ public:
using StorePipelineState = cutlass::PipelineState<ReuseSmemC ? StagesC : StagesD>;
struct SharedStorage {
struct TensorStorage : aligned_struct<128> {
cute::conditional_t<not is_source_supported,
detail::EmptyStorage<InternalElementC>,
array_aligned<InternalElementC, size(SmemLayoutC{})>> smem_C;
alignas(128) cute::conditional_t<ReuseSmemC,
detail::EmptyStorage<ElementD>,
array_aligned<ElementD, size(SmemLayoutD{})>> smem_D;
using FusionStorage = typename FusionCallbacks::SharedStorage;
alignas(128) FusionStorage thread;
} tensors;
using TensorStorage =
cute::conditional_t<not is_source_supported or ReuseSmemC, TensorStorageWithoutC, TensorStorageWithC>;
TensorStorage tensors;
using PipelineStorage = typename LoadPipeline::SharedStorage;
PipelineStorage pipeline;
@@ -192,7 +197,7 @@ public:
struct Params {
using TMA_C = decltype(make_tma_copy(
CopyOpG2S{},
make_tensor(static_cast<InternalElementC const*>(nullptr),
make_tensor(static_cast<SmemElementC const*>(nullptr),
repeat_like(StrideC{}, int32_t(0)), StrideC{}),
SmemLayoutC{}(_,_,0)));
using TMA_D = decltype(make_tma_copy(
@@ -316,21 +321,22 @@ public:
int thread_idx,
TensorStorage& shared_tensors) {
using namespace cute;
using _X = Underscore;
// Indexing variables
auto [M, N, K, L] = problem_shape_mnkl;
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
// Represent the full source tensor, slice to get the tile this CTA is currently responsible for
Tensor mC_mnl = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gC_mnl = local_tile(mC_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1,_X>{}); // (CTA_M,CTA_N,m,n,l)
Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (CTA_M,CTA_N)
Tensor mC = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gC = local_tile(mC, take<0,2>(CtaTileMNK{}), make_coord(m_coord,n_coord,l_coord)); // (CTA_M,CTA_N)
// Apply epilogue subtile, get matching smem tensor
auto ptr_sC = make_smem_ptr(shared_tensors.smem_C.data());
SmemElementC* ptr_sC = reinterpret_cast<SmemElementC*>(shared_tensors.smem_D.data());
if constexpr (not ReuseSmemC and is_source_supported) {
ptr_sC = shared_tensors.smem_C.data();
}
Tensor gC_epi = local_tile(gC, EpilogueTile{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sC_epi = make_tensor(ptr_sC, SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sC_epi = make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
// Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_)
ThrCopy thrblk_g2s = params.tma_load_c.get_slice(Int<0>{});
@@ -420,8 +426,9 @@ public:
int thread_idx,
TensorStorage& shared_tensors) {
using namespace cute;
using _X = Underscore;
using ElementAccumulator = typename AccEngine::value_type;
using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::ElementCompute;
using ElementCompute = cute::conditional_t<cute::is_void_v<ElementCompute_>,ElementAccumulator,ElementCompute_>;
static_assert(is_rmem<AccEngine>::value, "Accumulator must be RF resident.");
static_assert(rank(AccLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)");
@@ -439,16 +446,22 @@ public:
auto epi_tile_n = size<1>(EpilogueTile{});
// Represent the full output tensor, slice to get the tile this CTA is responsible for
Tensor mD_mnl = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gD_mnl = local_tile(mD_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1,_X>{}); // (CTA_M,CTA_N,m,n,l)
Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (CTA_M,CTA_N)
// Apply epilogue subtiling, construct corresponding pipelined smem tensors
auto ptr_sC = make_smem_ptr(shared_tensors.smem_C.data());
auto ptr_sD = make_smem_ptr(shared_tensors.smem_D.data());
Tensor mD = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gD = local_tile(mD, take<0,2>(CtaTileMNK{}), make_coord(m_coord,n_coord,l_coord)); // (CTA_M,CTA_N)
// Apply epilogue subtiling
Tensor gD_epi = local_tile(gD, EpilogueTile{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sC_epi = make_tensor(ptr_sC, SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sD_epi = make_tensor(ptr_sD, SmemLayoutD{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
// Construct the corresponding pipelined smem tensors
SmemElementC* ptr_sC = reinterpret_cast<SmemElementC*>(shared_tensors.smem_D.data());
if constexpr (not ReuseSmemC and is_source_supported) {
ptr_sC = shared_tensors.smem_C.data();
}
ElementD* ptr_sD = shared_tensors.smem_D.data();
Tensor sC_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sD_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sD), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
// Get the smallest tiled copy we can use to retile the accumulators
using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>;
@@ -458,14 +471,11 @@ public:
TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom<CopyOpR2S,ElementD>{}, tiled_copy_C_atom);
ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx);
Tensor tRS_rAcc = thread_r2s.retile_S(accumulators); // ((R2S,R2S_V),MMA_M,MMA_N)
Tensor tRS_sD = conditional_return<ReuseSmemC>(
thread_r2s.partition_D(recast<ElementD>(sC_epi)), // (R2S,R2S_M,R2S_N,PIPE_C)
thread_r2s.partition_D(sD_epi) ); // (R2S,R2S_M,R2S_N,PIPE_D)
Tensor tRS_sD = thread_r2s.partition_D(sD_epi); // (R2S,R2S_M,R2S_N,PIPE_D)
// Allocate register tensors
auto tRS_rD_shape = take<0,3>(shape(thread_r2s.partition_S(sD_epi)));
Tensor tRS_rC = make_tensor<InternalElementC>(tRS_rD_shape); // (R2S,R2S_M,R2S_N)
Tensor tRS_rD = make_tensor<ElementD >(tRS_rD_shape); // (R2S,R2S_M,R2S_N)
// Allocate D registers
Layout tRS_rD_layout = make_layout(take<0,3>(shape(thread_r2s.partition_S(sD_epi))));
Tensor tRS_rD = make_tensor<ElementD>(tRS_rD_layout); // (R2S,R2S_M,R2S_N)
// Vectorized fragment view
constexpr int FragmentSize = DispatchPolicy::FragmentSize;
@@ -474,16 +484,23 @@ public:
CUTE_STATIC_ASSERT(size<0>(tRS_rAcc) % FragmentSize == 0, "Fragment size does not vectorize properly");
// (t)hread-partition for (s)mem to (r)egister copy (tSR_)
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R,InternalElementC>{}, tiled_copy_C_atom);
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R, SmemElementC>{}, tiled_copy_C_atom);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C)
Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N)
Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C)
Layout tSR_rC_layout = thread_s2r.retile_D(tRS_rD).layout(); // (S2R,S2R_M,S2R_N)
// Allocate C registers
// If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type
// to eliminate some redundant pack+unpack instruction sequences for sub-word types
constexpr bool IsDirectS2R = cute::is_same_v<CopyOpS2R,DefaultCopy>
&& decltype(max_common_vector(tSR_rC_layout, tSR_sC.layout()))::value <= 1;
using RegisterElementC = cute::conditional_t<IsDirectS2R, ElementCompute, SmemElementC>;
Tensor tRS_rC = make_tensor<RegisterElementC>(tRS_rD_layout); // (R2S,R2S_M,R2S_N)
Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N)
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{});
Tensor bSG_sD = conditional_return<ReuseSmemC>(
thrblk_s2g.partition_S(recast<ElementD>(sC_epi)), // (S2G,S2G_M,S2G_N,PIPE_C)
thrblk_s2g.partition_S(sD_epi) ); // (S2G,S2G_M,S2G_N,PIPE_D)
Tensor bSG_sD = thrblk_s2g.partition_S(sD_epi); // (S2G,S2G_M,S2G_N,PIPE_D)
Tensor bSG_gD = thrblk_s2g.partition_D(gD_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N)
CUTE_STATIC_ASSERT(mma_tile_m == epi_tile_m, "EPI_TILE_M must equal MMA_TILE_M");
include/cutlass/epilogue/fusion/callbacks.hpp
+2
View File
@@ -62,6 +62,7 @@ struct FusionCallbacksTraits {
using Operation = T;
using CtaTile_MNK = void;
using EpilogueTile_MN = void;
using ElementCompute = void;
};
template <
@@ -78,6 +79,7 @@ struct FusionCallbacksTraits<
using Operation = Operation_;
using CtaTile_MNK = CtaTile_MNK_;
using EpilogueTile_MN = EpilogueTile_MN_;
using ElementCompute = typename Operation::ElementCompute;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/fusion/operations.hpp
+3 -5
View File
@@ -60,7 +60,7 @@ struct FusionOperation {
using ElementBias = void;
static constexpr int AlignmentBias = 0;
static constexpr bool IsPerRowBiasSupported = false;
template <class> using ActivationFn = void;
using ActivationFn = void;
static constexpr bool IsEltActSupported = false;
using ElementAux = void;
@@ -108,8 +108,7 @@ template<
>
struct LinCombEltAct
: LinearCombination<ElementOutput_, ElementCompute_, ElementScalar_, RoundStyle_> {
template <class T>
using ActivationFn = ActivationFn_<T>;
using ActivationFn = ActivationFn_<ElementCompute_>;
static constexpr bool IsEltActSupported = true;
};
@@ -142,8 +141,7 @@ template<
struct LinCombPerRowBiasEltAct
: LinCombPerRowBias<ElementOutput_, ElementCompute_,
ElementBias_, ElementScalar_, AlignmentBias_, RoundStyle_> {
template <class T>
using ActivationFn = ActivationFn_<T>;
using ActivationFn = ActivationFn_<ElementCompute_>;
static constexpr bool IsEltActSupported = true;
};
include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp
+63 -25
View File
@@ -217,6 +217,9 @@ struct FusionCallbacks<
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
{ // unary op: activation(beta * C + (alpha * acc))
@@ -230,7 +233,7 @@ struct FusionCallbacks<
}, // end binary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args: activation
activation // unary args: activation
}; // end unary op
}
};
@@ -258,7 +261,7 @@ using Sm90LinCombPerRowBias =
Sm90EVT<Sm90Compute<multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ScalarBroadcast<ElementScalar>, // alpha
Sm90AccFetch, // acc
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,_0>, AlignmentBias> // bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias
>
>;
@@ -293,7 +296,10 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
operator typename Impl::Arguments() const {
return
@@ -303,7 +309,7 @@ struct FusionCallbacks<
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
@@ -373,7 +379,13 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
@@ -384,12 +396,12 @@ struct FusionCallbacks<
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args : activation
activation // unary args : activation
}; // end unary op
}
};
@@ -461,10 +473,9 @@ struct FusionCallbacks<
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
> {
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
using Impl =
Sm90LinCombPerRowBiasEltActAux<
CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>, SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
using Operation =
@@ -478,7 +489,15 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
ElementAux* aux_ptr = nullptr;
StrideAux dAux = {};
@@ -492,14 +511,14 @@ struct FusionCallbacks<
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{aux_ptr, dAux} // unary args : store
}, // end unary op
{} // unary args : activation
activation // unary args : activation
}; // end unary op
}
};
@@ -528,7 +547,7 @@ using Sm90PerRowLinCombPerRowBias =
Sm90EVT<Sm90Compute<multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,_0>, AlignmentScalar>, // alpha
Sm90AccFetch, // acc
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,_0>, AlignmentBias> // bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias
>
>;
@@ -591,7 +610,13 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
@@ -600,14 +625,14 @@ struct FusionCallbacks<
{beta_ptr, beta}, // leaf args : beta
{}, // leaf args : C
{ // ternary op : alpha * acc + bias
{alpha_ptr, alpha}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{alpha_ptr, alpha}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args : activation
activation // unary args : activation
}; // end unary op
}
};
@@ -650,7 +675,7 @@ using Sm90ScaledLinCombPerRowBias =
Sm90EVT<Sm90Compute<multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,_0>, 3>, // scale_a * scale_b * alpha
Sm90AccFetch, // acc
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,_0>, AlignmentBias> // bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias
>
>;
@@ -728,7 +753,12 @@ struct FusionCallbacks<
ElementScalar const* scale_c_ptr = nullptr;
ElementScalar const* scale_d_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
@@ -742,14 +772,14 @@ struct FusionCallbacks<
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{scale_a, scale_b, alpha},
{scale_a_ptr, scale_b_ptr, alpha_ptr}
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args : activation
activation // unary args : activation
}, // end unary op
{{scale_d},
{scale_d_ptr}
@@ -855,10 +885,10 @@ struct FusionCallbacks<
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
> {
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
using Impl =
Sm90ScaledLinCombPerRowBiasEltActAmaxAux<
CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>,
SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
using Operation =
@@ -885,9 +915,17 @@ struct FusionCallbacks<
ElementScalar scale_aux = ElementScalar(1);
ElementScalar const* scale_aux_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
ElementAmax* amax_D_ptr = nullptr;
ElementAmax* amax_aux_ptr = nullptr;
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
ElementAux* aux_ptr = nullptr;
StrideAux dAux = {};
@@ -905,9 +943,9 @@ struct FusionCallbacks<
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{scale_a, scale_b, alpha},
{scale_a_ptr, scale_b_ptr, alpha_ptr}
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
@@ -924,7 +962,7 @@ struct FusionCallbacks<
{ // unary op : reduce(activation(Z))
{ // unary op : activation(Z)
{}, // leaf args : Z
{} // unary args : activation
activation // unary args : activation
}, // end unary op
{amax_D_ptr_} // unary args : reduce
}, // end unary op
include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
+121
View File
@@ -99,7 +99,9 @@ struct Sm90Compute : Sm90VisitorImpl<> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -118,6 +120,123 @@ struct Sm90Compute : Sm90VisitorImpl<> {
};
// partial specialization for compute fns that define an Arguments member, e.g. activation hyperparameters
template<
template <class> class ComputeFn,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle
>
struct Sm90Compute<
ComputeFn,
ElementOutput,
ElementCompute,
RoundStyle,
cute::void_t<typename ComputeFn<ElementCompute>::Arguments>
> {
struct SharedStorage { };
using Arguments = typename ComputeFn<ElementCompute>::Arguments;
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
CUTLASS_DEVICE bool
is_producer_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_C_load_needed() const {
return false;
}
CUTLASS_HOST_DEVICE
Sm90Compute() { }
CUTLASS_HOST_DEVICE
Sm90Compute(Params const& params, SharedStorage const& shared_storage)
: params(params) {}
Params const params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
get_producer_load_callbacks(
ProblemShapeMNKL problem_shape_mnkl,
TileShapeMNK tile_shape_mnk,
TileCoordMNKL tile_coord_mnkl,
EpilogueTile epi_tile,
int thread_idx) {
return EmptyProducerLoadCallbacks{};
}
struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
CUTLASS_DEVICE
ConsumerStoreCallbacks(Params const& params)
: params(params) {}
Params const& params;
template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
CUTLASS_DEVICE Array<ElementOutput, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n,
Array<ElementInputs, FragmentSize> const&... frg_inputs) {
return transform_apply(cute::make_tuple(frg_inputs...),
[&] (auto&& frg_input) {
using ElementInput = typename cute::remove_cvref_t<decltype(frg_input)>::Element;
using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
return convert_input(frg_input);
},
[&] (auto&&... cvt_frg_inputs) {
using ComputeOutput = ComputeFn<Array<ElementCompute, FragmentSize>>;
using ConvertOutput = NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize, RoundStyle>;
ComputeOutput compute_output{};
ConvertOutput convert_output{};
return convert_output(compute_output(cvt_frg_inputs..., params));
}
);
}
};
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
>
CUTLASS_DEVICE auto
get_consumer_store_callbacks(
ProblemShapeMNKL problem_shape_mnkl,
TileShapeMNK tile_shape_mnk,
TileCoordMNKL tile_coord_mnkl,
EpilogueTile epi_tile,
TiledCopy tiled_copy,
int thread_idx,
SrcTensor const& tCrC) {
return ConsumerStoreCallbacks(params);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Performance Optimized Specializations
@@ -215,7 +334,9 @@ struct Sm90TreeVisitor<
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
+45 -15
View File
@@ -71,7 +71,9 @@ struct Sm90AccFetch : Sm90VisitorImpl<> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -129,7 +131,9 @@ struct Sm90SrcFetch : Sm90VisitorImpl<> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -181,7 +185,8 @@ struct Sm90AuxLoad {
cute::conditional_t<is_m_major, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
struct SharedStorage {
alignas(128) array_aligned<Element, size(SmemLayout{})> smem_aux;
alignas(cutlass::detail::alignment_for_swizzle(SmemLayout{}))
array_aligned<Element, size(SmemLayout{})> smem_aux;
};
struct Arguments {
@@ -222,9 +227,9 @@ struct Sm90AuxLoad {
Sm90AuxLoad() { }
CUTLASS_HOST_DEVICE
Sm90AuxLoad(Params const& params, SharedStorage& shared_storage)
Sm90AuxLoad(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params),
smem_aux(shared_storage.smem_aux.data()) { }
smem_aux(const_cast<Element*>(shared_storage.smem_aux.data())) { }
Params const* params_ptr;
Element* smem_aux;
@@ -273,7 +278,9 @@ struct Sm90AuxLoad {
};
template <
class TileShapeMNK
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL
>
CUTLASS_DEVICE auto
get_producer_load_callbacks(
@@ -284,8 +291,9 @@ struct Sm90AuxLoad {
int thread_idx) {
auto [M, N, K, L] = problem_shape_mnkl;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mAux = params_ptr->tma_load_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gAux = sm90_tensor_to_cta_tile(mAux, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
Tensor gAux = local_tile(mAux, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor gAux_epi = local_tile(gAux, epi_tile, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE)
@@ -339,7 +347,9 @@ struct Sm90AuxLoad {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class TiledCopy,
class SrcTensor
>
@@ -363,7 +373,8 @@ struct Sm90AuxLoad {
make_tiled_copy_S(Copy_Atom<CopyOpS2R,Element>{}, tiled_copy),
make_tiled_copy_D(Copy_Atom<CopyOpS2R,Element>{}, tiled_copy)
);
Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE)
Tensor sAux_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(smem_aux), SmemLayout{})); // (EPI_TILE_M,EPI_TILE_N,PIPE)
auto tSR_sAux = tiled_s2r.get_slice(thread_idx).partition_S(sAux_epi); // (S2R,S2R_M,S2R_N,PIPE)
@@ -378,6 +389,7 @@ struct Sm90AuxLoad {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Scalar broadcast
// Supports reduction over multiple broadcasts to support fusions such as fp8 scaling factors
template<
class Element,
class StrideMNL = Stride<_0,_0,_0>,
@@ -387,7 +399,8 @@ template<
struct Sm90ScalarBroadcast {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0, _0>>) || // scalar broadcast, e.g. alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>)); // batched scalar broadcast, e.g. per-batch alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0, _1>>) || // batched scalar broadcast, e.g. per-batch alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>));
struct SharedStorage { };
@@ -419,7 +432,7 @@ struct Sm90ScalarBroadcast {
Sm90ScalarBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90ScalarBroadcast(Params const& params, SharedStorage& shared_storage)
Sm90ScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) {
// Get the scalar for non-batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,_0>>) {
@@ -431,7 +444,9 @@ struct Sm90ScalarBroadcast {
Params const* params_ptr;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -442,7 +457,9 @@ struct Sm90ScalarBroadcast {
EpilogueTile epi_tile,
int thread_idx) {
// Get the scalar for batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
if constexpr (
cute::is_same_v<StrideMNL, Stride<_0,_0,_1>> ||
cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
update_scalar(l_coord);
}
@@ -470,7 +487,9 @@ struct Sm90ScalarBroadcast {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -486,7 +505,9 @@ struct Sm90ScalarBroadcast {
SrcTensor const& tCrC) {
// Get the scalar for batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
if constexpr (
cute::is_same_v<StrideMNL, Stride<_0,_0,_1>> ||
cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
update_scalar(l_coord);
}
@@ -541,7 +562,7 @@ struct Sm90RowBroadcast {
// Accumulator doesn't distribute row elements evenly amongst threads so we must buffer in smem
struct SharedStorage {
array_aligned<Element, size<1>(CtaTileShapeMNK{}) * Stages> smem_row;
alignas(16) array_aligned<Element, size<1>(CtaTileShapeMNK{}) * Stages> smem_row;
};
struct Arguments {
@@ -562,9 +583,9 @@ struct Sm90RowBroadcast {
Sm90RowBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90RowBroadcast(Params const& params, SharedStorage& shared_storage)
Sm90RowBroadcast(Params const& params, SharedStorage const& shared_storage)
: params(params),
smem_row(shared_storage.smem_row.data()) { }
smem_row(const_cast<Element*>(shared_storage.smem_row.data())) { }
Params params;
Element* smem_row;
@@ -613,7 +634,9 @@ struct Sm90RowBroadcast {
};
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -625,8 +648,9 @@ struct Sm90RowBroadcast {
int thread_idx) {
auto [M, N, K, L] = problem_shape_mnkl;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
Tensor gRow = sm90_tensor_to_cta_tile(mRow, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
Tensor gRow = local_tile(mRow, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor sRow = make_tensor(make_smem_ptr(smem_row), // (CTA_M,CTA_N,PIPE)
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{}), Stages),
make_stride(_0{},_1{},size<1>(CtaTileShapeMNK{})));
@@ -680,7 +704,9 @@ struct Sm90RowBroadcast {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -757,13 +783,15 @@ struct Sm90ColBroadcast {
Sm90ColBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90ColBroadcast(Params const& params, SharedStorage& shared_storage)
Sm90ColBroadcast(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -819,7 +847,9 @@ struct Sm90ColBroadcast {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp
+32 -12
View File
@@ -83,7 +83,8 @@ struct Sm90AuxStore {
cute::conditional_t<is_m_major, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
struct SharedStorage {
alignas(128) array_aligned<Element, size(SmemLayout{})> smem_aux;
alignas(cutlass::detail::alignment_for_swizzle(SmemLayout{}))
array_aligned<Element, size(SmemLayout{})> smem_aux;
};
struct Arguments {
@@ -125,9 +126,9 @@ struct Sm90AuxStore {
Sm90AuxStore() { }
CUTLASS_HOST_DEVICE
Sm90AuxStore(Params const& params, SharedStorage& shared_storage)
Sm90AuxStore(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params),
smem_aux(shared_storage.smem_aux.data()) { }
smem_aux(const_cast<Element*>(shared_storage.smem_aux.data())) { }
Params const* params_ptr;
Element* smem_aux;
@@ -143,7 +144,9 @@ struct Sm90AuxStore {
}
template <
class TileShapeMNK
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL
>
CUTLASS_DEVICE auto
get_producer_load_callbacks(
@@ -233,7 +236,9 @@ struct Sm90AuxStore {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class TiledCopy,
class SrcTensor
>
@@ -248,14 +253,16 @@ struct Sm90AuxStore {
SrcTensor const& tCrC) {
auto [M, N, K, L] = problem_shape_mnkl;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mAux = params_ptr->tma_store_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gAux = sm90_tensor_to_cta_tile(mAux, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
Tensor gAux = local_tile(mAux, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor tC_gAux = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
gAux, epi_tile, tiled_copy, thread_idx);
Tensor tC_rAux = make_tensor<Element>(take<0,3>(shape(tC_gAux))); // (CPY,CPY_M,CPY_N)
Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE)
Tensor sAux_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(smem_aux), SmemLayout{})); // (EPI_TILE_M,EPI_TILE_N,PIPE)
Tensor gAux_epi = local_tile(gAux, epi_tile, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
auto tiled_r2s = conditional_return<ReferenceSrc>(
@@ -297,8 +304,8 @@ template <
struct Sm90ScalarReduction {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0, _0>>) || // scalar reduction, e.g. tensor max element
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>)); // batched scalar reduction, e.g. per-batch max element
(cute::is_same_v<StrideMNL, Stride<_0,_0, _1>>) || // batched scalar reduction, e.g. per-batch max element
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>));
struct SharedStorage { };
struct Arguments {
@@ -329,13 +336,15 @@ struct Sm90ScalarReduction {
Sm90ScalarReduction() { }
CUTLASS_HOST_DEVICE
Sm90ScalarReduction(Params const& params, SharedStorage& shared_storage)
Sm90ScalarReduction(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params const params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -417,7 +426,9 @@ struct Sm90ScalarReduction {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -502,13 +513,15 @@ struct Sm90RowReduction {
Sm90RowReduction() { }
CUTLASS_HOST_DEVICE
Sm90RowReduction(Params const& params, SharedStorage& shared_storage)
Sm90RowReduction(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -619,7 +632,9 @@ struct Sm90RowReduction {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -707,13 +722,15 @@ struct Sm90ColReduction {
Sm90ColReduction() { }
CUTLASS_HOST_DEVICE
Sm90ColReduction(Params const& params, SharedStorage& shared_storage)
Sm90ColReduction(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -765,10 +782,11 @@ struct Sm90ColReduction {
Array frg_I = convert_input(frg_input);
Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
Tensor tCcCol_mn = tCcCol(_,_,_,epi_m,epi_n);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
if (elem_less(tCcCol(i), residue_mn)) {
if (elem_less(tCcCol_mn(i), residue_mn)) {
ElementCompute& tCrCol_vmn = tCrCol_mn(epi_v * FragmentSize + i);
tCrCol_vmn = reduce_input(tCrCol_vmn, frg_I[i]);
}
@@ -808,7 +826,9 @@ struct Sm90ColReduction {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp
+48 -56
View File
@@ -51,34 +51,12 @@ using cute::tuple;
namespace detail {
// Convenience aliases
using ProblemShapeMNKL = tuple<int,int,int,int>;
using TileCoordMNKL = tuple<int,int,int,int>;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Partitioning Helpers
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class Engine, class LayoutMNL,
class TileShapeMNK
>
CUTLASS_HOST_DEVICE
constexpr auto
sm90_tensor_to_cta_tile(
Tensor<Engine, LayoutMNL> mT, // (M,N,L)
TileShapeMNK tile_shape_mnk, // (CTA_M,CTA_N,CTA_K)
TileCoordMNKL tile_coord_mnkl) {
using _X = Underscore;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mT_mnl = local_tile(mT, tile_shape_mnk, make_coord(_,_,_), Step<_1,_1,_X>{}); // (CTA_M,CTA_N)
return mT_mnl(_,_,m,n,l);
}
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class CtaTileMN,
@@ -106,6 +84,7 @@ template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class Engine, class LayoutMNL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy
>
@@ -118,7 +97,8 @@ sm90_partition_for_epilogue(
EpilogueTile epi_tile, // (EPI_TILE_M,EPI_TILE_N)
TiledCopy tiled_copy,
int thread_idx) {
Tensor cT = sm90_tensor_to_cta_tile(mT, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
auto [m, n, k, l] = tile_coord_mnkl;
Tensor cT = local_tile(mT, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor tCcT =
sm90_partition_for_epilogue<ReferenceSrc>(cT, epi_tile, tiled_copy, thread_idx); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
@@ -156,7 +136,7 @@ struct Sm90VisitorImplBase {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops(transform_apply(tuple<Ops...>{}, params, shared_storage,
[] (auto&& op, auto const& op_params, auto&& op_storage) {
using Op = cute::remove_cvref_t<decltype(op)>;
@@ -262,7 +242,9 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase<Ops...> {
// Producer load callbacks factory
// All operations must redefine this, but most can just dispatch to the base impl
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -363,7 +345,9 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase<Ops...> {
// All operations must redefine this
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -446,7 +430,9 @@ struct Sm90TreeVisitor : Sm90VisitorImpl<ChildOps..., NodeOp> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -516,7 +502,9 @@ struct Sm90SplitTreeVisitor : Sm90VisitorImpl<InputTree, AuxOutTrees..., OutputT
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -613,7 +601,9 @@ struct Sm90TopologicalVisitor : Sm90VisitorImpl<Ops...> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -651,9 +641,11 @@ namespace detail {
template <class Op0>
struct Sm90VisitorImplBase<Op0> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
};
// Retain tuple for SharedStorage because empty structs have 1B alignment
// tuples use multiple inheritance, avoids this problem
using SharedStorage = tuple<
typename Op0::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -675,9 +667,9 @@ struct Sm90VisitorImplBase<Op0> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0)
Op0(params.op_0, get<0>(shared_storage))
}) {}
tuple<Op0> ops;
@@ -686,10 +678,10 @@ struct Sm90VisitorImplBase<Op0> {
template <class Op0, class Op1>
struct Sm90VisitorImplBase<Op0, Op1> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
typename Op1::SharedStorage op_1;
};
using SharedStorage = tuple<
typename Op0::SharedStorage,
typename Op1::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -714,10 +706,10 @@ struct Sm90VisitorImplBase<Op0, Op1> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0),
Op1(params.op_1, shared_storage.op_1)
Op0(params.op_0, get<0>(shared_storage)),
Op1(params.op_1, get<1>(shared_storage))
}) {}
tuple<Op0, Op1> ops;
@@ -726,11 +718,11 @@ struct Sm90VisitorImplBase<Op0, Op1> {
template <class Op0, class Op1, class Op2>
struct Sm90VisitorImplBase<Op0, Op1, Op2> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
typename Op1::SharedStorage op_1;
typename Op2::SharedStorage op_2;
};
using SharedStorage = tuple<
typename Op0::SharedStorage,
typename Op1::SharedStorage,
typename Op2::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -758,11 +750,11 @@ struct Sm90VisitorImplBase<Op0, Op1, Op2> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0),
Op1(params.op_1, shared_storage.op_1),
Op2(params.op_2, shared_storage.op_2)
Op0(params.op_0, get<0>(shared_storage)),
Op1(params.op_1, get<1>(shared_storage)),
Op2(params.op_2, get<2>(shared_storage))
}) {}
tuple<Op0, Op1, Op2> ops;
@@ -771,12 +763,12 @@ struct Sm90VisitorImplBase<Op0, Op1, Op2> {
template <class Op0, class Op1, class Op2, class Op3>
struct Sm90VisitorImplBase<Op0, Op1, Op2, Op3> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
typename Op1::SharedStorage op_1;
typename Op2::SharedStorage op_2;
typename Op3::SharedStorage op_3;
};
using SharedStorage = tuple<
typename Op0::SharedStorage,
typename Op1::SharedStorage,
typename Op2::SharedStorage,
typename Op3::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -807,12 +799,12 @@ struct Sm90VisitorImplBase<Op0, Op1, Op2, Op3> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0),
Op1(params.op_1, shared_storage.op_1),
Op2(params.op_2, shared_storage.op_2),
Op3(params.op_3, shared_storage.op_3)
Op0(params.op_0, get<0>(shared_storage)),
Op1(params.op_1, get<1>(shared_storage)),
Op2(params.op_2, get<2>(shared_storage)),
Op3(params.op_3, get<3>(shared_storage))
}) {}
tuple<Op0, Op1, Op2, Op3> ops;
include/cutlass/epilogue/thread/activation.h
+119 -250
View File
@@ -49,38 +49,6 @@ namespace cutlass {
namespace epilogue {
namespace thread {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct LinearCombinationGenericParams {
T alpha; ///< scales accumulators
T beta; ///< scales source tensor
T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams():
alpha(T(1)),
beta(T(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T alpha,
T beta = T(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T const *alpha_ptr,
T const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Identity operator
@@ -92,13 +60,6 @@ struct Identity {
T operator()(T value) const {
return value;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T, int N>
@@ -107,12 +68,59 @@ struct Identity<Array<T, N> > {
Array<T, N> operator()(Array<T, N> const &value) const {
return value;
}
};
using Params = LinearCombinationGenericParams<T>;
/// Scale operator
template <typename T>
struct Scale {
struct Arguments {
T scale = T(1);
};
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
T operator()(T const& value, T const& scale) const {
multiplies<T> mul;
return mul(scale, value);
}
CUTLASS_HOST_DEVICE
T operator()(T const& value, Arguments const& args = Arguments()) const {
return this->operator()(value, args.scale);
}
};
template <typename T, int N>
struct Scale<Array<T, N>> {
using Arguments = typename Scale<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& values, T const& scale) const {
multiplies<Array<T, N>> mul;
return mul(scale, values);
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& values, Arguments const& args = Arguments()) const {
return this->operator()(values, args.scale);
}
};
/// Specialization to compose other activations with a defined unary operator
/// e.g. Scale<Identity<T>>
template <template <class> class Activation, typename T>
struct Scale<Activation<T>> {
using Arguments = typename Scale<T>::Arguments;
CUTLASS_HOST_DEVICE
T operator()(T const &value, decltype(Arguments{}.scale) const& scale) const {
multiplies<T> mul;
Activation<T> act;
return mul(scale, act(value));
}
CUTLASS_HOST_DEVICE
T operator()(T const& value, Arguments const& args = Arguments()) const {
return this->operator()(value, args.scale);
}
};
@@ -134,14 +142,6 @@ struct ReLu {
return mx(value, T(0));
}
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T>
@@ -162,90 +162,87 @@ struct ReLu<Array<T, N>> {
maximum<Array<T, N>> mx;
return mx(frag, T(0));
}
};
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
// Generic clamp
template <typename T>
struct Clamp {
struct Arguments {
T lower_bound = cutlass::platform::numeric_limits<T>::min();
T upper_bound = cutlass::platform::numeric_limits<T>::max();
};
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &frag, Params const &params_) const {
return this->operator()(frag);
T operator()(T const& value, T const& lower_bound, T const& upper_bound) const {
maximum<T> mx;
minimum<T> mn;
return mn(mx(value, lower_bound), upper_bound);
}
CUTLASS_HOST_DEVICE
T operator()(T const& value, Arguments const& args = Arguments()) const {
return this->operator()(value, args.lower_bound, args.upper_bound);
}
};
template <typename T, int N>
struct Clamp<Array<T,N>> {
using Arguments = typename Clamp<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T,N> operator()(Array<T,N> const& values, T const& lower_bound, T const& upper_bound) const {
maximum<Array<T,N>> mx;
minimum<Array<T,N>> mn;
return mn(mx(values, lower_bound), upper_bound);
}
CUTLASS_HOST_DEVICE
Array<T,N> operator()(Array<T,N> const& values, Arguments const& args = Arguments()) const {
return this->operator()(values, args.lower_bound, args.upper_bound);
}
};
// Leaky Relu operator
template <typename T>
struct LeakyReLU {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
// Methods
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
struct Arguments {
T leaky_alpha = T(0);
};
CUTLASS_HOST_DEVICE
T operator()(T const &value, T const & alpha_recip) const {
T res = value > T(0) ? value : value * alpha_recip;
T operator()(T const& value, T const& leaky_alpha) const {
T res = value > T(0) ? value : value * leaky_alpha;
return res;
}
CUTLASS_HOST_DEVICE
T operator()(T const &value, Params const &params_) const {
this->operator()(value, params_.leaky_alpha);
T operator()(T const& value, Arguments const& args = Arguments()) const {
this->operator()(value, args.leaky_alpha);
}
};
template <typename T, int N>
struct LeakyReLU<Array<T, N> > {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
// Methods
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
};
using Arguments = typename LeakyReLU<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, T const & alpha_recip) const {
Array<T, N> operator()(Array<T, N> const& values, T const& leaky_alpha) const {
Array<T, N> y;
LeakyReLU<T> leaky_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < int(value.size()); ++i) {
y[i] = leaky_op(value[i], alpha_recip);
for (int i = 0; i < int(values.size()); ++i) {
y[i] = leaky_op(values[i], leaky_alpha);
}
return y;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value, params_.leaky_alpha);
Array<T, N> operator()(Array<T, N> const& values, Arguments const& args = Arguments()) const {
return this->operator()(values, args.leaky_alpha);
}
};
@@ -253,15 +250,8 @@ struct LeakyReLU<Array<T, N> > {
template <typename T>
struct Tanh {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
return fast_tanh(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
T operator()(T const &value) const {
return fast_tanh(value);
}
};
@@ -279,13 +269,6 @@ struct Tanh<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <int N>
@@ -296,14 +279,6 @@ struct Tanh<Array<half_t, N>> {
Array<T, N> operator()(Array<T, N> const& z) const {
fast_tanh_op<Array<T, N>> tanh;
return tanh(z);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
@@ -311,15 +286,8 @@ struct Tanh<Array<half_t, N>> {
template <typename T>
struct Sigmoid {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
return T(1) / (T(1) + fast_exp(-scalar));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
T operator()(T const &value) const {
return T(1) / (T(1) + fast_exp(-value));
}
};
@@ -337,13 +305,6 @@ struct Sigmoid<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <int N>
@@ -368,13 +329,6 @@ struct Sigmoid<Array<half_t, N>> {
fast_exp(neg(z))));
#endif
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &z, Params const &params_) const {
return this->operator()(z);
}
};
// SiLu (swish) operator introduced by Elfwing et al. in the following paper
@@ -385,16 +339,9 @@ struct Sigmoid<Array<half_t, N>> {
template <typename T>
struct SiLu {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
T operator()(T const &value) const {
Sigmoid<T> sigmoid;
return scalar * sigmoid(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
return value * sigmoid(value);
}
};
@@ -406,13 +353,6 @@ struct SiLu<Array<T, N>> {
multiplies<Array<T, N>> mul;
return mul(value, sigmoid_op(value));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
// Hardswish operator introduced by Howard et al. in the following paper
@@ -429,13 +369,6 @@ struct HardSwish {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 / T(6);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <>
@@ -449,13 +382,6 @@ struct HardSwish<float> {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 * 0.16666667f;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <typename T, int N>
@@ -472,13 +398,6 @@ struct HardSwish<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <int N>
@@ -494,13 +413,6 @@ struct HardSwish<Array<half_t, N> > {
return mul(mul(mn(mx(add(value, T(3)), T(0)), T(6)), value), T(0.16666667f));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
//
@@ -516,48 +428,27 @@ struct HardSwish<Array<half_t, N> > {
template <typename T>
struct GELU {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
return T(cutlass::constants::half<T>() * scalar *
(cutlass::constants::one<T>() + (T)erff((float)(scalar * cutlass::constants::half_root_two<T>()))));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
T operator()(T const &value) const {
return T(cutlass::constants::half<T>() * value *
(cutlass::constants::one<T>() + (T)erff((float)(value * cutlass::constants::half_root_two<T>()))));
}
};
template <>
struct GELU<float> {
CUTLASS_HOST_DEVICE
float operator()(float const &scalar) const {
return cutlass::constants::half<float>() * scalar *
(cutlass::constants::one<float>() + erff(scalar * cutlass::constants::half_root_two<float>() ));
}
using Params = LinearCombinationGenericParams<float>;
CUTLASS_HOST_DEVICE
float operator()(float const &scalar, Params const &params_) const {
return this->operator()(scalar);
float operator()(float const &value) const {
return cutlass::constants::half<float>() * value *
(cutlass::constants::one<float>() + erff(value * cutlass::constants::half_root_two<float>() ));
}
};
template <>
struct GELU<double> {
CUTLASS_HOST_DEVICE
double operator()(double const &scalar) const {
return cutlass::constants::half<double>() * scalar *
(cutlass::constants::one<double>() + erf( scalar * cutlass::constants::half_root_two<double>() ));
}
using Params = LinearCombinationGenericParams<double>;
CUTLASS_HOST_DEVICE
double operator()(double const &scalar, Params const &params_) const {
return this->operator()(scalar);
double operator()(double const &value) const {
return cutlass::constants::half<double>() * value *
(cutlass::constants::one<double>() + erf( value * cutlass::constants::half_root_two<double>() ));
}
};
@@ -575,15 +466,11 @@ struct GELU<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T>
using ScaledGELU = Scale<GELU<T>>;
// GELU operator implemented using the Taylor series approximation
template <typename T>
struct GELU_taylor {
@@ -597,13 +484,6 @@ struct GELU_taylor {
return T(cutlass::constants::half<T>() * z *
(cutlass::constants::one<T>() + fast_tanh(k0 * z * (cutlass::constants::one<T>() + k1 * z * z))));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <int N>
@@ -630,13 +510,6 @@ struct GELU_taylor<Array<half_t, N> > {
return y;
}
using Params = LinearCombinationGenericParams<half_t>;
CUTLASS_HOST_DEVICE
Array<half_t, N> operator()(Array<half_t, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T, int N>
@@ -654,15 +527,11 @@ struct GELU_taylor<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T>
using ScaledGELU_taylor = Scale<GELU_taylor<T>>;
/// Computes backwards pass for GELU operator assuming d_t is the layer gradient and
/// z is computed from the forward pass.
template <typename T>
include/cutlass/epilogue/thread/linear_combination_generic.h
+60 -3
View File
@@ -49,6 +49,51 @@ namespace thread {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <class Activation, class = void>
struct GenericActivationTraits {
static constexpr bool IsArgumentsNeeded = false;
struct Arguments {};
};
template <class Activation>
struct GenericActivationTraits<Activation, decltype(typename Activation::Arguments(), void())> {
static constexpr bool IsArgumentsNeeded = true;
using Arguments = typename Activation::Arguments;
};
template <typename T>
struct LinearCombinationGenericParams {
T alpha; ///< scales accumulators
T beta; ///< scales source tensor
T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams():
alpha(T(1)),
beta(T(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T alpha,
T beta = T(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T const *alpha_ptr,
T const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Applies a linear combination operator followed by an activation function to an array of elements.
///
/// D = activation(alpha * accumulator + beta * source + uniform)
@@ -84,7 +129,11 @@ public:
static FloatRoundStyle const kRound = Round;
/// Host-constructable parameters structure
using Params = typename ActivationFunctor<FragmentCompute>::Params;
struct Params
: LinearCombinationGenericParams<ElementCompute>,
GenericActivationTraits<ActivationFunctor<ElementCompute>>::Arguments {
using LinearCombinationGenericParams<ElementCompute>::LinearCombinationGenericParams;
};
private:
@@ -161,7 +210,11 @@ public:
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator, intermediate); // D = alpha * Accum + X
}
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
if constexpr (GenericActivationTraits<ActivationFunctor<ElementCompute>>::IsArgumentsNeeded) {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
} else {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate);
}
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
@@ -192,7 +245,11 @@ public:
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator); // D = alpha * Accum
}
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
if constexpr (GenericActivationTraits<ActivationFunctor<FragmentCompute>>::IsArgumentsNeeded) {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
} else {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate);
}
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
include/cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h
+495
View File
@@ -0,0 +1,495 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Functor performing elementwise operations used by epilogues.
*/
#pragma once
/////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/epilogue/threadblock/epilogue_base.h"
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace threadblock {
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Epilogue operator
template <
typename DefaultEpilogue, ///< Default Epilogue Descriptor
typename FusionCallbacks_, ///< The called fusion callbacks
int Stages = 2, ///< Software pipeline stages for epilogue
int IterationsUnroll = true ///< Used to reduce binary size when epilogue op is large
>
class EpilogueWithVisitorCallbacks :
public EpilogueBase<
typename DefaultEpilogue::Shape,
typename DefaultEpilogue::WarpMmaOperator::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::AccumulatorFragmentIterator,
typename DefaultEpilogue::WarpTileIterator,
typename DefaultEpilogue::Padding,
DefaultEpilogue::kFragmentsPerIteration>,
public EpilogueBaseStreamK<
typename DefaultEpilogue::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::WarpMmaOperator,
typename DefaultEpilogue::AccumulatorFragmentIterator>
{
public:
static_assert(Stages <= 2, "Sm80 EVT only support upto 2 Stages.");
// Whether the epilogue is pipelined
static bool constexpr Pipelined = Stages > 1;
using FusionCallbacks = FusionCallbacks_;
using OutputTileIterator = typename DefaultEpilogue::OutputTileIterator;
// Number of epilogue iterations.
// Each iteration processes a 8xThreadblockTile::kN output tile
static const int kIterations = OutputTileIterator::kIterations;
using Base = EpilogueBase<
typename DefaultEpilogue::Shape,
typename DefaultEpilogue::WarpMmaOperator::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::AccumulatorFragmentIterator,
typename DefaultEpilogue::WarpTileIterator,
typename DefaultEpilogue::Padding,
DefaultEpilogue::kFragmentsPerIteration>;
using BaseStreamK = EpilogueBaseStreamK<
typename DefaultEpilogue::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::WarpMmaOperator,
typename DefaultEpilogue::AccumulatorFragmentIterator>;
static int const kPartitionsK = DefaultEpilogue::kPartitionsK;
using AccumulatorFragmentIterator = typename DefaultEpilogue::AccumulatorFragmentIterator;
using WarpTileIterator = typename DefaultEpilogue::WarpTileIterator;
using SharedLoadIterator = typename DefaultEpilogue::SharedLoadIterator;
/// The complete warp-level accumulator tile
using AccumulatorTile = typename Base::AccumulatorTile;
/// Accumulator element
using ElementAccumulator = typename WarpTileIterator::Element;
struct OutputOp{
using ElementAccumulator = ElementAccumulator;
using Params = typename FusionCallbacks::Arguments;
};
/// Fragment type used by the accumulator tile's fragment iterator
using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment;
// Output access size
static int const kElementsPerAccess = DefaultEpilogue::kElementsPerAccess;
/// Array type used by output functor
using AccumulatorAccessType = Array<
typename WarpTileIterator::Element, kElementsPerAccess>;
static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK;
static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles;
using Params = typename FusionCallbacks::Params;
static size_t constexpr kSmemStageOffset = sizeof(Base::SharedStorage) / sizeof(ElementAccumulator);
static int constexpr kAccumulatorFragmentCount = AccumulatorTile::kElements / (kIterations * AccumulatorAccessType::kElements) / kPartitionsK;
struct SharedStorage {
typename Base::SharedStorage acc_smem[Stages];
typename FusionCallbacks::SharedStorage callback_smem;
};
private:
/// Loads fragment from shared memory aligned with output tensor
SharedLoadIterator shared_load_iterator_;
FusionCallbacks fusion_callbacks;
public:
/// Constructor
CUTLASS_DEVICE
EpilogueWithVisitorCallbacks(
const Params &params_callbacks, ///< Epilogue Visitor params
SharedStorage &shared_storage, ///< Shared storage object
int thread_idx, ///< ID of a thread within the threadblock
int warp_idx, ///< ID of warp within threadblock
int lane_idx ///< Id of thread within warp
):
Base(shared_storage.acc_smem[0], thread_idx, warp_idx, lane_idx),
BaseStreamK(thread_idx),
shared_load_iterator_(shared_storage.acc_smem[0].reference(), thread_idx),
fusion_callbacks(params_callbacks, shared_storage.callback_smem)
{ }
/// Aggregates the accumulator sets shared by peer blocks in the global workspace,
/// performing epilogue computations, writing to output
template <class ProblemShape>
CUTLASS_DEVICE
void reduce(
int peer_idx_begin,
int peer_idx_end,
int reduce_fragment_idx,
void *element_workspace,
cutlass::gemm::GemmCoord threadblock_tile_offset,
ProblemShape problem_shape,
int thread_idx)
{
auto callbacks = fusion_callbacks.get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
);
callbacks.begin_epilogue();
// Reduce peer accumulator fragments into one fragment
AccumulatorFragment accum_fragment;
BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace);
// Store fragment to shared memory
this->warp_tile_iterator_.store(accum_fragment);
__syncthreads();
callbacks.begin_step(reduce_fragment_idx);
// Load fragment from shared memory
typename SharedLoadIterator::Fragment aligned_accum_fragment;
shared_load_iterator_.load(aligned_accum_fragment);
// Add fragments shared by other k partitions
if (kPartitionsK > 1)
{
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
typename SharedLoadIterator::Fragment aligned_addend_fragment;
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_addend_fragment);
aligned_accum_fragment = add_fragments(aligned_accum_fragment, aligned_addend_fragment);
}
}
//
// Iterate over output fragment
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const*>(&aligned_accum_fragment);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
callbacks.begin_row(row_idx);
}
callbacks.visit(
reduce_fragment_idx,
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
callbacks.end_row(row_idx);
}
}
callbacks.end_step(reduce_fragment_idx);
callbacks.end_epilogue();
}
/// Streams the result to global memory
template <class ProblemShape>
CUTLASS_DEVICE
void operator()(
AccumulatorTile const &accumulators,
cutlass::gemm::GemmCoord threadblock_tile_offset,
ProblemShape problem_shape,
int thread_idx
) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
auto callbacks = fusion_callbacks.get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
);
callbacks.begin_epilogue();
//
// Iterator over warp-level accumulator fragment
//
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
//
// Iterate over accumulator tile
//
if constexpr(Pipelined){
__syncthreads();
//
// Pipeline Prologue
//
size_t warp_iterator_offset = kSmemStageOffset;
size_t smem_iterator_offset = kSmemStageOffset;
callbacks.begin_step(0);
acc2smem_source_needed<cutlass::make_index_sequence<kIterations>>::push(
0, accum_fragment_iterator, this->warp_tile_iterator_);
this->warp_tile_iterator_.add_pointer_offset(warp_iterator_offset);
warp_iterator_offset = -warp_iterator_offset;
//
// Pipeline Loop
//
#pragma unroll(IterationsUnroll ? kIterations : 1)
for (int iter_idx = 1; iter_idx < kIterations + 1; ++iter_idx) {
__syncthreads();
// Skip the load for epilogue
if (iter_idx < kIterations) {
callbacks.begin_step(iter_idx);
acc2smem_source_needed<cutlass::make_index_sequence<kIterations>>::push(
iter_idx, accum_fragment_iterator, this->warp_tile_iterator_);
this->warp_tile_iterator_.add_pointer_offset(warp_iterator_offset);
warp_iterator_offset = -warp_iterator_offset;
}
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment[0]);
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
if (kPartitionsK > 1) {
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_accum_fragment[i]);
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
}
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
shared_load_iterator_.add_pointer_offset(smem_iterator_offset);
smem_iterator_offset = -smem_iterator_offset;
//
// Iterate over output fragments
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
callbacks.begin_row(row_idx);
}
callbacks.visit(
iter_idx-1,
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
callbacks.end_row(row_idx);
}
}
//
// Conclude the step
//
callbacks.end_step(iter_idx-1);
}
} else {
#pragma unroll(IterationsUnroll ? kIterations : 1)
for (int iter_idx = 0; iter_idx < kIterations; ++iter_idx) {
//
// Load the source
//
callbacks.begin_step(iter_idx);
//
// Convert and store fragment
//
__syncthreads();
acc2smem_source_needed<cutlass::make_index_sequence<kIterations>>::push(
iter_idx, accum_fragment_iterator, this->warp_tile_iterator_);
__syncthreads();
//
// Load fragments from shared memory
//
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment[0]);
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
if (kPartitionsK > 1) {
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_accum_fragment[i]);
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
}
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
//
// Iterate over output fragments
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment[0]);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
callbacks.begin_row(row_idx);
}
callbacks.visit(
iter_idx,
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
callbacks.end_row(row_idx);
}
}
//
// Conclude the step
//
callbacks.end_step(iter_idx);
}
}
callbacks.end_epilogue();
}
private:
template<class Seq>
struct acc2smem_source_needed;
template <size_t... Seq>
struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
template<int Advance>
CUTLASS_DEVICE
static void helper(AccumulatorFragmentIterator accum_fragment_iterator,
WarpTileIterator &warp_tile_iterator) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Advance; i++) {
++accum_fragment_iterator;
}
typename AccumulatorFragmentIterator::Fragment accum_fragment;
accum_fragment_iterator.load(accum_fragment);
warp_tile_iterator.store(accum_fragment);
}
CUTLASS_DEVICE
static void push(size_t pos,
AccumulatorFragmentIterator const &iterator_begin,
WarpTileIterator &warp_tile_iterator) {
int dummy[] = {(pos == Seq) && (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
}
};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp
+433
View File
@@ -0,0 +1,433 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Visitor tree operation base implementation to enable composable fusions
for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using cute::tuple;
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
template <class... Ops>
struct VisitorImpl2x: fusion::detail::Sm90VisitorImplBase<Ops...> {
using fusion::detail::Sm90VisitorImplBase<Ops...>::Sm90VisitorImplBase;
using fusion::detail::Sm90VisitorImplBase<Ops...>::ops;
template <class CallbacksTuple>
struct Callbacks {
// Callbacks can store non-persistent variables (e.g. tensors) or copies of persistent variables
CallbacksTuple callbacks_tuple;
/// Called at the start of the epilogue just before iterating over accumulator slices
CUTLASS_DEVICE void
begin_epilogue() {
for_each(callbacks_tuple,
[] (auto& callbacks) {
callbacks.begin_epilogue();
}
);
}
/// Called at the start of one step before starting accumulator exchange
CUTLASS_DEVICE void
begin_step(int step_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.begin_step(step_idx);
}
);
}
/// Called at the start of a row
CUTLASS_DEVICE void
begin_row(int row_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.begin_row(row_idx);
}
);
}
/// Called after accumulators have been exchanged for each accumulator vector
template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInputs, FragmentSize> const&... frg_inputs) // depends on the N-naryness of the op
= delete; // Must be implemented for each operation
/// Called at the start of a row
CUTLASS_DEVICE void
end_row(int row_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.end_row(row_idx);
}
);
}
/// Called after all accumulator elements have been visited
CUTLASS_DEVICE void
end_step(int step_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.end_step(step_idx);
}
);
}
/// Called after all steps have been completed
CUTLASS_DEVICE void
end_epilogue() {
for_each(callbacks_tuple,
[] (auto& callbacks) {
callbacks.end_epilogue();
}
);
}
};
// Callbacks factory
// All operations must redefine this
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return transform_apply(ops,
[&] (auto& op) {
return op.get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape);
},
[] (auto&&... callbacks) {
auto callbacks_tuple = cute::make_tuple(callbacks...);
return Callbacks<decltype(callbacks_tuple)>{callbacks_tuple};
}
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Convenience aliases
using EmptyCallbacks = VisitorImpl2x<>::Callbacks<cute::tuple<>>;
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace detail
using namespace detail;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Tree visitor
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template <class NodeOp, class... ChildOps>
struct TreeVisitor2x : VisitorImpl2x<ChildOps..., NodeOp> {
using VisitorImpl2x<ChildOps..., NodeOp>::VisitorImpl2x;
template<class CallbacksImpl>
struct Callbacks : CallbacksImpl {
CUTLASS_DEVICE
Callbacks(CallbacksImpl&& impl)
: CallbacksImpl(cute::forward<CallbacksImpl>(impl)) {}
using CallbacksImpl::callbacks_tuple;
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
constexpr int Rm1 = sizeof...(ChildOps);
return cute::detail::tapply(callbacks_tuple,
[&] (auto& child_callbacks) {
return child_callbacks.visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc);
},
[&] (auto&&... frg_inputs) {
return get<Rm1>(callbacks_tuple).visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc, frg_inputs...);
},
make_seq<Rm1>{}
);
}
};
// Callbacks factory
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks<
decltype(VisitorImpl2x<ChildOps..., NodeOp>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
))>(
VisitorImpl2x<ChildOps..., NodeOp>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
)
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class ElementCompute,
class EdgeTuple,
class... Ops
>
struct TopologicalVisitor2x : VisitorImpl2x<Ops...> {
static_assert(is_static_v<EdgeTuple>);
static_assert(rank(EdgeTuple{}) == sizeof...(Ops));
static_assert(sizeof...(Ops) > 1);
using VisitorImpl2x<Ops...>::VisitorImpl2x;
template<class CallbacksImpl>
struct Callbacks : CallbacksImpl {
CUTLASS_DEVICE
Callbacks(CallbacksImpl&& impl)
: CallbacksImpl(cute::forward<CallbacksImpl>(impl)) {}
using CallbacksImpl::callbacks_tuple;
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
constexpr int Rm1 = sizeof...(Ops) - 1;
auto frg_compute_tuple = cute::repeat<Rm1>(Array<ElementCompute, FragmentSize>{});
return cute::detail::tapply(EdgeTuple{}, callbacks_tuple, frg_compute_tuple,
// Visit the first R-1 ops in topological order
[&] (auto&& edge_seq, auto& callbacks, auto& frg_compute) {
frg_compute = cute::detail::apply(frg_compute_tuple,
// Compute the current op with children inputs
[&] (auto const&... frg_inputs) {
auto frg_output = callbacks.visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc, frg_inputs...);
using ElementOutput = typename decltype(frg_output)::Element;
using ConvertOutput = NumericArrayConverter<ElementCompute, ElementOutput, FragmentSize>;
ConvertOutput convert_output{};
return convert_output(frg_output);
},
// Get inputs in the sequence given by the children indices of the current op
edge_seq
);
return frg_compute;
},
// Visit the last op
[&] (auto const&...) {
return cute::detail::apply(frg_compute_tuple,
// Compute the last op with children inputs
[&] (auto const&... frg_inputs) {
return get<Rm1>(callbacks_tuple).visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc, frg_inputs...);
},
// Get inputs in the sequence given by the children indices of the last op
get<Rm1>(EdgeTuple{})
);
},
// Transform to visit R-1 ops, apply to visit last op
make_seq<Rm1>{}
);
}
};
// Callbacks factory
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks<decltype(
VisitorImpl2x<Ops...>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
))>(
VisitorImpl2x<Ops...>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
)
);
}
};
template <class NodeOp, class... ChildOps>
using Sm80EVT = TreeVisitor2x<NodeOp, ChildOps...>;
template<
class ElementCompute,
class EdgeTuple,
class... Ops
>
using Sm80TopologicalVisitor = TopologicalVisitor2x<ElementCompute, EdgeTuple, Ops...>;
using X = Underscore;
/////////////////////////////////////////////////////////////////////////////////////////////////
// OutputTileThreadLayout translate the CUTLASS 2.X OutputTileOptimalThreadMap into cute layout
// used by CUTLASS 3.X Epilogue
template <
typename ThreadblockShape_,
typename WarpShape_,
typename Element_,
int ElementsPerAccess,
int Stages_=1
>
struct OutputTileThreadLayout: DefaultThreadMapTensorOp<
ThreadblockShape_,
WarpShape_,
ThreadblockShape_::kK/WarpShape_::kK,
Element_,
ElementsPerAccess>::Type {
using Base = typename DefaultThreadMapTensorOp<
ThreadblockShape_,
WarpShape_,
ThreadblockShape_::kK/WarpShape_::kK,
Element_,
ElementsPerAccess>::Type;
using Base::Base;
// Software pipeline stages in epilogue
static_assert(Stages_ <= 2, "Sm80 EVT only support upto 2 Stages.");
static const int Stages = Stages_;
using ThreadShape = cute::Shape<
cute::Int<Base::Detail::kAccessWidth>, // lane col idx
cute::Int<Base::Detail::kAccessRows>, // lane row idx
cute::Int<Base::Detail::kWarpsRemainingForRows>, // warp row idx
cute::Int<Base::Shape::kGroup>, // group idx
cute::Int<Base::Shape::kCluster> // cluster idx
>;
using Shape = typename Base::Shape;
using Count = typename Base::Count;
using ThreadMapShape = cute::Shape<
// Column
Int<Base::kElementsPerAccess>, // vector
Int<Base::Detail::kAccessWidth>, // lane_col_coord
Int<Base::Iterations::kColumn>, // iteration::column
// Row
Int<Base::Detail::kAccessRows>, // lane_row_coord
Int<Base::Iterations::kRow>, // iterations in row
Int<Base::Detail::kWarpsRemainingForRows>, // warp_row_coord
Int<Count::kRow>, // iteration::row
Int<Count::kGroup>, // iteration::group
Int<Shape::kGroup>, // group_coord
Int<Count::kCluster>, // iteration::cluster
Int<Shape::kCluster> // cluster_coord
>;
// The shape of CTA Tile
using CtaShapeMNL = cute::Shape<
Int<
Shape::kRow * Count::kRow *
Shape::kGroup * Count::kGroup *
Shape::kCluster * Count::kCluster
>,
Int<Shape::kColumn * Count::kColumn>,
_1
>;
static const int kElementsPerAccess = ElementsPerAccess;
//
// Methods
//
CUTLASS_DEVICE
static auto tid2coord(int thread_idx) {
return make_layout(ThreadShape{})[thread_idx];
}
template <class TensorInput>
CUTLASS_DEVICE
static auto partition(TensorInput &&xT, int thread_idx, gemm::GemmCoord threadblock_tile_offset) {
// (BLK_M,BLK_N)
Tensor bCxT = local_tile(
xT, CtaShapeMNL{}, make_coord(_,_,_), Step<_1,_1, X>{}
)(_,_,threadblock_tile_offset.m(),threadblock_tile_offset.n(),threadblock_tile_offset.k());
auto [lane_col_coord, lane_row_coord, warp_row_coord, group_coord, cluster_coord] = tid2coord(thread_idx);
// transform to column-major
Tensor bCxT_nm = make_tensor(
std::forward<decltype(bCxT)>(bCxT).data(), make_layout(get<1>(bCxT.layout()), get<0>(bCxT.layout()))
).compose(make_layout(ThreadMapShape{}));
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
return bCxT_nm(_,lane_col_coord,_,lane_row_coord,_,warp_row_coord,_,_,group_coord,_,cluster_coord);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_compute.hpp
+109
View File
@@ -0,0 +1,109 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Visitor tree compute operations for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using namespace detail;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// N-nary Elementwise Compute Operation
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
template <class> class ComputeFn,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class = void
>
struct VisitorCompute : VisitorImpl2x<> {
using VisitorImpl2x<>::VisitorImpl2x;
struct Callbacks : EmptyCallbacks {
template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
CUTLASS_DEVICE Array<ElementOutput, FragmentSize>
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInputs, FragmentSize> const&... frg_inputs) {
return transform_apply(cute::make_tuple(frg_inputs...),
[&] (auto&& frg_input) {
using ElementInput = typename cute::remove_cvref_t<decltype(frg_input)>::Element;
using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
return convert_input(frg_input);
},
[&] (auto&&... cvt_frg_inputs) {
using ComputeOutput = ComputeFn<Array<ElementCompute, FragmentSize>>;
using ConvertOutput = NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize, RoundStyle>;
ComputeOutput compute_output{};
ConvertOutput convert_output{};
return convert_output(compute_output(cvt_frg_inputs...));
}
);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp
+559
View File
@@ -0,0 +1,559 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Visitor tree load operations for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
#include "cute/tensor.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using namespace detail;
using X = Underscore;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Elementwise Fetch Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// returns accumulator
struct VisitorAccFetch : VisitorImpl2x<> {
using VisitorImpl2x<>::VisitorImpl2x;
struct Callbacks : EmptyCallbacks {
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE Array<ElementAccumulator, FragmentSize>
visit(int iter_idx, int row_idx, int column_idx, int frg_idx, Array<ElementAccumulator, FragmentSize> const& frg_acc) {
return frg_acc;
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks{};
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Broadcast Load Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// Scalar broadcast
template<
class Element,
class StrideMNL = Stride<_0,_0,_0>,
int BroadcastCount = 1,
template <class> class ReductionFn = multiplies
>
struct VisitorScalarBroadcast {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0,_0>>) || // scalar broadcast, e.g. alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0,_1>>) ||
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>)); // batched scalar broadcast, e.g. per-batch alpha
struct SharedStorage { };
struct Arguments {
Element scalars[BroadcastCount] = {};
Element const* scalar_ptrs[BroadcastCount] = {};
StrideMNL dScalar = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
CUTLASS_HOST_DEVICE
VisitorScalarBroadcast() { }
CUTLASS_HOST_DEVICE
VisitorScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) {
// Get the scalar for non-batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,_0>>) {
update_scalar();
}
}
Element scalar;
Params const* params_ptr;
struct Callbacks: EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(Element scalar)
: scalar(scalar) {}
Element scalar;
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Array<Element, FragmentSize> frg_scalar;
frg_scalar.fill(scalar);
return frg_scalar;
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
// Get the scalar for batched broadcast
if constexpr (
cute::is_same_v<StrideMNL, Stride<_0,_0,_1>> ||
cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
update_scalar(threadblock_tile_offset.k());
}
return Callbacks(scalar);
}
private:
CUTLASS_DEVICE void
update_scalar(int l_coord = 0) {
int l_offset = l_coord * size<2>(params_ptr->dScalar);
if (params_ptr->scalar_ptrs[0] != nullptr) {
scalar = params_ptr->scalar_ptrs[0][l_offset];
} else {
// batch stride is ignored for nullptr fallback
scalar = params_ptr->scalars[0];
}
// Do reduction over multiple broadcasts if necessary
ReductionFn<Element> reduction_fn;
CUTLASS_PRAGMA_UNROLL
for (int i = 1; i < BroadcastCount; ++i) {
if (params_ptr->scalar_ptrs[i] != nullptr) {
scalar = reduction_fn(scalar, params_ptr->scalar_ptrs[i][l_offset]);
} else {
// batch stride is ignored for nullptr fallback
scalar = reduction_fn(scalar, params_ptr->scalars[i]);
}
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Elementwise Load Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class ThreadMap,
class Element,
class StrideMNL
>
struct VisitorAuxLoad{
struct Arguments {
Element* ptr_aux = nullptr;
Element null_default = Element(0);
StrideMNL dAux = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
// Software pipeline stages
static const int Stages = ThreadMap::Stages;
struct SharedStorage {};
// Global load type
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
static int constexpr VecLength = sizeof(VecType) / sizeof(Element);
CUTLASS_HOST_DEVICE
VisitorAuxLoad() { }
CUTLASS_HOST_DEVICE
VisitorAuxLoad(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gAux,
RTensor&& tC_rAux,
CTensor&& tC_cAux,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gAux(cute::forward<GTensor>(tC_gAux)),
tC_rAux(cute::forward<RTensor>(tC_rAux)),
tC_cAux(cute::forward<CTensor>(tC_cAux)),
problem_shape(problem_shape),
params_ptr(params_ptr) { }
GTensor tC_gAux;
RTensor tC_rAux;
CTensor tC_cAux;
Params const* params_ptr;
ProblemShape problem_shape;
CUTLASS_DEVICE void
begin_step(int step_idx) {
clear(tC_rAux(_,_,_,step_idx%Stages));
auto src_v = filter(tC_gAux(_,_,_,step_idx));
auto coord_v = filter(tC_cAux(_,_,_,step_idx));
auto dst_v = filter(tC_rAux(_,_,_,step_idx%Stages));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(src_v); ++i) {
bool guard = elem_less(coord_v(i), problem_shape);
cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const*)&src_v(i), guard);
}
}
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Tensor tC_rAux_frg = recast<Array<Element, FragmentSize>>(coalesce(tC_rAux(_,_,_,iter_idx%Stages)));
return tC_rAux_frg(frg_idx);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mAux = make_tensor(
make_gmem_ptr(params_ptr->ptr_aux),
problem_shape,
params_ptr->dAux); // (M,N,L)
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
Tensor tC_gAux = recast<VecType>(
group_modes<3,6>(ThreadMap::partition(mAux, thread_idx, threadblock_tile_offset)));
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, Stages
Tensor tC_rAux = make_tensor<VecType>(
make_layout(flatten(make_shape(take<0,3>(tC_gAux.shape()), Int<Stages>{}))));
// Generate the pred tensor
Tensor cAux = make_identity_tensor(mAux.shape());
Tensor tC_cAux = local_partition(
group_modes<3,6>(ThreadMap::partition(cAux, thread_idx, threadblock_tile_offset)),
Shape<Int<VecLength>>{},
(_0{})
);
return Callbacks<
decltype(tC_gAux), decltype(tC_rAux),
decltype(tC_cAux), ProblemShape>(
cute::move(tC_gAux),
cute::move(tC_rAux),
cute::move(tC_cAux),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Row vector broadcast
template<
class ThreadMap,
class Element,
class StrideMNL
>
struct VisitorRowBroadcast {
struct Arguments {
Element const* ptr_row = nullptr;
Element null_default = Element(0);
StrideMNL dRow = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage {};
// Global load type
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
static int constexpr VecLength = sizeof(VecType) / sizeof(Element);
CUTLASS_HOST_DEVICE
VisitorRowBroadcast() { }
CUTLASS_HOST_DEVICE
VisitorRowBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gRow,
RTensor&& tC_rRow,
CTensor&& tC_cRow,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gRow(cute::forward<GTensor>(tC_gRow)),
tC_rRow(cute::forward<RTensor>(tC_rRow)),
tC_cRow(cute::forward<CTensor>(tC_cRow)),
n(get<1>(problem_shape)),
params_ptr(params_ptr) { }
GTensor tC_gRow;
RTensor tC_rRow;
CTensor tC_cRow;
Params const* params_ptr;
int n;
CUTLASS_DEVICE void
begin_epilogue() {
clear(tC_rRow);
auto src_v = filter(tC_gRow);
auto coord_v = filter(tC_cRow);
auto dst_v = filter(tC_rRow);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(src_v); ++i) {
bool guard = get<1>(coord_v(i)) < n;
cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const*)&src_v(i), guard);
}
}
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Tensor rRow_frg = recast<Array<Element, FragmentSize>>(coalesce(tC_rRow));
return rRow_frg(column_idx);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mRow = make_tensor(
make_gmem_ptr(params_ptr->ptr_row),
problem_shape,
params_ptr->dRow);
// VECTOR, FRAGMENT_COLUMN
Tensor tC_gRow = recast<VecType>(
ThreadMap::partition(mRow, thread_idx, threadblock_tile_offset)
)(_,_,_0{},_0{},_0{},_0{});
Tensor tC_rRow = make_tensor_like(tC_gRow);
// Generate the pred tensor
Tensor cRow = make_identity_tensor(mRow.shape());
Tensor tC_cRow = local_partition(
ThreadMap::partition(cRow, thread_idx, threadblock_tile_offset)(_,_,_0{},_0{},_0{},_0{}),
Shape<Int<VecLength>>{},
(_0{})
);
return Callbacks<
decltype(tC_gRow), decltype(tC_rRow),
decltype(tC_cRow), ProblemShape>(
cute::move(tC_gRow),
cute::move(tC_rRow),
cute::move(tC_cRow),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Column vector broadcast
template<
class ThreadMap,
class Element,
class StrideMNL = Stride<_1,_0,_0>
>
struct VisitorColBroadcast {
struct Arguments {
Element const* ptr_col = nullptr;
Element null_default = Element(0);
StrideMNL dCol = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage { };
CUTLASS_HOST_DEVICE
VisitorColBroadcast() { }
CUTLASS_HOST_DEVICE
VisitorColBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gCol,
RTensor&& tC_rCol,
CTensor&& tC_cCol,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gCol(cute::forward<GTensor>(tC_gCol)),
tC_rCol(cute::forward<RTensor>(tC_rCol)),
tC_cCol(cute::forward<CTensor>(tC_cCol)),
m(get<0>(problem_shape)),
params_ptr(params_ptr) { }
GTensor tC_gCol;
RTensor tC_rCol;
CTensor tC_cCol;
Params const* params_ptr;
int m;
CUTLASS_DEVICE void
begin_epilogue() {
clear(tC_rCol);
Tensor pred = make_tensor<bool>(shape(tC_gCol));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tC_cCol(i)) < m;
}
copy_if(pred, tC_gCol, tC_rCol);
}
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Array<Element, FragmentSize> frg_col;
frg_col.fill(tC_rCol(row_idx,iter_idx));
return frg_col;
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mCol = make_tensor(
make_gmem_ptr(params_ptr->ptr_col),
problem_shape,
params_ptr->dCol);
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
Tensor tC_gCol = group_modes<1,4>(
ThreadMap::partition(mCol, thread_idx, threadblock_tile_offset)(_0{},_0{},_,_,_,_));
Tensor tC_rCol = make_tensor_like(tC_gCol);
// Generate the pred tensor
Tensor cCol = make_identity_tensor(mCol.shape());
Tensor tC_cCol = group_modes<1,4>(
ThreadMap::partition(cCol, thread_idx, threadblock_tile_offset)(_0{},_0{},_,_,_,_));
return Callbacks<
decltype(tC_gCol), decltype(tC_rCol),
decltype(tC_cCol), ProblemShape>(
cute::move(tC_gCol),
cute::move(tC_rCol),
cute::move(tC_cCol),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp
+781
View File
@@ -0,0 +1,781 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Visitor tree store operations for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using namespace detail;
using X = Underscore;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Elementwise Store Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class ThreadMap,
class Element,
FloatRoundStyle RoundStyle,
class StrideMNL
>
struct VisitorAuxStore{
struct Arguments {
Element* ptr_aux = nullptr;
StrideMNL dAux = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage {};
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
static int constexpr VecLength = sizeof(VecType) / sizeof(Element);
CUTLASS_HOST_DEVICE
VisitorAuxStore() { }
CUTLASS_HOST_DEVICE
VisitorAuxStore(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gAux,
RTensor&& tC_rAux,
CTensor&& tC_cAux,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gAux(cute::forward<GTensor>(tC_gAux)),
tC_rAux(cute::forward<RTensor>(tC_rAux)),
tC_cAux(cute::forward<CTensor>(tC_cAux)),
problem_shape(problem_shape),
params_ptr(params_ptr) { }
GTensor tC_gAux;
RTensor tC_rAux;
CTensor tC_cAux;
Params const* params_ptr;
ProblemShape problem_shape;
CUTLASS_DEVICE void
begin_step(int step_idx) {
clear(tC_rAux);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
using ConvertInput = NumericArrayConverter<Element, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
Tensor tC_rAux_frg = recast<Array<Element, FragmentSize>>(coalesce(tC_rAux));
tC_rAux_frg(frg_idx) = convert_input(frg_input);
return frg_input;
}
CUTLASS_DEVICE void
end_step(int step_idx) {
auto src_v = filter(tC_rAux);
auto coord_v = filter(tC_cAux(_,_,_,step_idx));
auto dst_v = filter(tC_gAux(_,_,_,step_idx));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(src_v); ++i) {
bool guard = elem_less(coord_v(i), problem_shape);
cutlass::arch::global_store<VecType, sizeof(VecType)>(src_v(i), (void*)&dst_v(i), guard);
}
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mAux = make_tensor(
make_gmem_ptr(params_ptr->ptr_aux),
problem_shape,
params_ptr->dAux); // (M,N,L)
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
Tensor tC_gAux = recast<VecType>(group_modes<3,6>(ThreadMap::partition(mAux, thread_idx, threadblock_tile_offset)));
Tensor tC_rAux = make_tensor_like(take<0,3>(tC_gAux));
// Generate the pred tensor
Tensor cAux = make_identity_tensor(mAux.shape());
Tensor tC_cAux = local_partition(
group_modes<3,6>(ThreadMap::partition(cAux, thread_idx, threadblock_tile_offset)),
Shape<Int<VecLength>>{},
(_0{})
);
return Callbacks<
decltype(tC_gAux), decltype(tC_rAux),
decltype(tC_cAux), ProblemShape>(
cute::move(tC_gAux),
cute::move(tC_rAux),
cute::move(tC_cAux),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Reduction Store Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// Helper functions
template <
template <class> class ReduceFn,
int kThreads, class T>
CUTLASS_DEVICE
void intra_warp_row_reduce(T& value) {
using ReduceInput = ReduceFn<T>;
ReduceInput reduce_input{};
constexpr int kHalfThreads = kThreads >> 1;
CUTLASS_PRAGMA_UNROLL
for (int i = kHalfThreads; i > 0; i >>= 1) {
value = reduce_input(value, __shfl_xor_sync(0xFFFFFFFF, value, i));
}
}
template <
template <class> class ReduceFn,
FloatRoundStyle RoundStyle,
class ElementCompute,
class ElementFragment, int FragmentSize>
CUTLASS_DEVICE
void fragment_reduce(ElementCompute& value, Array<ElementFragment, FragmentSize> const& frg) {
using ReduceInput = ReduceFn<ElementCompute>;
ReduceInput reduce_input{};
using ConvertInput = NumericConverter<ElementCompute, ElementFragment, RoundStyle>;
ConvertInput convert_input{};
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
value = reduce_input(value, convert_input(frg[i]));
}
}
template<
template <class> class AtomicReduceFn,
FloatRoundStyle RoundStyle,
class ElementCompute,
class ElementOutput>
CUTLASS_DEVICE
void atomic_reduce(ElementOutput* ptr, ElementCompute const& value) {
using ReduceOutput = AtomicReduceFn<ElementOutput>;
using ConvertOutput = NumericConverter<ElementOutput, ElementCompute, RoundStyle>;
ReduceOutput reduce_output{};
ConvertOutput convert_output{};
reduce_output(ptr, convert_output(value));
}
// Col vector reduction
template <
template <class> class RegReduceFn,
template <class> class AtomicReduceFn,
class ThreadMap,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class StrideMNL = Stride<_1,_0,_0>
>
struct VisitorColReduction {
struct Arguments {
ElementOutput* ptr_col = nullptr;
ElementCompute reduction_identity = 0;
StrideMNL dCol = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage { };
CUTLASS_HOST_DEVICE
VisitorColReduction() { }
CUTLASS_HOST_DEVICE
VisitorColReduction(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gCol,
CTensor&& tC_cCol,
ProblemShape problem_shape,
Params const* params_ptr,
int thread_idx
):
tC_gCol(cute::forward<GTensor>(tC_gCol)),
tC_cCol(cute::forward<CTensor>(tC_cCol)),
m(get<0>(problem_shape)),
n(get<1>(problem_shape)),
params_ptr(params_ptr) {
// The partial reduction results of each warp are further
// reduced to the first thread in each row.
// Only the first thread in each row is the writing thread
is_writing_thread = thread_idx % ThreadMap::Detail::kAccessWidth == 0;
}
GTensor tC_gCol;
CTensor tC_cCol;
Params const* params_ptr;
int m;
int n;
int curr_iter_idx;
bool is_writing_thread;
ElementCompute reduction_accum;
CUTLASS_DEVICE void
begin_row(int row_idx) {
reduction_accum = ElementCompute(params_ptr->reduction_identity);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
curr_iter_idx = iter_idx;
int coord_n = get<1>(tC_cCol(column_idx, row_idx, iter_idx));
if (coord_n < n) {
fragment_reduce<RegReduceFn, RoundStyle>(reduction_accum, frg_input);
}
// Intra-warp reduction
if (column_idx + 1 == ThreadMap::Iterations::kColumn) {
intra_warp_row_reduce<RegReduceFn, ThreadMap::Detail::kAccessWidth>(reduction_accum);
}
return frg_input;
}
CUTLASS_DEVICE auto
end_row(int row_idx) {
bool guard = get<0>(tC_cCol(_0{}, row_idx,curr_iter_idx)) < m;
if (guard && is_writing_thread) {
atomic_reduce<AtomicReduceFn, RoundStyle>(&tC_gCol(row_idx,curr_iter_idx), reduction_accum);
}
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mCol = make_tensor(
make_gmem_ptr(params_ptr->ptr_col),
problem_shape,
params_ptr->dCol);
// FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tC_gCol = group_modes<1,4>(
ThreadMap::partition(mCol, thread_idx, threadblock_tile_offset)(_0{},_0{},_,_,_,_));
// Generate the pred tensor
Tensor cCol = make_identity_tensor(mCol.shape());
// FRAGMENT_COL, FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tC_cCol = group_modes<2,5>(
ThreadMap::partition(cCol, thread_idx, threadblock_tile_offset)(_0{},_,_,_,_,_));
return Callbacks<
decltype(tC_gCol), decltype(tC_cCol),
ProblemShape>(
cute::move(tC_gCol),
cute::move(tC_cCol),
problem_shape,
params_ptr,
thread_idx
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Row vector reduction
template <
template <class> class RegReduceFn,
template <class> class AtomicReduceFn,
class ThreadMap,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class StrideMNL = Stride<_0,_1,_0>
>
struct VisitorRowReduction {
struct Arguments {
ElementOutput* ptr_row = nullptr;
ElementCompute reduction_identity = 0;
StrideMNL dRow = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
using SharedStorageShape = decltype(select<0,1,2,3,5,8,10>(typename ThreadMap::ThreadMapShape{}));
struct SharedStorage {
AlignedArray<ElementCompute, size(SharedStorageShape{}), 16> reduction;
};
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<ElementOutput>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
CUTLASS_HOST_DEVICE
VisitorRowReduction() { }
CUTLASS_HOST_DEVICE
VisitorRowReduction(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params),
smem_reduce(const_cast<ElementCompute*>(shared_storage.reduction.data())) { }
Params const* params_ptr;
ElementCompute* smem_reduce;
template <
class RTensorR2S, class STensorR2S, class CTensorR2S,
class STensorS2R, class RTensorS2R, class CTensorS2R,
class GTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
// R->S
RTensorR2S&& tRS_rSrc,
STensorR2S&& tRS_sRows,
CTensorR2S&& tRS_cSrc,
// S->R
STensorS2R&& tSR_sRows,
RTensorS2R&& tSR_rRows,
CTensorS2R&& tSR_cRows,
// R->G
GTensor&& tC_gRow,
CTensor&& tC_cRow,
ProblemShape problem_shape,
Params const* params_ptr
):
// R->S
tRS_rSrc(cute::forward<RTensorR2S>(tRS_rSrc)),
tRS_sRows(cute::forward<STensorR2S>(tRS_sRows)),
tRS_cSrc(cute::forward<CTensorR2S>(tRS_cSrc)),
// S->R
tSR_sRows(cute::forward<STensorS2R>(tSR_sRows)),
tSR_rRows(cute::forward<RTensorS2R>(tSR_rRows)),
tSR_cRows(cute::forward<CTensorS2R>(tSR_cRows)),
// R->G
tC_gRow(cute::forward<GTensor>(tC_gRow)),
tC_cRow(cute::forward<CTensor>(tC_cRow)),
m(get<0>(problem_shape)),
n(get<1>(problem_shape)),
params_ptr(params_ptr) { }
// R->S
RTensorR2S tRS_rSrc;
STensorR2S tRS_sRows;
CTensorR2S tRS_cSrc;
// S->R
STensorS2R tSR_sRows;
RTensorS2R tSR_rRows;
CTensorS2R tSR_cRows;
// R->G
GTensor tC_gRow;
CTensor tC_cRow;
Params const* params_ptr;
int n;
int m;
CUTLASS_DEVICE void
begin_epilogue() {
fill(tRS_rSrc, params_ptr->reduction_identity);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
Tensor tRS_rRow_frg = recast<Array<ElementCompute, FragmentSize>>(coalesce(tRS_rSrc));
int coord_m = get<0>(tRS_cSrc(column_idx,row_idx,iter_idx));
if (coord_m < m)
reduction(tRS_rRow_frg[column_idx], convert_input(frg_input));
return frg_input;
}
CUTLASS_DEVICE void
end_epilogue() {
//
// Store the partially reduced value to SMEM
//
// Guard against uses of the existing SMEM tile
__syncthreads();
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tRS_rSrc); ++i) {
copy_vec<VecType>(filter(tRS_rSrc), filter(tRS_sRows));
}
__syncthreads();
//
// Now, threads are assigned several columns of the output. They fetch over all rows from
// the compacted SMEM tile and perform a reduction.
//
fill(tSR_rRows, params_ptr->reduction_identity);
using ReduceInputReg = RegReduceFn<ElementCompute>;
ReduceInputReg reduce_input_reg{};
CUTLASS_PRAGMA_UNROLL
for (int j = 0; j < size(tSR_rRows); ++j) {
if (get<0>(tSR_cRows(j)) < get<1>(typename ThreadMap::CtaShapeMNL{}) && get<1>(tC_cRow(j)) < n) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tSR_sRows) / size(tSR_rRows); ++i) {
tSR_rRows(j) = reduce_input_reg(tSR_rRows(j), tSR_sRows(i + j * size(tSR_sRows) / size(tSR_rRows)));
}
atomic_reduce<AtomicReduceFn, RoundStyle>(&tC_gRow(j), tSR_rRows(j));
}
}
}
private:
template <int FragmentSize>
CUTLASS_DEVICE ElementCompute
reduction(Array<ElementCompute, FragmentSize>& reduce_buffer, Array<ElementCompute, FragmentSize> const& result) {
using ReduceInput = RegReduceFn<ElementCompute>;
ReduceInput reduce_input{};
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
reduce_buffer[i] = reduce_input(reduce_buffer[i], result[i]);
}
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mRow = make_tensor(
make_gmem_ptr(params_ptr->ptr_row),
problem_shape,
params_ptr->dRow);
//
// Step 1: reduce fragment input (Src) into tRS_rSrc
//
// VECTOR,FRAGMENT_COL
Tensor tRS_rSrc = make_tensor<ElementCompute>(select<0,2>(typename ThreadMap::ThreadMapShape{}));
Tensor cSrc = make_identity_tensor(mRow.shape());
// FRAGMENT_COLUMN, FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tRS_cSrc = group_modes<2,5>(ThreadMap::partition(cSrc, thread_idx, threadblock_tile_offset)(_0{},_,_,_,_,_));
//
// Step 2: copy the partial results in tRS_rSrc to sRows in shared memory
//
// VECTOR,ACCESS_WIDTH,FRAGMENT_COL,ACCESS_ROWS,WARPS_PER_ROW,GROUPS,CLUSTERS
Tensor sRows = make_tensor(
make_smem_ptr(smem_reduce), SharedStorageShape{}
);
auto [lane_col_coord, lane_row_coord, warp_row_coord, group_coord, cluster_coord] = ThreadMap::tid2coord(thread_idx);
Tensor tRS_sRows = sRows(_,lane_col_coord,_,lane_row_coord,warp_row_coord,group_coord,cluster_coord);
//
// Step 3: copy the partial results in sRows to tSR_sRow for reduction
//
// VECTOR*ACCESS_WIDTH*FRAGMENT_COL,ACCESS_ROWS*WARPS_PER_ROW*GROUPS*CLUSTERS
Tensor sRows_nm = coalesce(group_modes<1,5>(group_modes<0,3>(sRows)), Shape<_1,_1>{});
// SMEM_ROW/THREADS,ACCESS_ROWS*WARPS_PER_ROW*GROUPS*CLUSTERS
Tensor tSR_sRows = local_partition(sRows_nm, Shape<Int<ThreadMap::kThreads>,_1>{}, thread_idx);
// SMEM_ROW/THREADS
Tensor tSR_rRows = make_tensor_like(tSR_sRows(_,_0{}));
// Coord
Tensor cRows_nm = make_identity_tensor(sRows_nm.shape());
Tensor tSR_cRows = local_partition(cRows_nm, Shape<Int<ThreadMap::kThreads>,_1>{}, thread_idx)(_,_0{});
//
// Step 4: atomically reduce the results to global memory
//
Tensor tC_gRow = local_partition(
// Cta tile
local_tile(
mRow, typename ThreadMap::CtaShapeMNL{}, make_coord(_,_,_),Step<_1,_1, X>{}
)(_,_,threadblock_tile_offset.m(),threadblock_tile_offset.n(),threadblock_tile_offset.k()),
// Partition to threads
Shape<_1,Int<ThreadMap::kThreads>>{}, thread_idx
)(_0{},_);
Tensor cRow = make_identity_tensor(mRow.shape());
Tensor tC_cRow = local_partition(
// Cta tile
local_tile(
cRow, typename ThreadMap::CtaShapeMNL{}, make_coord(_,_,_), Step<_1,_1, X>{}
)(_,_,threadblock_tile_offset.m(),threadblock_tile_offset.n(),threadblock_tile_offset.k()),
// Partition to threads
Shape<_1,Int<ThreadMap::kThreads>>{}, thread_idx
)(_0{},_);
return Callbacks<
decltype(tRS_rSrc), decltype(tRS_sRows),
decltype(tRS_cSrc), decltype(tSR_sRows),
decltype(tSR_rRows), decltype(tSR_cRows),
decltype(tC_gRow), decltype(tC_cRow),
ProblemShape>(
// R->S
cute::move(tRS_rSrc),
cute::move(tRS_sRows),
cute::move(tRS_cSrc),
// S->R
cute::move(tSR_sRows),
cute::move(tSR_rRows),
cute::move(tSR_cRows),
// R->G
cute::move(tC_gRow),
cute::move(tC_cRow),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Scalar reduction
template <
template <class> class RegReduceFn,
template <class> class AtomicReduceFn,
class ThreadMap,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class StrideMNL = Stride<_0,_0,_0>
>
struct VisitorScalarReduction {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0, _0>>) || // scalar reduction, e.g. tensor max element
(cute::is_same_v<StrideMNL, Stride<_0,_0, _1>>) || // batched scalar reduction, e.g. per-batch max element
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>));
struct Arguments {
ElementOutput* ptr_scalar = nullptr;
ElementCompute reduction_identity = 0;
StrideMNL dScalar = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage { };
CUTLASS_HOST_DEVICE
VisitorScalarReduction(){ };
CUTLASS_HOST_DEVICE
VisitorScalarReduction(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class CTensor, class GTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
CTensor&& tC_cSrc,
GTensor&& tC_gScalar,
ProblemShape problem_shape,
Params const* params_ptr,
int thread_idx
):
tC_cSrc(cute::forward<CTensor>(tC_cSrc)),
tC_gScalar(cute::forward<GTensor>(tC_gScalar)),
problem_shape(problem_shape),
params_ptr(params_ptr) {
// The partial reduction results of each warp are further
// reduced to this first thread.
// Only the first thread of each warp is the writing thread
is_writing_thread = thread_idx % ThreadMap::kWarpSize == 0;
}
GTensor tC_gScalar;
CTensor tC_cSrc;
Params const* params_ptr;
ProblemShape problem_shape;
bool is_writing_thread;
ElementCompute reduction_accum;
CUTLASS_DEVICE void
begin_epilogue() {
reduction_accum = ElementCompute(params_ptr->reduction_identity);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
auto coord = tC_cSrc(column_idx, row_idx, iter_idx);
if (elem_less(coord, problem_shape)) {
fragment_reduce<RegReduceFn, RoundStyle>(reduction_accum, frg_input);
}
return frg_input;
}
CUTLASS_DEVICE auto
end_epilogue() {
// Intra-warp reduction
intra_warp_row_reduce<RegReduceFn, ThreadMap::kWarpSize>(reduction_accum);
// Atomically reduce to global memory
atomic_reduce<AtomicReduceFn, RoundStyle>(&tC_gScalar(_0{},_0{}), reduction_accum);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor cSrc = make_identity_tensor(problem_shape);
// FRAGMENT_COL, FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tC_cSrc = group_modes<2,5>(
ThreadMap::partition(cSrc, thread_idx, threadblock_tile_offset)(_0{},_,_,_,_,_)
);
Tensor mScalar = make_tensor(
make_gmem_ptr(params_ptr->ptr_scalar),
problem_shape,
params_ptr->dScalar
);
Tensor tC_gScalar = mScalar(_,_,threadblock_tile_offset.k());
return Callbacks<
decltype(tC_cSrc), decltype(tC_gScalar),
ProblemShape>(
cute::move(tC_cSrc),
cute::move(tC_gScalar),
problem_shape,
params_ptr,
thread_idx
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitors.hpp
+38
View File
@@ -0,0 +1,38 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Higher-level header file includes all the CUTLASS 2x visitors
*/
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
#include "cutlass/epilogue/threadblock/fusion/visitor_load.hpp"
#include "cutlass/epilogue/threadblock/fusion/visitor_store.hpp"
#include "cutlass/epilogue/threadblock/fusion/visitor_compute.hpp"
include/cutlass/epilogue/threadblock/predicated_tile_iterator_params.h
+10
View File
@@ -32,6 +32,16 @@
\brief
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by this unit test: `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
include/cutlass/fast_math.h
+50 -10
View File
@@ -28,6 +28,15 @@
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
@@ -43,7 +52,7 @@
#include "cutlass/array.h"
#include "cutlass/uint128.h"
#include "cutlass/coord.h"
#include "cutlass/numeric_types.h"
#include "cutlass/half.h"
/**
* \file
@@ -148,15 +157,23 @@ CUTLASS_HOST_DEVICE dividend_t round_nearest(dividend_t dividend, divisor_t divi
return ((dividend + divisor - 1) / divisor) * divisor;
}
template <typename value_t>
CUTLASS_HOST_DEVICE
constexpr
value_t abs_for_integer(value_t a) {
return ((a > 0) ? a : -a);
}
/**
* Greatest common divisor
*/
template <typename value_t>
CUTLASS_HOST_DEVICE constexpr value_t gcd(value_t a, value_t b) {
CUTLASS_HOST_DEVICE
CUTLASS_CONSTEXPR_IF_CXX17
value_t gcd(value_t a, value_t b) {
for (;;) {
if (a == 0) return b;
if (a == 0) return cutlass::abs_for_integer(b);
b %= a;
if (b == 0) return a;
if (b == 0) return cutlass::abs_for_integer(a);
a %= b;
}
}
@@ -165,21 +182,44 @@ CUTLASS_HOST_DEVICE constexpr value_t gcd(value_t a, value_t b) {
* Least common multiple
*/
template <typename value_t>
CUTLASS_HOST_DEVICE constexpr value_t lcm(value_t a, value_t b) {
CUTLASS_HOST_DEVICE
CUTLASS_CONSTEXPR_IF_CXX17
value_t lcm(value_t a, value_t b) {
value_t temp = gcd(a, b);
return temp ? (a / temp * b) : 0;
return temp ? (cutlass::abs_for_integer(a) / temp * cutlass::abs_for_integer(b)) : 0;
}
/**
* Greatest common divisor
*/
template <typename value_t>
CUTLASS_HOST_DEVICE
constexpr
value_t gcd_cxx11(value_t a, value_t b) {
return (a == 0 || b == 0) ? cutlass::abs_for_integer(a | b) : gcd_cxx11(b, a % b);
}
/**
* Least common multiple
*/
template <typename value_t>
CUTLASS_HOST_DEVICE
constexpr
value_t lcm_cxx11(value_t a, value_t b) {
return gcd_cxx11(a, b) ? (cutlass::abs_for_integer(a) / gcd_cxx11(a, b) * cutlass::abs_for_integer(b)) : 0;
}
/// Returns the smallest value in the half-open range [a, a+b) that is a multiple of b
CUTLASS_HOST_DEVICE
constexpr int round_up(int a, int b) {
CUTLASS_CONSTEXPR_IF_CXX17
int round_up(int a, int b) {
return ((a + b - 1) / b) * b;
}
/// Returns the ceiling of (a / b)
CUTLASS_HOST_DEVICE
constexpr int ceil_div(int a, int b) {
CUTLASS_CONSTEXPR_IF_CXX17
int ceil_div(int a, int b) {
return (a + b - 1) / b;
}
@@ -639,12 +679,12 @@ struct Max {
};
CUTLASS_HOST_DEVICE
constexpr int const_min(int a, int b) {
CUTLASS_CONSTEXPR_IF_CXX17 int const_min(int a, int b) {
return (b < a ? b : a);
}
CUTLASS_HOST_DEVICE
constexpr int const_max(int a, int b) {
CUTLASS_CONSTEXPR_IF_CXX17 int const_max(int a, int b) {
return (b > a ? b : a);
}
include/cutlass/float8.h
+16
View File
@@ -33,6 +33,16 @@
\brief Defines a class for using IEEE half-precision floating-point types in host or
device code.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
// FP8 types are available starting CUDA 11.8+
@@ -447,6 +457,9 @@ struct alignas(1) float_e4m3_t : float8_base<FloatEncoding::E4M3> {
/// Constructor inheritance
using Base::Base;
/// Default constructor
float_e4m3_t() = default;
#ifdef CUDA_FP8_ENABLED
/// Conversion from CUDA's FP8 type
CUTLASS_HOST_DEVICE
@@ -649,6 +662,9 @@ struct alignas(1) float_e5m2_t : float8_base<FloatEncoding::E5M2> {
/// Constructor inheritance
using Base::Base;
/// Default constructor
float_e5m2_t() = default;
#ifdef CUDA_FP8_ENABLED
/// Conversion from CUDA's FP8 type
CUTLASS_HOST_DEVICE
include/cutlass/functional.h
+10 -1
View File
@@ -33,12 +33,21 @@
This is inspired by the Standard Library's <functional> header.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/half.h"
#include "cutlass/tfloat32.h"
#include "cutlass/bfloat16.h"
#if defined(CUTLASS_ARCH_WMMA_ENABLED)
#include <mma.h>
include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl
+1 -1
View File
@@ -59,7 +59,7 @@ compute_stage_count_or_override(StageCount<stages> stage_count) {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template<int CapacityBytes, class ElementA, class ElementB, class TileShapeMNK, int stages>
constexpr int
compute_stage_count_or_override(cute::integral_constant<int, stages> stage_count) {
compute_stage_count_or_override(cute::Int<stages> stage_count) {
return stages;
}
include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp
+50 -46
View File
@@ -36,6 +36,7 @@
#include "cutlass/gemm/gemm.h"
#include "cutlass/detail/dependent_false.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/detail/layout.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
@@ -193,12 +194,18 @@ struct CollectiveMma<
"SmemLayoutB K must be 128bytes to be transposed.");
static_assert(!transform::collective::detail::use_universal_transposition<InternalSmemLayoutAtomB, InternalElementB>(),
"Warp specialized ARF kernels have not supported universal B transposition yet.");
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{});
static_assert(SmemAlignmentA >= 128 and SmemAlignmentB >= 128, "Require at least 128B alignment");
struct SharedStorage
{
struct TensorStorage : cute::aligned_struct<256> {
cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>, 256> smem_A;
cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>, 256> smem_B;
struct TensorStorage : cute::aligned_struct<cute::max(SmemAlignmentA, SmemAlignmentB)> {
cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>, SmemAlignmentA> smem_A;
cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>, SmemAlignmentB> smem_B;
} tensors;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
@@ -222,14 +229,14 @@ struct CollectiveMma<
using TMA_A = decltype(make_tma_copy(
GmemTiledCopyA{},
make_tensor(static_cast<InternalElementA const*>(nullptr), repeat_like(InternalStrideA{}, int32_t(0)), InternalStrideA{}),
SmemLayoutA{}(_,_,0),
SmemLayoutA{}(_,_,cute::Int<0>{}),
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any
// Assumption: StrideB is congruent with Problem_NK
using TMA_B = decltype(make_tma_copy(
GmemTiledCopyB{},
make_tensor(static_cast<InternalElementB const*>(nullptr), repeat_like(InternalStrideB{}, int32_t(0)), InternalStrideB{}),
SmemLayoutB{}(_,_,0),
SmemLayoutB{}(_,_,cute::Int<0>{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any
TMA_A tma_load_a;
@@ -353,8 +360,10 @@ struct CollectiveMma<
int lane_predicate = cute::elect_one_sync();
if (warp_idx_in_warp_group == 0 and lane_predicate) {
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_N,BLK_K,PIPE)
//
// Prepare the TMA loads for A and B
@@ -466,12 +475,15 @@ struct CollectiveMma<
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE)
// If TransposeB, GMMA will read from transposed B layout SMEM
Tensor gmma_sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor gmma_sB_position_dependent = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor gmma_sB = as_position_independent_swizzle_tensor(gmma_sB_position_dependent); // (BLK_N,BLK_K,PIPE)
//
// Define C accumulators and A/B partitioning
@@ -483,7 +495,7 @@ struct CollectiveMma<
// Allocate fragments and descriptors
Tensor tCsA = thread_mma.partition_A(sA);
Tensor tCrA = thread_mma.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCsB = thread_mma.partition_B(gmma_sB); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCsB = thread_mma.partition_B(gmma_sB_position_dependent); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
//
@@ -523,29 +535,29 @@ struct CollectiveMma<
cute::bool_constant<TransposeB>{});
warpgroup_fence_operand(accum);
ConsumerToken barrier_token = {BarrierStatus::WaitAgain};
// first k tile
{
pipeline.consumer_wait(smem_pipe_read);
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
pipeline.consumer_wait(smem_pipe_read, barrier_token);
int read_stage = smem_pipe_read.index();
++smem_pipe_read;
bool skip_wait = (pipeline.consumer_try_wait(smem_pipe_read) == BarrierStatus::WaitDone);
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// copy smem->rmem for A operand
copy(smem_tiled_copy_A, tCsA(_,_,0,read_stage), tCrA_copy_view(_,_,0));
// transpose B operand in SMEM
transpose(sB, gmma_sB, read_stage, 0);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
if (k_block == 0) {
transpose(sB, gmma_sB, read_stage, 1);
transpose.synchronize();
}
transpose.synchronize(k_block);
transpose(sB, gmma_sB, read_stage, k_block + 1);
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
@@ -555,15 +567,12 @@ struct CollectiveMma<
warpgroup_wait<2>();
if (k_tile_count - 1 > 0) {
if (!skip_wait) {
pipeline.consumer_wait(smem_pipe_read);
}
--k_tile_count;
if (k_tile_count > 0) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
}
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
@@ -572,9 +581,12 @@ struct CollectiveMma<
warpgroup_wait<2>();
}
if (k_tile_count == 0) {
return;
}
warpgroup_fence_operand(accum);
// Mainloop GMMAs
--k_tile_count;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 1; --k_tile_count) {
@@ -583,30 +595,26 @@ struct CollectiveMma<
//
int read_stage = smem_pipe_read.index();
++smem_pipe_read;
bool skip_wait = (pipeline.consumer_try_wait(smem_pipe_read) == BarrierStatus::WaitDone);
warpgroup_fence_operand(accum);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
if (k_block == 0) {
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
}
if (k_block == size<2>(tCrA) - 1) {
if (!skip_wait) {
pipeline.consumer_wait(smem_pipe_read);
}
pipeline.consumer_wait(smem_pipe_read, barrier_token);
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
// transpose B operand in SMEM
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
} else {
}
else {
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
// transpose B operand in SMEM
if (k_block < 2) {
transpose.synchronize(k_block); // make transpose of k_block available
}
if (k_block == 0) {
transpose(sB, gmma_sB, read_stage, 1);
}
transpose.synchronize(k_block); // make transpose of k_block available
transpose(sB, gmma_sB, read_stage, k_block + 1);
}
warpgroup_arrive();
@@ -627,7 +635,7 @@ struct CollectiveMma<
warpgroup_fence_operand(accum);
if (k_tile_count > 0) {
{
//
// Compute on k_tile
//
@@ -635,16 +643,14 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
warpgroup_fence_operand(accum);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
if (k_block < 2) {
transpose.synchronize(k_block); // make k_block transpose available
}
if (k_block == 0) {
transpose(sB, gmma_sB, read_stage, 1);
}
transpose.synchronize(k_block); // make k_block transpose available
transpose(sB, gmma_sB, read_stage, k_block + 1);
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
@@ -663,8 +669,6 @@ struct CollectiveMma<
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
warpgroup_commit_batch();
warpgroup_wait<2>();
warpgroup_fence_operand(accum);
}
warpgroup_fence_operand(accum);
include/cutlass/gemm/device/gemm_universal_base.h
-14
View File
@@ -157,14 +157,6 @@ protected:
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error " << cudaGetErrorString(cudart_result));
return Status::kErrorInternal;
}
cudart_result = cudaFuncSetAttribute(
Kernel2<GemmKernel>,
cudaFuncAttributePreferredSharedMemoryCarveout, 100); // 100% shared memory
if (cudart_result != cudaSuccess) {
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error " << cudaGetErrorString(cudart_result));
return Status::kErrorInternal;
}
}
// Update SM occupancy member
@@ -228,12 +220,6 @@ public:
{
CUTLASS_TRACE_HOST("GemmUniversalBase::can_implement()");
// Initialize static kernel and device properties, if necessary.
Status result = init_device_props();
if (result != Status::kSuccess) {
return result;
}
dim3 grid = get_grid_shape(args);
if (!(grid.y <= std::numeric_limits<uint16_t>::max() &&
include/cutlass/gemm/gemm.h
+3 -28
View File
@@ -37,40 +37,15 @@
#include "cutlass/coord.h"
#include "cutlass/gemm_coord.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/gemm/gemm_enumerated_types.h"
#include "cute/layout.hpp"
#include "cutlass/detail/layout.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// GEMM operand enumeration: D = A * B + C
enum class Operand {
kA, /// A multiplicand
kB, /// B multiplicand
kC, /// Source accumulator
kD /// Destination accumulator
};
/////////////////////////////////////////////////////////////////////////////////////////////////
enum class GemmUniversalMode {
kGemm,
kGemmSplitKParallel,
kBatched,
kArray,
kInvalid
};
////////////////////////////////////////////////////////////////////////////////
/// Some options for clearing shared memory
enum class SharedMemoryClearOption {
kNone, ///< SMEM is in don't-care state
kZfill, ///< Kernels fill out of bounds accesses with zeros
kClearLastStage ///< Last SMEM stage is explicitly cleared. Mainloop uses 'kNone'
};
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/gemm_enumerated_types.h
+89
View File
@@ -0,0 +1,89 @@
/***************************************************************************************************
* Copyright (c) 2017 - 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.
*
**************************************************************************************************/
/*! \file
\brief Defines common types used for all GEMM-like operators.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/coord.h"
#include "cutlass/gemm_coord.h"
#include "cutlass/layout/matrix.h"
namespace cutlass {
namespace gemm {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// GEMM operand enumeration: D = A * B + C
enum class Operand {
kA, /// A multiplicand
kB, /// B multiplicand
kC, /// Source accumulator
kD /// Destination accumulator
};
/////////////////////////////////////////////////////////////////////////////////////////////////
enum class GemmUniversalMode {
kGemm,
kGemmSplitKParallel,
kBatched,
kArray,
kInvalid
};
////////////////////////////////////////////////////////////////////////////////
/// Some options for clearing shared memory
enum class SharedMemoryClearOption {
kNone, ///< SMEM is in don't-care state
kZfill, ///< Kernels fill out of bounds accesses with zeros
kClearLastStage ///< Last SMEM stage is explicitly cleared. Mainloop uses 'kNone'
};
/////////////////////////////////////////////////////////////////////////
} // namespace gemm
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/default_gemm_universal_with_visitor.h
+157
View File
@@ -0,0 +1,157 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief
Default configuration for a GEMM with fused epilogue visitor callbacks
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"
#include "cutlass/gemm/kernel/gemm_universal_with_visitor.h"
#include "cutlass/gemm/kernel/gemm_universal_with_visitor_streamk.h"
#include "cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Complex elementwise transformation on A operand
ComplexTransform TransformA,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Complex elementwise transformation on B operand
ComplexTransform TransformB,
/// Access granularity of B matrix in units of elements
int kAlignmentB,
/// Element type for C and D matrix operands
typename ElementC_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Access granularity of C matrix in unit of elements
int kAlignmentC,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Element type for epilogue computation
typename ElementEpilogue,
/// Operator class tag
typename OperatorClass,
/// Tag indicating architecture to tune for
typename ArchTag,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Warp-level tile size (concept: GemmShape)
typename InstructionShape,
/// Epilogue output operator
typename FusionCallbacks,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Operation performed by GEMM
typename Operator,
/// Number of stages used in the pipelined epilogue
int EpilogueStages = 1
>
struct DefaultGemmWithVisitor {
using GemmBase = typename DefaultGemmUniversal<
ElementA_, LayoutA_, TransformA, kAlignmentA,
ElementB_, LayoutB_, TransformB, kAlignmentB,
ElementC_, LayoutC_, ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
epilogue::thread::LinearCombination<
ElementC_, kAlignmentC,
ElementAccumulator, ElementEpilogue
>,
ThreadblockSwizzle,
Stages,
Operator
>::GemmKernel;
// Define epilogue
using Epilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<
typename GemmBase::Epilogue,
FusionCallbacks,
EpilogueStages
>;
/// GemmWithVisitor without StreamkFeature member type
template <class SwizzleT, class Enable = void>
class SelectBase :
public GemmWithEpilogueVisitor<
typename GemmBase::Mma,
Epilogue,
SwizzleT>
{};
/// GemmWIthVisitor with StreamkFeature member type
template <class SwizzleT>
class SelectBase<SwizzleT, typename SwizzleT::StreamkFeature> :
public GemmWithEpilogueVisitorStreamk<
typename GemmBase::Mma,
Epilogue,
SwizzleT>
{};
/// Select kernel by ThreadblockSwizzle's support for StreamkFeature
using GemmKernel = SelectBase<ThreadblockSwizzle>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/gemm_grouped_problem_visitor.h
-1
View File
@@ -39,7 +39,6 @@
#include "cutlass/gemm/gemm.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/gemm/kernel/grouped_problem_visitor.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
include/cutlass/gemm/kernel/gemm_universal.h
+8 -8
View File
@@ -47,7 +47,6 @@
#include "cutlass/layout/matrix.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/kernel/params_universal_base.h"
#include "cutlass/trace.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -346,8 +345,8 @@ public:
output_op = args.epilogue;
}
};
};
/// Shared memory storage structure
union SharedStorage {
@@ -465,13 +464,14 @@ public:
/// Executes one GEMM
CUTLASS_DEVICE
void operator()(
Params const &params,
SharedStorage &shared_storage)
{
// Compute threadblock location
void operator()(Params const &params, SharedStorage &shared_storage) {
ThreadblockSwizzle threadblock_swizzle;
run_with_swizzle(params, shared_storage, threadblock_swizzle);
}
/// Executes one GEMM with an externally-provided swizzling function
CUTLASS_DEVICE
void run_with_swizzle(Params const &params, SharedStorage &shared_storage, ThreadblockSwizzle& threadblock_swizzle) {
cutlass::gemm::GemmCoord threadblock_tile_offset =
threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
include/cutlass/gemm/kernel/gemm_universal_with_visitor.h
+321
View File
@@ -0,0 +1,321 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Gemm kernel with an epilogue defined under the epilogue visitor concept
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/gemm/kernel/gemm_universal.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Gemm that compute the epilogue visitor functor
template <
typename Mma, ///! Threadblock-scoped matrix multiply-accumulate
typename Epilogue, ///! Epilogue
typename ThreadblockSwizzle_ ///! Threadblock swizzling function
>
class GemmWithEpilogueVisitor: GemmUniversal<Mma,Epilogue, ThreadblockSwizzle_> {
public:
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Base = GemmUniversal<Mma,Epilogue, ThreadblockSwizzle>;
using Base::Base;
using FusionCallbacks = typename Epilogue::FusionCallbacks;
using ElementA = typename Base::ElementA;
using LayoutA = typename Base::LayoutA;
using ElementB = typename Base::ElementB;
using LayoutB = typename Base::LayoutB;
using ElementC = typename Base::ElementC;
using LayoutC = typename Base::LayoutC;
using ThreadblockShape = typename Mma::Shape;
//
// Structures
//
using SharedStorage = typename Base::SharedStorage;
using Arguments = typename Base::Arguments;
//
// Structure for precomputing values in host memory and passing to kernels
//
/// Parameters structure
struct Params : UniversalParamsBase<
ThreadblockSwizzle,
ThreadblockShape,
ElementA,
ElementB,
ElementC,
LayoutA,
LayoutB>
{
using ParamsBase = UniversalParamsBase<
ThreadblockSwizzle,
ThreadblockShape,
ElementA,
ElementB,
ElementC,
LayoutA,
LayoutB>;
//
// Data members
//
cute::Shape<int32_t,int32_t,int32_t> problem_shape;
typename Mma::IteratorA::Params params_A;
typename Mma::IteratorB::Params params_B;
typename FusionCallbacks::Params output_op;
void * ptr_A;
void * ptr_B;
int64_t batch_stride_A;
int64_t batch_stride_B;
int * ptr_gather_A_indices;
int * ptr_gather_B_indices;
//
// Host dispatch API
//
/// Default constructor
Params() = default;
/// Constructor
Params(
Arguments const &args, /// GEMM application arguments
int device_sms, /// Number of SMs on the device
int sm_occupancy) /// Kernel SM occupancy (in thread blocks)
:
ParamsBase(args, device_sms, sm_occupancy),
params_A(args.lda ? make_Coord_with_padding<LayoutA::kStrideRank>(args.lda) : args.stride_a),
params_B(args.ldb ? make_Coord_with_padding<LayoutB::kStrideRank>(args.ldb) : args.stride_b),
output_op(FusionCallbacks::to_underlying_arguments(args.problem_size, args.epilogue, nullptr /*workspace*/)),
problem_shape({args.problem_size.m(), args.problem_size.n(), args.batch_count}),
ptr_A(const_cast<void *>(args.ptr_A)),
ptr_B(const_cast<void *>(args.ptr_B)),
batch_stride_A(args.batch_stride_A),
batch_stride_B(args.batch_stride_B),
ptr_gather_A_indices(const_cast<int *>(args.ptr_gather_A_indices)),
ptr_gather_B_indices(const_cast<int *>(args.ptr_gather_B_indices))
{
// Raise error on unsupported modes
assert(args.mode != GemmUniversalMode::kGemmSplitKParallel && "Sm80 EVT does not support SplitKParallel.");
assert(!(args.mode == GemmUniversalMode::kGemm && this->grid_tiled_shape.k() > 1 )
&& "Sm80 EVT does not support SplitKSerial.");
assert(args.mode != GemmUniversalMode::kArray && "Sm80 EVT does not support Array Gemm.");
}
/// Lightweight update given a subset of arguments.
void update(Arguments const &args)
{
CUTLASS_TRACE_HOST("GemmUniversalwithVisitor::Params::update()");
// Update input pointers
ptr_A = const_cast<void *>(args.ptr_A);
ptr_B = const_cast<void *>(args.ptr_B);
batch_stride_A = args.batch_stride_A;
batch_stride_B = args.batch_stride_B;
this->batch_stride_D = args.batch_stride_D;
ptr_gather_A_indices = const_cast<int *>(args.ptr_gather_A_indices);
ptr_gather_B_indices = const_cast<int *>(args.ptr_gather_B_indices);
output_op = FusionCallbacks::to_underlying_arguments(args.problem_size, args.epilogue, nullptr /*workspace*/);
problem_shape = make_shape(args.problem_size.m(), args.problem_size.n(), args.batch_count);
}
};
public:
//
// Device-only API
//
// Factory invocation
CUTLASS_DEVICE
static void invoke(
Params const &params,
SharedStorage &shared_storage)
{
GemmWithEpilogueVisitor op;
op(params, shared_storage);
}
/// Executes one GEMM
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage) {
ThreadblockSwizzle threadblock_swizzle;
run_with_swizzle(params, shared_storage, threadblock_swizzle);
}
/// Executes one GEMM with an externally-provided swizzling function
CUTLASS_DEVICE
void run_with_swizzle(Params const &params, SharedStorage &shared_storage, ThreadblockSwizzle& threadblock_swizzle) {
cutlass::gemm::GemmCoord threadblock_tile_offset =
threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
// Early exit if CTA is out of range
if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
return;
}
int offset_k = 0;
int problem_size_k = params.problem_size.k();
ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
//
// Fetch pointers based on mode.
//
if (params.mode == GemmUniversalMode::kGemm) {
if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
}
offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
}
else if (params.mode == GemmUniversalMode::kBatched) {
ptr_A += threadblock_tile_offset.k() * params.batch_stride_A;
ptr_B += threadblock_tile_offset.k() * params.batch_stride_B;
}
__syncthreads();
// Compute initial location in logical coordinates
cutlass::MatrixCoord tb_offset_A{
threadblock_tile_offset.m() * Mma::Shape::kM,
offset_k,
};
cutlass::MatrixCoord tb_offset_B{
offset_k,
threadblock_tile_offset.n() * Mma::Shape::kN
};
// Compute position within threadblock
int thread_idx = threadIdx.x;
// Construct iterators to A and B operands
typename Mma::IteratorA iterator_A(
params.params_A,
ptr_A,
{params.problem_size.m(), problem_size_k},
thread_idx,
tb_offset_A,
params.ptr_gather_A_indices);
typename Mma::IteratorB iterator_B(
params.params_B,
ptr_B,
{problem_size_k, params.problem_size.n()},
thread_idx,
tb_offset_B,
params.ptr_gather_B_indices);
// Broadcast the warp_id computed by lane 0 to ensure dependent code
// is compiled as warp-uniform.
int warp_idx = canonical_warp_idx_sync();
int lane_idx = threadIdx.x % 32;
//
// Main loop
//
// Construct thread-scoped matrix multiply
Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
typename Mma::FragmentC accumulators;
accumulators.clear();
// Compute threadblock-scoped matrix multiply-add
int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
// Compute threadblock-scoped matrix multiply-add
mma(
gemm_k_iterations,
accumulators,
iterator_A,
iterator_B,
accumulators);
//
// Epilogue
//
threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
Epilogue epilogue(
params.output_op,
shared_storage.epilogue,
thread_idx,
warp_idx,
lane_idx);
// Execute the epilogue operator to update the destination tensor.
epilogue(accumulators, threadblock_tile_offset, params.problem_shape, thread_idx);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/gemm_universal_with_visitor_streamk.h
+892
View File
@@ -0,0 +1,892 @@
/***************************************************************************************************
* Copyright (c) 2017 - 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.
*
**************************************************************************************************/
/*! \file
\brief Gemm kernel with an epilogue defined under the epilogue visitor concept with streamk.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/complex.h"
#include "cutlass/barrier.h"
#include "cutlass/block_striped.h"
#include "cutlass/trace.h"
#include "cutlass/gemm/kernel/gemm_universal_streamk.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate
typename Epilogue_, ///! Epilogue
typename ThreadblockSwizzle_ ///! Threadblock mapping function
>
class GemmWithEpilogueVisitorStreamk {
public:
using Base = GemmUniversalStreamk<Mma_, Epilogue_, ThreadblockSwizzle_>;
//
// Types and constants
//
using Mma = Mma_;
using Epilogue = Epilogue_;
using FusionCallbacks = typename Epilogue::FusionCallbacks;
using EpilogueOutputOp = typename Epilogue::OutputOp;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using ElementA = typename Mma::IteratorA::Element;
using LayoutA = typename Mma::IteratorA::Layout;
using ElementB = typename Mma::IteratorB::Element;
using LayoutB = typename Mma::IteratorB::Layout;
using ElementC = typename Epilogue::OutputTileIterator::Element;
using LayoutC = typename Epilogue::OutputTileIterator::Layout;
/// The per-thread tile of raw accumulators
using AccumulatorTile = typename Mma::FragmentC;
static ComplexTransform const kTransformA = Mma::kTransformA;
static ComplexTransform const kTransformB = Mma::kTransformB;
using Operator = typename Mma::Operator;
using OperatorClass = typename Mma::Operator::OperatorClass;
using ThreadblockShape = typename Mma::Shape;
using WarpShape = typename Mma::Operator::Shape;
using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
using ArchTag = typename Mma::ArchTag;
static int const kStages = Mma::kStages;
static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
/// Warp count (concept: GemmShape)
using WarpCount = typename Mma::WarpCount;
static int const kThreadCount = 32 * WarpCount::kCount;
/// Workspace bytes per thread block
static size_t const kWorkspaceBytesPerBlock =
__NV_STD_MAX(
kThreadCount * sizeof(AccumulatorTile),
Epilogue::kWorkspaceBytesPerBlock);
/// Block-striped reduction utility
using BlockStripedReduceT = BlockStripedReduce<kThreadCount, AccumulatorTile>;
//
// Structures
//
using Arguments = typename Base::Arguments;
/// Parameters structure
struct Params
{
public:
//
// Data members
//
cute::Shape<int32_t,int32_t,int32_t> problem_shape;
void * ptr_A;
void * ptr_B;
typename Mma::IteratorA::Params params_A;
typename Mma::IteratorB::Params params_B;
int64_t batch_stride_A;
int64_t batch_stride_B;
GemmUniversalMode mode;
ThreadblockSwizzle block_mapping;
void *barrier_workspace;
void *partials_workspace;
typename FusionCallbacks::Params output_op;
void * ptr_D;
void * ptr_C;
typename Epilogue::OutputTileIterator::Params params_D;
typename Epilogue::OutputTileIterator::Params params_C;
int64_t batch_stride_D;
int64_t batch_stride_C;
protected:
//
// Host-only dispatch-utilities
//
/// Pad the given allocation size up to the nearest cache line
static size_t cacheline_align_up(size_t size)
{
static const int CACHELINE_SIZE = 128;
return (size + CACHELINE_SIZE - 1) / CACHELINE_SIZE * CACHELINE_SIZE;
}
/// Get the workspace size needed for barrier
size_t get_barrier_workspace_size() const
{
// For atomic reduction, each SK-block needs a synchronization flag. For parallel reduction,
// each reduction block needs its own synchronization flag.
int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region();
int num_flags = fast_max(sk_blocks, block_mapping.reduction_blocks);
return cacheline_align_up(sizeof(typename Barrier::T) * num_flags);
}
/// Get the workspace size needed for intermediate partial sums
size_t get_partials_workspace_size() const
{
int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region();
return cacheline_align_up(kWorkspaceBytesPerBlock * sk_blocks);
}
public:
//
// Host dispatch API
//
/// Default constructor
Params() = default;
/// Constructor
Params(
Arguments const &args, /// GEMM application arguments
int device_sms, /// Number of SMs on the device
int sm_occupancy) /// Kernel SM occupancy (in thread blocks)
:
problem_shape({args.problem_size.m(), args.problem_size.n(), args.batch_count}),
params_A(args.lda ? make_Coord_with_padding<LayoutA::kStrideRank>(args.lda) : args.stride_a),
params_B(args.ldb ? make_Coord_with_padding<LayoutB::kStrideRank>(args.ldb) : args.stride_b),
params_C(args.ldc ? make_Coord_with_padding<LayoutC::kStrideRank>(args.ldc) : args.stride_c),
params_D(args.ldd ? make_Coord_with_padding<LayoutC::kStrideRank>(args.ldd) : args.stride_d),
output_op(FusionCallbacks::to_underlying_arguments(args.problem_size, args.epilogue, nullptr /*workspace*/)),
mode(args.mode),
ptr_A(const_cast<void *>(args.ptr_A)),
ptr_B(const_cast<void *>(args.ptr_B)),
ptr_C(const_cast<void *>(args.ptr_C)),
ptr_D(args.ptr_D),
batch_stride_A(args.batch_stride_A),
batch_stride_B(args.batch_stride_B),
batch_stride_C(args.batch_stride_C),
batch_stride_D(args.batch_stride_D),
barrier_workspace(nullptr),
partials_workspace(nullptr)
{
// Number of SMs to make available for StreamK decomposition
int avail_sms = (args.avail_sms == -1) ?
device_sms :
fast_min(args.avail_sms, device_sms);
// Initialize the block mapping structure
block_mapping = ThreadblockSwizzle(
typename ThreadblockSwizzle::template KernelTraits<GemmWithEpilogueVisitorStreamk>(),
args.mode,
args.problem_size,
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
args.batch_count,
sm_occupancy,
device_sms,
avail_sms);
}
/// Returns the workspace size (in bytes) needed for these parameters
size_t get_workspace_size() const
{
return
get_barrier_workspace_size() +
get_partials_workspace_size();
}
/// Assign and initialize the specified workspace buffer. Assumes
/// the memory allocated to workspace is at least as large as get_workspace_size().
Status init_workspace(
void *workspace,
cudaStream_t stream = nullptr)
{
uint8_t *ptr = static_cast<uint8_t*>(workspace);
// Establish partials workspace
partials_workspace = nullptr;
size_t partials_workspace_bytes = get_partials_workspace_size();
if (partials_workspace_bytes > 0)
{
if (!workspace) {
return Status::kErrorWorkspaceNull;
}
partials_workspace = ptr;
ptr += partials_workspace_bytes;
}
// Establish barrier workspace
barrier_workspace = nullptr;
size_t barrier_workspace_bytes = get_barrier_workspace_size();
if (barrier_workspace_bytes > 0)
{
if (!workspace) {
return Status::kErrorWorkspaceNull;
}
barrier_workspace = ptr;
ptr += barrier_workspace_bytes;
}
// Zero-initialize barrier workspace
if (barrier_workspace)
{
size_t barrier_workspace_bytes = get_barrier_workspace_size();
CUTLASS_TRACE_HOST(" Initialize " << barrier_workspace_bytes << " barrier bytes");
cudaError_t result = cudaMemsetAsync(
barrier_workspace,
0,
barrier_workspace_bytes,
stream);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(" cudaMemsetAsync() returned error " << cudaGetErrorString(result));
return Status::kErrorInternal;
}
}
return Status::kSuccess;
}
/// Returns the GEMM volume in thread block tiles
cutlass::gemm::GemmCoord get_tiled_shape() const
{
return block_mapping.tiled_shape();
}
/// Returns the total number of thread blocks to launch
int get_grid_blocks() const
{
dim3 grid_dims = get_grid_dims();
return grid_dims.x * grid_dims.y * grid_dims.z;
}
/// Returns the grid extents in thread blocks to launch
dim3 get_grid_dims() const
{
return block_mapping.get_grid_dims();
}
/// Lightweight update given a subset of arguments.
void update(Arguments const &args)
{
CUTLASS_TRACE_HOST("GemmUniversalStreamK::Params::update()");
// Update input/output pointers
ptr_A = const_cast<void *>(args.ptr_A);
ptr_B = const_cast<void *>(args.ptr_B);
ptr_C = const_cast<void *>(args.ptr_C);
ptr_D = args.ptr_D;
batch_stride_A = args.batch_stride_A;
batch_stride_B = args.batch_stride_B;
batch_stride_C = args.batch_stride_C;
batch_stride_D = args.batch_stride_D;
output_op = FusionCallbacks::to_underlying_arguments(args.problem_size, args.epilogue, nullptr /*workspace*/);
problem_shape = make_shape(args.problem_size.m(), args.problem_size.n(), args.batch_count);
}
};
struct TileWorkDesc: Base::TileWorkDesc {
int k_end;
CUTLASS_DEVICE
bool tile_finished(Params const &params)
{
return (k_end == params.block_mapping.problem_size.k());
}
};
// using TileWorkDesc = typename Base::TileWorkDesc;
using SharedStorage = typename Base::SharedStorage;
protected:
//
// Data members
//
/// GEMM problem parameters
Params params;
/// Shared storage reference
SharedStorage &shared_storage;
/// ID within the threadblock
int thread_idx;
/// ID of warp
int warp_idx;
/// ID of each thread within a warp
int lane_idx;
/// Threadblock scoped epilogue
Epilogue epilogue;
public:
//
// Host-only dispatch API
//
/// Determines whether the GEMM problem size satisfies this kernel's
/// alignment requirements
static Status can_implement(
cutlass::gemm::GemmCoord const & problem_size)
{
return Base::can_implement(problem_size);
}
/// Determines whether the GEMM problem satisfies this kernel's
/// alignment requirements
static Status can_implement(Arguments const &args) {
return can_implement(args.problem_size);
}
protected:
//
// Device-only utility methods
//
/// Iterator for fetching tile fragments from A
CUTLASS_DEVICE
typename Mma::IteratorA init_iterator_A(
TileWorkDesc &tile_work,
GemmUniversalMode mode)
{
// The input A matrix
ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
// Update input pointers based on batched/array mode
if (mode == GemmUniversalMode::kBatched) {
ptr_A += tile_work.tiled_coord.k() * params.batch_stride_A;
}
if (mode == GemmUniversalMode::kArray) {
ptr_A = static_cast<ElementA * const *>(params.ptr_A)[tile_work.tiled_coord.k()];
}
int m_begin = tile_work.tiled_coord.m() * Mma::Shape::kM;
int m_end = params.block_mapping.problem_size.m();
return Mma::IteratorA(
params.params_A,
ptr_A,
{ m_end, tile_work.k_end },
threadIdx.x,
{ m_begin, tile_work.k_begin });
}
/// Iterator for fetching tile fragments from B
CUTLASS_DEVICE
typename Mma::IteratorB init_iterator_B(
TileWorkDesc &tile_work,
GemmUniversalMode mode)
{
// The input B matrix
ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
// Update input pointers based on batched/array mode
if (mode == GemmUniversalMode::kBatched) {
ptr_B += tile_work.tiled_coord.k() * params.batch_stride_B;
}
if (mode == GemmUniversalMode::kArray) {
ptr_B = static_cast<ElementB * const *>(params.ptr_B)[tile_work.tiled_coord.k()];
}
int n_begin = tile_work.tiled_coord.n() * Mma::Shape::kN;
int n_end = params.block_mapping.problem_size.n();
return Mma::IteratorB(
params.params_B,
ptr_B,
{ tile_work.k_end, n_end },
threadIdx.x,
{ tile_work.k_begin, n_begin });
}
CUTLASS_DEVICE
void init_dp_tile_work(
TileWorkDesc &tile_work,
int tile_idx)
{
// The linear tile index
tile_work.tile_idx = tile_idx;
// The first global-scoped MAC-iteration this threadblock will perform for this tile
tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile();
// The number of MAC-iterations this threadblock will perform for this tile
tile_work.k_iters_remaining = params.block_mapping.iters_per_tile();
// The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile
tile_work.k_begin = 0;
// The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile
tile_work.k_end = params.block_mapping.problem_size.k();
// The location of this tile (in threadblock-tile coordinates) in the output matrix
tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx);
}
CUTLASS_DEVICE
void init_sk_tile_work(
TileWorkDesc &tile_work,
int tile_idx,
int block_iter_begin,
int block_iter_end)
{
// The linear tile index
tile_work.tile_idx = tile_idx;
// The first global-scoped MAC-iteration for this tile
int tile_iter_begin = tile_idx * params.block_mapping.iters_per_tile();
// The first global-scoped MAC-iteration this threadblock will perform for this tile
tile_work.iter_begin = max(block_iter_begin, tile_iter_begin);
// The first tile-scoped MAC-iteration this threadblock will perform for this tile
int k_iter_begin = tile_work.iter_begin - tile_iter_begin;
// The last (one past) tile-scoped MAC-iteration this threadblock will perform for this tile
int k_iter_end = block_iter_end - tile_iter_begin;
// The number of MAC-iterations this threadblock will perform for this tile
tile_work.k_iters_remaining = k_iter_end - k_iter_begin;
// The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile
tile_work.k_begin = k_iter_begin * Mma::Shape::kK;
// The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile
tile_work.k_end = min(
params.block_mapping.problem_size.k(), // extent of k domain
(k_iter_end * Mma::Shape::kK)); // extent of the threadblock's global iteration assignment
// The location of this tile (in threadblock-tile coordinates) in the output matrix
tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx);
}
/// Share accumulators with peers
CUTLASS_DEVICE
void share_accumulators(
AccumulatorTile const &accumulator_tile,
int block_idx,
int first_block_idx)
{
AccumulatorTile *accum_tile_workspace = reinterpret_cast<AccumulatorTile *>(params.partials_workspace);
int accum_tile_offset = first_block_idx * kThreadCount;
if (block_idx == first_block_idx)
{
// First peer initializes the workspace partials
BlockStripedReduceT::store(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx);
}
else
{
// Subsequent peers atomically accumulate into the workspace partials
if (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic)
{
// Non-deterministic reduction order: wait for the first peer to have initialized the partials before we add to them
Barrier::wait_lt(params.barrier_workspace, thread_idx, first_block_idx, 1);
}
else
{
// Turnstile reduction order: wait until the previous peer has written
int wait_count = block_idx - first_block_idx;
Barrier::wait_eq(params.barrier_workspace, thread_idx, first_block_idx, wait_count);
}
// Perform reduction in workspace
BlockStripedReduceT::reduce(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx);
}
// Signal our arrival
Barrier::arrive_inc(params.barrier_workspace, thread_idx, first_block_idx);
}
/// Acquire accumulators from peers
CUTLASS_DEVICE
void acquire_accumulators(
AccumulatorTile &accumulator_tile,
int block_idx,
int first_block_idx)
{
AccumulatorTile *accum_tile_workspace = reinterpret_cast<AccumulatorTile *>(params.partials_workspace);
// Wait for arrival
int num_carry_in = block_idx - first_block_idx;
Barrier::wait_eq_reset(params.barrier_workspace, thread_idx, first_block_idx, num_carry_in);
// Load and add peer-partials accumulator tile to local accumulator tile
int accum_tile_offset = first_block_idx * kThreadCount;
BlockStripedReduceT::load_add(accumulator_tile, accum_tile_workspace + accum_tile_offset, thread_idx);
}
/// Perform epilogue computations and output
CUTLASS_DEVICE
void do_epilogue(
TileWorkDesc &tile_work,
AccumulatorTile &accumulator_tile)
{
cutlass::gemm::GemmCoord threadblock_tile_offset{
tile_work.tiled_coord.m(),
tile_work.tiled_coord.n(),
tile_work.tiled_coord.k()
};
// Execute the epilogue operator to update the destination tensor.
epilogue(
accumulator_tile,
threadblock_tile_offset,
params.problem_shape,
thread_idx);
}
CUTLASS_DEVICE
void separate_reduction(int reduce_idx)
{
int peer_idx_begin, peer_idx_last, reduce_tile_idx, reduce_fragment_idx;
// Reduce by sk-tile (every tile contributed to by one or more blocks)
reduce_tile_idx = reduce_idx / Epilogue::kAccumulatorFragments;
reduce_fragment_idx = reduce_idx % Epilogue::kAccumulatorFragments;
int iter_tile_first = reduce_tile_idx * params.block_mapping.iters_per_tile();
int iter_tile_last = iter_tile_first + params.block_mapping.iters_per_tile() - 1;
peer_idx_begin = params.block_mapping.get_sk_block_idx(iter_tile_first);
peer_idx_last = params.block_mapping.get_sk_block_idx(iter_tile_last);
// Wait for peers to complete
int peer_idx_end = peer_idx_last + 1;
int num_peers = peer_idx_end - peer_idx_begin;
Barrier::wait_eq_reset(
params.barrier_workspace,
thread_idx,
(reduce_tile_idx * Epilogue::kAccumulatorFragments) + reduce_fragment_idx,
num_peers);
/// The location of this tile (in threadblock-tile coordinates) in the output matrix
GemmCoord tiled_coord = params.block_mapping.get_tile_offset(reduce_tile_idx);
// Execute the epilogue operator to update the destination tensor.
epilogue.reduce(
peer_idx_begin,
peer_idx_end,
reduce_fragment_idx,
params.partials_workspace,
tiled_coord,
params.problem_shape,
thread_idx);
}
CUTLASS_DEVICE
void process_tile(
TileWorkDesc tile_work,
int block_idx,
int dp_start_block_idx,
int block_iter_begin)
{
// Initialize input iterators
typename Mma::IteratorA iterator_A = init_iterator_A(tile_work, params.mode);
typename Mma::IteratorB iterator_B = init_iterator_B(tile_work, params.mode);
// Initialize accumulators
AccumulatorTile accumulator_tile;
accumulator_tile.clear();
// Initialize MMA abstraction
Mma mma(
shared_storage.main_loop,
thread_idx,
warp_idx,
lane_idx);
// Perform this tile's range of multiply-accumulate (MAC) iterations
mma(tile_work.k_iters_remaining, accumulator_tile, iterator_A, iterator_B, accumulator_tile);
if ((ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) ||
(params.block_mapping.reduction_blocks == 0) ||
(block_idx >= dp_start_block_idx))
{
//
// Cooperative SK peer reduction or DP block
//
int first_block_idx = params.block_mapping.get_first_block_idx(tile_work.tile_idx, block_idx);
if (!tile_work.tile_finished(params)) {
// Non "finishing" SK blocks must share their partial accumulator sums through global scratch workspace
share_accumulators(accumulator_tile, block_idx, first_block_idx);
}
else
{
// DP blocks and "finishing" SK blocks must perform epilogue operations and write the output tile
if (!tile_work.tile_started())
{
// A "finishing" SK block must first aggregate its accumulator partial sums with those shared by peer threadblocks
acquire_accumulators(accumulator_tile, block_idx, first_block_idx);
}
do_epilogue(tile_work, accumulator_tile);
}
}
else
{
//
// Separate peer reduction
//
// Share accumulator partial sums with peer threadblock(s) through scratch workspace
epilogue.share(block_idx, params.partials_workspace, accumulator_tile, tile_work.tile_started());
// Signal arrival
Barrier::arrive_range_inc(
params.barrier_workspace,
thread_idx,
tile_work.tile_idx * Epilogue::kAccumulatorFragments,
Epilogue::kAccumulatorFragments);
}
}
/// Executes one GEMM
CUTLASS_DEVICE
void gemm()
{
// Initialize block's iteration range
int tile_idx = 0;
int block_iter_begin = 0;
int block_iters_remaining = 0;
int block_idx = params.block_mapping.get_block_idx();
int sk_padding_start_block_idx = params.block_mapping.sk_regions() * params.block_mapping.sk_blocks_per_region();
int dp_start_block_idx = params.block_mapping.sk_waves * params.block_mapping.avail_sms;
int reduce_start_block_idx = dp_start_block_idx + params.block_mapping.dp_blocks;
int grid_padding_start_block_idx = reduce_start_block_idx + params.block_mapping.reduction_blocks;
// Initialize tile work descriptor
TileWorkDesc tile_work;
bool dp_block = (block_idx >= dp_start_block_idx) && (block_idx < reduce_start_block_idx);
bool sk_block = (block_idx < sk_padding_start_block_idx);
bool reduce_block = (block_idx >= reduce_start_block_idx) &&
(block_idx < grid_padding_start_block_idx) &&
(ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kMixed);
if (dp_block)
{
// This is a DP block
int dp_block_idx = block_idx - dp_start_block_idx;
int first_dp_tile = (params.block_mapping.cohort_raster) ? 0 : params.block_mapping.sk_tiles;
// Blocks in first DP wave get configured number of tiles
tile_idx = first_dp_tile + dp_block_idx;
int tile_allottment = params.block_mapping.dp_first_wave_tiles;
// Blocks in subsequent DP waves get 1 tile
if (dp_block_idx >= params.block_mapping.avail_sms) {
tile_allottment = 1;
tile_idx += (params.block_mapping.dp_first_wave_tiles - 1) * params.block_mapping.avail_sms;
}
block_iters_remaining = params.block_mapping.iters_per_tile() * tile_allottment;
init_dp_tile_work(tile_work, tile_idx);
// DP blocks exit if out of bounds or overlap an SK tile (only possible during cohort rasterization, where dp_first_wave_tiles must be 1)
if ((tile_idx < params.block_mapping.sk_tiles) ||
(tile_work.tiled_coord.m() >= params.block_mapping.tiled_shape().m()) ||
(tile_work.tiled_coord.n() >= params.block_mapping.tiled_shape().n()))
{
return;
}
}
else if (sk_block)
{
// This is a SK block
int block_iter_end;
params.block_mapping.get_iter_extents(block_idx, block_iter_begin, block_iter_end);
block_iters_remaining = block_iter_end - block_iter_begin;
tile_idx = params.block_mapping.get_sk_tile_idx(block_iter_end - 1);
init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining);
}
else
{
if (reduce_block)
{
// This is a reduction threadblock
int reduce_block_idx = block_idx - reduce_start_block_idx;
separate_reduction(reduce_block_idx);
}
return;
}
// Iteration-processing loop body
CUTLASS_PRAGMA_NO_UNROLL
while (true)
{
// Perform this block's share of work for this tile
process_tile(
tile_work,
block_idx,
dp_start_block_idx,
block_iter_begin);
block_iters_remaining -= tile_work.k_iters_remaining;
if (block_iters_remaining == 0)
{
break;
}
// Continue to next tile
__syncthreads();
if (block_idx >= dp_start_block_idx)
{
// DP block consume their tiles at stride
tile_idx += params.block_mapping.avail_sms;
init_dp_tile_work(tile_work, tile_idx);
}
else
{
// SK blocks consume their tiles in backwards order
tile_idx--;
init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining);
}
}
}
public:
//
// Device-only API
//
// Factory invocation
CUTLASS_DEVICE
static void invoke(
Params const &params,
SharedStorage &shared_storage)
{
GemmWithEpilogueVisitorStreamk op(params, shared_storage);
op();
}
CUTLASS_DEVICE
GemmWithEpilogueVisitorStreamk(
Params const &params,
SharedStorage &shared_storage)
:
params(params),
shared_storage(shared_storage),
thread_idx(threadIdx.x),
warp_idx(__shfl_sync(0xffffffff, threadIdx.x / 32, 0)), // broadcast the warp_id computed by lane 0 to ensure dependent code
lane_idx(threadIdx.x % 32),
epilogue(
params.output_op,
shared_storage.epilogue,
thread_idx,
warp_idx,
lane_idx)
{}
/// Executes one GEMM
CUTLASS_DEVICE
void operator()()
{
// Generic SK code path
gemm();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/params_universal_base.h
+38 -35
View File
@@ -51,7 +51,8 @@ namespace kernel {
namespace util {
template <class LayoutA, class LayoutB>
static inline bool
CUTLASS_HOST_DEVICE
static bool
is_continous_k_aligned(GemmCoord problem_size, size_t alignmentA, size_t alignmentB) {
return (std::is_same<LayoutA, layout::RowMajor>::value && (problem_size.k() % alignmentA) == 0) ||
(std::is_same<LayoutB, layout::ColumnMajor>::value && (problem_size.k() % alignmentB) == 0);
@@ -149,42 +150,9 @@ struct UniversalParamsBase
batch_stride_D(args.batch_stride_D),
semaphore(nullptr)
{
ThreadblockSwizzle swizzle;
// Get GEMM volume in thread block tiles
grid_tiled_shape = swizzle.get_tiled_shape(
args.problem_size,
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
args.batch_count);
swizzle_log_tile = swizzle.get_log_tile(grid_tiled_shape);
// Determine extent of K-dimension assigned to each block
gemm_k_size = args.problem_size.k();
if (args.mode == GemmUniversalMode::kGemm || args.mode == GemmUniversalMode::kGemmSplitKParallel)
{
static const uint32_t CACHELINE_BYTES = 128;
static const size_t element_bytes_a = sizeof(ElementA);
static const size_t element_bytes_b = sizeof(ElementB);
static const size_t cacheline_elements_a = CACHELINE_BYTES / element_bytes_a;
static const size_t cacheline_elements_b = CACHELINE_BYTES / element_bytes_b;
const bool cacheline_alignment_needed =
util::is_continous_k_aligned<LayoutA, LayoutB>(problem_size, cacheline_elements_a, cacheline_elements_b);
int const kAlignK = const_max(
const_max(128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value),
cacheline_alignment_needed ? const_max(cacheline_elements_a, cacheline_elements_b) : 1);
gemm_k_size = round_up(ceil_div(args.problem_size.k(), args.batch_count), kAlignK);
if (gemm_k_size) {
grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
}
}
init_grid_tiled_shape();
}
/// Returns the workspace size (in bytes) needed for this problem geometry
size_t get_workspace_size() const
{
@@ -259,6 +227,41 @@ struct UniversalParamsBase
return ThreadblockSwizzle().get_grid_shape(grid_tiled_shape);
}
private:
CUTLASS_HOST_DEVICE
void init_grid_tiled_shape() {
// Get GEMM volume in thread block tiles
grid_tiled_shape = ThreadblockSwizzle::get_tiled_shape(
problem_size,
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
batch_count);
swizzle_log_tile = ThreadblockSwizzle::get_log_tile(grid_tiled_shape);
// Determine extent of K-dimension assigned to each block
gemm_k_size = problem_size.k();
if (mode == GemmUniversalMode::kGemm || mode == GemmUniversalMode::kGemmSplitKParallel)
{
static const uint32_t CACHELINE_BYTES = 128;
static const size_t element_bytes_a = sizeof(ElementA);
static const size_t element_bytes_b = sizeof(ElementB);
static const size_t cacheline_elements_a = CACHELINE_BYTES / element_bytes_a;
static const size_t cacheline_elements_b = CACHELINE_BYTES / element_bytes_b;
const bool cacheline_alignment_needed =
util::is_continous_k_aligned<LayoutA, LayoutB>(problem_size, cacheline_elements_a, cacheline_elements_b);
int const kAlignK = const_max(
const_max(128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value),
cacheline_alignment_needed ? const_max(cacheline_elements_a, cacheline_elements_b) : 1);
gemm_k_size = round_up(ceil_div(problem_size.k(), batch_count), kAlignK);
if (gemm_k_size) {
grid_tiled_shape.k() = ceil_div(problem_size.k(), gemm_k_size);
}
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/sm70_gemm.hpp
+1 -1
View File
@@ -79,7 +79,7 @@ public:
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
"SM70 kernel does not support specializing the tile scheduler.");
using TileScheduleTag = TileScheduler_;
using TileSchedulerTag = TileScheduler_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape,
cute::Shape<cute::Int<1>, cute::Int<1>, cute::Int<1>>>::Scheduler;
include/cutlass/gemm/kernel/sm90_gemm_tma.hpp
+2 -2
View File
@@ -105,14 +105,14 @@ public:
using StrideC = typename CollectiveEpilogue::StrideC;
using ElementD = typename CollectiveEpilogue::ElementD;
using StrideD = typename CollectiveEpilogue::StrideD;
using EpilogueArguments = typename CollectiveEpilogue::Params;
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
using EpilogueParams = typename CollectiveEpilogue::Params;
static_assert(cute::is_same_v<ElementAccumulator, typename CollectiveEpilogue::ElementAccumulator>,
"Mainloop and epilogue do not agree on accumulator value type.");
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
"TMA kernel does not support specializing the tile scheduler.");
using TileScheduleTag = TileScheduler_;
using TileSchedulerTag = TileScheduler_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
+1 -3
View File
@@ -99,7 +99,7 @@ public:
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
"TMA warp-specialized kernel does not support specializing the tile scheduler.");
using TileScheduleTag = TileScheduler_;
using TileSchedulerTag = TileScheduler_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
@@ -357,8 +357,6 @@ public:
// Get pipeline iterators and increments from tensor shapes
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
auto k_tile_count = size<2>(gA);
auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape);
[[maybe_unused]] auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape);
// Wait for all thread blocks in the Cluster
cluster_wait_fn();
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+8 -11
View File
@@ -97,7 +97,7 @@ public:
static_assert(ArchTag::kMinComputeCapability >= 90);
using TileScheduleTag = TileScheduler_;
using TileSchedulerTag = TileScheduler_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
@@ -238,6 +238,7 @@ public:
if constexpr (!std::is_const_v<decltype(args.max_swizzle_size)>) {
args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
}
args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
@@ -390,8 +391,6 @@ public:
// Get pipeline stage increments from tensor shapes
auto k_tile_count = size<3>(gA_mkl);
auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape);
auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape);
TileScheduler scheduler{params.scheduler};
auto work_tile_info = scheduler.get_current_work();
@@ -417,15 +416,13 @@ public:
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
// Slice with our work tile coordinates to construct mainloop tensor views
Tensor gA_presplit = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gB_presplit = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
// Split operands A and B along the K dimension according to work_tile_info
Tensor gA = TileScheduler::split_MK(gA_presplit, work_tile_info); // (BLK_N,BLK_K,k_split_iters)
Tensor gB = TileScheduler::split_NK(gB_presplit, work_tile_info); // (BLK_N,BLK_K,k_split_iters)
auto work_k_tile_count = size<2>(gA);
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA_presplit));
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
auto work_k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info);
auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<2>(gA)), shape<2>(gA));
collective_mainloop.load(
mainloop_pipeline,
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+3 -2
View File
@@ -99,7 +99,7 @@ public:
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
"Ping-pong kernel only supports the default scheduler.");
using TileScheduleTag = TileScheduler_;
using TileSchedulerTag = TileScheduler_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
@@ -240,6 +240,7 @@ public:
if constexpr (!std::is_const_v<decltype(args.max_swizzle_size)>) {
args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
}
args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
@@ -282,7 +283,7 @@ public:
// Kernel level shared memory storage
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
int thread_idx = int(threadIdx.x);
int lane_idx = canonical_lane_idx();
int warp_idx = canonical_warp_idx_sync();
include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp
+41 -184
View File
@@ -31,7 +31,9 @@
#pragma once
#include "cutlass/fast_math.h"
#include "cutlass/gemm_coord.hpp"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cute/layout.hpp"
#include "cute/tensor.hpp"
#include "cute/arch/cluster_sm90.hpp"
@@ -57,30 +59,15 @@ public:
bool is_valid_tile = false;
};
//
// Methods
//
enum class RasterOrder {
AlongM,
AlongN
};
using Params = PersistentTileSchedulerSm90Params;
using RasterOrder = typename Params::RasterOrder;
using RasterOrderOptions = typename Params::RasterOrderOptions;
struct Arguments {
int max_swizzle_size = 1;
RasterOrderOptions raster_order = RasterOrderOptions::Heuristic;
};
struct Params {
FastDivmodU64 divmod_cluster_shape_major_{};
FastDivmodU64 divmod_cluster_shape_minor_{};
FastDivmodU64 divmod_batch_{};
FastDivmodU64 divmod_cluster_blk_major_{};
uint64_t blocks_per_problem_ = 0;
int32_t log_swizzle_size_ = 0;
RasterOrder raster_order_ = RasterOrder::AlongN;
};
// Sink scheduler params as a member
Params scheduler_params;
@@ -102,40 +89,18 @@ public:
static_assert(cute::is_static<TileShape>::value);
static_assert(cute::is_static<ClusterShape>::value);
// Round up to nearest multiple of cluster dim along each mode
auto [problem_blocks_m, problem_blocks_n, problem_blocks_l] = get_tiled_cta_shape_mnl(
problem_shape_mnkl, tile_shape, cluster_shape);
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape, cluster_shape);
// Round up to nearest multiple of swizzle_size along each mode
auto log_swizzle_size = get_log_swizzle_size(problem_blocks_m, problem_blocks_n, arguments.max_swizzle_size);
problem_blocks_m = round_up(problem_blocks_m, (1 << log_swizzle_size) * cute::size<0>(cluster_shape));
problem_blocks_n = round_up(problem_blocks_n, (1 << log_swizzle_size) * cute::size<1>(cluster_shape));
Params params;
params.initialize(
problem_blocks,
to_gemm_coord(cluster_shape),
hw_info,
arguments.max_swizzle_size,
arguments.raster_order
);
RasterOrder raster_order;
raster_order = get_rasterization_order(problem_shape_mnkl, tile_shape);
if (raster_order == RasterOrder::AlongN) {
return {
FastDivmodU64(cute::size<1>(cluster_shape)),
FastDivmodU64(cute::size<0>(cluster_shape)),
FastDivmodU64(problem_blocks_m * problem_blocks_n),
FastDivmodU64(problem_blocks_n / cute::size<1>(cluster_shape)),
problem_blocks_m * problem_blocks_n * problem_blocks_l,
log_swizzle_size,
raster_order
};
}
else {
return {
FastDivmodU64(cute::size<0>(cluster_shape)),
FastDivmodU64(cute::size<1>(cluster_shape)),
FastDivmodU64(problem_blocks_m * problem_blocks_n),
FastDivmodU64(problem_blocks_m / cute::size<0>(cluster_shape)),
problem_blocks_m * problem_blocks_n * problem_blocks_l,
log_swizzle_size,
raster_order
};
}
return params;
}
CUTLASS_HOST_DEVICE
@@ -146,10 +111,10 @@ public:
// like blockIdx and gridDim, with __CUDA_ARCH__.
#if defined(__CUDA_ARCH__)
if (params_.raster_order_ == RasterOrder::AlongN) {
current_work_linear_idx_ = static_cast<uint64_t>(int(blockIdx.x) + (int(blockIdx.y) * int(gridDim.x)));
current_work_linear_idx_ = uint64_t(blockIdx.x) + uint64_t(blockIdx.y) * uint64_t(gridDim.x);
}
else {
current_work_linear_idx_ = static_cast<uint64_t>((int(blockIdx.x) * int(gridDim.y)) + int(blockIdx.y));
current_work_linear_idx_ = uint64_t(blockIdx.x) * uint64_t(gridDim.y) + uint64_t(blockIdx.y);
}
#else
CUTLASS_ASSERT(false && "This line should never be reached");
@@ -187,7 +152,7 @@ public:
// MSVC requires protecting use of CUDA-specific nonstandard syntax,
// like blockIdx and gridDim, with __CUDA_ARCH__.
#if defined(__CUDA_ARCH__)
current_work_linear_idx_ += static_cast<uint64_t>(int(gridDim.x) * int(gridDim.y) * int(gridDim.z)) * advance_count;
current_work_linear_idx_ += uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z) * uint64_t(advance_count);
#else
CUTLASS_ASSERT(false && "This line should never be reached");
#endif
@@ -246,35 +211,14 @@ public:
CUTLASS_HOST_DEVICE static
dim3
get_tiled_cta_shape_mnl(ProblemShapeMNKL problem_shape_mnkl, BlockShape cta_shape, ClusterShape cluster_shape) {
// Across M and N is our Cluster tile, so we must round up the blocks to the nearest whole number of Cluster tiles
auto cta_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shape_mnkl), cute::shape<0>(cta_shape)));
auto cta_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shape_mnkl), cute::shape<1>(cta_shape)));
// Round up to nearest multiple of cluster dim along each mode
int problem_blocks_m = round_up(cta_m, cute::size<0>(cluster_shape));
int problem_blocks_n = round_up(cta_n, cute::size<1>(cluster_shape));
// Cluster tile does not span the batch mode, so no extra rounding up required for it
int problem_blocks_l = int(cute::size<3>(problem_shape_mnkl));
return {uint32_t(problem_blocks_m), uint32_t(problem_blocks_n), uint32_t(problem_blocks_l)};
}
CUTLASS_HOST_DEVICE
static int32_t
get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) {
int min_cta_dim = min(problem_ctas_m, problem_ctas_n);
if (max_swizzle_size >= 8 && min_cta_dim >= 6) {
return 3;
}
else if (max_swizzle_size >= 4 && min_cta_dim >= 3) {
return 2;
}
else if (max_swizzle_size >= 2 && min_cta_dim >= 2) {
return 1;
}
else {
return 0;
}
return Params::get_tiled_cta_shape_mnl(
to_gemm_coord(problem_shape_mnkl),
to_gemm_coord(cluster_shape),
cta_m, cta_n
);
}
// Given the inputs, computes the physical grid we should launch.
@@ -289,111 +233,17 @@ public:
Arguments arguments,
bool truncate_by_problem_size=true) {
int const sm_count = hw_info.sm_count;
CUTLASS_TRACE_HOST("get_grid_shape(): Persistent schedule grid plan using SM count = " << sm_count);
auto problem_shape_mnkl = cute::append<4>(problem_shape_mnk, cute::Int<1>{});
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, cta_shape, cluster_shape);
// Compute the total number of output tiles our problem has
auto problem_shape_MNKL = cute::append<4>(problem_shape_mnk, cute::Int<1>{});
auto [problem_blocks_m, problem_blocks_n, problem_blocks_l] =
get_tiled_cta_shape_mnl(problem_shape_MNKL, cta_shape, cluster_shape);
// Round up to nearest multiple of swizzle_size along each mode
auto swizzle_size = 1 << get_log_swizzle_size(problem_blocks_m, problem_blocks_n, arguments.max_swizzle_size);
problem_blocks_m = round_up(problem_blocks_m, swizzle_size * cute::size<0>(cluster_shape));
problem_blocks_n = round_up(problem_blocks_n, swizzle_size * cute::size<1>(cluster_shape));
int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks_l;
RasterOrder raster_order;
raster_order = get_rasterization_order(problem_shape_mnk, cta_shape);
dim3 launch_grid;
if (raster_order == RasterOrder::AlongN) {
launch_grid = dim3(cute::size<0>(cluster_shape), 1, 1);
}
else {
launch_grid = dim3(1, cute::size<1>(cluster_shape), 1);
}
auto possibly_truncate = [&](int x, int y) {
if (truncate_by_problem_size) {
return std::min(x, y);
}
else {
return x;
}
};
// The else path is generic, however, we can avoid some divs if we know cluster size is 1
if constexpr (size(cluster_shape) == 1) {
if (raster_order == RasterOrder::AlongN) {
launch_grid.y = possibly_truncate(sm_count, problem_blocks_total);
}
else {
launch_grid.x = possibly_truncate(sm_count, problem_blocks_total);
}
}
else {
/*
* Optimal grid size calculation is based on
* GH100: 8 GPCs, 72 TPCs (9 TPCs/GPC), 2 SMs/TPC, 144 SMs per full GPU
* Hence, maximum SMs per GPC = 18
*/
constexpr int max_sm_per_gpc = 18;
// Provided SM count could possibly be less than the assumed maximum SMs per GPC
int const min_num_gpc = sm_count < max_sm_per_gpc ? 1 : sm_count / max_sm_per_gpc;
int const max_cta_occupancy_per_gpc = max_sm_per_gpc - (max_sm_per_gpc % size(cluster_shape));
int cta_per_device = min_num_gpc * max_cta_occupancy_per_gpc;
// The calculation below allows for larger grid size launch for different GPUs.
int const num_gpc_residual = sm_count < max_sm_per_gpc ? 0 : sm_count % max_sm_per_gpc;
int const max_cta_occupancy_per_residual_gpc = num_gpc_residual - (num_gpc_residual % size(cluster_shape));
cta_per_device += max_cta_occupancy_per_residual_gpc;
cta_per_device = sm_count < cta_per_device ? sm_count : cta_per_device;
if (raster_order == RasterOrder::AlongN) {
launch_grid.y = possibly_truncate(
cta_per_device / cute::size<0>(cluster_shape),
problem_blocks_total / cute::size<0>(cluster_shape));
}
else {
launch_grid.x = possibly_truncate(
cta_per_device / cute::size<1>(cluster_shape),
problem_blocks_total / cute::size<1>(cluster_shape));
}
}
return launch_grid;
}
template <class ProblemShapeMNKL, class BlockShape>
CUTLASS_HOST_DEVICE static RasterOrder get_rasterization_order(ProblemShapeMNKL problem_shape_mnkl, BlockShape cta_shape) {
auto tiles_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shape_mnkl), cute::shape<0>(cta_shape)));
auto tiles_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shape_mnkl), cute::shape<1>(cta_shape)));
if (tiles_n > tiles_m) {
return RasterOrder::AlongM;
}
return RasterOrder::AlongN;
}
// Splits an input tensor with MxK according to the splitting configuration specified by work_tile_info.
// Since the basic tile scheduler does not split output tiles, this method is a no-op.
template<class Engine, class Layout>
CUTLASS_DEVICE
static auto
split_MK(cute::Tensor<Engine, Layout> const& tensor, WorkTileInfo const&) {
return tensor;
}
// Splits an input tensor with NxK tiles according to the splitting configuration specified by work_tile_info.
// Since the basic tile scheduler does not split output tiles, this method is a no-op.
template<class Engine, class Layout>
CUTLASS_DEVICE
static auto
split_NK(cute::Tensor<Engine, Layout> const& tensor, WorkTileInfo const&) {
return tensor;
return Params::get_grid_shape(
problem_blocks,
to_gemm_coord(cluster_shape),
hw_info,
arguments.max_swizzle_size,
arguments.raster_order,
/* truncate_by_problem_size = */true
);
}
// Returns whether the block assigned this work should compute the epilogue for the corresponding
@@ -441,6 +291,13 @@ public:
// space of the output tile assigned to the work unit.
return cute::size(cute::ceil_div(cute::get<2>(problem_shape), cute::get<2>(tile_shape)));
}
CUTLASS_HOST_DEVICE
static uint32_t
get_work_k_tile_start(WorkTileInfo const&) {
// All work units returned by this scheduler start from K tile 0
return 0u;
}
};
} // namespace cutlass::gemm::kernel::detail
include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp
+115 -416
View File
@@ -34,7 +34,6 @@
#include "cutlass/barrier.h"
#include "cutlass/block_striped.h"
#include "cutlass/fast_math.h"
#include "cutlass/workspace.hpp"
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
#include "cutlass/kernel_hardware_info.hpp"
#include "cute/layout.hpp"
@@ -54,22 +53,21 @@ class PersistentTileSchedulerSm90StreamK {
private:
using UnderlyingScheduler = PersistentTileSchedulerSm90;
public:
using RasterOrder = UnderlyingScheduler::RasterOrder;
private:
using UnderlyingArguments = typename UnderlyingScheduler::Arguments;
using UnderlyingParams = typename UnderlyingScheduler::Params;
uint64_t current_work_linear_idx_ = 0;
// Minimum number of k iterations that can be assigned to a stream-K unit
static constexpr uint32_t min_iters_per_sk_unit_ = 2;
public:
using RasterOrder = UnderlyingScheduler::RasterOrder;
using RasterOrderOptions = UnderlyingScheduler::RasterOrderOptions;
// Use a dummy barrier manager to simply get the type used to store the barrier
using BarrierType = typename NamedBarrierManager<1>::T;
public:
struct WorkTileInfo {
int32_t M_idx = 0;
int32_t N_idx = 0;
@@ -91,27 +89,32 @@ public:
bool is_final_split = true;
};
using Params = PersistentTileSchedulerSm90StreamKParams;
using ReductionMode = Params::ReductionMode;
struct Arguments {
Arguments() = default;
Arguments(Arguments const&) = default;
Arguments(Arguments&&) = default;
CUTLASS_HOST_DEVICE
Arguments&
Arguments&
operator=(Arguments const& args) {
splits = args.splits;
return *this;
}
CUTLASS_HOST_DEVICE
Arguments&
operator=(Arguments&& args) noexcept {
splits = args.splits;
raster_order = args.raster_order;
return *this;
}
CUTLASS_HOST_DEVICE
CUTLASS_HOST_DEVICE
Arguments&
operator=(Arguments&& args) noexcept {
splits = args.splits;
raster_order = args.raster_order;
return *this;
}
CUTLASS_HOST_DEVICE
Arguments(int splits_) : splits(splits_) {}
// The splitting factor to be used in a split-K decomposition of the problem.
@@ -119,48 +122,8 @@ public:
// is bypassed in favor of a split-K decomposition.
int splits = 1;
const int max_swizzle_size = 1;
};
struct Params {
FastDivmodU64 divmod_cluster_shape_major_{};
FastDivmodU64 divmod_cluster_shape_minor_{};
FastDivmodU64 divmod_batch_{};
FastDivmodU64 divmod_k_{};
FastDivmodU64 divmod_cluster_blk_major_{};
int32_t log_swizzle_size_ = 0;
uint64_t units_per_problem_ = 0;
RasterOrder raster_order_ = RasterOrder::AlongN;
ClusterShape cluster_shape_{};
// The splitting factor to be used in a split-K decomposition of the problem.
// If this is set to a value greater than 1, stream-K decomposition logic
// is bypassed in favor of a split-K decomposition.
uint32_t splits_ = 1;
// Number of tiled k iterations required to compute a single output tile.
uint32_t k_tiles_per_output_tile_ = 0;
// Number of stream-K or split-K work units that compute an extra k iteration.
// This is done to handle residuals in dividing up the k iteration space.
// For stream-K, since the actual assignment of work to stream-K units will be done
// at the granularity of a cluster, we store only the number of big clusters.
uint32_t big_units_ = 0;
// Workspace for holding partial accumulators to be reduced across stream-K/split-K units
void* reduction_workspace_ = nullptr;
// Number of tiles covered by stream-K work units
uint32_t sk_tiles_ = 0;
// Number of work units computing stream-K tiles
uint32_t sk_units_ = 0;
// Number of tiled k iterations computed by each stream-K work unit. This
// can potentially cover more than one output tile.
uint32_t k_tiles_per_sk_unit_ = 0;
RasterOrderOptions raster_order = RasterOrderOptions::Heuristic;
ReductionMode reduction_mode = ReductionMode::Deterministic;
};
// Sink scheduler params as a member
@@ -173,7 +136,7 @@ public:
template <class ProblemShape>
static Params
to_underlying_arguments(
ProblemShape problem_shape_mnkl,
ProblemShape problem_shape,
TileShape tile_shape,
ClusterShape cluster_shape,
KernelHardwareInfo const& hw_info,
@@ -183,143 +146,23 @@ public:
static_assert(cute::is_static<TileShape>::value);
static_assert(cute::is_static<ClusterShape>::value);
// Round up to nearest multiple of cluster dim along each mode
auto [problem_blocks_m, problem_blocks_n, problem_blocks_l] = get_tiled_cta_shape_mnl(
problem_shape_mnkl, tile_shape, cluster_shape);
auto problem_shape_mnkl = cute::append<4>(problem_shape, cute::Int<1>{});
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape, cluster_shape);
uint32_t k_tile_per_output_tile = cute::size(cute::ceil_div(cute::shape<2>(problem_shape_mnkl), cute::shape<2>(TileShape{})));
uint64_t output_tiles = problem_blocks_m * problem_blocks_n * problem_blocks_l;
// Number of k tile iterations in each output tile
uint32_t k_tiles_per_output_tile = (cute::size<2>(problem_shape_mnkl) + cute::size<2>(tile_shape) - 1) /
cute::size<2>(tile_shape);
UnderlyingArguments underlying_args;
underlying_args.max_swizzle_size = 1;
UnderlyingParams underlying_params = UnderlyingScheduler::to_underlying_arguments(
problem_shape_mnkl, tile_shape, cluster_shape, hw_info, underlying_args, workspace);
void* reduction_workspace = nullptr;
if (workspace != nullptr) {
// Reduction workspace is at the beginning of the workspace. Lock workspace follows.
reduction_workspace = workspace;
}
if (args.splits > 1) {
// Short circuit to basic split-K decomposition
// Don't split by more than the available number of SMs
auto splits = args.splits > hw_info.sm_count ? hw_info.sm_count : args.splits;
// Don't split by more than the K tile iterations
//
// splits is almost certainly nonnegative here (e.g., hw_info.sm_count,
// despite being an int, is a count), so it can safely be converted to unsigned
// in the comparison to avoid a signed-unsigned comparison warning-as-error.
splits = static_cast<decltype(k_tiles_per_output_tile)>(splits) > k_tiles_per_output_tile ? k_tiles_per_output_tile : splits;
return get_params_basic(
underlying_params, problem_blocks_m, problem_blocks_n, problem_blocks_l, cluster_shape,
splits, k_tiles_per_output_tile, reduction_workspace);
}
// Calculate the maximum number of blocks from clusters of shape cluster_shape that we
// can fit within sm_count SMs.
dim3 grid = get_grid_shape(problem_shape_mnkl, tile_shape, cluster_shape, hw_info, args);
uint64_t ctas_per_wave = grid.x * grid.y;
// The number of output tiles to be computed in stream-K and data-parallel fashion, respectively.
uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, k_tiles_per_output_tile);
uint64_t dp_tiles = output_tiles - sk_tiles;
// Calculate the number of work units covering the data-parallel and stream-K tiles.
// A "work unit" is a single index in the linearized ID space used by the scheduler.
// We distinguish it from a "block," which is typically tied to a hardware unit
// (e.g., the callers into this scheduler will be persistent thread blocks).
// A work unit can encompass multiple output tiles worth of work (as will be the
// case for stream-K blocks).
// Since splitting is not required for data-parallel tiles, only one data-parallel unit
// is needed per data-parallel tile.
uint64_t dp_units = dp_tiles;
// Number of k iterations computed by the stream-K units as a whole
uint64_t k_tiles_sk_total = k_tiles_per_output_tile * sk_tiles;
// If there are stream-K tiles to compute and a sufficiently large number of k iterations
// across them, they will be covered by a single wave of persistent threadblocks. Thus, there
// will be as many work units as there are threadblocks in a single wave.
//
// When the total k iterations across stream-K tiles is too small to justify distributing
// across an entire wave of blocks, we instead distribute the iterations over a smaller
// set of blocks.
// Calculate the number of stream-K units that would be needed if each stream-K unit
// computed the minimum allowable k iterations. Truncate this to be in units of clusters.
uint64_t min_sized_sk_units = (k_tiles_sk_total / min_iters_per_sk_unit_);
min_sized_sk_units = (min_sized_sk_units / cute::size(cluster_shape)) * cute::size(cluster_shape);
uint64_t sk_units = min(ctas_per_wave, min_sized_sk_units);
if (sk_units == 0) {
// Short circuit to basic data-parallel decomposition
return get_params_basic(
underlying_params, problem_blocks_m, problem_blocks_n, problem_blocks_l, cluster_shape,
1, k_tiles_per_output_tile, reduction_workspace);
}
// If the number of stream-K units is a multiple of the number of stream-K tiles, then
// the problem can leverage a basic split-K decomposition for the stream-K tiles.
if (sk_tiles < sk_units && sk_units % sk_tiles == 0) {
// Short circuit to basic split-K decomposition
uint32_t sk_splits = static_cast<uint32_t>(sk_units / sk_tiles);
return get_params_basic(
underlying_params, problem_blocks_m, problem_blocks_n, problem_blocks_l, cluster_shape,
sk_splits, k_tiles_per_output_tile, reduction_workspace);
}
// Number of k iterations computed per stream-K units
uint64_t k_tiles_per_sk_unit = k_tiles_sk_total / sk_units;
// Number of stream-K units that need to compute extra iterations in order to cover
// the residual k iterations. This assumes that each such unit computes one additional
// iteration.
uint64_t sk_big_units = k_tiles_sk_total - (k_tiles_per_sk_unit * sk_units);
// The division below is guaranteed to be exact because sk_big_units is guaranteed
// to be a multiple of cluster_size (cute::size(cluster_shape)). This is useful because
// it allows us to use a block's linearized cluster ID to determine whether it is
// a big block. The reasoning behind this guarnatee is explained as follows:
// sk_big_units = k_tiles_sk_total - (k_tiles_per_sk_unit * sk_units);
//
// - k_tiles_sk_total is a multiple of cluster_size because it is the product
// of number of tail tiles and the number of k iterations per tile. Because
// both the number of output tiles and number of available SMs are rounded
// to be multiples of cluster shape, the number of tail tiles
// (output_tiles % avail_sms) is a multpile of cluster_size.
//
// - sk_units is a multiple of cluster_size because it is either blocks_per_wave
// or 0, and blocks_per_wave is a multiple of the cluster_size due to the grid-planning
// logic rounding to multiples of cluster dimensions
uint64_t sk_big_units_per_cluster = sk_big_units / cute::size(cluster_shape);
return {
underlying_params.divmod_cluster_shape_major_,
underlying_params.divmod_cluster_shape_minor_,
underlying_params.divmod_batch_,
FastDivmodU64(problem_blocks_m * problem_blocks_n), // Static k-splitting divmod. Unused for stream-K.
underlying_params.divmod_cluster_blk_major_,
underlying_params.log_swizzle_size_,
static_cast<uint32_t>(dp_units + sk_units),
underlying_params.raster_order_,
cluster_shape,
1, // Static k-splitting factor. Unused for stream-K.
k_tiles_per_output_tile,
static_cast<uint32_t>(sk_big_units_per_cluster),
reduction_workspace,
sk_tiles,
static_cast<uint32_t>(sk_units),
static_cast<uint32_t>(k_tiles_per_sk_unit)
};
Params params;
params.initialize(
problem_blocks,
k_tile_per_output_tile,
to_gemm_coord(cluster_shape),
hw_info,
args.splits,
args.max_swizzle_size,
args.raster_order,
args.reduction_mode,
workspace
);
return params;
}
CUTLASS_HOST_DEVICE
@@ -328,10 +171,10 @@ public:
CUTLASS_HOST_DEVICE
PersistentTileSchedulerSm90StreamK(Params const& params_) : scheduler_params(params_) {
if (params_.raster_order_ == RasterOrder::AlongN) {
current_work_linear_idx_ = static_cast<uint64_t>(int(blockIdx.x) + (int(blockIdx.y) * int(gridDim.x)));
current_work_linear_idx_ = uint64_t(blockIdx.x) + uint64_t(blockIdx.y) * uint64_t(gridDim.x);
}
else {
current_work_linear_idx_ = static_cast<uint64_t>((int(blockIdx.x) * int(gridDim.y)) + int(blockIdx.y));
current_work_linear_idx_ = uint64_t(blockIdx.x) * uint64_t(gridDim.y) + uint64_t(blockIdx.y);
}
}
@@ -397,7 +240,7 @@ public:
CUTLASS_DEVICE
void
advance_to_next_work(uint32_t advance_count = 1) {
current_work_linear_idx_ += static_cast<uint64_t>(int(gridDim.x) * int(gridDim.y) * int(gridDim.z)) * advance_count;
current_work_linear_idx_ += uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z) * uint64_t(advance_count);
}
// Given the inputs, computes the total number of output blocks this problem will compute over
@@ -420,17 +263,16 @@ public:
KernelHardwareInfo hw_info,
Arguments arguments) {
UnderlyingArguments underlying_args;
underlying_args.max_swizzle_size = 1;
// Call into the underlying get_grid_shape method, but do not allow the grid shape returned
// to be truncated based on the number of output tiles in the problem.
return UnderlyingScheduler::get_grid_shape(
problem_shape,
tile_shape,
cluster_shape,
auto problem_shape_mnkl = cute::append<4>(problem_shape, cute::Int<1>{});
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape, cluster_shape);
return Params::get_grid_shape(
problem_blocks,
to_gemm_coord(cluster_shape),
hw_info,
underlying_args,
/*truncate_by_problem_size=*/false);
arguments.max_swizzle_size,
arguments.raster_order
);
}
// Returns whether fixup is needed for `work_tile_info`.
@@ -501,7 +343,8 @@ public:
// note that, in the split-K case, the units_per_problem_ member of Params will be
// the total number of output tiles multiplied by the number of splits.
auto reduction_tiles = params.splits_ > 1 ? (params.units_per_problem_ / params.splits_) : params.sk_tiles_;
auto reduction_workspace_size = get_reduction_workspace_size<ElementAccumulator>(reduction_tiles);
auto reduction_workspace_size = Params::get_reduction_workspace_size(
reduction_tiles, to_gemm_coord(TileShape{}), sizeof_bits<ElementAccumulator>::value);
BarrierType* lock_workspace = reinterpret_cast<BarrierType*>(
reinterpret_cast<uint8_t*>(params.reduction_workspace_) + reduction_workspace_size);
@@ -511,8 +354,14 @@ public:
BlockStripedReduceT::store(reduction_workspace_array, *accumulator_array, barrier_group_thread_idx);
}
else {
// Wait until the preceding split added its accumulators
BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
if (params.reduction_mode_ == ReductionMode::Deterministic) {
// Wait until the preceding split added its accumulators
BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
}
else {
// Wait unitl the first split has stored its accumulators
BarrierManager::wait_lt(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, 1);
}
// Perform reduction in workspace
BlockStripedReduceT::reduce(reduction_workspace_array, *accumulator_array, barrier_group_thread_idx);
@@ -531,22 +380,6 @@ public:
}
}
// Splits an input tensor with MxK according to the splitting configuration specified by work_tile_info
template<class Engine, class Layout>
CUTLASS_DEVICE
static auto
split_MK(cute::Tensor<Engine, Layout> const& tensor, WorkTileInfo const& work_tile_info) {
return split<Engine, Layout, 0>(tensor, work_tile_info);
}
// Splits an input tensor with NxK tiles according to the splitting configuration specified by work_tile_info
template<class Engine, class Layout>
CUTLASS_DEVICE
static auto
split_NK(cute::Tensor<Engine, Layout> const& tensor, WorkTileInfo const& work_tile_info) {
return split<Engine, Layout, 1>(tensor, work_tile_info);
}
// Returns whether the block assigned this work should compute the epilogue for the corresponding
// output tile. For the case of stream-K, this should only occur if the work is marked as the final split.
CUTLASS_HOST_DEVICE
@@ -564,14 +397,14 @@ public:
if (params.raster_order_ == RasterOrder::AlongN) {
return
(tiles_mn * work_tile_info.L_idx) +
(params.divmod_cluster_shape_major_.divisor *
(params.divmod_cluster_shape_major_.divisor *
params.divmod_cluster_blk_major_.divisor * work_tile_info.M_idx) +
work_tile_info.N_idx;
}
else {
return
(tiles_mn * work_tile_info.L_idx) +
(params.divmod_cluster_shape_major_.divisor *
(params.divmod_cluster_shape_major_.divisor *
params.divmod_cluster_blk_major_.divisor * work_tile_info.N_idx) +
work_tile_info.M_idx;
}
@@ -582,16 +415,16 @@ public:
uint64_t cta_per_grid_dim;
uint64_t cluster_dim_idx;
if (params.raster_order_ == RasterOrder::AlongN) {
uint64_t block_idx_m = (work_tile_info.M_idx - cta_m_in_cluster) / cute::size<0>(params.cluster_shape_);
uint64_t block_idx_m = (work_tile_info.M_idx - cta_m_in_cluster) / params.divmod_cluster_shape_minor_.divisor;
uint64_t block_idx_n = work_tile_info.N_idx;
cta_per_grid_dim = (params.divmod_cluster_shape_major_.divisor *
cta_per_grid_dim = (params.divmod_cluster_shape_major_.divisor *
params.divmod_cluster_blk_major_.divisor * block_idx_m) + block_idx_n;
cluster_dim_idx = cta_m_in_cluster;
}
else {
uint64_t block_idx_m = work_tile_info.M_idx;
uint64_t block_idx_n = (work_tile_info.N_idx - cta_n_in_cluster) / cute::size<1>(params.cluster_shape_);
cta_per_grid_dim = (params.divmod_cluster_shape_major_.divisor *
uint64_t block_idx_n = (work_tile_info.N_idx - cta_n_in_cluster) / params.divmod_cluster_shape_minor_.divisor;
cta_per_grid_dim = (params.divmod_cluster_shape_major_.divisor *
params.divmod_cluster_blk_major_.divisor * block_idx_n) + block_idx_m;
cluster_dim_idx = cta_n_in_cluster;
}
@@ -609,13 +442,27 @@ public:
KernelHardwareInfo const& hw_info,
uint32_t mma_warp_groups) {
int barrier_workspace_size = 0;
int reduction_workspace_size = 0;
auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
get_workspace_component_sizes<ProblemShape, ElementAccumulator>(
args, problem_shape, barrier_workspace_size, reduction_workspace_size, hw_info, mma_warp_groups);
ClusterShape cluster_shape;
TileShape tile_shape;
return barrier_workspace_size + reduction_workspace_size;
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape, cluster_shape);
uint32_t k_tile_per_output_tile = cute::size(cute::ceil_div(cute::shape<2>(problem_shape_mnkl), cute::shape<2>(TileShape{})));
return Params::get_workspace_size(
problem_blocks,
k_tile_per_output_tile,
to_gemm_coord(tile_shape),
to_gemm_coord(cluster_shape),
hw_info,
args.splits,
args.max_swizzle_size,
args.raster_order,
mma_warp_groups,
sizeof_bits<BarrierType>::value,
sizeof_bits<ElementAccumulator>::value
);
}
template <class ProblemShape, class ElementAccumulator>
@@ -628,26 +475,29 @@ public:
KernelHardwareInfo const& hw_info,
uint32_t mma_warp_groups) {
#if !defined(__CUDACC_RTC__)
int barrier_workspace_size = 0;
int reduction_workspace_size = 0;
auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
get_workspace_component_sizes<ProblemShape, ElementAccumulator>(
args, problem_shape, barrier_workspace_size, reduction_workspace_size, hw_info, mma_warp_groups);
ClusterShape cluster_shape;
TileShape tile_shape;
if (barrier_workspace_size > 0) {
if (workspace == nullptr) {
return Status::kErrorWorkspaceNull;
}
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape, cluster_shape);
uint32_t k_tile_per_output_tile = cute::size(cute::ceil_div(cute::shape<2>(problem_shape_mnkl), cute::shape<2>(TileShape{})));
// Only the barrier workspace needs to be cleared for stream-K.
// Barrier workspace follows reduction workspace.
uint8_t* barrier_workspace = reinterpret_cast<uint8_t*>(workspace) + reduction_workspace_size;
return zero_workspace(static_cast<void*>(barrier_workspace), barrier_workspace_size, stream);
}
return Status::kSuccess;
#endif
return Params::initialize_workspace(
workspace,
stream,
problem_blocks,
k_tile_per_output_tile,
to_gemm_coord(tile_shape),
to_gemm_coord(cluster_shape),
hw_info,
args.splits,
args.max_swizzle_size,
args.raster_order,
mma_warp_groups,
sizeof_bits<BarrierType>::value,
sizeof_bits<ElementAccumulator>::value
);
}
template <class ProblemShape>
@@ -657,162 +507,10 @@ public:
return work_tile_info.k_tile_count;
}
private:
// Splits a tensor using the splitting configuration specified by work_tile_info using
// a MN shape detemined by TileDim0.
template <class Engine, class Layout, int TileDim0>
CUTLASS_DEVICE
static auto
split(cute::Tensor<Engine, Layout> const& tensor, WorkTileInfo const& work_tile_info) {
using namespace cute;
// Divide input tensor into `splits` chunks along the k dimension
auto div_shape = make_shape(size<TileDim0>(TileShape{}), size<2>(TileShape{}), work_tile_info.splits);
auto split = zipped_divide(tensor, div_shape);
// Index into the split tensor at the work tile's split index
auto indexed = split(make_coord(make_coord(_, _, work_tile_info.K_idx), make_coord(0, 0, _)));
// Construct a layout for the indexed tensor. The main purpose of this new layout is to
// override the k extent to support cases in which the split computes a number of iterations
// not equal to total_k_tiles / splits. A common example of this is in stream-K is when a
// unit computes the final 20 of the total 32 k iterations of the output tile. In this case,
// set splits = 32 and the split index (K_idx) to 11. The zipped divide above results in each
// of the splits computing only one k iteration.
auto overridden_shape = make_shape(size<0>(indexed.layout()), size<1>(indexed.layout()), work_tile_info.k_tile_count);
auto layout = make_layout(overridden_shape, tensor.stride());
return make_tensor(indexed.data(), layout);
}
// Returns the number of stream-K tiles that will be computed amongst `output_tiles` total
// output tiles on a device with `ctas_per_wave` CTAs in each wave.
CUTLASS_HOST_DEVICE
static uint32_t
get_num_sk_tiles(uint64_t output_tiles, uint64_t ctas_per_wave, uint32_t k_tiles_per_output_tile) {
uint32_t full_waves = static_cast<uint32_t>(output_tiles / ctas_per_wave);
uint32_t total_waves = static_cast<uint32_t>((output_tiles + ctas_per_wave - 1) / ctas_per_wave);
if (full_waves == total_waves || k_tiles_per_output_tile == 1) {
// All tiles will be data-parallel tiles if there is either no quantization
// or if there is no work to be split.
return 0;
}
//
// The final wave is not full. Perform some stream-K work.
//
// Rudimentary heuristic: prefer data-parallel decomposition if we have more than
// one wave and the tail wave is more than half full. This is subject to change.
if (full_waves != 0) {
uint64_t tail_tiles = output_tiles - (full_waves * ctas_per_wave);
if (tail_tiles >= (ctas_per_wave / 2)) {
return 0;
}
}
// If there is wave quantization, assign the first two waves worth of tiles to be
// covered by stream-K work and the remainder to be data-parallel. Since we know
// that full_waves == total_waves - 1 in this case, the number of data-parallel
// waves is simply full_waves-1 (unless full_waves == 0).
uint32_t dp_waves = full_waves > 0 ? full_waves - 1 : 0;
uint64_t dp_tiles = dp_waves * ctas_per_wave;
return static_cast<uint32_t>(output_tiles - dp_tiles);
}
// Calculates the size of the workspace needed for holding reduction barriers
CUTLASS_HOST_DEVICE
static int
get_barrier_workspace_size(uint64_t num_tiles, uint32_t mma_warp_groups) {
auto workspace_bits = num_tiles * mma_warp_groups * sizeof_bits<BarrierType>::value;
return bits_to_bytes(static_cast<int>(workspace_bits));
}
// Calculates the size of the workspace needed for holding partial outputs from splits
template <class ElementAccumulator>
CUTLASS_HOST_DEVICE
static int
get_reduction_workspace_size(uint64_t num_tiles) {
auto output_tile_size = cute::size<0>(TileShape{}) * cute::size<1>(TileShape{});
auto workspace_bits = sizeof_bits<ElementAccumulator>::value * output_tile_size * num_tiles;
return bits_to_bytes(static_cast<int>(workspace_bits));
}
template <class ProblemShape, class ElementAccumulator>
static void
get_workspace_component_sizes(
Arguments const& args,
ProblemShape problem_shape,
int& barrier_workspace_size,
int& reduction_workspace_size,
KernelHardwareInfo const& hw_info,
uint32_t mma_warp_groups) {
// Workspace is needed only for output tiles that will be split. Thus, we first determine the number
// of output tiles that will be split, and then calculate the workspace needed to cover these.
auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
ClusterShape cluster_shape;
auto [problem_blocks_m, problem_blocks_n, problem_blocks_l] = get_tiled_cta_shape_mnl(
problem_shape_mnkl, TileShape{}, cluster_shape);
uint64_t output_tiles = problem_blocks_m * problem_blocks_n * problem_blocks_l;
if (args.splits > 1) {
// Basic split-K variant requires workspace for all output tiles
barrier_workspace_size = get_barrier_workspace_size(output_tiles, mma_warp_groups);
reduction_workspace_size = get_reduction_workspace_size<ElementAccumulator>(output_tiles);
}
else {
int sm_count = hw_info.sm_count;
if (sm_count <= 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
}
uint32_t k_tiles_per_output_tile = (cute::size<2>(problem_shape_mnkl) + cute::size<2>(TileShape{}) - 1) /
cute::size<2>(TileShape{});
dim3 grid = get_grid_shape(problem_shape_mnkl, TileShape{}, cluster_shape, {0, sm_count}, args);
uint64_t ctas_per_wave = grid.x * grid.y;
uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, k_tiles_per_output_tile);
barrier_workspace_size = get_barrier_workspace_size(sk_tiles, mma_warp_groups);
reduction_workspace_size = get_reduction_workspace_size<ElementAccumulator>(sk_tiles);
}
}
// Constructs parameters for either a basic data-parallel or basic split-K decomposition of the problem
static Params
get_params_basic(
UnderlyingParams const& underlying_params,
uint32_t blocks_m,
uint32_t blocks_n,
uint32_t blocks_l,
ClusterShape cluster_shape,
uint32_t splits,
uint32_t k_tiles_per_output_tile,
void* reduction_workspace) {
uint32_t big_units = k_tiles_per_output_tile % splits;
return {
underlying_params.divmod_cluster_shape_major_,
underlying_params.divmod_cluster_shape_minor_,
FastDivmodU64(blocks_m * blocks_n * splits),
FastDivmodU64(blocks_m * blocks_n),
underlying_params.divmod_cluster_blk_major_,
underlying_params.log_swizzle_size_,
blocks_m * blocks_n * blocks_l * splits,
underlying_params.raster_order_,
cluster_shape,
splits,
k_tiles_per_output_tile,
big_units,
reduction_workspace
};
get_work_k_tile_start(WorkTileInfo const& work_tile_info) {
return work_tile_info.K_idx;
}
// Sets the current stream-K work to compute within work_tile_info. If new_unit is true, work_tile_info
@@ -840,7 +538,8 @@ private:
//
// To do so, we divide up the linearized stream-K units into clusters and share the same K
// offsets for work within clusters.
auto cluster_linear_work_idx = linear_idx / size(params.cluster_shape_);
auto cluster_size = params.divmod_cluster_shape_major_.divisor * params.divmod_cluster_shape_minor_.divisor;
auto cluster_linear_work_idx = linear_idx / cluster_size;
// Determine the starting k iteration computed by this stream-K work unit
uint32_t unit_iter_start = params.k_tiles_per_sk_unit_ * cluster_linear_work_idx;
@@ -890,16 +589,16 @@ private:
uint32_t true_tile_iter_end = true_tile_iter_start + params.k_tiles_per_output_tile_;
// Bring the linearized tile ID back into the space of tiles, rather than clusters
true_tile_id *= size(params.cluster_shape_);
true_tile_id *= cluster_size;
auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
// The final linearized tile ID is in units of the cluster dimension over which we rasterize.
if (params.raster_order_ == RasterOrder::AlongN) {
true_tile_id += cta_n_in_cluster * cute::size<0>(params.cluster_shape_);
true_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
}
else {
true_tile_id += cta_m_in_cluster * cute::size<1>(params.cluster_shape_);
true_tile_id += cta_m_in_cluster * params.divmod_cluster_shape_minor_.divisor;
}
// The unit's starting k iteration in the current tile is either the starting
@@ -925,7 +624,7 @@ private:
params.divmod_cluster_shape_major_,
params.divmod_cluster_shape_minor_,
params.divmod_cluster_blk_major_,
params.log_swizzle_size_,
params.log_swizzle_size_,
params.raster_order_);
//
include/cutlass/gemm/kernel/tile_scheduler_params.h
+1005
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/threadblock/threadblock_swizzle.h
+36 -36
View File
@@ -62,10 +62,10 @@ struct GemmIdentityThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
/// *Gemm* problem size: gemm(M, N, K)
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
GemmCoord problem_size,
GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
return GemmCoord(
(problem_size.m() + tile_size.m() - 1) / tile_size.m(),
@@ -76,11 +76,11 @@ struct GemmIdentityThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
/// *ImplicitGemm* Conv2d problem size: conv_operator(NPQK, NHWC, KRSC)
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
cutlass::conv::Operator conv_operator,
cutlass::conv::Conv2dProblemSize const &problem_size,
GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
gemm::GemmCoord implicit_gemm_problem_size =
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
@@ -92,11 +92,11 @@ struct GemmIdentityThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
/// *ImplicitGemm* Conv3d problem size: conv_operator(NZPQK, NDHWC, KTRSC)
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
cutlass::conv::Operator conv_operator,
cutlass::conv::Conv3dProblemSize const &problem_size,
GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
gemm::GemmCoord implicit_gemm_problem_size =
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
@@ -107,7 +107,7 @@ struct GemmIdentityThreadblockSwizzle {
/// Computes CUDA grid dimensions given a size in units of logical tiles
CUTLASS_HOST_DEVICE
dim3 get_grid_shape(GemmCoord tiled_shape) const {
static dim3 get_grid_shape(GemmCoord tiled_shape) {
int tile = 1 << get_log_tile(tiled_shape);
return dim3(tiled_shape.m() * tile, (tiled_shape.n() + tile - 1) / tile, tiled_shape.k());
}
@@ -129,7 +129,7 @@ struct GemmIdentityThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(int log_tile) const {
static GemmCoord get_tile_offset(int log_tile) {
int block_idx_x = RematerializeBlockIdxX();
int block_idx_y = RematerializeBlockIdxY();
int block_idx_z = RematerializeBlockIdxZ();
@@ -141,7 +141,7 @@ struct GemmIdentityThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(GemmCoord tiled_shape) const {
static GemmCoord get_tile_offset(GemmCoord tiled_shape) {
int const kTile = N;
int block_idx_x = RematerializeBlockIdxX();
@@ -168,10 +168,10 @@ struct GemmHorizontalThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
GemmCoord problem_size,
GemmCoord tile_size,
int split_k_slices) const {
int split_k_slices) {
return GemmCoord(
(problem_size.m() + tile_size.m() - 1) / tile_size.m(),
@@ -181,7 +181,7 @@ struct GemmHorizontalThreadblockSwizzle {
/// Computes CUDA grid dimensions given a size in units of logical tiles
CUTLASS_HOST_DEVICE
dim3 get_grid_shape(GemmCoord tiled_shape) const {
static dim3 get_grid_shape(GemmCoord tiled_shape) {
return dim3(tiled_shape.n(), tiled_shape.m(), tiled_shape.k());
}
@@ -193,7 +193,7 @@ struct GemmHorizontalThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(GemmCoord tiled_shape) const {
static GemmCoord get_tile_offset(GemmCoord tiled_shape) {
return GemmCoord{
RematerializeBlockIdxY(),
RematerializeBlockIdxX(),
@@ -209,10 +209,10 @@ struct GemmBatchedIdentityThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
GemmCoord problem_size,
GemmCoord tile_size,
int batch_count) const {
int batch_count) {
return GemmCoord(
(problem_size.m() + tile_size.m() - 1) / tile_size.m(),
@@ -222,7 +222,7 @@ struct GemmBatchedIdentityThreadblockSwizzle {
/// Computes CUDA grid dimensions given a size in units of logical tiles
CUTLASS_HOST_DEVICE
dim3 get_grid_shape(GemmCoord tiled_shape) const {
static dim3 get_grid_shape(GemmCoord tiled_shape) {
return dim3(tiled_shape.m(), tiled_shape.n(), tiled_shape.k());
}
@@ -234,7 +234,7 @@ struct GemmBatchedIdentityThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(GemmCoord tiled_shape) const {
static GemmCoord get_tile_offset(GemmCoord tiled_shape) {
return GemmCoord{
RematerializeBlockIdxX(),
RematerializeBlockIdxY(),
@@ -244,7 +244,7 @@ struct GemmBatchedIdentityThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(int log_tile) const {
static GemmCoord get_tile_offset(int log_tile) {
int block_idx_x = RematerializeBlockIdxX();
int block_idx_y = RematerializeBlockIdxY();
int block_idx_z = RematerializeBlockIdxZ();
@@ -256,7 +256,7 @@ struct GemmBatchedIdentityThreadblockSwizzle {
/// Gets the batch index
CUTLASS_DEVICE
int get_batch_idx() const {
static int get_batch_idx() {
return RematerializeBlockIdxZ();
}
};
@@ -271,10 +271,10 @@ struct GemmSplitKIdentityThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
GemmCoord problem_size,
GemmCoord tile_size,
int partitions) const {
int partitions) {
return GemmCoord(
(problem_size.m() + tile_size.m() - 1) / tile_size.m(),
@@ -299,14 +299,14 @@ struct GemmSplitKIdentityThreadblockSwizzle {
/// Computes CUDA grid dimensions given a size in units of logical tiles
CUTLASS_HOST_DEVICE
dim3 get_grid_shape(GemmCoord tiled_shape) const {
static dim3 get_grid_shape(GemmCoord tiled_shape) {
int tile = 1 << get_log_tile(tiled_shape);
return dim3(tiled_shape.m() * tile, (tiled_shape.n() + tile - 1) / tile, tiled_shape.k());
}
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(int log_tile) const {
static GemmCoord get_tile_offset(int log_tile) {
int block_idx_x = RematerializeBlockIdxX();
int block_idx_y = RematerializeBlockIdxY();
int block_idx_z = RematerializeBlockIdxZ();
@@ -318,7 +318,7 @@ struct GemmSplitKIdentityThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(GemmCoord tiled_shape) const {
static GemmCoord get_tile_offset(GemmCoord tiled_shape) {
int const kTile = N;
int block_idx_x = RematerializeBlockIdxX();
@@ -342,10 +342,10 @@ struct GemmSplitKHorizontalThreadblockSwizzle {
/// Returns the shape of the problem in units of logical tiles
CUTLASS_HOST_DEVICE
GemmCoord get_tiled_shape(
static GemmCoord get_tiled_shape(
GemmCoord problem_size,
GemmCoord tile_size,
int partitions) const {
int partitions) {
return GemmCoord(
(problem_size.m() + tile_size.m() - 1) / tile_size.m(),
@@ -355,7 +355,7 @@ struct GemmSplitKHorizontalThreadblockSwizzle {
/// Computes CUDA grid dimensions given a size in units of logical tiles
CUTLASS_HOST_DEVICE
dim3 get_grid_shape(GemmCoord tiled_shape) const {
static dim3 get_grid_shape(GemmCoord tiled_shape) {
return dim3(tiled_shape.n(), tiled_shape.m(), tiled_shape.k());
}
@@ -367,7 +367,7 @@ struct GemmSplitKHorizontalThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(int log_tile) const {
static GemmCoord get_tile_offset(int log_tile) {
return GemmCoord{
RematerializeBlockIdxY(),
RematerializeBlockIdxX(),
@@ -377,7 +377,7 @@ struct GemmSplitKHorizontalThreadblockSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
GemmCoord get_tile_offset(GemmCoord tiled_shape) const {
static GemmCoord get_tile_offset(GemmCoord tiled_shape) {
return GemmCoord{
RematerializeBlockIdxY(),
RematerializeBlockIdxX(),
@@ -393,9 +393,9 @@ struct GemvBatchedStridedThreadblockDefaultSwizzle {
/// Returns the shape of the problem in units of logical tiles
CUTLASS_HOST_DEVICE
BatchedGemmCoord get_tiled_shape(
static BatchedGemmCoord get_tiled_shape(
BatchedGemmCoord problem_size,
BatchedGemmCoord tile_size) const {
BatchedGemmCoord tile_size) {
return BatchedGemmCoord(
1, // M is always 1
@@ -406,7 +406,7 @@ struct GemvBatchedStridedThreadblockDefaultSwizzle {
/// Computes CUDA grid dimensions given a size in units of logical tiles
CUTLASS_HOST_DEVICE
dim3 get_grid_shape(BatchedGemmCoord tiled_shape) const {
static dim3 get_grid_shape(BatchedGemmCoord tiled_shape) {
return dim3(tiled_shape.n(), tiled_shape.batch(), tiled_shape.k());
}
@@ -418,7 +418,7 @@ struct GemvBatchedStridedThreadblockDefaultSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
BatchedGemmCoord get_tile_offset(int log_tile) const {
static BatchedGemmCoord get_tile_offset(int log_tile) {
return BatchedGemmCoord{
0, // M is always 1
RematerializeBlockIdxX(),
@@ -429,7 +429,7 @@ struct GemvBatchedStridedThreadblockDefaultSwizzle {
/// Obtains the threadblock offset (in units of threadblock-scoped tiles)
CUTLASS_DEVICE
BatchedGemmCoord get_tile_offset() const {
static BatchedGemmCoord get_tile_offset() {
return BatchedGemmCoord{
0, // M is always 1
RematerializeBlockIdxX(),
@@ -440,13 +440,13 @@ struct GemvBatchedStridedThreadblockDefaultSwizzle {
/// Gets the batch tile index
CUTLASS_DEVICE
int get_batch_tile_idx() const {
static int get_batch_tile_idx() {
return RematerializeBlockIdxY();
}
/// Gets the absolute batch index
CUTLASS_DEVICE
int get_batch_idx() const {
static int get_batch_idx() {
return RematerializeBlockDimY()*RematerializeBlockIdxY() + RematerializeThreadIdxY();
}
};
include/cutlass/gemm/threadblock/threadblock_swizzle_streamk.h
+11 -1
View File
@@ -32,13 +32,23 @@
\brief Implements streamk threadblock mapping blockIdx to GEMM problems.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/platform/platform.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/gemm_enumerated_types.h"
#include "cutlass/conv/conv2d_problem_size.h"
#include "cutlass/conv/conv3d_problem_size.h"
#include "cutlass/gemm/threadblock/index_remat.h"
include/cutlass/gemm_coord.hpp
+65
View File
@@ -0,0 +1,65 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*! \file
\brief Utilities to convert a CuTe tuple to a GemmCoord or BatchedGemmCoord
*/
#pragma once
#include "cutlass/gemm_coord.h"
namespace cutlass {
namespace gemm {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <class Tuple>
CUTLASS_HOST_DEVICE
auto
to_gemm_coord(Tuple tuple) {
static_assert(cute::rank(tuple) <= 4, "Can only convert tuples of rank <= 4.");
if constexpr (cute::rank(tuple) <= 3) {
auto tuple_mnk = cute::append<3>(tuple, cute::Int<0>{});
return GemmCoord(cute::size<0>(tuple_mnk), cute::size<1>(tuple_mnk), cute::size<2>(tuple_mnk));
}
else {
return BatchedGemmCoord(cute::size<0>(tuple), cute::size<1>(tuple), cute::size<2>(tuple), cute::size<3>(tuple));
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/half.h
+11 -1
View File
@@ -33,6 +33,17 @@
\brief Defines a class for using IEEE half-precision floating-point types in host or
device code.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#ifndef CUTLASS_ENABLE_F16C
@@ -158,7 +169,6 @@ public:
///////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
///////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/integer_subbyte.h
+14
View File
@@ -33,6 +33,17 @@
\brief Defines a class for using integer types smaller than one byte in host or
device code.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#if defined(__CUDACC_RTC__)
@@ -41,6 +52,8 @@
#include <cstdint>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_size.h"
#include "cutlass/platform/platform.h"
namespace cutlass {
@@ -238,3 +251,4 @@ struct numeric_limits<cutlass::uint1b_t> {
} // namespace platform
} // namespace cutlass
include/cutlass/kernel_hardware_info.h
+86
View File
@@ -0,0 +1,86 @@
/***************************************************************************************************
* 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
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by this unit test: `cutlass_test_unit_core_cpp11`.
*/
#if !defined(__CUDACC_RTC__)
#include "cuda_runtime.h"
#include "cutlass/trace.h"
#endif
namespace cutlass {
struct KernelHardwareInfo {
//
// Data members
//
int device_id = 0;
int sm_count = 0;
//
// Methods
//
#if !defined(__CUDACC_RTC__)
static inline int
query_device_multiprocessor_count(int device_id = 0) {
cudaError_t result = cudaGetDevice(&device_id);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(
" cudaGetDevice() returned error "
<< cudaGetErrorString(result));
return 0;
}
int multiprocessor_count;
result = cudaDeviceGetAttribute(&multiprocessor_count,
cudaDevAttrMultiProcessorCount, device_id);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(
" cudaDeviceGetAttribute() returned error "
<< cudaGetErrorString(result));
return 0;
}
return multiprocessor_count;
}
#endif
};
} // namespace cutlass
include/cutlass/kernel_hardware_info.hpp
+3 -44
View File
@@ -30,47 +30,6 @@
**************************************************************************************************/
#pragma once
#if !defined(__CUDACC_RTC__)
#include "cuda_runtime.h"
#include "cutlass/trace.h"
#endif
namespace cutlass {
struct KernelHardwareInfo {
//
// Data members
//
int device_id = 0;
int sm_count = 0;
//
// Methods
//
#if !defined(__CUDACC_RTC__)
static int
query_device_multiprocessor_count(int device_id = 0) {
cudaError_t result = cudaGetDevice(&device_id);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(
" cudaGetDevice() returned error "
<< cudaGetErrorString(result));
return 0;
}
int multiprocessor_count;
result = cudaDeviceGetAttribute(&multiprocessor_count,
cudaDevAttrMultiProcessorCount, device_id);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(
" cudaDeviceGetAttribute() returned error "
<< cudaGetErrorString(result));
return 0;
}
return multiprocessor_count;
}
#endif
};
} // namespace cutlass
// Simply import .h version of header so as to avoid breaking any existing CUTLASS builds
// after .hpp was changed to .h
#include "cutlass/kernel_hardware_info.h"
include/cutlass/layout/matrix.h
+14
View File
@@ -37,6 +37,17 @@
Layout functions must implement all members in the public interface of IdentityTensorLayout<>
defined in cutlass/tensor_ref.h.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by this unit test: `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
@@ -538,6 +549,7 @@ public:
/// Inverse of layout function, mapping linear offset to logical coordinate
CUTLASS_HOST_DEVICE
MatrixCoord inverse(LongIndex offset) const {
CUTLASS_UNUSED(offset);
return MatrixCoord(0, 0);
}
@@ -796,6 +808,7 @@ public:
/// Inverse of layout function, mapping linear offset to logical coordinate
CUTLASS_HOST_DEVICE
MatrixCoord inverse(LongIndex offset) const {
CUTLASS_UNUSED(offset);
return MatrixCoord(0, 0);
}
@@ -901,6 +914,7 @@ public:
/// Inverse of layout function, mapping linear offset to logical coordinate
CUTLASS_HOST_DEVICE
MatrixCoord inverse(LongIndex offset) const {
CUTLASS_UNUSED(offset);
return MatrixCoord(0, 0);
}
include/cutlass/layout/pitch_linear.h
+11
View File
@@ -31,6 +31,17 @@
/*! \file
\brief Defines layout functions used by TensorRef and derived classes for pitch-linear memory.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by this unit test: `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
include/cutlass/layout/vector.h
+1
View File
@@ -78,6 +78,7 @@ public:
/// Helper returns a layout to a tightly packed tensor
CUTLASS_HOST_DEVICE
static PackedVectorLayout packed(TensorCoord const &size) {
CUTLASS_UNUSED(size);
return PackedVectorLayout();
}
include/cutlass/numeric_size.h
+93
View File
@@ -0,0 +1,93 @@
/***************************************************************************************************
* Copyright (c) 2017 - 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.
*
**************************************************************************************************/
/*!
\file
\brief Top-level include for all CUTLASS numeric types.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines the size of an element in bits
template <typename T>
struct sizeof_bits {
static int const value = int(sizeof(T) * 8);
};
template <typename T>
struct sizeof_bits<T const>: sizeof_bits<T> {};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Definitions for 1-bit binary and 4-bit integer types
//
/// 1-bit binary type
using bin1_t = bool;
/// Defines the size of an element in bits - specialized for bin1_t
template <>
struct sizeof_bits<bin1_t> {
static int const value = 1;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Returns the number of bytes required to hold a specified number of bits
CUTLASS_HOST_DEVICE
constexpr int
bits_to_bytes(int bits) {
return (bits + 7) / 8;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/numeric_types.h
+12 -35
View File
@@ -32,49 +32,26 @@
\file
\brief Top-level include for all CUTLASS numeric types.
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/numeric_size.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines the size of an element in bits
template <typename T>
struct sizeof_bits {
static int const value = int(sizeof(T) * 8);
};
template <typename T>
struct sizeof_bits<T const>: sizeof_bits<T> {};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Definitions for 1-bit binary and 4-bit integer types
//
/// 1-bit binary type
using bin1_t = bool;
/// Defines the size of an element in bits - specialized for bin1_t
template <>
struct sizeof_bits<bin1_t> {
static int const value = 1;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Returns the number of bytes required to hold a specified number of bits
CUTLASS_HOST_DEVICE
constexpr int
bits_to_bytes(int bits) {
return (bits + 7) / 8;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
template <size_t... Seq>
struct index_sequence;
include/cutlass/pipeline/sm90_pipeline.hpp
+47 -8
View File
@@ -219,7 +219,6 @@ public :
ThreadCategory role = ThreadCategory::NonParticipant;
uint32_t is_leader = 0;
uint32_t num_consumers = 0;
cute::tuple<int, int> active_warps = {0, 0};
};
// Constructor
@@ -232,7 +231,7 @@ public :
int warp_idx = canonical_warp_idx();
int lane_predicate = cute::elect_one_sync();
auto cluster_shape = ClusterShape{};
if (warp_idx == cute::get<0>(params.active_warps) && lane_predicate == 1) {
if (warp_idx == 0 && lane_predicate == 1) {
// Barrier FULL init
for (int i = 0; i < Stages; ++i) {
full_barrier_ptr_[i].init(1);
@@ -350,6 +349,11 @@ public :
return consumer_try_wait(state.index(), state.phase(), skip_wait);
}
CUTLASS_DEVICE
ConsumerToken consumer_test_wait(PipelineState state, uint32_t skip_wait = false) {
return consumer_test_wait(state.index(), state.phase(), skip_wait);
}
CUTLASS_DEVICE
void consumer_wait(PipelineState state) {
consumer_wait(state.index(), state.phase());
@@ -440,6 +444,15 @@ private :
return {static_cast<BarrierStatus>(barrier_status)};
}
CUTLASS_DEVICE
ConsumerToken consumer_test_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
if (skip_wait) {
return {BarrierStatus::WaitDone};
}
uint32_t barrier_status = full_barrier_ptr_[stage].test_wait(phase);
return {static_cast<BarrierStatus>(barrier_status)};
}
// Wait for producer to commit transactions (done by TMA)
CUTLASS_DEVICE
void consumer_wait(uint32_t stage, uint32_t phase) {
@@ -621,7 +634,6 @@ public :
uint32_t producer_arv_count = 1;
uint32_t consumer_arv_count = 1;
uint32_t dst_blockid = cute::block_rank_in_cluster();
cute::tuple<int, int> active_warps = {0, 0};
};
// Constructor
@@ -636,7 +648,7 @@ public :
// Barrier FULL, EMPTY init
// Init is done only by thread 0 of the block
if (warp_idx == cute::get<0>(params.active_warps) && lane_predicate == 1) {
if (warp_idx == 0 && lane_predicate == 1) {
for (int i = 0; i < Stages; ++i) {
full_barrier_ptr_[i].init(params.producer_arv_count);
empty_barrier_ptr_[i].init(params.consumer_arv_count);
@@ -708,6 +720,11 @@ public :
return consumer_try_wait(state.index(), state.phase(), skip_wait);
}
CUTLASS_DEVICE
ConsumerToken consumer_test_wait(PipelineState state, uint32_t skip_wait = false) {
return consumer_test_wait(state.index(), state.phase(), skip_wait);
}
CUTLASS_DEVICE
void consumer_wait(PipelineState state, ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) {
consumer_wait(state.index(), state.phase(), barrier_token);
@@ -764,6 +781,15 @@ private:
return {static_cast<BarrierStatus>(barrier_status)};
}
CUTLASS_DEVICE
ConsumerToken consumer_test_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
if (skip_wait) {
return {BarrierStatus::WaitDone};
}
uint32_t barrier_status = full_barrier_ptr_[stage].test_wait(phase);
return {static_cast<BarrierStatus>(barrier_status)};
}
CUTLASS_DEVICE
void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) {
if (barrier_token == BarrierStatus::WaitAgain) {
@@ -809,7 +835,6 @@ public :
uint32_t producer_arv_count = 1;
uint32_t consumer_arv_count = 1;
uint32_t dst_blockid = cute::block_rank_in_cluster();
cute::tuple<int, int> active_warps = {0, 0};
};
// Default assumption when only storage is passed is :
@@ -831,7 +856,7 @@ public :
// Barrier FULL, EMPTY init
// Init is done only by thread 0 of the block
if (warp_idx == cute::get<0>(params.active_warps) && lane_predicate == 1) {
if (warp_idx == 0 && lane_predicate == 1) {
for (int i = 0; i < Stages; ++i) {
full_barrier_ptr_[i].init(params.producer_arv_count);
empty_barrier_ptr_[i].init(params.consumer_arv_count);
@@ -897,6 +922,11 @@ public :
return consumer_try_wait(state.index(), state.phase(), skip_wait);
}
CUTLASS_DEVICE
ConsumerToken consumer_test_wait(PipelineState state, uint32_t skip_wait = false) {
return consumer_test_wait(state.index(), state.phase(), skip_wait);
}
CUTLASS_DEVICE
void consumer_wait(PipelineState state, ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) {
consumer_wait(state.index(), state.phase(), barrier_token);
@@ -947,6 +977,15 @@ private:
return {static_cast<BarrierStatus>(barrier_status)};
}
CUTLASS_DEVICE
ConsumerToken consumer_test_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
if (skip_wait) {
return {BarrierStatus::WaitDone};
}
uint32_t barrier_status = full_barrier_ptr_[stage].test_wait(phase);
return {static_cast<BarrierStatus>(barrier_status)};
}
CUTLASS_DEVICE
void consumer_wait(uint32_t stage, uint32_t phase) {
uint32_t done = full_barrier_ptr_[stage].test_wait(phase);
@@ -968,6 +1007,7 @@ private:
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
//
// Barrier to ensure an Ordered Sequence between
@@ -989,7 +1029,6 @@ public :
struct Params {
uint32_t group_id;
uint32_t group_size;
cute::tuple<int, int> active_warps = {0, 0};
};
private :
@@ -1020,7 +1059,7 @@ public:
// Barrier FULL, EMPTY init
// Init is done only by the one elected thread of the block
if (warp_idx == cute::get<0>(params.active_warps) && lane_predicate == 1) {
if (warp_idx == 0 && lane_predicate == 1) {
for (int d = 0; d < Depth; ++d) {
for (int l = 0; l < Length; ++l) {
barrier_ptr_[d * Length + l].init(params.group_size);
include/cutlass/platform/platform.h
+23 -2
View File
@@ -95,6 +95,16 @@
* counterparts (or trivially find-and-replace their occurrences in code text).
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
*/
//-----------------------------------------------------------------------------
// Dependencies
//-----------------------------------------------------------------------------
@@ -354,7 +364,7 @@ using std::nullptr_t;
// Conditional metaprogramming <type_traits>
//-----------------------------------------------------------------------------
#if defined(__CUDACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1600))
#if defined(__CUDACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201700L)) || (defined(_MSC_VER) && (_MSC_VER < 1600))
/// std::enable_if (true specialization)
template <bool C, typename T = void>
@@ -385,14 +395,16 @@ using std::conditional;
#endif
#if (201703L <=__cplusplus)
/// std::conditional_t
using CUTLASS_STL_NAMESPACE::conditional_t;
#endif
//-----------------------------------------------------------------------------
// Const/volatility specifiers <type_traits>
//-----------------------------------------------------------------------------
#if defined(__CUDACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
#if defined(__CUDACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201703L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
/// std::remove_const (non-const specialization)
template <typename T>
@@ -432,6 +444,8 @@ using std::remove_cv;
#endif
#if (201703L <=__cplusplus)
/// std::remove_cv_t
using CUTLASS_STL_NAMESPACE::remove_cv_t;
/// std::remove_reference_t
@@ -449,6 +463,9 @@ struct remove_cvref {
template <class T>
using remove_cvref_t = typename remove_cvref<T>::type;
#endif
//-----------------------------------------------------------------------------
// Type relationships <type_traits>
//-----------------------------------------------------------------------------
@@ -613,11 +630,15 @@ using std::is_trivially_copyable;
#endif
#if (201703L <=__cplusplus)
/// std::is_unsigned_v
using CUTLASS_STL_NAMESPACE::is_integral_v;
/// std::is_unsigned_v
using CUTLASS_STL_NAMESPACE::is_unsigned_v;
#endif
//-----------------------------------------------------------------------------
// bit_cast <bit>
//-----------------------------------------------------------------------------

Some files were not shown because too many files have changed in this diff Show More