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

3.6.0 update (#2005)

Browse Source 
* 3.6.0 update

* doc and swap stuff

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-12-25 01:34:40 -05:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
parent e1cd8c7866
commit 3d261a5974
258 changed files with 10863 additions and 3883 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/cooperative_copy.hpp
+31 -25
View File
@@ -51,19 +51,14 @@ naive_cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
auto N = size(src);
if (tid < N) {
uint32_t upper_bound = (N / NumThreads) * NumThreads;
CUTE_UNROLL
for (uint32_t i = 0; i < upper_bound; i += NumThreads) { // All in-bounds
dst[tid + i] = src[tid + i];
}
if (N % NumThreads != 0) { // Likely static condition
uint32_t final_idx = tid + upper_bound;
if (final_idx < N) { // Final in-bounds
dst[final_idx] = src[final_idx];
}
}
auto N = size(dst);
auto R = N % Int<NumThreads>{};
if (R > 0 && tid < R) { // Likely static condition && Residue in-bounds
dst[tid] = src[tid];
}
CUTE_UNROLL
for (uint32_t i = uint32_t(R); i < uint32_t(N); i += NumThreads) { // All in-bounds
dst[tid + i] = src[tid + i];
}
}
@@ -117,12 +112,14 @@ heuristic_permutation(Tensor<AEngine, ALayout> const& a,
//
template <uint32_t NumThreads, uint32_t MaxVecBits,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
class DstEngine, class DstLayout,
class CopyPolicy = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
Tensor<DstEngine, DstLayout> & dst,
CopyPolicy const& cpy = {})
{
// Assumes the shapes are static, can generalize/fallback
CUTE_STATIC_ASSERT_V(is_static<decltype(shape(src))>{} && is_static<decltype(shape(dst))>{});
@@ -283,23 +280,28 @@ cooperative_copy(uint32_t const& tid,
// If we're using all threads (static) or the tid is in-range (dynamic)
if (vec_thrs == NumThreads or tid < vec_thrs) {
return copy_if(TrivialPredTensor{}, recast<VecType const>(src_v), recast<VecType>(dst_v));
auto src_c = recast<VecType const>(src_v);
auto dst_c = recast<VecType>(dst_v);
return copy(cpy, src_c, dst_c);
}
}
}
// Default max-vectorization size to value_type size
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
class DstEngine, class DstLayout,
class CopyPolicy = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
Tensor<DstEngine, DstLayout> & dst,
CopyPolicy const& cpy = {})
{
constexpr uint32_t MaxVecBits = sizeof_bits_v<typename SrcEngine::value_type>;
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst);
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst, cpy);
}
//
@@ -308,26 +310,30 @@ cooperative_copy(uint32_t const& tid,
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
class DstEngine, class DstLayout,
class CopyPolicy = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
Tensor<DstEngine, DstLayout> && dst,
CopyPolicy const& cpy = {})
{
return cooperative_copy<NumThreads>(tid, src, dst);
return cooperative_copy<NumThreads>(tid, src, dst, cpy);
}
template <uint32_t NumThreads, uint32_t MaxVecBits,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
class DstEngine, class DstLayout,
class CopyPolicy = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
Tensor<DstEngine, DstLayout> && dst,
CopyPolicy const& cpy = {})
{
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst);
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst, cpy);
}
} // end namespace cute
include/cute/algorithm/cooperative_gemm.hpp
+320 -247
View File
@@ -50,31 +50,115 @@ namespace cute
namespace detail {
// Predicated Cooperative GEMM
template <class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
class ALoadTransformOp, class BLoadTransformOp,
class CLoadTransformOp, class CStoreTransformOp,
__CUTE_REQUIRES(ALayout::rank == 2 && is_smem<TA>::value &&
BLayout::rank == 2 && is_smem<TB>::value &&
CLayout::rank == 2 && is_smem<TC>::value)>
// Slow fallback path:
template<typename ... Args,
typename Alpha, typename TRC, typename RCLayout,
typename Beta, class TSC, typename CLayout, typename SCLayout,
typename CLoadTransformOp, typename CStoreTransformOp>
CUTE_HOST_DEVICE
void
cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
Alpha const& alpha,
Tensor<TA, ALayout> sA,
Tensor<TB, BLayout> sB,
Beta const& beta,
Tensor<TC, CLayout> sC,
ALoadTransformOp const& sA_load_op, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op) // transforms results before they are stored to C
epilogue_predication(ThrMMA<Args...> const& thr_mma,
Alpha const& alpha,
Tensor<TRC, RCLayout> & tCrC,
Beta const& beta,
Tensor<TSC, CLayout> & sC,
Tensor<TSC, SCLayout> & tCsC,
CLoadTransformOp const& sC_load_op, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op) // transforms results before they are stored to C
{
using TypeA = typename TA::value_type;
using TypeB = typename TB::value_type;
using TypeC = typename TC::value_type;
using InputTypeC = typename TSC::value_type;
using ComputeTypeC = typename ThrMMA<Args...>::ValTypeC;
CUTE_STATIC_ASSERT(CUTE_STL_NAMESPACE::is_same_v<ComputeTypeC, typename TRC::value_type>);
// Create coordinate tensors for the problem
Tensor cC = make_identity_tensor(shape(sC)); // (M,N) -> (m,n)
// Repeat partitioning with thr_mma
Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) -> (m,n)
const bool isBetaZero = [&] () {
if constexpr (is_complex<Beta>::value) {
return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
}
else {
return beta == Int<0>{};
}
CUTE_GCC_UNREACHABLE;
} ();
// Custom axpby_if for now
CUTE_UNROLL
for (int i = 0; i < size(tCrC); ++i)
{
if (elem_less(tCcC(i), shape(sC)))
{
tCsC(i) = sC_store_op(isBetaZero ? alpha * tCrC(i)
: alpha * tCrC(i) +
beta * static_cast<ComputeTypeC>(sC_load_op(tCsC(i))));
}
}
}
template<class Alpha, class TRC, class RCLayout,
class Beta, class TSC, class SCLayout,
class CLoadTransformOp, class CStoreTransformOp,
class SmemCopyOpC>
CUTE_HOST_DEVICE
void
epilogue_no_predication(Alpha const& alpha,
Tensor<TRC, RCLayout> & tCrC,
Beta const& beta,
Tensor<TSC, SCLayout> & tCsC,
CLoadTransformOp const& sC_load_op, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op, // transforms results before they are stored to C
SmemCopyOpC const& sC_copy_op)
{
using InputTypeC = typename TSC::value_type;
using ComputeTypeC = typename TRC::value_type;
const bool isBetaZero = [&] () {
if constexpr (is_complex<Beta>::value) {
return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
}
else {
return beta == Int<0>{};
}
CUTE_GCC_UNREACHABLE;
} ();
Tensor tCrDi = make_fragment_like(tCsC);
Tensor tCrD = make_fragment_like(tCrC);
if(!isBetaZero) {
copy(sC_copy_op, tCsC, tCrDi);
// Transform C on/after load
cute::transform(tCrDi, tCrD, sC_load_op);
}
// C = alpha * (A * B) + beta * C
axpby(alpha, tCrC, beta, tCrD);
// Transform C before/on store
cute::transform(tCrD, tCrDi, sC_store_op);
copy(sC_copy_op, tCrDi, tCsC);
}
// Predicated Cooperative GEMM
template <class... Args,
class TA, class ALayout, class TB, class BLayout,
class TC, class RCLayout,
class ALoadTransformOp, class BLoadTransformOp>
CUTE_HOST_DEVICE
void
cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
Tensor<TA, ALayout> const& sA,
Tensor<TB, BLayout> const& sB,
Tensor<TC, RCLayout> & tCrC,
ALoadTransformOp const& sA_load_op, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op) // transforms B values before use in GEMM
{
using InputTypeA = typename TA::value_type;
using InputTypeB = typename TB::value_type;
using InputTypeC = typename TC::value_type;
using ComputeTypeA = typename ThrMMA<Args...>::ValTypeA;
using ComputeTypeB = typename ThrMMA<Args...>::ValTypeB;
using ComputeTypeC = typename ThrMMA<Args...>::ValTypeC;
//
// MMA Partitioning
@@ -83,22 +167,18 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
// Partition the sA, sB, and sC tiles across the threads for the MMA
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K)
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K)
Tensor tCsC = thr_mma.partition_C(sC); // (MMA,MMA_M,MMA_N)
// Create register tensors for the MMA to operate on
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K)
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K)
Tensor tCrC = thr_mma.make_fragment_C(tCsC); // (MMA,MMA_M,MMA_N)
#if 0
if (thread0()) {
print(" sA: "); print( sA); print("\n");
print(" sB: "); print( sB); print("\n");
print(" sC: "); print( sC); print("\n");
print(thr_mma);
print("tCsA: "); print(tCsA); print("\n");
print("tCsB: "); print(tCsB); print("\n");
print("tCsC: "); print(tCsC); print("\n");
print("tCrA: "); print(tCrA); print("\n");
print("tCrB: "); print(tCrB); print("\n");
print("tCrC: "); print(tCrC); print("\n");
@@ -154,23 +234,20 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
for (int m = 0; m < size<1>(tCrA); ++m) { // Copy MMA_M
CUTE_UNROLL
for (int i = 0; i < size<0>(tCrA); ++i) { // Copy MMA_I
tCrA(i,m,0) = (tCpA(i,m) && (0 < K_BLOCK_MAX-1 || elem_less(get<1>(tCcA(i,m,0)), shape<1>(sA)))) ? sA_load_op(tCsA(i,m,0)) : TypeA{};
tCrA(i,m,0) = (tCpA(i,m) && (0 < K_BLOCK_MAX-1 || elem_less(get<1>(tCcA(i,m,0)), shape<1>(sA)))) ? static_cast<ComputeTypeA>(sA_load_op(tCsA(i,m,0))) : ComputeTypeA{};
}
}
CUTE_UNROLL
for (int n = 0; n < size<1>(tCrB); ++n) { // Copy MMA_N
CUTE_UNROLL
for (int i = 0; i < size<0>(tCrB); ++i) { // Copy MMA_I
tCrB(i,n,0) = (tCpB(i,n) && (0 < K_BLOCK_MAX-1 || elem_less(get<1>(tCcB(i,n,0)), shape<1>(sB)))) ? sB_load_op(tCsB(i,n,0)) : TypeB{};
tCrB(i,n,0) = (tCpB(i,n) && (0 < K_BLOCK_MAX-1 || elem_less(get<1>(tCcB(i,n,0)), shape<1>(sB)))) ? static_cast<ComputeTypeB>(sB_load_op(tCsB(i,n,0))) : ComputeTypeB{};
}
}
//
// MAINLOOP
//
// Clear accumulators
clear(tCrC);
CUTE_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block)
{
@@ -185,138 +262,80 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
for (int m = 0; m < size<1>(tCrA); ++m) { // Copy MMA_M
CUTE_UNROLL
for (int i = 0; i < size<0>(tCrA); ++i) { // Copy MMA_I
tCrA(i,m,k_next) = (tCpA(i,m) && (k_next < K_BLOCK_MAX-1 || elem_less(get<1>(tCcA(i,m,k_next)), shape<1>(sA)))) ? sA_load_op(tCsA(i,m,k_next)) : TypeA{};
tCrA(i,m,k_next) = (tCpA(i,m) && (k_next < K_BLOCK_MAX-1 || elem_less(get<1>(tCcA(i,m,k_next)), shape<1>(sA)))) ? static_cast<ComputeTypeA>(sA_load_op(tCsA(i,m,k_next))) : ComputeTypeA{};
}
}
CUTE_UNROLL
for (int n = 0; n < size<1>(tCrB); ++n) { // Copy MMA_N
CUTE_UNROLL
for (int i = 0; i < size<0>(tCrB); ++i) { // Copy MMA_I
tCrB(i,n,k_next) = (tCpB(i,n) && (k_next < K_BLOCK_MAX-1 || elem_less(get<1>(tCcB(i,n,k_next)), shape<1>(sB)))) ? sB_load_op(tCsB(i,n,k_next)) : TypeB{};
tCrB(i,n,k_next) = (tCpB(i,n) && (k_next < K_BLOCK_MAX-1 || elem_less(get<1>(tCcB(i,n,k_next)), shape<1>(sB)))) ? static_cast<ComputeTypeB>(sB_load_op(tCsB(i,n,k_next))) : ComputeTypeB{};
}
}
}
// GEMM on k_block in registers
gemm(thr_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
}
//
// Epilogue
//
// Create coordinate tensors for the problem
Tensor cC = make_identity_tensor(shape(sC)); // (M,N) -> (m,n)
// Repeat partitioning with thr_mma
Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) -> (m,n)
const bool isBetaZero = (beta == Beta{});
// Custom axpby_if for now
CUTE_UNROLL
for (int i = 0; i < size(tCrC); ++i)
{
if (elem_less(tCcC(i), shape(sC)))
{
tCsC(i) = sC_store_op(isBetaZero ? alpha * static_cast<TypeC>(tCrC(i))
: alpha * static_cast<TypeC>(tCrC(i)) +
beta * static_cast<TypeC>(sC_load_op(tCsC(i))));
}
}
}
// Slow fallback path
template <class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
class ALoadTransformOp, class BLoadTransformOp,
class CLoadTransformOp, class CStoreTransformOp,
__CUTE_REQUIRES(ALayout::rank == 2 && is_smem<TA>::value &&
BLayout::rank == 2 && is_smem<TB>::value &&
CLayout::rank == 2 && is_smem<TC>::value)>
CUTE_HOST_DEVICE
void
cooperative_gemm_predication(uint32_t thread_idx,
TiledMMA<Args...> const& tiled_mma,
Alpha const& alpha,
Tensor<TA, ALayout> sA,
Tensor<TB, BLayout> sB,
Beta const& beta,
Tensor<TC, CLayout> sC,
ALoadTransformOp const& sA_load_op, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op) // transforms results before they are stored to C
{
// ThrMMA
auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
cooperative_gemm_predication(thr_mma, alpha, sA, sB, beta, sC, sA_load_op, sB_load_op, sC_load_op, sC_store_op);
}
// Unpredicated Cooperative GEMM
template <class SmemCopyOpA, class SmemCopyOpB, class SmemCopyOpC,
class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
template <class... Args,
class TA, class ALayout, class TB, class BLayout,
class TC, class CLayout,
class ALoadTransformOp, class BLoadTransformOp,
class CLoadTransformOp, class CStoreTransformOp,
__CUTE_REQUIRES(ALayout::rank == 2 && is_smem<TA>::value &&
BLayout::rank == 2 && is_smem<TB>::value &&
CLayout::rank == 2 && is_smem<TC>::value)>
class SmemCopyOpA, class SmemCopyOpB>
CUTE_HOST_DEVICE
void
cooperative_gemm_no_predication(uint32_t thread_idx,
TiledMMA<Args...> const& tiled_mma,
Alpha const& alpha,
Tensor<TA, ALayout> sA,
Tensor<TB, BLayout> sB,
Beta const& beta,
Tensor<TC, CLayout> sC,
ALoadTransformOp const& sA_load_op, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op) // transforms results before they are stored to C
cooperative_gemm_no_predication(uint32_t thread_idx,
ThrMMA<Args...> const& thr_mma,
Tensor<TA, ALayout> const& sA,
Tensor<TB, BLayout> const& sB,
Tensor<TC, CLayout> & tCrC,
ALoadTransformOp const& sA_load_op, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op, // transforms B values before use in GEMM
SmemCopyOpA const& sA_copy_op,
SmemCopyOpB const& sB_copy_op)
{
using TypeA = typename TA::value_type;
using TypeB = typename TB::value_type;
using TypeC = typename TC::value_type;
using InputTypeA = typename TA::value_type;
using InputTypeB = typename TB::value_type;
using InputTypeC = typename TC::value_type;
using ComputeTypeA = typename ThrMMA<Args...>::ValTypeA;
using ComputeTypeB = typename ThrMMA<Args...>::ValTypeB;
using ComputeTypeC = typename ThrMMA<Args...>::ValTypeC;
// ThrMMA
auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
//
// MMA Partitioning
//
Tensor tCsC = thr_mma.partition_C(sC);
// Create register tensors for the MMA to operate on
Tensor tCrA = thr_mma.partition_fragment_A(sA); // (MMA,MMA_M,MMA_K)
Tensor tCrB = thr_mma.partition_fragment_B(sB); // (MMA,MMA_N,MMA_K)
Tensor tCrC = thr_mma.make_fragment_C(tCsC); // (MMA,MMA_M,MMA_N)
using CopyOpAType = SmemCopyOpA;
using CopyOpBType = SmemCopyOpB;
auto smem_tiled_copy_A = make_tiled_copy_A(Copy_Atom<CopyOpAType, TypeA>{}, thr_mma);
auto smem_tiled_copy_A = make_tiled_copy_A(Copy_Atom<CopyOpAType, InputTypeA>{}, thr_mma);
auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx);
Tensor tCsA = smem_thr_copy_A.partition_S(sA);
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K
Tensor tCrAi = make_fragment_like(tCsA);
Tensor tCrAi_copy_view = smem_thr_copy_A.retile_D(tCrAi);
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrAi_copy_view)); // CPY_M
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrAi_copy_view)); // CPY_K
auto smem_tiled_copy_B = make_tiled_copy_B(Copy_Atom<CopyOpBType, TypeB>{}, thr_mma);
auto smem_tiled_copy_B = make_tiled_copy_B(Copy_Atom<CopyOpBType, InputTypeB>{}, thr_mma);
auto smem_thr_copy_B = smem_tiled_copy_B.get_thread_slice(thread_idx);
Tensor tCsB = smem_thr_copy_B.partition_S(sB);
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // CPY_N
CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCrB_copy_view)); // CPY_K
Tensor tCrBi = make_fragment_like(tCsB);
Tensor tCrBi_copy_view = smem_thr_copy_B.retile_D(tCrBi);
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrBi_copy_view)); // CPY_N
CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCrBi_copy_view)); // CPY_K
#if 0
if (thread0()) {
print(" sA: "); print(sA); print("\n");
print(" sB: "); print(sB); print("\n");
print(" sC: "); print(sC); print("\n");
print(thr_mma); print("\n");
print("tCsC: "); print(tCsC); print("\n");
print("tCrA: "); print(tCrA); print("\n");
print("tCrB: "); print(tCrB); print("\n");
print("tCrC: "); print(tCrC); print("\n");
@@ -333,15 +352,12 @@ cooperative_gemm_no_predication(uint32_t thread_idx,
// PREFETCH
//
copy(smem_tiled_copy_A, tCsA(_,_,Int<0>{}), tCrA_copy_view(_,_,Int<0>{}));
copy(smem_tiled_copy_B, tCsB(_,_,Int<0>{}), tCrB_copy_view(_,_,Int<0>{}));
copy(smem_tiled_copy_A, tCsA(_,_,Int<0>{}), tCrAi_copy_view(_,_,Int<0>{}));
copy(smem_tiled_copy_B, tCsB(_,_,Int<0>{}), tCrBi_copy_view(_,_,Int<0>{}));
//
// MAINLOOP
//
// Clear accumulators
clear(tCrC);
constexpr int K_BLOCK_MAX = size<2>(tCrA);
CUTE_UNROLL
@@ -352,132 +368,178 @@ cooperative_gemm_no_predication(uint32_t thread_idx,
{
// Load the next k_block
int k_next = k_block + 1; // statically unrolled
copy(smem_tiled_copy_A, tCsA(_,_,k_next), tCrA_copy_view(_,_,k_next));
copy(smem_tiled_copy_B, tCsB(_,_,k_next), tCrB_copy_view(_,_,k_next));
copy(smem_tiled_copy_A, tCsA(_,_,k_next), tCrAi_copy_view(_,_,k_next));
copy(smem_tiled_copy_B, tCsB(_,_,k_next), tCrBi_copy_view(_,_,k_next));
}
// Transform A and B, relying on the compiler to remove in case of identity ops
cute::transform(tCrA(_,_,k_block), sA_load_op);
cute::transform(tCrB(_,_,k_block), sB_load_op);
cute::transform(tCrAi(_,_,k_block), tCrA(_,_,k_block), sA_load_op);
cute::transform(tCrBi(_,_,k_block), tCrB(_,_,k_block), sB_load_op);
// GEMM on k_block in registers
gemm(thr_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
}
//
// Epilogue
//
auto isBetaZero = [&] () {
if constexpr (is_complex<Beta>::value) {
return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
}
else {
return beta == Int<0>{};
}
CUTE_GCC_UNREACHABLE;
} ();
using CopyOpCType = SmemCopyOpC;
Tensor tCrD = thr_mma.make_fragment_C(tCsC);
if(!isBetaZero) {
copy(CopyOpCType{}, tCsC, tCrD);
// Transform C on/after load
cute::transform(tCrD, sC_load_op);
}
// C = alpha * (A * B) + beta * C
axpby(alpha, tCrC, beta, tCrD);
// Transform C before/on store
cute::transform(tCrD, sC_store_op);
copy(CopyOpCType{}, tCrD, tCsC);
}
} // end namespace detail
template <class SmemCopyOpA, class SmemCopyOpB, class SmemCopyOpC,
class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
class ALoadTransformOp = cute::identity, class BLoadTransformOp = cute::identity,
class CLoadTransformOp = cute::identity, class CStoreTransformOp = cute::identity,
__CUTE_REQUIRES(ALayout::rank == 2 && is_smem<TA>::value &&
BLayout::rank == 2 && is_smem<TB>::value &&
CLayout::rank == 2 && is_smem<TC>::value)>
CUTE_HOST_DEVICE
void
cooperative_gemm(uint32_t thread_idx,
TiledMMA<Args...> const& tiled_mma,
Alpha const& alpha,
Tensor<TA, ALayout> sA,
Tensor<TB, BLayout> sB,
Beta const& beta,
Tensor<TC, CLayout> sC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op = {}, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op = {}) // transforms results before they are stored to C
{
CUTE_STATIC_ASSERT_V(size<0>(sA) == size<0>(sC)); // AM == CM
CUTE_STATIC_ASSERT_V(size<0>(sB) == size<1>(sC)); // BN == CN
CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // AK == BK
using TypeA = typename TA::value_type;
using TypeB = typename TB::value_type;
using TypeC = typename TC::value_type;
static_assert(is_convertible_v<decay_t<invoke_result_t<ALoadTransformOp, TypeA>>, TypeA>,
"ALoadTransformOp functor must accept value of type TA::value_type and return value convertible to type TA::value_type");
static_assert(is_convertible_v<decay_t<invoke_result_t<BLoadTransformOp, TypeB>>, TypeB>,
"BLoadTransformOp functor must accept value of type TB::value_type and return value convertible to type TB::value_type");
static_assert(is_convertible_v<decay_t<invoke_result_t<CLoadTransformOp, TypeC>>, TypeC>,
"CLoadTransformOp functor must accept value of type TC::value_type and return value convertible to type TC::value_type");
static_assert(is_convertible_v<decay_t<invoke_result_t<CStoreTransformOp, TypeC>>, TypeC>,
"CStoreTransformOp functor must accept value of type TC::value_type and return value convertible to type TC::value_type");
static constexpr bool compat = evenly_divides(make_shape(size<0>(sA), size<0>(sB), size<1>(sA)),
tile_shape(TiledMMA<Args...>{}));
if constexpr (compat) {
detail::cooperative_gemm_no_predication<SmemCopyOpA, SmemCopyOpB, SmemCopyOpC>(
thread_idx, tiled_mma, alpha, sA, sB, beta, sC,
sA_load_op, sB_load_op, sC_load_op, sC_store_op
);
} else {
detail::cooperative_gemm_predication(
thread_idx, tiled_mma, alpha, sA, sB, beta, sC,
sA_load_op, sB_load_op, sC_load_op, sC_store_op
);
}
}
// C passed as a shared memory tensor
// Epilogue included
template <class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
class ALoadTransformOp = cute::identity, class BLoadTransformOp = cute::identity,
class CLoadTransformOp = cute::identity, class CStoreTransformOp = cute::identity,
__CUTE_REQUIRES(ALayout::rank == 2 && is_smem<TA>::value &&
BLayout::rank == 2 && is_smem<TB>::value &&
CLayout::rank == 2 && is_smem<TC>::value)>
class SmemCopyOpA = DefaultCopy, class SmemCopyOpB = DefaultCopy,
class SmemCopyOpC = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_gemm(uint32_t thread_idx,
TiledMMA<Args...> const& tiled_mma,
Alpha const& alpha,
Tensor<TA, ALayout> const& sA,
Tensor<TB, BLayout> const& sB,
Beta const& beta,
Tensor<TC, CLayout> & sC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op = {}, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op = {}, // transforms results before they are stored to C
SmemCopyOpA const& sA_copy_op = {},
SmemCopyOpB const& sB_copy_op = {},
SmemCopyOpC const& sC_copy_op = {})
{
CUTE_STATIC_ASSERT_V(rank(sA) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(sB) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(sC) == Int<2>{});
CUTE_STATIC_ASSERT_V(size<0>(sA) == size<0>(sC)); // AM == CM
CUTE_STATIC_ASSERT_V(size<0>(sB) == size<1>(sC)); // BN == CN
CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // AK == BK
using InputTypeA = typename TA::value_type;
using InputTypeB = typename TB::value_type;
using InputTypeC = typename TC::value_type;
using ComputeTypeA = typename TiledMMA<Args...>::ValTypeA;
using ComputeTypeB = typename TiledMMA<Args...>::ValTypeB;
using ComputeTypeC = typename TiledMMA<Args...>::ValTypeC;
auto compat = evenly_divides(make_shape(size<0>(sA), size<0>(sB), size<1>(sA)),
tile_shape(TiledMMA<Args...>{}));
// ThrMMA
auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
Tensor tCsC = thr_mma.partition_C(sC); // (MMA,MMA_M,MMA_N) :: InputTypeC
Tensor tCrC = thr_mma.make_fragment_C(tCsC); // (MMA,MMA_M,MMA_N) :: ComputeTypeC
// Clear accumulators
clear(tCrC);
#if 0
if (thread0()) {
print(" sC: "); print(sC); print("\n");
print(" tCsC: "); print(tCsC); print("\n");
}
#endif
if constexpr (is_constant<true, decltype(compat)>::value) {
detail::cooperative_gemm_no_predication(
thread_idx, thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op, sA_copy_op, sB_copy_op
);
detail::epilogue_no_predication(
alpha, tCrC, beta, tCsC, sC_load_op, sC_store_op, sC_copy_op
);
} else {
detail::cooperative_gemm_predication(
thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op
);
detail::epilogue_predication(
thr_mma, alpha, tCrC, beta, sC, tCsC, sC_load_op, sC_store_op
);
}
}
// C already partitioned into registers on input
// It can be passed non-empty
// Epilogue not included
template <class... Args,
class TA, class ALayout, class TB, class BLayout,
class TC, class CLayout,
class ALoadTransformOp = cute::identity, class BLoadTransformOp = cute::identity,
class SmemCopyOpA = DefaultCopy, class SmemCopyOpB = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_gemm(uint32_t thread_idx,
TiledMMA<Args...> const& tiled_mma,
Tensor<TA, ALayout> const& sA,
Tensor<TB, BLayout> const& sB,
Tensor<TC, CLayout> & tCrC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
SmemCopyOpA const& sA_copy_op = {},
SmemCopyOpB const& sB_copy_op = {})
{
CUTE_STATIC_ASSERT_V(rank(sA) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(sB) == Int<2>{});
CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // AK == BK
using InputTypeA = typename TA::value_type;
using InputTypeB = typename TB::value_type;
using InputTypeC = typename TC::value_type;
using ComputeTypeA = typename TiledMMA<Args...>::ValTypeA;
using ComputeTypeB = typename TiledMMA<Args...>::ValTypeB;
using ComputeTypeC = typename TiledMMA<Args...>::ValTypeC;
// Check if input C fragment is compatible with thr_mma and problem size
using ref_c_frag = decltype(partition_shape_C(tiled_mma, make_shape(size<0>(sA), size<0>(sB))));
CUTE_STATIC_ASSERT_V(compatible(shape(ref_c_frag{}), shape(tCrC)));
auto compat = evenly_divides(make_shape(size<0>(sA), size<0>(sB), size<1>(sA)),
tile_shape(TiledMMA<Args...>{}));
// ThrMMA
auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
if constexpr (is_constant<true, decltype(compat)>::value) {
detail::cooperative_gemm_no_predication(
thread_idx, thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op, sA_copy_op, sB_copy_op
);
} else {
detail::cooperative_gemm_predication(
thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op
);
}
}
// Accept mutable temporaries
template <class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
class ALoadTransformOp = cute::identity, class BLoadTransformOp = cute::identity,
class CLoadTransformOp = cute::identity, class CStoreTransformOp = cute::identity,
class SmemCopyOpA = DefaultCopy, class SmemCopyOpB = DefaultCopy,
class SmemCopyOpC = DefaultCopy>
CUTE_HOST_DEVICE
void
cooperative_gemm(uint32_t thread_idx,
TiledMMA<Args...> const& tiled_mma,
Alpha const& alpha,
Tensor<TA, ALayout> sA,
Tensor<TB, BLayout> sB,
Beta const& beta,
Tensor<TC, CLayout> sC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op = {}, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op = {}) // transforms results before they are stored to C
TiledMMA<Args...> const& tiled_mma,
Alpha const& alpha,
Tensor<TA, ALayout> const& sA,
Tensor<TB, BLayout> const& sB,
Beta const& beta,
Tensor<TC, CLayout> && sC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op = {}, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op = {}, // transforms results before they are stored to C
SmemCopyOpA const& sA_copy_op = {},
SmemCopyOpB const& sB_copy_op = {},
SmemCopyOpC const& sC_copy_op = {})
{
using CopyOpA = AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<typename TA::value_type>>;
using CopyOpB = AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<typename TB::value_type>>;
using CopyOpC = AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<typename TC::value_type>>;
cooperative_gemm<CopyOpA, CopyOpB, CopyOpC>(
thread_idx, tiled_mma, alpha, sA, sB, beta, sC,
sA_load_op, sB_load_op, sC_load_op, sC_store_op
);
cooperative_gemm(thread_idx, tiled_mma, alpha, sA, sB, beta, sC,
sA_load_op, sB_load_op, sC_load_op, sC_store_op,
sA_copy_op, sB_copy_op, sC_copy_op);
}
// Legacy overload of cute::gemm for backwards-compatibility
@@ -485,27 +547,38 @@ template <class... Args,
class Alpha, class TA, class ALayout, class TB, class BLayout,
class Beta, class TC, class CLayout,
class ALoadTransformOp = cute::identity, class BLoadTransformOp = cute::identity,
class CLoadTransformOp = cute::identity, class CStoreTransformOp = cute::identity,
__CUTE_REQUIRES(ALayout::rank == 2 && is_smem<TA>::value &&
BLayout::rank == 2 && is_smem<TB>::value &&
CLayout::rank == 2 && is_smem<TC>::value)>
class CLoadTransformOp = cute::identity, class CStoreTransformOp = cute::identity>
CUTE_HOST_DEVICE
void
gemm(ThrMMA<Args...> const& thr_mma,
Alpha const& alpha,
Tensor<TA, ALayout> sA,
Tensor<TB, BLayout> sB,
Beta const& beta,
Tensor<TC, CLayout> sC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op = {}, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op = {}) // transforms results before they are stored to C
gemm(ThrMMA<Args...> const& thr_mma,
Alpha const& alpha,
Tensor<TA, ALayout> const& sA,
Tensor<TB, BLayout> const& sB,
Beta const& beta,
Tensor<TC, CLayout> & sC,
ALoadTransformOp const& sA_load_op = {}, // transforms A values before use in GEMM
BLoadTransformOp const& sB_load_op = {}, // transforms B values before use in GEMM
CLoadTransformOp const& sC_load_op = {}, // transforms C values before use in GEMM
CStoreTransformOp const& sC_store_op = {}) // transforms results before they are stored to C
{
CUTE_STATIC_ASSERT_V(rank(sA) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(sB) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(sC) == Int<2>{});
CUTE_STATIC_ASSERT_V(size<0>(sA) == size<0>(sC)); // AM == CM
CUTE_STATIC_ASSERT_V(size<0>(sB) == size<1>(sC)); // BN == CN
CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // AK == BK
Tensor tCsC = thr_mma.partition_C(sC); // (MMA,MMA_M,MMA_N)
Tensor tCrC = thr_mma.make_fragment_C(tCsC); // (MMA,MMA_M,MMA_N)
// Goes directly to the slow path to avoid getting thread_idx from thr_mma
detail::cooperative_gemm_predication(
thr_mma, alpha, sA, sB, beta, sC,
sA_load_op, sB_load_op, sC_load_op, sC_store_op
thr_mma, sA, sB, sC, sA_load_op, sB_load_op
);
detail::epilogue_predication(
thr_mma, alpha, tCrC, beta, sC, tCsC, sC_load_op, sC_store_op
);
}
include/cute/algorithm/copy.hpp
+307 -144
View File
@@ -38,79 +38,6 @@
namespace cute
{
//
// Accept mutable temporaries
//
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy(src, dst);
}
template <class VecType,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_vec<VecType>(src, dst);
}
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_aligned(src, dst);
}
template <class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(pred, src, dst);
}
template <class CopyPolicy,
class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(CopyPolicy const& copy_policy,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(copy_policy, pred, src, dst);
}
template <class CopyPolicy,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(CopyPolicy const& copy_policy,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy(copy_policy, src, dst);
}
//
// copy_if -- Predicated Copy
//
@@ -124,12 +51,13 @@ copy_if(PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
auto copy_op = select_elementwise_copy(src, dst);
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
CUTE_UNROLL
for (int i = 0; i < size(src); ++i) {
for (int i = 0; i < size(dst); ++i) {
if (pred(i)) {
copy_op.copy(src(i), dst(i));
dst(i) = static_cast<DstType>(static_cast<SrcType>(src(i)));
}
}
}
@@ -138,17 +66,6 @@ copy_if(PrdTensor const& pred,
// copy_if -- Predicated CopyAtom
//
namespace detail {
// Trait that detects if atom's traits has a member function with(bool)
template <class, class Enable = void>
constexpr bool has_with_bool = false;
template <class T>
constexpr bool has_with_bool<T, cute::void_t<decltype(declval<typename T::Traits>().with(declval<bool>()))>> = true;
} // end namespace detail
template <class... CopyArgs,
class PredTensor,
class SrcEngine, class SrcLayout,
@@ -161,73 +78,90 @@ copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
Tensor<DstEngine, DstLayout> & dst) // (V,Rest...)
{
static_assert(SrcLayout::rank == DstLayout::rank, "CopyAtom rank-mismatch.");
auto has_with_bool = cute::is_valid([](auto t)->void_t<decltype(declval<typename decltype(t)::Traits>().with(true))>{}, copy_atom);
if constexpr (SrcLayout::rank == 1) { // Dispatch the copy
copy_atom.call(src, dst);
if constexpr (has_with_bool) {
copy_atom.with(pred()).call(src, dst);
} else {
if (pred()) { copy_atom.call(src, dst); }
}
} else { // Loop over all but the first mode
constexpr int R = SrcLayout::rank;
Tensor src_v = group_modes<1,R>(src);
Tensor dst_v = group_modes<1,R>(dst);
CUTE_UNROLL
for (int i = 0; i < size<1>(src_v); ++i) {
// If copy traits can be transformed with a predicate value, do it, otherwise branch here
if constexpr (detail::has_with_bool<Copy_Atom<CopyArgs...>>) {
for (int i = 0; i < size<1>(dst_v); ++i) {
if constexpr (has_with_bool) {
copy_atom.with(pred(i)).call(src_v(_,i), dst_v(_,i));
} else {
if (pred(i)) {
copy_atom.call(src_v(_,i), dst_v(_,i));
}
if (pred(i)) { copy_atom.call(src_v(_,i), dst_v(_,i)); }
}
}
}
}
//
// copy_vec -- attempt vectorized copy with VecType
// copy_if -- AutoCopyAsync
//
template <class VecType,
template <class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
copy_if(AutoCopyAsync const& cpy,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
static_assert(sizeof_bits_v<VecType> >= 8 && sizeof_bits_v<VecType> % 8 == 0,
"Expected a vectorization type of at least a byte.");
using SrcElemWithConst = remove_reference_t<typename SrcEngine::reference>;
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
if constexpr (cute::is_same<SrcType, DstType>::value &&
sizeof_bits_v<VecType> > sizeof_bits_v<DstType>)
{
// Preserve volatility of Src/Dst types.
using SrcVecType = conditional_t<is_volatile_v<typename SrcEngine::element_type>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<typename DstEngine::element_type>, VecType volatile, VecType >;
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
#if 0
if (thread0()) {
print("copy_vec<%db> -- vectorizing copy:\n", int(sizeof_bits_v<VecType>));
print(" "); print(src); print(" => "); print(src_v); print("\n");
print(" "); print(dst); print(" => "); print(dst_v); print("\n");
auto copy_op = []() {
#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
if constexpr (is_gmem<SrcEngine>::value && is_smem<DstEngine>::value &&
sizeof(SrcType) == sizeof(DstType)) {
if constexpr (is_const_v<SrcElemWithConst> && sizeof(SrcType) == 16) {
return SM80_CP_ASYNC_CACHEGLOBAL<SrcType,DstType>{};
} else if constexpr (sizeof(SrcType) == 4 || sizeof(SrcType) == 8 || sizeof(SrcType) == 16) {
return SM80_CP_ASYNC_CACHEALWAYS<SrcType,DstType>{};
} else {
return UniversalCopy<SrcType,DstType>{};
}
} else {
return UniversalCopy<SrcType,DstType>{};
}
#endif
return copy_if(TrivialPredTensor{}, src_v, dst_v);
} else {
#if 0
if (thread0()) {
print("copy_vec<%db> -- NOT vectorizing copy:\n", int(sizeof_bits_v<VecType>));
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
CUTE_GCC_UNREACHABLE;
#else
return UniversalCopy<SrcType,DstType>{};
#endif
}();
return copy_if(TrivialPredTensor{}, src, dst);
CUTE_UNROLL
for (int i = 0; i < size(dst); ++i) {
if (pred(i)) {
copy_op.copy(src(i), dst(i));
}
}
}
//
// copy -- AutoCopyAsync
//
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(AutoCopyAsync const& cpy,
Tensor<SrcEngine, SrcLayout> const& src, // (V,Rest...)
Tensor<DstEngine, DstLayout> & dst) // (V,Rest...)
{
copy_if(cpy, TrivialPredTensor{}, src, dst);
}
//
// copy -- CopyAtom
//
@@ -238,15 +172,56 @@ template <class... CopyArgs,
CUTE_HOST_DEVICE
void
copy(Copy_Atom<CopyArgs...> const& copy_atom,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
Tensor<SrcEngine, SrcLayout> const& src, // (V,Rest...)
Tensor<DstEngine, DstLayout> & dst) // (V,Rest...)
{
return copy_if(copy_atom, TrivialPredTensor{}, src, dst);
static_assert(SrcLayout::rank == DstLayout::rank, "CopyAtom rank-mismatch.");
if constexpr (SrcLayout::rank == 1) { // Dispatch the copy
copy_atom.call(src, dst);
} else { // Loop over all but the first mode
constexpr int R = SrcLayout::rank;
Tensor src_v = group_modes<1,R>(src);
Tensor dst_v = group_modes<1,R>(dst);
if constexpr (is_static<decltype(shape(src_v))>::value && is_static<decltype(shape(dst_v))>::value) {
CUTE_STATIC_ASSERT_V(size<1>(src_v) == size<1>(dst_v));
// AutoFilter on the Rest-mode
auto dst_null = nullspace(layout<1>(dst_v));
Tensor dst_n = zipped_divide(dst_v, make_tile(shape<0>(dst_v), dst_null)); // ((V, NLL), (_1, Rest))
Tensor src_n = zipped_divide(src_v, make_tile(shape<0>(src_v), dst_null)); // ((V, NLL), (_1, Rest))
CUTE_STATIC_ASSERT_V(size<1>(src_n) == size<1>(dst_n));
CUTE_STATIC_ASSERT_V((cosize<0,1>(dst_n.layout()) == Int<1>{}), "Nullspace definition error");
CUTE_STATIC_ASSERT_V((cosize<0,1>(src_n.layout()) == Int<1>{}), "Error: Ambiguous scatter detected in copy");
CUTE_STATIC_ASSERT_V((size<1,0>(dst_n) == Int<1>{}));
CUTE_STATIC_ASSERT_V((size<1,0>(src_n) == Int<1>{}));
Tensor dst_c = dst_n(make_coord(_,Int<0>{}),make_coord(Int<0>{},_)); // (V, Rest)
Tensor src_c = src_n(make_coord(_,Int<0>{}),make_coord(Int<0>{},_)); // (V, Rest)
CUTE_STATIC_ASSERT_V(size<1>(src_c) == size<1>(dst_c));
CUTE_STATIC_ASSERT_V(shape<0>(dst_c) == shape<0>(dst));
CUTE_STATIC_ASSERT_V(shape<0>(src_c) == shape<0>(src));
CUTE_UNROLL
for (int i = 0; i < size<1>(dst_c); ++i) {
copy_atom.call(src_c(_,i), dst_c(_,i));
}
} else {
CUTE_UNROLL
for (int i = 0; i < size<1>(dst_v); ++i) {
copy_atom.call(src_v(_,i), dst_v(_,i));
}
}
}
}
//////////////////////////////////////////
// Special Auto-Vectorizing Overloads
//////////////////////////////////////////
////////////////////////////////////////////////////////
// Special Auto-Vectorizing, Auto-Filtering Overloads //
////////////////////////////////////////////////////////
// Specialization for AutoVectorizingCopyAssumedAlignment<MaxVecBits>
template <int MaxVecBits, class... Args,
@@ -258,30 +233,67 @@ copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
constexpr int vec_elem = decltype(max_common_vector(src, dst))::value;
constexpr int common_elem = CUTE_STATIC_V(max_common_vector(src, dst));
constexpr int align_bits = CUTE_STATIC_V(gcd(max_alignment(src), max_alignment(dst), Int<MaxVecBits>{}));
static_assert(is_integral<decltype(Int<common_elem>{} * sizeof_bits_v<typename SrcEngine::value_type>)>::value, "Error: Attempting a subbit copy!");
constexpr int vec_bits = gcd(common_elem * sizeof_bits_v<typename SrcEngine::value_type>, align_bits);
constexpr int max_align_src = decltype(max_alignment(src.layout()))::value;
constexpr int max_align_dst = decltype(max_alignment(dst.layout()))::value;
constexpr int max_align = gcd(vec_elem, max_align_src, max_align_dst);
if constexpr (common_elem > 1 && ((vec_bits % 8) == 0)) {
// If more than one element vectorizes to 8bits or more, then recast and copy
using VecType = uint_bit_t<vec_bits>;
// Preserve volatility
using SrcVecType = conditional_t<is_volatile_v<typename SrcEngine::element_type>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<typename DstEngine::element_type>, VecType volatile, VecType >;
constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
constexpr int vec_bits = gcd(src_bits * max_align, MaxVecBits);
// Recast
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
if constexpr (vec_elem > 1 && vec_bits >= 8) {
// If more than one element vectorizes to 8bits or more, then copy_vec
#if 0
if (thread0()) {
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", common_elem, vec_bits);
print(" "); print(src); print(" => "); print(src_v); print("\n");
print(" "); print(dst); print(" => "); print(dst_v); print("\n");
}
#endif
return copy_vec<uint_bit_t<vec_bits>>(src, dst);
return copy_if(TrivialPredTensor{}, src_v, dst_v);
} else {
return copy_if(TrivialPredTensor{}, src, dst);
}
}
template <class Base>
struct AutoFilter {
Base const& base;
CUTE_HOST_DEVICE AutoFilter(Base const& b) : base(b) {}
};
// Specialization for AutoFilter
template <class CopyOp,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(AutoFilter<CopyOp> const& copy_op,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
if constexpr (is_constant<true, decltype(size(src) == size(dst))>::value) {
auto dst_null = nullspace(dst.layout());
Tensor dst_n = zipped_divide(dst, dst_null);
Tensor src_n = zipped_divide(src, dst_null);
CUTE_STATIC_ASSERT_V(cosize<0>(dst_n.layout()) == Int<1>{}, "Nullspace definition error");
CUTE_STATIC_ASSERT_V(cosize<0>(src_n.layout()) == Int<1>{}, "Error: Ambiguous scatter detected in copy");
copy(copy_op.base, src_n(Int<0>{},_), dst_n(Int<0>{},_));
} else {
copy(copy_op.base, src, dst);
}
}
// Auto-vectorizing copy for static layouts
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -292,7 +304,11 @@ copy(Tensor<SrcEngine, SrcLayout> const& src,
{
if constexpr (is_static<SrcLayout>::value && is_static<DstLayout>::value) {
// Assume Tensors with static layouts (e.g. registers) have pointers that are 128b aligned
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
return copy(AutoFilter(AutoVectorizingCopyWithAssumedAlignment<128>{}), src, dst);
} else
if constexpr (is_static<decltype(shape(src))>::value && is_static<decltype(shape(dst))>::value) {
// Tensors with static shapes can be filtered, but do not assume that dynamic layouts are aligned.
return copy(AutoFilter(AutoVectorizingCopyWithAssumedAlignment<8>{}), src, dst);
} else {
// Do not assume that dynamic layouts are aligned.
return copy(AutoVectorizingCopyWithAssumedAlignment<8>{}, src, dst);
@@ -307,7 +323,12 @@ void
copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
if constexpr (is_static<decltype(shape(src))>::value && is_static<decltype(shape(dst))>::value) {
// Tensors with static shapes can be filtered
return copy(AutoFilter(AutoVectorizingCopyWithAssumedAlignment<128>{}), src, dst);
} else {
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
}
}
// Specializaton for Atom AutoVectorizingCopyAssumedAlignment
@@ -379,4 +400,146 @@ copy(Copy_Atom<Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...>, CA_Args...> const&
}
#endif // #if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)
//
// Decay TiledCopy to CopyAtom
//
template <class CopyAtom, class TV, class Tiler,
class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(TiledCopy<CopyAtom, TV, Tiler> const& tiled_copy,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy_if(static_cast<CopyAtom const&>(tiled_copy), pred, src, dst);
}
template <class CopyAtom, class TV, class Tiler,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(TiledCopy<CopyAtom, TV, Tiler> const& tiled_copy,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(static_cast<CopyAtom const&>(tiled_copy), src, dst);
}
template <class TiledCopy, class ThrIdx,
class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(ThrCopy<TiledCopy, ThrIdx> const& thr_copy,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst) = delete;
template <class TiledCopy, class ThrIdx,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(ThrCopy<TiledCopy, ThrIdx> const& thr_copy,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst) = delete;
//
// Catch uncaught policies
//
template <class CopyPolicy,
class PredTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(CopyPolicy const& cpy,
PredTensor const& prd,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
static_assert(dependent_false<CopyPolicy>, "Unrecognized CopyPolicy.");
}
template <class CopyPolicy,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(CopyPolicy const& cpy,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
static_assert(dependent_false<CopyPolicy>, "Unrecognized CopyPolicy.");
}
//
// Accept mutable temporaries
//
template <class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(pred, src, dst);
}
template <class CopyPolicy,
class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(CopyPolicy const& copy_policy,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(copy_policy, pred, src, dst);
}
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy(src, dst);
}
template <class CopyPolicy,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(CopyPolicy const& copy_policy,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy(copy_policy, src, dst);
}
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_aligned(src, dst);
}
} // end namespace cute
include/cute/arch/copy.hpp
+13 -13
View File
@@ -39,7 +39,7 @@ namespace cute
{
//
// Direct Copy for any type
// Direct Copy for any specific types
//
template <class S, class D = S>
@@ -48,21 +48,15 @@ struct UniversalCopy
using SRegisters = S[1];
using DRegisters = D[1];
template <class S_, class D_>
CUTE_HOST_DEVICE static constexpr void
copy(S_ const& src,
D_ & dst)
{
dst = static_cast<D>(static_cast<S>(src));
}
// Sanity
static_assert(sizeof_bits_v<S> >= 8);
static_assert(sizeof_bits_v<D> >= 8);
// Accept mutable temporaries
template <class S_, class D_>
CUTE_HOST_DEVICE static constexpr void
copy(S_ const& src,
D_ && dst)
copy(S const& src,
D & dst)
{
UniversalCopy<S,D>::copy(src, dst);
dst = src;
}
};
@@ -92,6 +86,12 @@ using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<128>;
using DefaultCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
//
// Copy policy automatically selecting between
// UniversalCopy and cp.async , based on type and memory space.
//
struct AutoCopyAsync {};
//
// Global memory prefetch into L2
//
include/cute/arch/mma_sm80.hpp
+97 -99
View File
@@ -2040,6 +2040,103 @@ struct SM80_16x8x64_S32U4U4S32_TN_SATURATE
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 8x8x128 TN
struct SM80_8x8x128_S32U1U1S32_TN_ANDPOPC
{
using DRegisters = uint32_t[2];
using ARegisters = uint32_t[1];
using BRegisters = uint32_t[1];
using CRegisters = uint32_t[2];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1,
uint32_t const& a0,
uint32_t const& b0,
uint32_t const& c0, uint32_t const& c1)
{
#if defined(CUTE_ARCH_MMA_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.and.popc "
"{%0, %1},"
"{%2},"
"{%3},"
"{%4, %5};\n"
: "=r"(d0), "=r"(d1)
: "r"(a0),
"r"(b0),
"r"(c0), "r"(c1));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x128_S32U1U1S32_TN_ANDPOPC without CUTE_ARCH_MMA_SM80_ENABLED");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 16x8x128 TN
struct SM80_16x8x128_S32U1U1S32_TN_ANDPOPC
{
using DRegisters = uint32_t[4];
using ARegisters = uint32_t[2];
using BRegisters = uint32_t[1];
using CRegisters = uint32_t[4];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3,
uint32_t const& a0, uint32_t const& a1,
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)
asm volatile(
"mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.and.popc "
"{%0, %1, %2, %3},"
"{%4, %5},"
"{%6},"
"{%7, %8, %9, %10};\n"
: "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)
: "r"(a0), "r"(a1),
"r"(b0),
"r"(c0), "r"(c1), "r"(c2), "r"(c3));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x128_S32U1U1S32_TN_ANDPOPC without CUTE_ARCH_MMA_SM80_ENABLED");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 16x8x256 TN
struct SM80_16x8x256_S32U1U1S32_TN_ANDPOPC
{
using DRegisters = uint32_t[4];
using ARegisters = uint32_t[4];
using BRegisters = uint32_t[2];
using CRegisters = uint32_t[4];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3,
uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
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_B1_AND_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.and.popc "
"{%0, %1, %2, %3},"
"{%4, %5, %6, %7},"
"{%8, %9},"
"{%10, %11, %12, %13};\n"
: "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3),
"r"(b0), "r"(b1),
"r"(c0), "r"(c1), "r"(c2), "r"(c3));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x256_S32U1U1S32_TN_ANDPOPC without CUTE_ARCH_MMA_SM80_ENABLED");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 8x8x128 TN
@@ -2141,103 +2238,4 @@ struct SM80_16x8x256_S32U1U1S32_TN_XORPOPC
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 8x8x128 TN
struct SM80_8x8x128_S32U1U1S32_TN_ANDPOPC
{
using DRegisters = uint32_t[2];
using ARegisters = uint32_t[1];
using BRegisters = uint32_t[1];
using CRegisters = uint32_t[2];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1,
uint32_t const& a0,
uint32_t const& b0,
uint32_t const& c0, uint32_t const& c1)
{
#if defined(CUTE_ARCH_MMA_B1_AND_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.and.popc "
"{%0, %1},"
"{%2},"
"{%3},"
"{%4, %5};\n"
: "=r"(d0), "=r"(d1)
: "r"(a0),
"r"(b0),
"r"(c0), "r"(c1));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x128_S32U1U1S32_TN_ANDPOPC without CUTE_ARCH_MMA_SM80_ENABLED");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 16x8x128 TN
struct SM80_16x8x128_S32U1U1S32_TN_ANDPOPC
{
using DRegisters = uint32_t[4];
using ARegisters = uint32_t[2];
using BRegisters = uint32_t[1];
using CRegisters = uint32_t[4];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3,
uint32_t const& a0, uint32_t const& a1,
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_B1_AND_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.and.popc "
"{%0, %1, %2, %3},"
"{%4, %5},"
"{%6},"
"{%7, %8, %9, %10};\n"
: "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)
: "r"(a0), "r"(a1),
"r"(b0),
"r"(c0), "r"(c1), "r"(c2), "r"(c3));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x128_S32U1U1S32_TN_ANDPOPC without CUTE_ARCH_MMA_SM80_ENABLED");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// MMA 16x8x256 TN
struct SM80_16x8x256_S32U1U1S32_TN_ANDPOPC
{
using DRegisters = uint32_t[4];
using ARegisters = uint32_t[4];
using BRegisters = uint32_t[2];
using CRegisters = uint32_t[4];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3,
uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
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_B1_AND_SM80_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.and.popc "
"{%0, %1, %2, %3},"
"{%4, %5, %6, %7},"
"{%8, %9},"
"{%10, %11, %12, %13};\n"
: "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3),
"r"(b0), "r"(b1),
"r"(c0), "r"(c1), "r"(c2), "r"(c3));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x256_S32U1U1S32_TN_ANDPOPC without CUTE_ARCH_MMA_SM80_ENABLED");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cute
include/cute/atom/copy_atom.hpp
+5 -5
View File
@@ -100,16 +100,16 @@ struct Copy_Atom<Copy_Traits<Args...>, CopyInternalType>
if constexpr (is_constant<NumValSrc, decltype(size(src))>::value ||
is_constant<NumValDst, decltype(size(dst))>::value) {
// Dispatch to unpack to execute instruction
return copy_unpack(*this, src, dst);
} else
if constexpr (is_tuple<decltype(shape(src))>::value &&
is_tuple<decltype(shape(dst))>::value) {
return copy_unpack(static_cast<Traits const&>(*this), src, dst);
} else if constexpr (is_tuple<decltype(shape(src))>::value &&
is_tuple<decltype(shape(dst))>::value) {
// If the size of the src/dst doesn't match the instruction,
// recurse this rank-1 layout by peeling off the mode
// ((A,B,C,...)) -> (A,B,C,...)
return copy(*this, tensor<0>(src), tensor<0>(dst));
} else {
static_assert(dependent_false<SEngine>, "No instruction match and no recursion possible.");
static_assert(dependent_false<SEngine>,
"CopyAtom: Src/Dst partitioning does not match the instruction requirement.");
}
}
include/cute/atom/copy_traits.hpp
+18 -15
View File
@@ -92,23 +92,29 @@ struct Copy_Traits<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>>
using RefLayout = SrcLayout;
};
// Extract a CPY_Op from a CPY_Traits
template <class CPY_Traits>
struct CPY_Op {};
template <class CPY_Op_Arg, class... Args>
struct CPY_Op<Copy_Traits<CPY_Op_Arg, Args...>> {
using type = CPY_Op_Arg;
};
//
// Generic copy_unpack for common argument-based Copy_Traits
//
template <class CopyOp, class... Args,
template <class AnyCPYTraits,
class SEngine, class SLayout,
class DEngine, class DLayout>
CUTE_HOST_DEVICE constexpr
void
copy_unpack(Copy_Traits<CopyOp,Args...> const&,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> & dst)
copy_unpack(AnyCPYTraits const&,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> & dst)
{
// Specializations can generalize on these checks
//static_assert(is_smem<TS>::value, "Expected smem for this Copy_Traits<CopyOp>");
//static_assert(is_rmem<TD>::value, "Expected rmem for this Copy_Traits<CopyOp>");
using CopyOp = typename CPY_Op<AnyCPYTraits>::type;
using RegistersSrc = typename CopyOp::SRegisters;
using RegistersDst = typename CopyOp::DRegisters;
using RegTypeSrc = typename remove_extent<RegistersSrc>::type;
@@ -129,18 +135,15 @@ copy_unpack(Copy_Traits<CopyOp,Args...> const&,
rD, make_int_sequence<RegNumDst>{});
}
//
// Accept mutable temporaries
//
template <class CopyOp, class... Args,
template <class AnyCPYTraits,
class SEngine, class SLayout,
class DEngine, class DLayout>
CUTE_HOST_DEVICE constexpr
void
copy_unpack(Copy_Traits<CopyOp,Args...> const& traits,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> && dst)
copy_unpack(AnyCPYTraits const& traits,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> && dst)
{
copy_unpack(traits, src, dst);
}
include/cute/atom/copy_traits_sm80.hpp
+19 -46
View File
@@ -51,13 +51,6 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<S,D>>
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Construct a zfill variant with a given predicate value
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
with(bool pred) const {
return {pred};
}
};
template <class S, class D>
@@ -73,13 +66,6 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<S,D>>
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Construct a zfill variant with a given predicate value
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
with(bool pred) const {
return {pred};
}
};
template <class S, class D>
@@ -96,8 +82,15 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Predicate value that determines whether to load or zfill
bool pred = false;
// Predicate value: true = load, false = zfill
bool pred = true;
// Construct a zfill variant with a given predicate value
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
with(bool pred) const {
return {pred};
}
// Overload copy_unpack for zfill variant to pass the predicate into the op
template <class TS, class SLayout,
@@ -137,8 +130,15 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// Predicate value that determines whether to load or zfill
bool pred = false;
// Predicate value: true = load, false = zfill
bool pred = true;
// Construct a zfill variant with a given predicate value
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
with(bool pred) const {
return {pred};
}
// Overload copy_unpack for zfill variant to pass the predicate into the op
template <class TS, class SLayout,
@@ -164,31 +164,4 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// Element copy selector
template <class SrcTensor, class DstTensor>
CUTE_HOST_DEVICE constexpr
auto
select_elementwise_copy(SrcTensor const&, DstTensor const&)
{
using SrcType = typename SrcTensor::value_type;
using DstType = typename DstTensor::value_type;
#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
if constexpr (is_gmem<SrcTensor>::value && is_smem<DstTensor>::value &&
sizeof(SrcType) == sizeof(DstType) &&
(sizeof(SrcType) == 4 || sizeof(SrcType) == 8 || sizeof(SrcType) == 16))
{
return SM80_CP_ASYNC_CACHEALWAYS<SrcType,DstType>{};
} else {
return UniversalCopy<SrcType,DstType>{};
}
CUTE_GCC_UNREACHABLE;
#else
return UniversalCopy<SrcType,DstType>{};
#endif
}
}
} // end namespace cute
include/cute/atom/copy_traits_sm90_tma.hpp
+109 -74
View File
@@ -58,37 +58,31 @@ struct AuxTmaParams {
};
// Utility for unpacking TMA_LOAD arguments into a CopyOp
template <class CopyOp>
template <class CopyOp, class... Args>
struct TMA_LOAD_Unpack
{
template <class... Args,
class TS, class SLayout,
template <class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
copy_unpack(Copy_Traits<CopyOp, Args...> const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_smem<TD>::value, "SM90_TMA_LOAD requires the destination be shared memory.");
auto src_coord = src.data().coord_;
if constexpr (detail::is_prefetch<CopyOp>) {
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord, tuple_seq<decltype(src_coord)>{});
} else {
static_assert(is_smem<TD>::value, "SM90_TMA_LOAD requires the destination be shared memory.");
void* dst_ptr = cute::raw_pointer_cast(dst.data());
void* dst_ptr = cute::raw_pointer_cast(dst.data());
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr);
auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr);
#endif
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
make_tuple(dst_ptr), seq<0>{},
src_coord, tuple_seq<decltype(src_coord)>{});
}
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
make_tuple(dst_ptr), seq<0>{},
src_coord, tuple_seq<decltype(src_coord)>{});
}
};
@@ -131,7 +125,7 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {{}, {&tma_desc_, &tma_mbar, static_cast<uint64_t>(cache_hint)}};
return {&tma_desc_, &tma_mbar, static_cast<uint64_t>(cache_hint)};
}
// Construct an executable SM90_TMA_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
@@ -143,7 +137,7 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {{}, {new_tma_desc, &tma_mbar, static_cast<uint64_t>(cache_hint)}};
return {new_tma_desc, &tma_mbar, static_cast<uint64_t>(cache_hint)};
}
// Generate the TMA coord tensor
@@ -167,7 +161,7 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
// The executable SM90_TMA_LOAD with tma_desc and tma_mbar
template <class NumBitsPerTMA>
struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
: TMA_LOAD_Unpack<SM90_TMA_LOAD_OP>
: TMA_LOAD_Unpack<SM90_TMA_LOAD_OP, NumBitsPerTMA>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
@@ -183,12 +177,15 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
uint64_t*, // smem mbarrier
uint64_t // cache hint
> const opargs_;
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc, uint64_t* mbar, uint64_t cache)
: opargs_(desc, mbar, cache) {}
};
// The prefetch for SM90_TMA_LOAD with tma_desc
template <class NumBitsPerTMA, class... Args>
struct Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA, Args...>
: TMA_LOAD_Unpack<SM90_TMA_LOAD::PREFETCH>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
@@ -206,6 +203,19 @@ struct Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA, Args...>
CUTE_HOST_DEVICE
Copy_Traits(Copy_Traits<CopyArgs...> const& traits)
: opargs_({&traits.tma_desc_}) {}
template <class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
copy_unpack(Copy_Traits const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
auto src_coord = src.data().coord_;
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_LOAD::PREFETCH>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord, tuple_seq<decltype(src_coord)>{});
}
};
//////////////////////////////////////////////////////////////////////////////
@@ -246,7 +256,7 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
uint64_t& tma_load_mbar,
uint16_t const& multicast_mask,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
return {{}, {&tma_desc_, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)}};
return {&tma_desc_, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)};
}
// Construct an executable SM90_TMA_LOAD_MULTICAST_OP with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
@@ -257,7 +267,7 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
uint64_t& tma_load_mbar,
uint16_t const& multicast_mask,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
return {{}, {new_tma_desc, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)}};
return {new_tma_desc, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)};
}
// Generate the TMA coord tensor
@@ -281,7 +291,7 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
// The executable SM90_TMA_LOAD_MULTICAST with tma_desc and tma_mbar and multicast_mask
template <class NumBitsPerTMA>
struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
: TMA_LOAD_Unpack<SM90_TMA_LOAD_MULTICAST_OP>
: TMA_LOAD_Unpack<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
@@ -298,43 +308,17 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
uint16_t, // multicast mask
uint64_t // cache hint
> const opargs_;
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc, uint64_t* mbar, uint16_t mask, uint64_t hint)
: opargs_(desc, mbar, mask, hint) {}
};
//////////////////////////////////////////////////////////////////////////////
///////////////////////////// TMA_STORE //////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
// Utility for unpacking TMA_STORE arguments into a CopyOp
template <class CopyOp>
struct TMA_STORE_Unpack
{
template <class... Args,
class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
copy_unpack(Copy_Traits<CopyOp, Args...> const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_smem<TS>::value, "Expected smem src for SM90_TMA_STORE");
void const* const desc_ptr = traits.tma_desc_;
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = dst.data().coord_;
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr);
#endif
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_STORE>{},
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
struct SM90_TMA_STORE_OP : SM90_TMA_STORE {};
struct SM90_TMA_STORE_PTR : SM90_TMA_STORE {};
// The executable SM90_TMA_STORE with tma_desc
template <class NumBitsPerTMA, class AuxParams_>
@@ -369,6 +353,13 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Construct new TMA_STORE with (unsafe) swapped out TMA descriptor ptr (for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_STORE_PTR, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc) const {
return {new_tma_desc};
}
template <class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
@@ -393,19 +384,11 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
// Construct Copy_Traits executable (w/ swapped out TMA descriptor) for SM90_TMA_STORE (for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc) const {
return {{}, new_tma_desc};
}
};
// The executable SM90_TMA_STORE with tma_desc
// Same as SM90_TMA_STORE, but with an unsafe TMA Desc PTR instead
template <class NumBitsPerTMA>
struct Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
: TMA_STORE_Unpack<SM90_TMA_STORE_OP>
struct Copy_Traits<SM90_TMA_STORE_PTR, NumBitsPerTMA>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
@@ -417,6 +400,31 @@ struct Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
// SM90_TMA_STORE arguments
TmaDescriptor const* tma_desc_;
template <class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
copy_unpack(Copy_Traits const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_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
void const* const desc_ptr = traits.tma_desc_;
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = dst.data().coord_;
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr);
#endif
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_STORE_PTR>{},
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
//////////////////////////////////////////////////////////////////////////////
@@ -520,7 +528,7 @@ struct Copy_Traits<SM90_BULK_COPY_G2S, NumBitsPerTMA, OpArgs...>
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_BULK_COPY_G2S, NumBitsPerTMA, uint64_t*>
with(uint64_t& bulk_mbar) const {
return {{&bulk_mbar}};
return {&bulk_mbar};
}
template <class TS, class SLayout,
@@ -613,7 +621,7 @@ struct Copy_Traits<SM90_BULK_COPY_AUTO, OpArgs...>
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_BULK_COPY_AUTO, uint64_t*>
with(uint64_t& bulk_mbar) const {
return {{&bulk_mbar}};
return {&bulk_mbar};
}
};
@@ -1391,19 +1399,46 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
return cute::make_tuple(gresult, sresult);
}
// Explicit defaults for cta_coord and cta_layout
template <class... Args,
class SEngine, class SLayout,
class GEngine, class GLayout>
CUTE_DEVICE
auto
tma_partition(Copy_Atom<Args...> const& copy_atom,
Tensor<SEngine,SLayout> const& stensor, // SMEM Tensor (TMATile, Rest...)
Tensor<GEngine,GLayout> const& gtensor) // GMEM Tensor (TMATile, Rest...)
{
return tma_partition(copy_atom, Int<0>{}, Layout<_1,_0>{}, stensor, gtensor);
}
// TMA Multicast Masks Calculation
template <int Mode, class CtaLayout, class CtaCoord>
CUTE_HOST_DEVICE constexpr
auto
uint16_t
create_tma_multicast_mask(CtaLayout const& cta_layout_vmnk,
CtaCoord const& cta_coord_vmnk)
{
auto cta_coord_slicer = replace<Mode>(cta_coord_vmnk, _);
auto [cta_layout, elected_cta] = slice_and_offset(cta_coord_slicer, cta_layout_vmnk);
// Get the instruction code
uint16_t mcast_mask = 0;
for (int i = 0; i < size(cta_layout); ++i) {
mcast_mask |= uint16_t(1) << cta_layout(i);
if constexpr (rank_v<decltype(cta_layout)> == 1 and depth_v<decltype(cta_layout)> <= 1 and
not is_static<decltype(cta_layout)>::value) {
// Get the instruction code -- optimized for dynamic flat-rank-1 cta_layout
mcast_mask = uint16_t(1);
// Smear by stride<0> (may want to predicate on stride<0> mag?)
mcast_mask |= mcast_mask << (1*stride<0>(cta_layout));
mcast_mask |= mcast_mask << (2*stride<0>(cta_layout));
mcast_mask |= mcast_mask << (4*stride<0>(cta_layout));
mcast_mask |= mcast_mask << (8*stride<0>(cta_layout));
// Select shape<0>
mcast_mask &= (uint16_t(-1) >> (16 - shape<0>(cta_layout) * stride<0>(cta_layout)));
} else {
// Get the instruction code -- generic path
for (int i = 0; i < size(cta_layout); ++i) {
mcast_mask |= uint16_t(1) << cta_layout(i);
}
}
// Shift by the instruction's elected block rank (dynamic)
mcast_mask <<= elected_cta;
include/cute/atom/mma_atom.hpp
+22 -27
View File
@@ -250,12 +250,12 @@ struct TiledMMA : MMA_Atom
auto t_tensor = logical_divide(ctensor, t_tile); // (PermM,PermN)
// Tile the tensor for the Atom
auto a_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})),
auto c_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})),
make_layout(size<1>(AtomShape_MNK{})));
auto a_tensor = zipped_divide(t_tensor, a_tile); // ((AtomM,AtomN),(RestM,RestN))
auto c_tensor = zipped_divide(t_tensor, c_tile); // ((AtomM,AtomN),(RestM,RestN))
// Transform the Atom mode from (M,K) to (Thr,Val)
auto tv_tensor = a_tensor.compose(AtomLayoutC_TV{},_); // ((ThrV,FrgV),(RestM,RestN))
auto tv_tensor = c_tensor.compose(AtomLayoutC_TV{},_); // ((ThrV,FrgV),(RestM,RestN))
// Tile the tensor for the C-threads
auto thr_tile = make_tile(_,
@@ -604,16 +604,15 @@ CUTE_HOST_DEVICE constexpr
auto
partition_shape_C(TiledMMA<Args...> const& mma, Shape_MN const& shape_MN)
{
constexpr int R = rank_v<Shape_MN>;
static_assert(R >= 2, "Must have at least rank-2");
auto atomMNK = typename TiledMMA<Args...>::AtomShape_MNK{};
auto thrVMNK = typename TiledMMA<Args...>::ThrLayoutVMNK{};
auto V = shape<1>(typename TiledMMA<Args...>::AtomLayoutC_TV{});
auto M = shape_div(size<0>(shape_MN), size<0>(atomMNK) * size<1>(thrVMNK));
auto N = shape_div(size<1>(shape_MN), size<1>(atomMNK) * size<2>(thrVMNK));
return cute::tuple_cat(make_shape(V,M,N), take<2,R>(shape_MN));
auto dummy = make_layout(shape(shape_MN));
auto dummy_tv = mma.thrfrg_C(dummy);
// Slice+rearrange like partition_C
auto dummy_v = dummy_tv(Int<0>{}, make_coord(_, repeat<rank(dummy)>(_)));
return shape(dummy_v);
}
template <class... Args, class Shape_MN>
CUTE_HOST_DEVICE constexpr
auto
@@ -632,14 +631,12 @@ CUTE_HOST_DEVICE constexpr
auto
partition_shape_A(TiledMMA<Args...> const& mma, Shape_MK const& shape_MK)
{
constexpr int R = rank_v<Shape_MK>;
static_assert(R >= 2, "Must have at least rank-2");
auto atomMNK = typename TiledMMA<Args...>::AtomShape_MNK{};
auto thrVMNK = typename TiledMMA<Args...>::ThrLayoutVMNK{};
auto V = shape<1>(typename TiledMMA<Args...>::AtomLayoutA_TV{});
auto M = shape_div(size<0>(shape_MK), size<0>(atomMNK) * size<1>(thrVMNK));
auto K = shape_div(size<1>(shape_MK), size<2>(atomMNK) * size<3>(thrVMNK));
return cute::tuple_cat(make_shape(V,M,K), take<2,R>(shape_MK));
auto dummy = make_layout(shape(shape_MK));
auto dummy_tv = mma.thrfrg_A(dummy);
// Slice+rearrange like partition_A
auto dummy_v = dummy_tv(Int<0>{}, make_coord(_, repeat<rank(dummy)>(_)));
return shape(dummy_v);
}
template <class... Args, class Shape_NK>
@@ -647,14 +644,12 @@ CUTE_HOST_DEVICE constexpr
auto
partition_shape_B(TiledMMA<Args...> const& mma, Shape_NK const& shape_NK)
{
constexpr int R = rank_v<Shape_NK>;
static_assert(R >= 2, "Must have at least rank-2");
auto atomMNK = typename TiledMMA<Args...>::AtomShape_MNK{};
auto thrVMNK = typename TiledMMA<Args...>::ThrLayoutVMNK{};
auto V = shape<1>(typename TiledMMA<Args...>::AtomLayoutB_TV{});
auto N = shape_div(size<0>(shape_NK), size<1>(atomMNK) * size<2>(thrVMNK));
auto K = shape_div(size<1>(shape_NK), size<2>(atomMNK) * size<3>(thrVMNK));
return cute::tuple_cat(make_shape(V,N,K), take<2,R>(shape_NK));
auto dummy = make_layout(shape(shape_NK));
auto dummy_tv = mma.thrfrg_B(dummy);
// Slice+rearrange like partition_B
auto dummy_v = dummy_tv(Int<0>{}, make_coord(_, repeat<rank(dummy)>(_)));
return shape(dummy_v);
}
//
include/cute/atom/mma_traits_sm80.hpp
+204 -3
View File
@@ -419,6 +419,203 @@ template <>
struct MMA_Traits<SM80_16x8x32_S32U8U8S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x32_S32U8U8S32_TN> {};
///////////////////////////////////////////////////////////////////////////////
/////////////////////////// s32 = s4 * s4 + s32 ///////////////////////////////
///////////////////////////////////////////////////////////////////////////////
template <>
struct MMA_Traits<SM80_8x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = int4b_t;
using ValTypeB = int4b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_8, _8, _32>;
using ThrID = Layout<_32>;
// (T32,V8) -> (M8,N32)
using ALayout = Layout<Shape <Shape < _4, _8>, Shape <_8>>,
Stride<Stride<_64, _1>, Stride<_8>>>;
using BLayout = Layout<Shape <Shape < _4, _8>, Shape <_8>>,
Stride<Stride<_64, _1>, Stride<_8>>>;
using CLayout = SM80_8x8_Row;
};
template <>
struct MMA_Traits<SM80_8x8x32_S32S4S4S32_TN_SATURATE>
: MMA_Traits<SM80_8x8x32_S32S4S4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = int4b_t;
using ValTypeB = int4b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16, _8, _32>;
using ThrID = Layout<_32>;
// (T32,V16) -> (M16,N32)
using ALayout = Layout<Shape <Shape < _4, _8>, Shape < _8, _2>>,
Stride<Stride<_128, _1>, Stride<_16, _8>>>;
// (T32,V8) -> (M8,N32)
using BLayout = Layout<Shape <Shape < _4, _8>, Shape <_8>>,
Stride<Stride<_32, _1>, Stride<_8>>>;
using CLayout = SM80_16x8_Row;
};
template <>
struct MMA_Traits<SM80_16x8x32_S32S4S4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x32_S32S4S4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x64_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = int4b_t;
using ValTypeB = int4b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16, _8, _64>;
using ThrID = Layout<_32>;
// (T32,V32) -> (M16,N64)
using ALayout = Layout<Shape <Shape < _4, _8>, Shape < _8, _2, _2>>,
Stride<Stride<_128, _1>, Stride<_16, _8, _512>>>;
// (T32,V16) -> (M8,N64)
using BLayout = Layout<Shape <Shape < _4, _8>, Shape <_8, _2>>,
Stride<Stride<_64, _1>, Stride<_8, _256>>>;
using CLayout = SM80_16x8_Row;
};
template <>
struct MMA_Traits<SM80_16x8x64_S32S4S4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x64_S32S4S4S32_TN> {};
///////////////////////////////////////////////////////////////////////////////
/////////////////////////// s32 = s4 * u4 + s32 ///////////////////////////////
///////////////////////////////////////////////////////////////////////////////
template <>
struct MMA_Traits<SM80_8x8x32_S32S4U4S32_TN>
: MMA_Traits<SM80_8x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = int4b_t;
using ValTypeB = uint4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_8x8x32_S32S4U4S32_TN_SATURATE>
: MMA_Traits<SM80_8x8x32_S32S4U4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x32_S32S4U4S32_TN>
: MMA_Traits<SM80_16x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = int4b_t;
using ValTypeB = uint4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_16x8x32_S32S4U4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x32_S32S4U4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x64_S32S4U4S32_TN>
: MMA_Traits<SM80_16x8x64_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = int4b_t;
using ValTypeB = uint4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_16x8x64_S32S4U4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x64_S32S4U4S32_TN> {};
///////////////////////////////////////////////////////////////////////////////
/////////////////////////// s32 = u4 * s4 + s32 ///////////////////////////////
///////////////////////////////////////////////////////////////////////////////
template <>
struct MMA_Traits<SM80_8x8x32_S32U4S4S32_TN>
: MMA_Traits<SM80_8x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = uint4b_t;
using ValTypeB = int4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_8x8x32_S32U4S4S32_TN_SATURATE>
: MMA_Traits<SM80_8x8x32_S32U4S4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x32_S32U4S4S32_TN>
: MMA_Traits<SM80_16x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = uint4b_t;
using ValTypeB = int4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_16x8x32_S32U4S4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x32_S32U4S4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x64_S32U4S4S32_TN>
: MMA_Traits<SM80_16x8x64_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = uint4b_t;
using ValTypeB = int4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_16x8x64_S32U4S4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x64_S32U4S4S32_TN> {};
///////////////////////////////////////////////////////////////////////////////
/////////////////////////// s32 = u4 * u4 + s32 ///////////////////////////////
///////////////////////////////////////////////////////////////////////////////
template <>
struct MMA_Traits<SM80_8x8x32_S32U4U4S32_TN>
: MMA_Traits<SM80_8x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = uint4b_t;
using ValTypeB = uint4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_8x8x32_S32U4U4S32_TN_SATURATE>
: MMA_Traits<SM80_8x8x32_S32U4U4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x32_S32U4U4S32_TN>
: MMA_Traits<SM80_16x8x32_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = uint4b_t;
using ValTypeB = uint4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_16x8x32_S32U4U4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x32_S32U4U4S32_TN> {};
template <>
struct MMA_Traits<SM80_16x8x64_S32U4U4S32_TN>
: MMA_Traits<SM80_16x8x64_S32S4S4S32_TN> {
using ValTypeD = int32_t;
using ValTypeA = uint4b_t;
using ValTypeB = uint4b_t;
using ValTypeC = int32_t;
};
template <>
struct MMA_Traits<SM80_16x8x64_S32U4U4S32_TN_SATURATE>
: MMA_Traits<SM80_16x8x64_S32U4U4S32_TN> {};
///////////////////////////////////////////////////////////////////////////////
/////////////////////////// s32 = b1 ^ b1 + s32 ///////////////////////////////
///////////////////////////////////////////////////////////////////////////////
@@ -440,9 +637,13 @@ struct MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC>
using CLayout = SM80_16x8_Row;
};
///////////////////////////////////////////////////////////////////////////////
/////////////////////////// s32 = b1 & b1 + s32 ///////////////////////////////
///////////////////////////////////////////////////////////////////////////////
template <>
struct MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_ANDPOPC>
:MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC> {};
: MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC> {};
template<>
struct MMA_Traits<SM80_8x8x128_S32U1U1S32_TN_XORPOPC>
@@ -455,7 +656,7 @@ struct MMA_Traits<SM80_8x8x128_S32U1U1S32_TN_XORPOPC>
using Shape_MNK = Shape<_8,_8,_128>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape<Shape<_4,_8>,_32>,
Stride<Stride<_256,_1>,_8>>;
Stride<Stride<_256,_1>,_8>>;
using BLayout = Layout<Shape<Shape<_4,_8>,_32>,
Stride<Stride<_256,_1>,_8>>;
using CLayout = SM80_8x8_Row;
@@ -472,7 +673,7 @@ struct MMA_Traits<SM80_16x8x128_S32U1U1S32_TN_XORPOPC>
using ValTypeA = cute::uint1b_t;
using ValTypeB = cute::uint1b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16,_8,_128>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape<Shape<_4,_8>,Shape<_32,_2>>,
include/cute/atom/mma_traits_sm90_gmma.hpp
-1
View File
@@ -1128,7 +1128,6 @@ struct MMA_Traits<SM90_64x32x16_F32F16F16_RS<tnspA, tnspB, scaleA, scaleB>>
////////////////////////////////////////////////////////////////////////////////////////////////////
template <
GMMA::Major tnspA,
GMMA::Major tnspB,
include/cute/atom/mma_traits_sm90_gmma_sparse.hpp
+1 -1
View File
@@ -7735,4 +7735,4 @@ struct MMA_Traits<SM90::GMMA::SPARSE::GMMA_64x256x64_F32E5M2E5M2_RS_TN<scaleA, s
#if defined(CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED)
#include "mma_traits_sm90_gmma_sparse_ext.hpp"
#endif
#endif
include/cute/config.hpp
+1 -1
View File
@@ -100,7 +100,7 @@
#if defined(_MSC_VER)
// Provides support for alternative operators 'and', 'or', and 'not'
# include <iso646.h>
# include <ciso646>
#endif // _MSC_VER
#if defined(__CUDACC_RTC__)
include/cute/container/array_subbyte.hpp
+29 -10
View File
@@ -100,20 +100,30 @@ public:
// Copy Ctor
CUTE_HOST_DEVICE constexpr
subbyte_reference(subbyte_reference const& other) {
*this = element_type(other);
subbyte_reference(subbyte_reference<value_type> const& other) {
*this = other.get();
}
CUTE_HOST_DEVICE constexpr
subbyte_reference(subbyte_reference<value_type const> const& other) {
*this = other.get();
}
// Copy Assignment
CUTE_HOST_DEVICE constexpr
subbyte_reference& operator=(subbyte_reference const& other) {
return *this = element_type(other);
subbyte_reference& operator=(subbyte_reference<value_type> const& other) {
return *this = other.get();
}
CUTE_HOST_DEVICE constexpr
subbyte_reference& operator=(subbyte_reference<value_type const> const& other) {
return *this = other.get();
}
// Assignment
template <class T_ = element_type>
CUTE_HOST_DEVICE constexpr
enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(element_type x)
enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(value_type x)
{
static_assert(is_same_v<T_, element_type>, "Do not specify template arguments!");
storage_type item = (reinterpret_cast<storage_type const&>(x) & BitMask);
@@ -149,11 +159,11 @@ public:
// Value
CUTE_HOST_DEVICE
element_type get() const
value_type get() const
{
if constexpr (is_same_v<bool, value_type>) { // Extract to bool -- potentially faster impl
return bool((*ptr_) & (BitMask << idx_));
} else { // Extract to element_type
} else { // Extract to value_type
// Extract from the current storage element
auto item = storage_type((ptr_[0] >> idx_) & BitMask);
@@ -165,13 +175,13 @@ public:
item |= storage_type((ptr_[1] & bit_mask_1) << straddle_bits);
}
return reinterpret_cast<element_type&>(item);
return reinterpret_cast<value_type&>(item);
}
}
// Extract to type element_type
// Extract to type value_type
CUTE_HOST_DEVICE constexpr
operator element_type() const {
operator value_type() const {
return get();
}
@@ -341,6 +351,14 @@ recast_ptr(subbyte_iterator<T> const& x) {
CUTE_GCC_UNREACHABLE;
}
// Dynamic pointers have unknown static alignment
template <class T>
CUTE_HOST_DEVICE constexpr
Int<0>
max_alignment(subbyte_iterator<T> const& x) {
return {};
}
template <class T>
CUTE_HOST_DEVICE void
print(subbyte_iterator<T> const& x) {
@@ -352,6 +370,7 @@ CUTE_HOST_DEVICE void
print(subbyte_reference<T> const& x) {
print(x.get());
}
//
// array_subbyte
// Statically sized array for non-byte-aligned data types
include/cute/layout.hpp
+1 -1
View File
@@ -1830,7 +1830,7 @@ recast_layout(Layout<Shape,Stride> const& layout)
return upcast<scale::num>(layout);
}
else {
static_assert(dependent_false<scale>, "Recast not supported.");
return downcast<scale::den>(upcast<scale::num>(layout));
}
CUTE_GCC_UNREACHABLE;
include/cute/layout_composed.hpp
+10 -1
View File
@@ -616,7 +616,7 @@ recast_layout(ComposedLayout<A,O,B> const& layout)
return upcast<scale::num>(layout);
}
else {
static_assert(dependent_false<scale>, "Recast not supported.");
return downcast<scale::den>(upcast<scale::num>(layout));
}
CUTE_GCC_UNREACHABLE;
}
@@ -631,6 +631,15 @@ max_alignment(ComposedLayout<A,O,B> const& layout)
return Int<1>{};
}
template <class A, class O, class B>
CUTE_HOST_DEVICE constexpr
auto
nullspace(ComposedLayout<A,O,B> const& layout)
{
// Do not attempt for general ComposedLayouts
return Layout<_1,_0>{};
}
//
// Display utilities
//
include/cute/numeric/integral_ratio.hpp
+36 -7
View File
@@ -154,13 +154,6 @@ operator*(C<c>, R<a,b>) {
return {};
}
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<c*b,a>::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)>
@@ -179,6 +172,13 @@ operator*(R<a,b>, C const& c) {
return c * R<a,b>::num / R<a,b>::den;
}
template <class C, auto a, auto b>
CUTE_HOST_DEVICE constexpr
auto
operator/(C const& c, R<a,b>) {
return c * R<b,a>{};
}
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
typename R<a*y+b*x, b*y>::type
@@ -200,6 +200,10 @@ operator+(C<c>, R<a,b>) {
return {};
}
/////////////////
// Comparisons //
/////////////////
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>
@@ -221,6 +225,31 @@ operator==(C<c>, R<a,b>) {
return {};
}
///////////////////////
// Special functions //
///////////////////////
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
typename R<gcd(a*y,b*x),b*x>::type
gcd(R<a,b>, R<x,y>) {
return {};
}
template <auto a, auto b, auto c>
CUTE_HOST_DEVICE constexpr
typename R<gcd(a,b*c),b*c>::type
gcd(R<a,b>, C<c>) {
return {};
}
template <auto c, auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<gcd(a,b*c),b*c>::type
gcd(C<c>, R<a,b>) {
return {};
}
template <auto a, auto b>
CUTE_HOST_DEVICE constexpr
typename R<abs(a),abs(b)>::type
include/cute/numeric/numeric_types.hpp
+1
View File
@@ -46,6 +46,7 @@ template <class T>
static constexpr auto sizeof_bits_v = sizeof_bits<T>::value;
using cutlass::bits_to_bytes;
using cutlass::bytes_to_bits;
using cutlass::is_subbyte;
include/cute/pointer.hpp
+8
View File
@@ -214,6 +214,14 @@ make_smem_ptr(void const* ptr) {
return make_smem_ptr(recast_ptr<T const>(ptr));
}
// nullptr_t overload for make_smem_ptr<float>(nullptr) disambiguation
template <class T>
CUTE_HOST_DEVICE constexpr
auto
make_smem_ptr(decltype(nullptr)) { // nullptr_t
return make_smem_ptr(recast_ptr<T>(nullptr));
}
// The smem tag is invariant over type-recast
template <class NewT, class P>
CUTE_HOST_DEVICE constexpr
include/cute/pointer_base.hpp
+19 -3
View File
@@ -30,9 +30,10 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/numeric_types.hpp> // cute::sizeof_bits
#include <cute/util/type_traits.hpp> // cute::declval, cute::void_t, etc
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/numeric_types.hpp> // cute::sizeof_bits
#include <cute/numeric/integral_constant.hpp> // Int<0>
#include <cute/util/type_traits.hpp> // cute::declval, cute::void_t, etc
namespace cute
{
@@ -115,6 +116,14 @@ raw_pointer_cast(T* ptr) {
return ptr;
}
// The statically-known alignment of a dynamic pointer is unknown
template <class T>
CUTE_HOST_DEVICE constexpr
Int<0>
max_alignment(T*) {
return {};
}
//
// A very simplified iterator adaptor.
// Derived classed may override methods, but be careful to reproduce interfaces exactly.
@@ -169,6 +178,13 @@ raw_pointer_cast(iter_adaptor<I,D> const& x) {
return raw_pointer_cast(x.ptr_);
}
template <class I, class D>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(iter_adaptor<I,D> const& x) {
return max_alignment(x.ptr_);
}
//
// counting iterator -- quick and dirty
//
include/cute/pointer_swizzle.hpp
+8
View File
@@ -147,6 +147,14 @@ recast_ptr(swizzle_ptr<SwizzleFn,P> const& ptr) {
return make_swizzle_ptr(recast_ptr<NewT>(ptr.get()), SwizzleFn{});
}
// The statically-known alignment of a swizzle pointer is the alignment of the swizzle function converted to bits
template <class SwizzleFn, class P>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(swizzle_ptr<SwizzleFn,P> const&) {
return Int<8>{} * max_alignment(SwizzleFn{});
}
//
// Display utilities
//
include/cute/swizzle_layout.hpp
+2 -2
View File
@@ -447,7 +447,7 @@ recast_layout(Swizzle<B,M,S> const& swizzle)
return upcast<scale::num>(swizzle);
}
else {
static_assert(dependent_false<scale>, "Recast not supported.");
return downcast<scale::den>(upcast<scale::num>(layout));
}
CUTE_GCC_UNREACHABLE;
}
@@ -457,7 +457,7 @@ CUTE_HOST_DEVICE constexpr
auto
max_alignment(Swizzle<B,M,S> const&)
{
return Int<1 << M>{};
return Int<(1 << M)>{};
}
template <int B, int M, int S, class Offset, class LayoutB>
include/cute/tensor_impl.hpp
+13
View File
@@ -84,6 +84,8 @@ struct ArrayEngine
};
// Specialization for sparse_elem<S,T> tensor allocation/iteration
// NOTE: This can and should be used for allocation of SMEM as well!
// Fuse these two ArrayEngines?
template <int S, class T, size_t N>
struct ArrayEngine<sparse_elem<S,T>, N>
{
@@ -858,6 +860,17 @@ max_common_layout(Tensor<SrcEngine,SrcLayout> const& a,
CUTE_GCC_UNREACHABLE;
}
/* Return the maximum (statically known) alignment of a Tensor in the number of bits
*/
template <class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(Tensor<Engine,Layout> const& t)
{
return gcd(max_alignment(t.data()),
max_alignment(t.layout()) * static_value<sizeof_bits<typename Engine::value_type>>());
}
//
// Key algebraic operations -- Composition, Divide, and Product
//
include/cute/util/debug.hpp
+2 -2
View File
@@ -123,7 +123,7 @@ bool
block([[maybe_unused]] int bid)
{
#if defined(__CUDA_ARCH__)
return blockIdx.x + blockIdx.y*gridDim.x + blockIdx.z*gridDim.x*gridDim.y == bid;
return blockIdx.x + blockIdx.y*gridDim.x + blockIdx.z*gridDim.x*gridDim.y == static_cast<unsigned int>(bid);
#else
return true;
#endif
@@ -134,7 +134,7 @@ bool
thread([[maybe_unused]] int tid, [[maybe_unused]] int bid)
{
#if defined(__CUDA_ARCH__)
return (threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.x*blockDim.y == tid) && block(bid);
return (threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.x*blockDim.y == static_cast<unsigned int>(tid)) && block(bid);
#else
return true;
#endif
include/cute/util/type_traits.hpp
+6
View File
@@ -141,9 +141,15 @@ using CUTE_STL_NAMESPACE::common_type_t;
using CUTE_STL_NAMESPACE::remove_pointer;
using CUTE_STL_NAMESPACE::remove_pointer_t;
using CUTE_STL_NAMESPACE::add_pointer;
using CUTE_STL_NAMESPACE::add_pointer_t;
using CUTE_STL_NAMESPACE::alignment_of;
using CUTE_STL_NAMESPACE::alignment_of_v;
using CUTE_STL_NAMESPACE::is_pointer;
using CUTE_STL_NAMESPACE::is_pointer_v;
// <utility>
using CUTE_STL_NAMESPACE::declval;
include/cutlass/arch/barrier.h
+93
View File
@@ -47,6 +47,99 @@ namespace cutlass {
namespace arch {
////////////////////////////////////////////////////////////////////////////////////////////////////
CUTLASS_DEVICE void fence_view_async_shared();
namespace detail { // namespace detail begin
// Single threaded versions that need to be called in an elect_one region
template<typename T, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array(T ptr, int arv_cnt) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
ptr[i].init(arv_cnt);
}
}
template<typename T, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array(uint64_t *ptr, int arv_cnt) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
T::init(&ptr[i], arv_cnt);
}
}
template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array_pair(FullBarrier full_barriers, EmptyBarrier empty_barriers, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
full_barriers[i].init(full_barrier_arv_cnt);
empty_barriers[i].init(empty_barrier_arv_cnt);
}
}
template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array_pair(uint64_t *full_barriers_ptr, uint64_t *empty_barriers_ptr, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
FullBarrier::init(&full_barriers_ptr[i], full_barrier_arv_cnt);
EmptyBarrier::init(&empty_barriers_ptr[i], empty_barrier_arv_cnt);
}
}
// Aligned versions that need to be call warp wide
template<typename T, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array_aligned(T ptr, int arv_cnt) {
if(cute::elect_one_sync()) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
ptr[i].init(arv_cnt);
}
}
}
template<typename T, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array_aligned(uint64_t *ptr, int arv_cnt) {
if(cute::elect_one_sync()) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
T::init(&ptr[i], arv_cnt);
}
}
}
template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array_pair_aligned(FullBarrier full_barriers, EmptyBarrier empty_barriers, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
if(cute::elect_one_sync()) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
full_barriers[i].init(full_barrier_arv_cnt);
empty_barriers[i].init(empty_barrier_arv_cnt);
}
}
}
template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
CUTLASS_DEVICE
void initialize_barrier_array_pair_aligned(uint64_t *full_barriers_ptr, uint64_t *empty_barriers_ptr, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
if(cute::elect_one_sync()) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Stages; i++) {
FullBarrier::init(&full_barriers_ptr[i], full_barrier_arv_cnt);
EmptyBarrier::init(&empty_barriers_ptr[i], empty_barrier_arv_cnt);
}
}
}
} // namespace detail end
// Enumerates the reserved named barriers to avoid potential conflicts
// This enum class specifies the NamedBarriers reserved by CUTLASS.
enum class ReservedNamedBarriers {
include/cutlass/arch/config.h
+4
View File
@@ -35,6 +35,8 @@
#pragma once
#include "cutlass/platform/platform.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM90
@@ -79,3 +81,5 @@
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/memory_sm75.h
+10 -9
View File
@@ -35,6 +35,7 @@
#pragma once
#include "cutlass/array.h"
#include "cutlass/detail/helper_macros.hpp"
#include "cutlass/layout/matrix.h"
#include "cute/arch/copy_sm75.hpp"
#include "cute/arch/util.hpp"
@@ -50,7 +51,7 @@ template <
/// .x1, .x2, or .x4
int MatrixCount
>
inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
CUTLASS_DEVICE void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
/////////////////////////////////////////////////////////////////////////////////////////////////
//
@@ -59,19 +60,19 @@ inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
/////////////////////////////////////////////////////////////////////////////////////////////////
/// CUTLASS helper to get SMEM pointer
inline __device__ unsigned cutlass_get_smem_pointer(void *ptr) {
CUTLASS_DEVICE unsigned cutlass_get_smem_pointer(void *ptr) {
return cute::cast_smem_ptr_to_uint(ptr);
}
/// CUTLASS helper to get SMEM pointer
inline __device__ unsigned cutlass_get_smem_pointer(void const *ptr) {
CUTLASS_DEVICE unsigned cutlass_get_smem_pointer(void const *ptr) {
return cutlass_get_smem_pointer(const_cast<void *>(ptr));
}
/////////////////////////////////////////////////////////////////////////////////////////////////
template <>
inline __device__ void ldsm<layout::RowMajor, 1>(
CUTLASS_DEVICE void ldsm<layout::RowMajor, 1>(
Array<unsigned, 1> & D,
void const* ptr) {
@@ -95,7 +96,7 @@ inline __device__ void ldsm<layout::RowMajor, 1>(
/////////////////////////////////////////////////////////////////////////////////////////////////
template <>
inline __device__ void ldsm<layout::RowMajor, 2>(
CUTLASS_DEVICE void ldsm<layout::RowMajor, 2>(
Array<unsigned, 2> & D,
void const* ptr) {
@@ -119,7 +120,7 @@ inline __device__ void ldsm<layout::RowMajor, 2>(
/////////////////////////////////////////////////////////////////////////////////////////////////
template <>
inline __device__ void ldsm<layout::RowMajor, 4>(
CUTLASS_DEVICE void ldsm<layout::RowMajor, 4>(
Array<unsigned, 4> & D,
void const* ptr) {
@@ -147,7 +148,7 @@ inline __device__ void ldsm<layout::RowMajor, 4>(
/////////////////////////////////////////////////////////////////////////////////////////////////
template <>
inline __device__ void ldsm<layout::ColumnMajor, 1>(
CUTLASS_DEVICE void ldsm<layout::ColumnMajor, 1>(
Array<unsigned, 1> & D,
void const* ptr) {
@@ -171,7 +172,7 @@ inline __device__ void ldsm<layout::ColumnMajor, 1>(
/////////////////////////////////////////////////////////////////////////////////////////////////
template <>
inline __device__ void ldsm<layout::ColumnMajor, 2>(
CUTLASS_DEVICE void ldsm<layout::ColumnMajor, 2>(
Array<unsigned, 2> & D,
void const* ptr) {
@@ -195,7 +196,7 @@ inline __device__ void ldsm<layout::ColumnMajor, 2>(
/////////////////////////////////////////////////////////////////////////////////////////////////
template <>
inline __device__ void ldsm<layout::ColumnMajor, 4>(
CUTLASS_DEVICE void ldsm<layout::ColumnMajor, 4>(
Array<unsigned, 4> & D,
void const* ptr) {
include/cutlass/arch/mma_sm70.h
-4
View File
@@ -33,11 +33,7 @@
*/
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm75.h
-4
View File
@@ -34,11 +34,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/arch/wmma.h"
include/cutlass/arch/mma_sm80.h
-4
View File
@@ -34,11 +34,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "mma.h"
include/cutlass/arch/mma_sm89.h
-4
View File
@@ -35,11 +35,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "mma.h"
include/cutlass/arch/mma_sm90.h
-4
View File
@@ -34,11 +34,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sparse_sm80.h
-4
View File
@@ -35,11 +35,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sparse_sm89.h
-4
View File
@@ -35,11 +35,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/simd.h
+2 -2
View File
@@ -34,8 +34,8 @@
#pragma once
#include "../array.h"
#include "../numeric_types.h"
#include "cutlass/arch/array.h"
#include "cutlass/arch/numeric_types.h"
namespace cutlass {
namespace arch {
include/cutlass/arch/synclog.hpp
+1 -1
View File
@@ -59,7 +59,7 @@ constexpr uint32_t synclog_cap = 1 << 26;
inline std::mutex synclog_mutex;
inline std::vector<uint32_t*> synclog_buf_list;
#if defined(__NVCC__) || (defined(__clang__) && defined(__CUDA__))
inline __device__ uint32_t* synclog_buf;
CUTLASS_DEVICE uint32_t* synclog_buf;
#endif
CUTLASS_DEVICE
include/cutlass/arch/wmma_sm70.h
-4
View File
@@ -34,11 +34,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/wmma_sm72.h
-4
View File
@@ -34,11 +34,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/wmma_sm75.h
-4
View File
@@ -34,11 +34,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/array.h
+4 -15
View File
@@ -2573,20 +2573,8 @@ Array<T, N> fma(Array<T, N> const &a, Array<T, N> const &b, T c) {
return op(a, b, c);
}
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/array_subbyte.h"
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
////////////////////////////////////////////////////////////////////////////////////////////////////
// AlignedArray
@@ -2606,9 +2594,10 @@ public:
};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/array_subbyte.h"
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/array_subbyte.h
+2
View File
@@ -554,6 +554,8 @@ private:
////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/blas3.h
+1 -1
View File
@@ -132,7 +132,7 @@ struct MantissaInBits<double> {
template <>
struct MantissaInBits<cutlass::complex<double>> {
static int constexpr bits = 30;
static double constexpr error = 1.0e-15;
static double constexpr error = 1.0e-14;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp
+95 -5
View File
@@ -189,7 +189,7 @@ private:
-problem_shape.dilation[NumSpatialDimensions-1-i] :
problem_shape.dilation[NumSpatialDimensions-1-i];
}
return make_im2col_tma_copy(
GmemTiledCopyA{},
tensor_a,
@@ -225,7 +225,7 @@ private:
auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
auto lower_srt = detail::compute_lower_srt(problem_shape);
return make_im2col_tma_copy(
GmemTiledCopyB{},
tensor_b,
@@ -372,6 +372,96 @@ public:
return false;
}
if (is_im2col_A || is_im2col_B) {
// Check valid corner values for TMA_LOAD_IM2COL, signed int ranging from [-corner_limit, corner_limit - 1]
constexpr int32_t corner_limit = 1 << (16 / NumSpatialDimensions - 1);
auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
for (int i = 0; i < problem_shape.RankS; ++i) {
implementable = implementable && lower_corner_whd[i] >= -corner_limit && lower_corner_whd[i] <= (corner_limit - 1);
}
auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
for (int i = 0; i < problem_shape.RankS; ++i) {
implementable = implementable && upper_corner_whd[i] >= -corner_limit && upper_corner_whd[i] <= (corner_limit - 1);
}
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Padding values don't meet requirements for TMA LOAD IM2COL.\n");
return false;
}
}
// Wgrad kernels don't support non-packed output strides, non-packed tensor A stride (linearized)
if constexpr (ConvOp == conv::Operator::kWgrad) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::ostringstream os;
#endif
const auto & input_shape = problem_shape.shape_A;
const auto & input_stride = problem_shape.stride_A;
implementable &= input_stride[ProblemShape::RankT - 1] == 1;
int input_shape_size = 1;
for (int i = ProblemShape::RankT - 2; i >= 0; --i) {
input_shape_size *= input_shape[i + 1];
implementable &= input_stride[i] == input_shape_size;
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
if (input_stride[i] != input_shape_size) {
os << "\n *** input_stride[" << i << "] = " << input_stride[i] << " != input_shape_size = " << input_shape_size << " ***";
}
#endif
}
if (!implementable) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
os << "\n input_shape_size: " << input_shape_size
<< "\n input_shape: " << input_shape
<< "\n input_stride: " << input_stride
<< "\n";
#endif
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Wgrad kernels don't support non-packed input strides.\n");
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST(os.str());
#endif
return false;
}
const auto & output_shape = problem_shape.shape_C;
const auto & output_stride = problem_shape.stride_C;
implementable &= output_stride[ProblemShape::RankT - 1] == 1;
int output_shape_size = 1;
for (int i = ProblemShape::RankT - 2; i >= 0; --i) {
output_shape_size *= output_shape[i + 1];
implementable &= output_stride[i] == output_shape_size;
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
if (output_stride[i] != output_shape_size) {
os << "\n *** output_stride[" << i << "] = " << output_stride[i] << " != output_shape_size = " << output_shape_size << " ***";
}
#endif
}
if (!implementable) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
os << "\n output_shape_size: " << input_shape_size
<< "\n output_shape: " << input_shape
<< "\n output_stride: " << input_stride
<< "\n";
#endif
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Wgrad kernels don't support non-packed output strides.\n");
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST(os.str());
#endif
return false;
}
}
// Conv kernels only support cross correlation mode currently.
implementable &= problem_shape.mode == cutlass::conv::Mode::kCrossCorrelation;
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Conv kernels only support cross correlation mode currently.\n");
return false;
}
if (problem_shape.groups > 1) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: This kernel does not support conv groups > 1.\n");
return false;
@@ -516,9 +606,9 @@ public:
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* Waits for all stages to either be released (all
* Consumer UNLOCKs), or if the stage was never used
* then would just be acquired since the phase was
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_producer_state);
@@ -645,7 +735,7 @@ public:
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/conv/convnd_problem_shape.hpp
+40
View File
@@ -319,6 +319,7 @@ struct ConvProblemShape {
// | ShapeB | KTRSC | KTRSC | NDHWC |
// | ShapeC | NZPQK | NDHWC | KTRSC |
//
// Input comes from calculate_xformed_act, which does NOT depend on ConvOp.
CUTLASS_HOST_DEVICE
constexpr void
set_shape_stride_ABC(
@@ -328,6 +329,31 @@ struct ConvProblemShape {
TensorStride stride_flt,
TensorExtent shape_xformed_act,
TensorStride stride_xformed_act) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
printf("*** set_shape_stride_ABC ***");
printf("\n shape_act: ");
print(shape_act);
printf("\n stride_act: ");
print(stride_act);
printf("\n shape_flt: ");
print(shape_flt);
printf("\n stride_flt: ");
print(stride_flt);
printf("\n shape_xformed_act: ");
print(shape_xformed_act);
printf("\n stride_xformed_act: ");
print(stride_xformed_act);
if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
printf("\n ConvOp: Fprop");
}
if constexpr (ConvOp == cutlass::conv::Operator::kDgrad) {
printf("\n ConvOp: Dgrad");
}
if constexpr (ConvOp == cutlass::conv::Operator::kWgrad) {
printf("\n ConvOp: Wgrad");
}
printf("\n");
#endif
if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
shape_A = shape_act;
@@ -353,6 +379,20 @@ struct ConvProblemShape {
shape_C = shape_flt;
stride_C = stride_flt;
}
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
printf("\n shape_A: ");
print(shape_A);
printf("\n stride_A: ");
print(stride_A);
printf("\n shape_B: ");
print(shape_B);
printf("\n stride_B: ");
print(stride_B);
printf("\n shape_C: ");
print(shape_C);
printf("\n stride_C: ");
print(stride_C);
#endif
}
// Get A extents.
include/cutlass/conv/kernel/direct_convolution.h
+2 -1
View File
@@ -40,6 +40,7 @@
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/platform/platform.h"
#include "cutlass/semaphore.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/layout/tensor.h"
@@ -155,7 +156,7 @@ struct DirectConvolutionParams {
swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
// Dynamic SMEM usage because stride and dilation are runtime params.
smem_size_ = (max(iterator_A.activation_size, int(sizeof(typename Epilogue::SharedStorage))) * kStages + iterator_B.filter_size);
smem_size_ = (cutlass::platform::max(iterator_A.activation_size, int(sizeof(typename Epilogue::SharedStorage))) * kStages + iterator_B.filter_size);
}
CUTLASS_HOST_DEVICE
include/cutlass/coord.h
+1 -1
View File
@@ -37,7 +37,7 @@
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#else
#include <stdint.h>
#include <cstdint>
#endif
#include "cutlass/cutlass.h"
include/cutlass/cuda_host_adapter.hpp
+7 -2
View File
@@ -85,7 +85,11 @@ namespace cutlass {
#if !defined(__CUDACC_RTC__)
#if ((__CUDACC_VER_MAJOR__ >= 12) || \
((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)))
#include <cudaTypedefs.h>
#endif // (__CUDACC_VERSION__ >= 11.8)
#include <driver_types.h>
#define CUTLASS_CUDA_DRIVER_STRINGIFY(tok) #tok
@@ -100,7 +104,8 @@ namespace cutlass {
#else // defined(CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL)
#if (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 5)
#if ((__CUDACC_VER_MAJOR__ >= 13) || \
((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 5))) \
#define CUTLASS_CUDA_DRIVER_WRAPPER_DECL(func, ver) \
template <typename... Args> \
@@ -138,7 +143,7 @@ namespace cutlass {
return reinterpret_cast<PFN_##func>(pfn)(args...); \
}
#endif // (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 5)
#endif // (__CUDACC_VERSION__ >= 12.5)
#endif // defined(CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL)
include/cutlass/detail/collective.hpp
+1
View File
@@ -31,6 +31,7 @@
#pragma once
#include "cute/container/tuple.hpp"
#include "cute/layout.hpp" // cute::size(shape)
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
include/cutlass/detail/collective/mixed_input_utils.hpp
+6 -6
View File
@@ -237,7 +237,7 @@ struct LayoutAwareConvertImpl<
}
};
// Specialization for UINT4 -> FPF16 with [02461357] value order
// Specialization for UINT4 -> FP16 with [02461357] value order
template <>
struct LayoutAwareConvertImpl<
cutlass::uint4b_t,
@@ -754,7 +754,6 @@ public:
cute::tuple<Ts...>& partitioned_extra_info,
int const k_block) {
static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
static_assert(cosize_v<LayoutIn> == cosize_v<LayoutOut>);
@@ -805,14 +804,15 @@ public:
{
auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_vm_(i));
auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2> &>(scales_pos_vm_(i));
constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
asm volatile(
"{\n"
" and .b32 %0, %2, %4 ;\n" \
" and .b32 %1, %3, %5 ;\n" \
" lop3 .b32 %0, %2, %4, %5, %6;\n" \
" xor .b32 %1, %3, %5; \n" \
"}\n"
: "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0x7F7F7F00), "n"(0x7F7F7F7F)
);
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut)
);
}
}
CUTLASS_PRAGMA_UNROLL
include/cutlass/detail/helper_macros.hpp
+8 -2
View File
@@ -57,6 +57,12 @@
#define CUTLASS_DEVICE inline
#endif
#if ! defined(_MSC_VER)
#define CUTLASS_LAMBDA_FUNC_INLINE __attribute__((always_inline))
#else
#define CUTLASS_LAMBDA_FUNC_INLINE [[msvc::forceinline]]
#endif
#define CUTLASS_HOST __host__
#define CUTLASS_GLOBAL __global__ static
@@ -74,11 +80,11 @@ CUTLASS_HOST_DEVICE void __CUTLASS_UNUSED(T const &)
#ifdef _MSC_VER
// Provides support for alternative operators 'and', 'or', and 'not'
#include <iso646.h>
#include <ciso646>
#endif // _MSC_VER
#if !defined(__CUDACC_RTC__)
#include <assert.h>
#include <cassert>
#endif
#if defined(__CUDA_ARCH__)
include/cutlass/detail/mainloop_fusion_helper_scale_factor.hpp
+75
View File
@@ -0,0 +1,75 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 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 Mainloop Fusion configs specific for scale factors
*/
#pragma once
#include <cute/util/type_traits.hpp> // cute::void_t
namespace cutlass::detail {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename CollectiveMainloop, typename = void>
struct ElementSFType {
using type = void;
};
template <typename CollectiveMainloop>
struct ElementSFType<CollectiveMainloop, cute::void_t<typename CollectiveMainloop::ElementSF>> {
using type = typename CollectiveMainloop::ElementSF;
};
template <typename CollectiveMainloop, typename = void>
struct LayoutSFAType {
using type = void;
};
template <typename CollectiveMainloop>
struct LayoutSFAType<CollectiveMainloop, cute::void_t<typename CollectiveMainloop::LayoutSFA>> {
using type = typename CollectiveMainloop::LayoutSFA;
};
template <typename CollectiveMainloop, typename = void>
struct LayoutSFBType {
using type = void;
};
template <typename CollectiveMainloop>
struct LayoutSFBType<CollectiveMainloop, cute::void_t<typename CollectiveMainloop::LayoutSFB>> {
using type = typename CollectiveMainloop::LayoutSFB;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::detail
include/cutlass/device_kernel.h
+4 -1
View File
@@ -34,8 +34,11 @@
#pragma once
#include <cutlass/detail/helper_macros.hpp> // CUTLASS_HOST_DEVICE
#include <cutlass/platform/platform.h> // uint64_t
// __grid_constant__ was introduced in CUDA 11.7.
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7)))
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7))) && !CUTLASS_CLANG_CUDA
# define CUTLASS_GRID_CONSTANT_SUPPORTED
#endif
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+2 -1
View File
@@ -422,7 +422,8 @@ struct CollectiveBuilder<
Schedule,
fusion::LinearCombination<ElementD,ElementCompute,ElementC_,ElementCompute,RoundStyle>,
cute::enable_if_t<cute::is_same_v<Schedule, NoSmemWarpSpecialized> ||
cute::is_same_v<Schedule, PtrArrayNoSmemWarpSpecialized> >> {
cute::is_same_v<Schedule, PtrArrayNoSmemWarpSpecialized> ||
cute::is_same_v<Schedule, PtrArrayNoSmemWarpSpecializedTransposed> >> {
// Passing void C disables source load
using ElementC = cute::conditional_t<cute::is_void_v<ElementC_>,
include/cutlass/epilogue/collective/default_epilogue_array.hpp
+24 -14
View File
@@ -86,7 +86,7 @@ public:
static const int kOutputAlignment = ThreadEpilogueOp::kCount;
using AlignmentType = typename cute::uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
static_assert(cute::is_same_v<EpilogueSchedule, PtrArrayNoSmemWarpSpecialized> || cute::is_same_v<EpilogueSchedule, PtrArrayDefault>, "Incompatible epilogue schedule.");
static_assert(cute::is_same_v<EpilogueSchedule, PtrArrayNoSmemWarpSpecialized> || cute::is_same_v<EpilogueSchedule, PtrArrayDefault> || cute::is_same_v<EpilogueSchedule, PtrArrayNoSmemWarpSpecializedTransposed>, "Incompatible epilogue schedule.");
static_assert(rank(InternalStrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]");
static_assert(rank(InternalStrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
@@ -198,20 +198,30 @@ public:
assert(0);
}
InternalStrideC stride_c;
InternalStrideD stride_d;
if constexpr (!cute::is_same_v<InternalStrideC, StrideC>) {
// If grouped gemm
if (epilogue_op.is_source_needed()) {
stride_c = detail::get_epilogue_stride<EpilogueSchedule>(params.dC[l_coord]);
auto [stride_c, stride_d] = [&, l = l_coord]() {
if constexpr (!cute::is_same_v<InternalStrideC, StrideC>) {
// If grouped gemm
if (epilogue_op.is_source_needed()) {
return make_tuple(
detail::get_epilogue_stride<EpilogueSchedule>(params.dC[l]),
detail::get_epilogue_stride<EpilogueSchedule>(params.dD[l])
);
}
else {
return make_tuple(
InternalStrideC{},
detail::get_epilogue_stride<EpilogueSchedule>(params.dD[l])
);
}
}
else {
return make_tuple(
detail::get_epilogue_stride<EpilogueSchedule>(params.dC),
detail::get_epilogue_stride<EpilogueSchedule>(params.dD)
);
}
stride_d = detail::get_epilogue_stride<EpilogueSchedule>(params.dD[l_coord]);
}
else {
stride_c = detail::get_epilogue_stride<EpilogueSchedule>(params.dC);
stride_d = detail::get_epilogue_stride<EpilogueSchedule>(params.dD);
}
}();
// Represent the full output tensor
ElementC const* ptr_C_l = nullptr;
if (epilogue_op.is_source_needed()) {
include/cutlass/epilogue/collective/detail.hpp
+13 -2
View File
@@ -157,7 +157,8 @@ struct EmptyStorage {
template<class EpilogueSchedule, class Stride>
CUTLASS_HOST_DEVICE
auto get_epilogue_stride(Stride stride){
if constexpr (cute::is_base_of_v<cutlass::gemm::EpilogueTransposed, EpilogueSchedule>) {
if constexpr (cute::is_base_of_v<cutlass::gemm::EpilogueTransposed, EpilogueSchedule>||
cute::is_base_of_v<cutlass::epilogue::PtrArrayNoSmemWarpSpecializedTransposed, EpilogueSchedule>) {
return cute::make_stride(cute::get<1>(stride), cute::get<0>(stride), cute::get<2>(stride));
}
else {
@@ -464,7 +465,7 @@ public:
tensormaps_fence_acquire([[maybe_unused]] cute::TmaDescriptor const* tensormap) { }
};
// SFINAE helpers for detecting beta/beta_ptr in EVT arguments.
// SFINAE helpers for detecting beta/beta_ptr/beta_ptr_array in EVT arguments.
template <class Arguments, class = void>
struct has_beta {
static constexpr bool value = false;
@@ -485,6 +486,16 @@ struct has_beta_ptr<Arguments, cute::void_t<decltype(Arguments{}.thread.beta_ptr
static constexpr bool value = true;
};
template <class Arguments, class = void>
struct has_beta_ptr_array {
static constexpr bool value = false;
};
template <class Arguments>
struct has_beta_ptr_array<Arguments, cute::void_t<decltype(Arguments{}.thread.beta_ptr_array)>> {
static constexpr bool value = true;
};
} // namespace detail
} // namespace collective
} // namespace epilogue
include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp
+25 -18
View File
@@ -328,7 +328,7 @@ public:
}
uint32_t transaction_bytes = TmaTransactionBytes;
typename Params::TMA_C tma_load_c = {};
typename Params::TMA_C tma_load_c{};
if constexpr (is_source_supported) {
ElementC const* ptr_C_first_batch = reinterpret_cast<ElementC const*>(args.ptr_C);
Tensor tensor_c = make_tensor(ptr_C_first_batch, make_layout(make_shape(init_M,init_N,init_L), append<3>(stride_c, _0{})));
@@ -409,7 +409,7 @@ public:
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_D>(cute::make_shape(M,N,L), InternalStrideD{});
}
if constexpr (not cute::is_void_v<ElementC>) {
if constexpr (is_source_supported) {
constexpr int tma_alignment_bits_C = cutlass::detail::get_input_alignment_bits<ElementC>();
constexpr int min_tma_aligned_elements_C = tma_alignment_bits_C / cutlass::sizeof_bits<ElementC>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_C>(cute::make_shape(M,N,L), InternalStrideC{});
@@ -432,13 +432,16 @@ public:
bool beta_implementable = true;
if constexpr (cute::is_void_v<ElementC>) {
if (cute::is_void_v<ElementC> || args.ptr_C == nullptr) {
if constexpr (detail::has_beta<Arguments>::value) {
beta_implementable = args.thread.beta == 0.0;
}
if constexpr (detail::has_beta_ptr<Arguments>::value) {
beta_implementable = beta_implementable && args.thread.beta_ptr == nullptr;
}
if constexpr (detail::has_beta_ptr_array<Arguments>::value) {
beta_implementable = beta_implementable && args.thread.beta_ptr_array == nullptr;
}
}
if (!beta_implementable) {
@@ -775,7 +778,7 @@ public:
tRS_rC,
thread_idx
};
auto cst_callbacks = fusion_callbacks.get_consumer_store_callbacks<RefSrc>(cst_args);
auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
bool is_producer_load_needed = fusion_callbacks.is_producer_load_needed();
bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed();
@@ -1017,7 +1020,7 @@ public:
Tensor gmem_tensormap = make_tensor(params.tensormaps, desc_layout); // (SMs, NumInputTensors)
if constexpr (IsLoad) {
if (not cute::is_void_v<ElementC>) {
if (is_source_supported) {
constexpr int C_tensormap_index = NumEpilogueWarpGroups;
Tensor pC_tensormap = make_tensor(params.tma_load_c.get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sC_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_C), Int<1>{}, Int<1>{});
@@ -1058,8 +1061,10 @@ public:
// Replacing global_address for the next batch
if constexpr (IsLoad) {
if constexpr (is_source_supported) {
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_C,
params.ptr_C[next_batch]);
if (params.ptr_C != nullptr) {
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_C,
params.ptr_C[next_batch]);
}
}
}
else if constexpr (is_destination_supported) {
@@ -1087,18 +1092,20 @@ public:
if constexpr (IsLoad) {
if constexpr (is_source_supported) {
ElementC const* ptr_C = nullptr;
Tensor tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), params.dC[next_group]));
if (params.dC != nullptr) {
ElementC const* ptr_C = nullptr;
Tensor tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), params.dC[next_group]));
cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c,
prob_shape, prob_stride);
// Convert strides to byte strides
for (uint64_t& stride : prob_stride) {
stride = (stride * sizeof_bits_v<ElementC>) / 8;
cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c,
prob_shape, prob_stride);
// Convert strides to byte strides
for (uint64_t& stride : prob_stride) {
stride = (stride * sizeof_bits_v<ElementC>) / 8;
}
cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_C,
prob_shape,
prob_stride);
}
cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_C,
prob_shape,
prob_stride);
}
}
else if constexpr (is_destination_supported) {
@@ -1166,7 +1173,7 @@ public:
void
tensormaps_fence_acquire(cute::TmaDescriptor const* tensormap) {
if constexpr (IsLoad) {
if constexpr (not cute::is_void_v<ElementC>) {
if constexpr (is_source_supported) {
cute::tma_descriptor_fence_acquire(tensormap);
}
}
include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp
+21 -7
View File
@@ -94,7 +94,7 @@ class CollectiveEpilogue<
SmemLayoutAtomD_,
CopyOpR2S_,
CopyAtomC_,
CopyOpR2R_,
CopyOpR2R_
> {
public:
//
@@ -136,6 +136,9 @@ private:
static_assert(not cute::is_void_v<NonVoidElementD>, "SmemElementD is void");
using NonVoidElementC = cute::conditional_t<not is_source_supported,NonVoidElementD,ElementC>; // prevents void ref breakages
using TmaElementD = cute::conditional_t<cute::is_same_v<NonVoidElementD, cutlass::complex<float>>, uint64_t, NonVoidElementD>;
using TmaElementC = cute::conditional_t<cute::is_same_v<NonVoidElementC, cutlass::complex<float>>, uint64_t, NonVoidElementC>;
using SmemElementC = typename cutlass::detail::get_unpacked_element_type<NonVoidElementC>::type;
using SmemElementD = typename cutlass::detail::get_unpacked_element_type<NonVoidElementD>::type;
@@ -239,14 +242,14 @@ public:
struct Params {
using TMA_C = decltype(make_tma_copy(
CopyOpG2S{},
make_tensor(make_gmem_ptr(static_cast<NonVoidElementC const*>(nullptr)),
make_tensor(make_gmem_ptr<TmaElementC const>(nullptr),
repeat_like(StrideC{}, int32_t(0)), StrideC{}),
take<0,2>(SmemLayoutC{}),
EpilogueTile{},
_1{}));
using TMA_D = decltype(make_tma_copy(
CopyOpS2G{},
make_tensor(make_gmem_ptr(static_cast<NonVoidElementD const*>(nullptr)),
make_tensor(make_gmem_ptr<TmaElementD>(nullptr),
repeat_like(StrideD{}, int32_t(0)), StrideD{}),
take<0,2>(SmemLayoutD{}),
EpilogueTile{},
@@ -273,9 +276,9 @@ public:
auto [M, N, K, L] = problem_shape_MNKL;
uint32_t transaction_bytes = TmaTransactionBytes;
typename Params::TMA_C tma_load_c = {};
typename Params::TMA_C tma_load_c{};
if constexpr (is_source_supported) {
Tensor tensor_c = make_tensor(make_gmem_ptr(args.ptr_C), make_layout(make_shape(M,N,L), args.dC));
Tensor tensor_c = make_tensor(make_gmem_ptr<TmaElementC const>(args.ptr_C), make_layout(make_shape(M,N,L), args.dC));
tma_load_c = make_tma_copy_C_sm90(
CopyOpG2S{},
tensor_c,
@@ -285,7 +288,7 @@ public:
typename Params::TMA_D tma_store_d;
if constexpr (is_destination_supported) {
Tensor tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), make_layout(make_shape(M,N,L), args.dD));
Tensor tensor_d = make_tensor(make_gmem_ptr<TmaElementD>(args.ptr_D), make_layout(make_shape(M,N,L), args.dD));
tma_store_d = make_tma_copy_C_sm90(
CopyOpS2G{},
tensor_d,
@@ -644,7 +647,18 @@ public:
// Absolute coordinate tensors (dynamic)
Tensor mD_crd = make_identity_tensor(make_shape(M,N)); // (M,N)
Tensor cD_mn = local_tile(mD_crd, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tRS_cD_mn = thread_r2s.partition_S(flat_divide(cD_mn, EpilogueTile{})); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
Tensor tRS_cD_mn = [&]() {
if constexpr (IsUseR2R) {
// (t)hread-partition for ConsumerStoreCallbacks.
TiledCopy tiled_cst = make_tiled_copy_S(Copy_Atom<CopyOpR2S,SmemElementC>{}, tiled_copy_C_atom);
ThrCopy thread_cst = tiled_cst.get_slice(thread_idx);
return thread_cst.partition_S(flat_divide(cD_mn, EpilogueTile{})); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
}
else {
return thread_r2s.partition_S(flat_divide(cD_mn, EpilogueTile{})); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
}
}();
// Relative coordinate tensors (static)
Tensor cD = make_counting_tensor(cD_mn.layout()); // (CTA_M,CTA_N)
Tensor tRS_cD = make_counting_tensor(tRS_cD_mn.layout()); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
include/cutlass/epilogue/dispatch_policy.hpp
+1
View File
@@ -50,6 +50,7 @@ struct EpilogueSimtVectorized {};
struct EpiloguePtrArraySimtVectorized {};
struct NoSmemWarpSpecialized {};
struct PtrArrayNoSmemWarpSpecialized {};
struct PtrArrayNoSmemWarpSpecializedTransposed {};
struct PtrArrayPlanarComplexNoSmemWarpSpecialized {};
struct TmaWarpSpecialized {};
struct TmaWarpSpecializedCooperative {};
include/cutlass/epilogue/fusion/operations.hpp
+144
View File
@@ -34,6 +34,7 @@
#include <cutlass/numeric_conversion.h>
#include <cutlass/layout/matrix.h>
#include <cute/numeric/numeric_types.hpp>
#include <cute/numeric/integral_constant.hpp> // cute::false_type
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -60,9 +61,12 @@ struct FusionOperation {
static constexpr int AlignmentScalar = 0;
static constexpr bool IsScaleFactorSupported = false;
static constexpr bool IsPerRowScaleSupported = false;
static constexpr bool IsPerColScaleSupported = false;
using ElementBias = void;
static constexpr int AlignmentBias = 0;
static constexpr bool IsPerRowBiasSupported = false;
static constexpr bool IsPerColBiasSupported = false;
static constexpr bool IsDePerRowBiasSupported = false;
using ActivationFn = void;
@@ -190,6 +194,24 @@ struct LinCombPerRowBiasEltAct
static constexpr bool IsEltActSupported = true;
};
// D = activation(alpha * acc + beta * C + per-column bias)
template<
template <class> class ActivationFn_,
class ElementOutput_,
class ElementCompute_,
class ElementBias_ = ElementOutput_,
class ElementSource_ = ElementOutput_,
class ElementScalar_ = ElementCompute_,
int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
>
struct LinCombPerColBiasEltAct
: LinCombPerColBias<ElementOutput_, ElementCompute_,
ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
using ActivationFn = ActivationFn_<ElementCompute_>;
static constexpr bool IsEltActSupported = true;
};
// D = activation(alpha * acc + beta * C + per-row bias)
// aux = alpha * acc + beta * C + per-row bias
template<
@@ -214,6 +236,30 @@ struct LinCombPerRowBiasEltActAux
static constexpr bool IsAuxOutSupported = true;
};
// D = activation(alpha * acc + beta * C + per-col bias)
// aux = alpha * acc + beta * C + per-col bias
template<
class GmemLayoutTagAux_,
template <class> class ActivationFn_,
class ElementOutput_,
class ElementCompute_,
class ElementAux_ = ElementOutput_,
class ElementBias_ = ElementOutput_,
class ElementSource_ = ElementOutput_,
class ElementScalar_ = ElementCompute_,
int AlignmentAux_ = 128 / cute::sizeof_bits_v<ElementAux_>,
int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
>
struct LinCombPerColBiasEltActAux
: LinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
using ElementAux = ElementAux_;
using GmemLayoutTagAux = GmemLayoutTagAux_;
static constexpr int AlignmentAux = AlignmentAux_;
static constexpr bool IsAuxOutSupported = true;
};
// D = activation(per-row alpha * acc + per-row beta * C + per-row bias)
template<
template <class> class ActivationFn_,
@@ -233,6 +279,46 @@ struct PerRowLinCombPerRowBiasEltAct
static constexpr bool IsPerRowScaleSupported = true;
};
// D = per-column alpha * per-row alpha * acc + beta * C
template<
class ElementOutput_,
class ElementCompute_,
class ElementSource_ = ElementCompute_,
class ElementScalar_ = ElementCompute_,
int AlignmentScalar_ = 128 / cute::sizeof_bits_v<ElementScalar_>,
FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
>
struct OuterProdLinComb : FusionOperation {
using ElementOutput = ElementOutput_;
using ElementCompute = ElementCompute_;
using ElementSource = ElementSource_;
using ElementScalar = ElementScalar_;
static constexpr int AlignmentScalar = AlignmentScalar_;
static constexpr auto RoundStyle = RoundStyle_;
static constexpr bool IsSourceSupported = true;
static constexpr bool IsPerRowScaleSupported = true;
static constexpr bool IsPerColScaleSupported = true;
};
// D = activation(per-col alpha * acc + per-col beta * C + per-column bias)
template<
template <class> class ActivationFn_,
class ElementOutput_,
class ElementCompute_,
class ElementBias_ = ElementOutput_,
class ElementSource_ = ElementOutput_,
class ElementScalar_ = ElementCompute_, // per-row alpha/beta
int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
int AlignmentScalar_ = 128 / cute::sizeof_bits_v<ElementScalar_>,
FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
>
struct PerColLinCombPerColBiasEltAct
: LinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
static constexpr int AlignmentScalar = AlignmentScalar_;
static constexpr bool IsPerColScaleSupported = true;
};
// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-row bias
// if D is fp8
// D = scale_d * activation(Z)
@@ -254,6 +340,27 @@ struct ScaledLinCombPerRowBiasEltAct
static constexpr bool IsScaleFactorSupported = true;
};
// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-col bias
// if D is fp8
// D = scale_d * activation(Z)
// else
// D = activation(Z)
template<
template <class> class ActivationFn_,
class ElementOutput_,
class ElementCompute_,
class ElementBias_ = ElementOutput_,
class ElementSource_ = ElementOutput_,
class ElementScalar_ = ElementCompute_,
int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
>
struct ScaledLinCombPerColBiasEltAct
: LinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
static constexpr bool IsScaleFactorSupported = true;
};
// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
// if D is fp8
// amax_d = max(abs(elements in activation(Z)))
@@ -291,6 +398,43 @@ struct ScaledLinCombPerRowBiasEltActAmaxAux
static constexpr bool IsAuxOutSupported = true;
};
// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
// if D is fp8
// amax_d = max(abs(elements in activation(Z)))
// D = scale_d * activation(Z)
// else
// D = activation(Z)
// if Aux is fp8
// amax_aux = max(abs(elements in Z))
// Aux = scale_aux * Z
// else
// Aux = Z
template<
class GmemLayoutTagAux_,
template <class> class ActivationFn_,
class ElementOutput_,
class ElementCompute_,
class ElementAux_ = ElementOutput_,
class ElementAmax_ = ElementCompute_,
class ElementBias_ = ElementOutput_,
class ElementSource_ = ElementOutput_,
class ElementScalar_ = ElementCompute_,
int AlignmentAux_ = 128 / cute::sizeof_bits_v<ElementAux_>,
int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
>
struct ScaledLinCombPerColBiasEltActAmaxAux
: ScaledLinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
using ElementAmax = ElementAmax_;
static constexpr bool IsAbsMaxSupported = true;
using ElementAux = ElementAux_;
using GmemLayoutTagAux = GmemLayoutTagAux_;
static constexpr int AlignmentAux = AlignmentAux_;
static constexpr bool IsAuxOutSupported = true;
};
// Z = Aux
// dY = alpha * acc + beta * C
// D = d_activation(dY, Z)
include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp
+901
View File
@@ -708,6 +708,105 @@ struct FusionCallbacks<
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = activation(alpha * acc + beta * C + per-column bias)
template<
int StagesC,
class CtaTileShapeMNK,
class EpilogueTile,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90LinCombPerColBiasEltAct =
Sm90EVT<Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>,
Sm90LinCombPerColBias<StagesC, CtaTileShapeMNK, EpilogueTile, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
>;
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementBias,
class ElementSource,
class ElementScalar,
int AlignmentBias,
FloatRoundStyle RoundStyle,
class CtaTileShapeMNK,
class EpilogueTile
>
struct FusionCallbacks<
epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
fusion::LinCombPerColBiasEltAct<
ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
>,
CtaTileShapeMNK,
EpilogueTile
> : Sm90LinCombPerColBiasEltAct<
StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
> {
using Impl =
Sm90LinCombPerColBiasEltAct<
StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
>;
using Operation =
fusion::LinCombPerColBiasEltAct<
ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
>;
struct Arguments {
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideAlpha = Stride<_0,_0,int64_t>;
using StrideBeta = Stride<_0,_0,int64_t>;
StrideAlpha dAlpha = {_0{}, _0{}, 0};
StrideBeta dBeta = {_0{}, _0{}, 0};
using StrideBias = Stride<_0,_1,int64_t>;
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
{ // unary op : activation(beta * C + (alpha * acc + bias))
{ // ternary op : beta * C + (alpha * acc + bias)
{{beta}, {beta_ptr}, {dBeta}}, // leaf args : beta
{}, // leaf args : C
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}, {dAlpha}}, // 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
activation // unary args : activation
}; // end unary op
}
};
// Ctor inheritance
using Impl::Impl;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = activation(alpha * acc + beta * C + per-row bias)
// Aux = alpha * acc + beta * C + per-row bias)
template<
@@ -832,6 +931,132 @@ struct FusionCallbacks<
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = activation(alpha * acc + beta * C + per_col bias)
// Aux = alpha * acc + beta * C + per_col bias)
template<
int StagesC,
class CtaTileShapeMNK,
class EpilogueTile,
int Stages,
class StrideAux,
class SmemLayoutAtom,
class CopyOpR2S,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementAux = ElementOutput,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90LinCombPerColBiasEltActAux =
Sm90EVT<Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>,
Sm90EVT<Sm90AuxStore<Stages, EpilogueTile, ElementAux, RoundStyle, StrideAux, SmemLayoutAtom, CopyOpR2S, AlignmentAux>,
Sm90LinCombPerColBias<StagesC, CtaTileShapeMNK, EpilogueTile, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
>
>;
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
class GmemLayoutTagAux,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementAux,
class ElementBias,
class ElementSource,
class ElementScalar,
int AlignmentAux,
int AlignmentBias,
FloatRoundStyle RoundStyle,
class CtaTileShapeMNK,
class EpilogueTile,
class SmemLayoutAtom,
class CopyOpR2S
>
struct FusionCallbacks<
epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
fusion::LinCombPerColBiasEltActAux<
GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute,
ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>,
CtaTileShapeMNK,
EpilogueTile,
SmemLayoutAtom,
CopyOpR2S
> : Sm90LinCombPerColBiasEltActAux<
StagesC, CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>, SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
> {
using Impl =
Sm90LinCombPerColBiasEltActAux<
StagesC, CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>, SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
using Operation =
fusion::LinCombPerColBiasEltActAux<
GmemLayoutTagAux, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
struct Arguments {
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideAlpha = Stride<_0,_0,int64_t>;
using StrideBeta = Stride<_0,_0,int64_t>;
StrideAlpha dAlpha = {_0{}, _0{}, 0};
StrideBeta dBeta = {_0{}, _0{}, 0};
using StrideBias = Stride<_0,_1,int64_t>;
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 = {};
operator typename Impl::Arguments() const {
return
{ // unary op : activation(store(beta * C + (alpha * acc + bias)))
{ // unary op : store(beta * C + (alpha * acc + bias))
{ // ternary op : beta * C + (alpha * acc + bias)
{{beta}, {beta_ptr}, {dBeta}}, // leaf args : beta
{}, // leaf args : C
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}, {dAlpha}}, // 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
{aux_ptr, dAux} // unary args : store
}, // end unary op
activation // unary args : activation
}; // end unary op
}
};
// Ctor inheritance
using Impl::Impl;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = per-row alpha * acc + per-row beta * C + per-row bias
template<
class CtaTileShapeMNK,
@@ -954,6 +1179,133 @@ struct FusionCallbacks<
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = per-col alpha * acc + per-col beta * C + per-column bias
template<
int StagesC,
class CtaTileShapeMNK,
class EpilogueTile,
class ElementOutput,
class ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
int AlignmentScalar = 128 / sizeof_bits_v<ElementScalar>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90PerColLinCombPerColBias =
Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + bias)
Sm90RowBroadcast<0, CtaTileShapeMNK, ElementScalar, ElementCompute, Stride<_0,bool,int64_t>, AlignmentScalar>, // beta, dynamic scalar/vector broadcast
Sm90SrcFetch<ElementSource>, // C
Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90RowBroadcast<0, CtaTileShapeMNK, ElementScalar, ElementCompute, Stride<_0,bool,int64_t>, AlignmentScalar>, // alpha, dynamic scalar/vector broadcast
Sm90AccFetch, // acc
Sm90RowBroadcast<0, CtaTileShapeMNK, ElementBias, ElementCompute, Stride<_0,_1,int64_t>, AlignmentBias> // bias
>
>;
// D = activation(per-col alpha * acc + per-col beta * C + per-column bias)
template<
int StagesC,
class CtaTileShapeMNK,
class EpilogueTile,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
int AlignmentScalar = 128 / sizeof_bits_v<ElementScalar>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90PerColLinCombPerColBiasEltAct =
Sm90EVT<Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>,
Sm90PerColLinCombPerColBias<StagesC, CtaTileShapeMNK, EpilogueTile, ElementCompute, ElementCompute,
ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle>
>;
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementBias,
class ElementSource,
class ElementScalar,
int AlignmentBias,
int AlignmentScalar,
FloatRoundStyle RoundStyle,
class CtaTileShapeMNK,
class EpilogueTile
>
struct FusionCallbacks<
epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
fusion::PerColLinCombPerColBiasEltAct<
ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
>,
CtaTileShapeMNK,
EpilogueTile
> : Sm90PerColLinCombPerColBiasEltAct<
StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
> {
using Impl =
Sm90PerColLinCombPerColBiasEltAct<
StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
>;
using Operation =
fusion::PerColLinCombPerColBiasEltAct<
ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
>;
struct Arguments {
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideAlpha = Stride<_0,bool,int64_t>;
using StrideBeta = Stride<_0,bool,int64_t>;
StrideAlpha dAlpha = {_0{}, bool(1), 0};
StrideBeta dBeta = {_0{}, bool(1), 0};
using StrideBias = Stride<_0,_1,int64_t>;
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
{ // unary op : activation(beta * C + (alpha * acc + bias))
{ // ternary op : beta * C + (alpha * acc + bias)
{beta_ptr, beta, dBeta}, // leaf args : beta
{}, // leaf args : C
{ // ternary op : alpha * acc + bias
{alpha_ptr, alpha, dAlpha}, // 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
activation // unary args : activation
}; // end unary op
}
};
// Ctor inheritance
using Impl::Impl;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
template <typename T>
@@ -1120,6 +1472,154 @@ struct FusionCallbacks<
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
template<
class CtaTileShapeMNK,
class ElementOutput,
class ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90ScaledLinCombPerColBias =
Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + bias)
Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,int64_t>, 2>, // scale_c * beta
Sm90SrcFetch<ElementSource>, // C
Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,int64_t>, 3>, // scale_a * scale_b * alpha
Sm90AccFetch, // acc
Sm90RowBroadcast<0, CtaTileShapeMNK, ElementBias, ElementCompute, Stride<_0,_1,int64_t>, AlignmentBias> // bias
>
>;
// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-col bias
// if D is fp8
// D = scale_d * activation(Z)
// else
// D = activation(Z)
template<
class CtaTileShapeMNK,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90ScaledLinCombPerColBiasEltAct =
Sm90EVT<Sm90Compute<detail::ScaleOutOp<ElementOutput>::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d
Sm90EVT<Sm90Compute<ActivationFn, ElementCompute, ElementCompute, RoundStyle>, // activation(Z)
// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-row bias
Sm90ScaledLinCombPerColBias<CtaTileShapeMNK, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
>,
Sm90ScalarBroadcast<ElementScalar> // scale_d
>;
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementBias,
class ElementSource,
class ElementScalar,
int AlignmentBias,
FloatRoundStyle RoundStyle,
class CtaTileShapeMNK,
class EpilogueTile
>
struct FusionCallbacks<
epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
fusion::ScaledLinCombPerColBiasEltAct<
ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
>,
CtaTileShapeMNK,
EpilogueTile
> : Sm90ScaledLinCombPerColBiasEltAct<
CtaTileShapeMNK, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
> {
using Impl =
Sm90ScaledLinCombPerColBiasEltAct<
CtaTileShapeMNK, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
>;
using Operation =
fusion::ScaledLinCombPerColBiasEltAct<
ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
>;
struct Arguments {
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
ElementScalar scale_a = ElementScalar(1);
ElementScalar scale_b = ElementScalar(1);
ElementScalar scale_c = ElementScalar(1);
ElementScalar scale_d = ElementScalar(1);
ElementScalar const* scale_a_ptr = nullptr;
ElementScalar const* scale_b_ptr = nullptr;
ElementScalar const* scale_c_ptr = nullptr;
ElementScalar const* scale_d_ptr = nullptr;
using StrideAlpha = Stride<_0,_0,int64_t>;
using StrideBeta = Stride<_0,_0,int64_t>;
StrideAlpha dAlpha = {_0{}, _0{}, 0};
StrideBeta dBeta = {_0{}, _0{}, 0};
using StrideBias = Stride<_0,_1,int64_t>;
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
{ // binary op : activation((scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)) * scale_d
{ // unary op : activation((scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias))
{ // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)
{{beta, scale_c},
{beta_ptr, scale_c_ptr},
{dBeta, {_0{}, _0{}, 0}}
}, // leaf args : (scale_c * beta)
{}, // leaf args : C
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{alpha, scale_a, scale_b},
{alpha_ptr, scale_a_ptr, scale_b_ptr},
{dAlpha, {_0{}, _0{}, 0}, {_0{}, _0{}, 0}}
}, // 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
activation // unary args : activation
}, // end unary op
{{scale_d},
{scale_d_ptr}
}, // leaf args : scale_d
{} // binary args : multiplies or first
}; // end binary op
}
};
// Ctor inheritance
using Impl::Impl;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
// if D is fp8
// amax_d = max(abs(elements in activation(Z)))
@@ -1440,6 +1940,326 @@ struct FusionCallbacks<
/////////////////////////////////////////////////////////////////////////////////////////////////
// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
// if D is fp8
// amax_d = max(abs(elements in activation(Z)))
// D = scale_d * activation(Z)
// else
// D = activation(Z)
// if Aux is fp8
// amax_aux = max(abs(elements in Z))
// Aux = scale_aux * Z
// else
// Aux = Z
// fp8 aux specialization
template<
class CtaTileShapeMNK,
class EpilogueTile,
int StagesD,
class StrideAux,
class SmemLayoutAtom,
class CopyOpR2S,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementAux = ElementOutput,
class ElementAmax = ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90ScaledLinCombPerColBiasEltActAmaxAuxFp8 =
Sm90SplitTreeVisitor<
// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
Sm90ScaledLinCombPerColBias<CtaTileShapeMNK, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>,
// D = activation(Z) * scale_d, amax_d = max(abs(elements in D))
Sm90EVT<Sm90Compute<detail::ScaleOutOp<ElementOutput>::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d
Sm90EVT<Sm90ScalarReduction<detail::amax, atomic_maximum, ElementAmax, ElementCompute, RoundStyle>, // amax_d
Sm90EVT<Sm90Compute<ActivationFn, ElementCompute, ElementCompute, RoundStyle>, // activation(Z)
Sm90SplitTreeFetch // Z
>
>,
Sm90ScalarBroadcast<ElementScalar> // scale_d
>,
// Aux = Z * scale_aux, amax_aux = max(abs(elements in Aux))
Sm90EVT<Sm90AuxStore<StagesD, EpilogueTile, ElementAux, RoundStyle, StrideAux, SmemLayoutAtom, CopyOpR2S, AlignmentAux>, // store(Aux)
Sm90EVT<Sm90Compute<cutlass::multiplies, ElementCompute, ElementCompute, RoundStyle>, // Z * scale_aux
Sm90EVT<Sm90ScalarReduction<detail::amax, atomic_maximum, ElementAmax, ElementCompute, RoundStyle>, // amax_aux
Sm90SplitTreeFetch // Z
>,
Sm90ScalarBroadcast<ElementScalar> // scale_aux
>
>
>;
// non-fp8 aux specialization
// lets us use some EVT specializations such as relu + uint1b_t aux
template<
class CtaTileShapeMNK,
class EpilogueTile,
int StagesD,
class StrideAux,
class SmemLayoutAtom,
class CopyOpR2S,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementAux = ElementOutput,
class ElementAmax = ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90ScaledLinCombPerColBiasEltActAmaxAuxNotFp8 =
// D = activation(Z) * scale_d, amax_d = max(abs(elements in D))
Sm90EVT<Sm90Compute<detail::ScaleOutOp<ElementOutput>::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d
Sm90EVT<Sm90ScalarReduction<detail::amax, atomic_maximum, ElementAmax, ElementCompute, RoundStyle>, // amax_d
Sm90EVT<Sm90Compute<ActivationFn, ElementCompute, ElementCompute, RoundStyle>, // activation(Z)
Sm90EVT<Sm90AuxStore<StagesD, EpilogueTile, ElementAux, RoundStyle, StrideAux, SmemLayoutAtom, CopyOpR2S, AlignmentAux>, // Aux = Z
// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
Sm90ScaledLinCombPerColBias<CtaTileShapeMNK, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
>
>
>,
Sm90ScalarBroadcast<ElementScalar> // scale_d
>;
// dispatcher
template<
class CtaTileShapeMNK,
class EpilogueTile,
int StagesD,
class StrideAux,
class SmemLayoutAtom,
class CopyOpR2S,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementAux = ElementOutput,
class ElementAmax = ElementCompute,
class ElementBias = ElementOutput,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90ScaledLinCombPerColBiasEltActAmaxAux = conditional_t<detail::is_fp8_v<ElementAux>,
Sm90ScaledLinCombPerColBiasEltActAmaxAuxFp8<
CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar,AlignmentAux, AlignmentBias, RoundStyle
>,
Sm90ScaledLinCombPerColBiasEltActAmaxAuxNotFp8<
CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>
>;
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
class GmemLayoutTagAux,
template <class> class ActivationFn,
class ElementOutput,
class ElementCompute,
class ElementAux,
class ElementAmax,
class ElementBias,
class ElementSource,
class ElementScalar,
int AlignmentAux,
int AlignmentBias,
FloatRoundStyle RoundStyle,
class CtaTileShapeMNK,
class EpilogueTile,
class SmemLayoutAtom,
class CopyOpR2S
>
struct FusionCallbacks<
epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
fusion::ScaledLinCombPerColBiasEltActAmaxAux<
GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute,
ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>,
CtaTileShapeMNK,
EpilogueTile,
SmemLayoutAtom,
CopyOpR2S
> : Sm90ScaledLinCombPerColBiasEltActAmaxAux<
CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>,
SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
> {
using Impl =
Sm90ScaledLinCombPerColBiasEltActAmaxAux<
CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>,
SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
using Operation =
fusion::ScaledLinCombPerColBiasEltActAmaxAux<
GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute,
ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
struct Arguments {
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
ElementScalar scale_a = ElementScalar(1);
ElementScalar scale_b = ElementScalar(1);
ElementScalar scale_c = ElementScalar(1);
ElementScalar scale_d = ElementScalar(1);
ElementScalar const* scale_a_ptr = nullptr;
ElementScalar const* scale_b_ptr = nullptr;
ElementScalar const* scale_c_ptr = nullptr;
ElementScalar const* scale_d_ptr = nullptr;
ElementScalar scale_aux = ElementScalar(1);
ElementScalar const* scale_aux_ptr = nullptr;
using StrideAlpha = Stride<_0,_0,int64_t>;
using StrideBeta = Stride<_0,_0,int64_t>;
StrideAlpha dAlpha = {_0{}, _0{}, 0};
StrideBeta dBeta = {_0{}, _0{}, 0};
using StrideBias = Stride<_0,_1,int64_t>;
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 = {};
operator typename Impl::Arguments() const {
// Only compute amax_d if D is fp8
ElementAmax* amax_D_ptr_ = nullptr;
if constexpr (detail::is_fp8_v<ElementOutput>) {
amax_D_ptr_ = amax_D_ptr;
}
// Aux is fp8 -> DAG arguments
if constexpr (detail::is_fp8_v<ElementAux>) {
typename Impl::Arguments args;
// always use structured binding to unpack DAG args since it may or may not be a tuple
auto& [Z_args, aux_args, D_args] = args;
Z_args =
{ // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)
{{beta, scale_c},
{beta_ptr, scale_c_ptr},
{dBeta, {_0{}, _0{}, 0}}
}, // leaf args : (scale_c * beta)
{}, // leaf args : C
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{alpha, scale_a, scale_b},
{alpha_ptr, scale_a_ptr, scale_b_ptr},
{dAlpha, {_0{}, _0{}, 0}, {_0{}, _0{}, 0}}
}, // 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
D_args =
{ // binary op : activation(Z) * scale_d or activation(Z)
{ // unary op : reduce(activation(Z))
{ // unary op : activation(Z)
{}, // leaf args : Z
activation // unary args : activation
}, // end unary op
{amax_D_ptr_} // unary args : reduce
}, // end unary op
{{scale_d},
{scale_d_ptr}
}, // leaf args : scale_d
{} // binary args : multiplies or first
}; // end binary op
aux_args =
{ // unary op : store(Aux)
{ // binary op : Z * scale_d or Z
{ // unary op : reduce(Z)
{}, // leaf args : Z
{amax_aux_ptr} // unary args : reduce
}, // end unary op
{{scale_aux},
{scale_aux_ptr}
}, // leaf args : scale_d
{} // binary args : multiplies
}, // end binary op
{aux_ptr, dAux} // unary args : store
}; // end unary op
return args;
}
// Aux is not fp8 -> Tree arguments
else {
return
{ // binary op : activation(Z) * scale_d or activation(Z)
{ // unary op : reduce(activation(Z))
{ // unary op : activation(Z)
{ // unary op : store(Z)
{ // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)
{{beta, scale_c},
{beta_ptr, scale_c_ptr},
{dBeta, {_0{}, _0{}, 0}}
}, // leaf args : (scale_c * beta)
{}, // leaf args : C
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{alpha, scale_a, scale_b},
{alpha_ptr, scale_a_ptr, scale_b_ptr},
{dAlpha, {_0{}, _0{}, 0}, {_0{}, _0{}, 0}}
}, // 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
{aux_ptr, dAux} // unary args : store
}, // end unary op
activation // unary args : activation
}, // end unary op
{amax_D_ptr_} // unary args : reduce
}, // end unary op
{{scale_d},{scale_d_ptr}}, // leaf args : scale_d
{} // binary args : multiplies or first
}; // end binary op
}
}
};
// Ctor inheritance
using Impl::Impl;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class CtaTileShapeMNK,
class EpilogueTile,
@@ -1679,6 +2499,87 @@ struct FusionCallbacks<
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = per-column alpha * per-row alpha * acc + beta * c
template<
class CtaTileShapeMNK,
class ElementOutput,
class ElementCompute,
class ElementSource = ElementOutput,
class ElementScalar = ElementCompute,
int AlignmentScalar = 128 / sizeof_bits_v<ElementScalar>, // Alignment of per-column and per-row scaling vectors
FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
>
using Sm90OuterProdLinComb =
Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // c(beta) * c(C) + c(alpha * acc)
Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,int>>, // beta
Sm90SrcFetch<ElementSource>, // C
Sm90EVT<Sm90Compute<multiplies, ElementCompute, ElementCompute, RoundStyle>, // c(alpha) * c(acc)
Sm90OuterProduct<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,int>, Stride<_0,_1,int>, AlignmentScalar>, // alpha_col * alpha_row
Sm90AccFetch // acc
>
>;
template <
int StagesC,
int StagesD,
int FragmentSize,
bool ReuseSmemC,
bool DelayTmaStore,
class ElementOutput,
class ElementCompute,
class ElementSource,
class ElementScalar,
int AlignmentScalar,
FloatRoundStyle RoundStyle,
class CtaTileShapeMNK,
class EpilogueTile
>
struct FusionCallbacks<
epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
OuterProdLinComb<ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle>,
CtaTileShapeMNK,
EpilogueTile
> : Sm90OuterProdLinComb<CtaTileShapeMNK, ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle> {
using Impl = Sm90OuterProdLinComb<CtaTileShapeMNK, ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle>;
using Operation = OuterProdLinComb<ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle>;
struct Arguments {
// Give a name and flat ordering to the fusion callback args
using StrideCol = Stride<_1,_0,int>;
using StrideRow = Stride<_0,_1,int>;
using StrideBeta = Stride<_0,_0,int>;
ElementScalar const* alpha_ptr_col = nullptr;
ElementScalar const* alpha_ptr_row = nullptr;
ElementScalar beta = static_cast<ElementScalar>(0);
ElementScalar const* beta_ptr = nullptr;
StrideCol dAlphaCol = {};
StrideRow dAlphaRow = {};
StrideBeta dBeta = {};
// Conversion to the args expected by the visitor implementation
// to_underlying_arguments will implicitly call this
operator typename Impl::Arguments() const {
return
{
{beta, beta_ptr, dBeta}, // leaf args : beta
{}, // leaf args : C
{
{ alpha_ptr_col, alpha_ptr_row, dAlphaCol, dAlphaRow }, // leaf args : alpha cols / rows
{}, // leaf args : acc
{}
},
{}
};
}
};
// Ctor inheritance
using Impl::Impl;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// D = softmax(top_k(alpha * acc + beta * C))
template<
int TopK,
include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
+7 -5
View File
@@ -266,8 +266,8 @@ struct Sm90TreeVisitor<
auto const& scale_op = get<0>(Impl::ops);
auto const& added_op = get<2>(Impl::ops);
if constexpr (detail::IsScalarBroadcast<InputScaleOp>::value && not is_void_v<ElementSource>) {
return (get<2>(scale_op.params_ptr->dScalar[0]) != 0 && scale_op.params_ptr->scalar_ptrs[0] != nullptr) ||
is_C_load_needed() ||
return (get<2>(scale_op.params_ptr->dScalar[0]) != 0 && scale_op.params_ptr->scalar_ptrs[0] != nullptr) ||
is_C_load_needed() ||
added_op.is_producer_load_needed();
}
else {
@@ -408,8 +408,9 @@ template <
>
struct Sm90TreeVisitor<
Sm90Compute<Activation, ElementOutput, ElementCompute, RoundStyle,
cute::enable_if_t<cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ReLu<ElementCompute>> ||
cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>> >>,
cute::enable_if_t<cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ReLu<ElementCompute>> ||
cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>> ||
cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ThresholdReLU<ElementCompute>> >>,
Sm90TreeVisitor<
Sm90AuxStore<
Stages,
@@ -503,7 +504,8 @@ struct Sm90TreeVisitor<
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
ElementCompute pre_relu = frg_compute[i];
if constexpr (cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>>) {
if constexpr (cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>> ||
cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ThresholdReLU<ElementCompute>>) {
frg_compute[i] = relu(frg_compute[i], params_compute);
}
else {
include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
+170 -4
View File
@@ -734,11 +734,12 @@ private:
// Supports reduction over multiple broadcasts to support fusions such as fp8 scaling factors
template<
class Element,
class StrideMNL = Stride<_0,_0,_0>,
class StrideMNL_ = Stride<_0,_0,_0>,
int BroadcastCount = 1,
template <class> class ReductionFn = multiplies
>
struct Sm90ScalarBroadcastPtrArray {
using StrideMNL = StrideMNL_;
static_assert(is_static_v<decltype(take<0,2>(StrideMNL{}))>); // batch stride can be dynamic or static
static_assert(take<0,2>(StrideMNL{}) == Stride<_0,_0>{});
@@ -780,8 +781,8 @@ struct Sm90ScalarBroadcastPtrArray {
CUTLASS_DEVICE bool
is_producer_load_needed() const {
// producer load is needed if Element is not void and we have multiple scalars
return !cute::is_void_v<Element> and size<2>(params_ptr->dScalar[0]) != 0;
// producer load is needed if Element is not void
return !cute::is_void_v<Element>;
}
CUTLASS_DEVICE bool
@@ -814,7 +815,7 @@ struct Sm90ScalarBroadcastPtrArray {
CUTLASS_DEVICE auto
get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
// Get the scalar for batched broadcast
if (get<2>(params_ptr->dScalar[0]) != 0) {
if (size<2>(params_ptr->dScalar[0]) != 0) {
auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
update_scalar(l_coord);
}
@@ -1377,6 +1378,171 @@ struct Sm90ColBroadcast {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Do outer product from the column and row loaded
//
template<
int Stages,
class CtaTileShapeMNK,
class ElementScalar,
class StrideColMNL_ = Stride<_1,_0,int64_t>, /// NOTE: Batched scaling untested for now
class StrideRowMNL_ = Stride<_0,_1,int64_t>,
int Alignment = 128 / sizeof_bits_v<ElementScalar>,
bool EnableNullptr = false // Fallback scalar broadcast for nullptr params
>
struct Sm90OuterProduct {
using StrideColMNL = StrideColMNL_;
using StrideRowMNL = StrideRowMNL_;
static_assert(Stages == 0, "OuterProduct doesn't support smem usage");
static_assert(Alignment * sizeof_bits_v<ElementScalar> % 128 == 0, "sub-16B alignment not supported yet");
static_assert(!EnableNullptr, "Nullptr fallback not implemented");
static_assert(is_static_v<decltype(take<0,2>(StrideColMNL{}))> &&
is_static_v<decltype(take<0,2>(StrideRowMNL{}))>, "Only batch stride can be dynamic");
static_assert(take<0,2>(StrideColMNL{}) == Stride<_1,_0>{} &&
take<0,2>(StrideRowMNL{}) == Stride<_0,_1>{}, "Row and column incorrectly formatted");
// Accumulator distributes col/row elements evenly amongst threads so we can just directly load from gmem
struct SharedStorage { };
struct Arguments {
ElementScalar const* ptr_col = nullptr;
ElementScalar const* ptr_row = nullptr;
StrideColMNL dCol = {};
StrideRowMNL dRow = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
template <class ProblemShape>
static bool
can_implement(ProblemShape const& problem_shape, Arguments const& args) {
return true;
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
return 0;
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
CUTLASS_DEVICE bool
is_producer_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_C_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_zero() const {
return false;
}
CUTLASS_HOST_DEVICE
Sm90OuterProduct() { }
CUTLASS_HOST_DEVICE
Sm90OuterProduct(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <class... Args>
CUTLASS_DEVICE auto
get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
return EmptyProducerLoadCallbacks{};
}
template<
class GTensorCol, class RTensorCol,
class GTensorRow, class RTensorRow
>
struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
CUTLASS_DEVICE
ConsumerStoreCallbacks(GTensorCol&& tCgCol, RTensorCol&& tCrCol,
GTensorRow&& tCgRow, RTensorRow&& tCrRow,
Params const& params)
: tCgCol(cute::forward<GTensorCol>(tCgCol))
, tCrCol(cute::forward<RTensorCol>(tCrCol))
, tCgRow(cute::forward<GTensorRow>(tCgRow))
, tCrRow(cute::forward<RTensorRow>(tCrRow))
, params(params) {}
GTensorCol tCgCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
RTensorCol tCrCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
GTensorRow tCgRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
RTensorRow tCrRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
Params const& params;
CUTLASS_DEVICE void
begin() {
// Filter so we don't issue redundant copies over stride-0 modes
copy(filter(tCgCol), filter(tCrCol));
copy(filter(tCgRow), filter(tCrRow));
}
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE Array<ElementScalar, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
Array<ElementScalar, FragmentSize> frg_colrow;
Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
Tensor tCrRow_mn = tCrRow(_,_,_,epi_m,epi_n);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
frg_colrow[i] = static_cast<ElementScalar>(tCrCol_mn(epi_v * FragmentSize + i) * tCrRow_mn(epi_v * FragmentSize + i));
}
return frg_colrow;
}
};
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class... Args
>
CUTLASS_DEVICE auto
get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
auto [M, N, K, L] = args.problem_shape_mnkl;
Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tCgRow = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
mRow, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
Tensor tCrRow = make_tensor_like(tCgRow); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
return ConsumerStoreCallbacks<
decltype(tCgCol), decltype(tCrCol),
decltype(tCgRow), decltype(tCrRow)
>(
cute::move(tCgCol), cute::move(tCrCol),
cute::move(tCgRow), cute::move(tCrRow),
params
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Batch matrix broadcast
include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp
+21 -33
View File
@@ -293,11 +293,11 @@ template <
class LayoutOrStrideMNL,
class SmemLayoutAtom, // Unused
class CopyOpR2S, // Unused
int Alignment,
int Alignment,
bool EnableNullptr
>
struct Sm90AuxStore<
0, EpilogueTile, Element, RoundStyle, LayoutOrStrideMNL,
0, EpilogueTile, Element, RoundStyle, LayoutOrStrideMNL,
SmemLayoutAtom, CopyOpR2S, Alignment, EnableNullptr
> {
using ElementAux = Element;
@@ -343,7 +343,7 @@ struct Sm90AuxStore<
CUTLASS_HOST_DEVICE
Sm90AuxStore(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
CUTLASS_DEVICE bool
@@ -381,7 +381,7 @@ struct Sm90AuxStore<
tC_cAux(cute::forward<CTensorR2G>(tC_cAux)),
problem_shape_mnl(problem_shape_mnl),
params_ptr(params_ptr) {}
GTensorR2G tC_gAux;
RTensor tC_rAux;
CTensorR2G tC_cAux;
@@ -414,7 +414,7 @@ struct Sm90AuxStore<
Tensor tC_cAux_mn = tC_cAux(_,_,_,epi_m,epi_n);
Tensor tC_cAux_vec = tensor<1>(zipped_divide(coalesce(tC_cAux_mn), MCL.compose(Int<V>{})));
Tensor tC_gAux_vec = recast<Array<Element, V>>(coalesce(tC_gAux(_,_,_,epi_m,epi_n)));
Tensor tC_rAux_vec = recast<Array<Element, V>>(coalesce(tC_rAux));
@@ -451,7 +451,7 @@ struct Sm90AuxStore<
// Predication support
Tensor coordAux = make_identity_tensor(shape(mAux));
Tensor tC_cAux = sm90_partition_for_epilogue<ReferenceSrc>(
coordAux, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
coordAux, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
return ConsumerStoreCallbacks<decltype(tC_gAux), decltype(tC_rAux), decltype(tC_cAux), decltype(problem_shape_mnl)>(
cute::move(tC_gAux),
@@ -703,7 +703,6 @@ public:
else if constexpr (FinalReduction) {
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
size_t tile_counters_offset = product(ceil_div(make_shape(size<>(M), size<>(N), L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute);
tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -753,19 +752,18 @@ public:
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
#if !defined(CUTLASS_SKIP_REDUCTION_INIT)
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
if constexpr (IsAtomic) {
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
Layout mRow_layout = make_layout(make_shape(size<>(M),size<>(N),size<>(L)), args.dRow);
if (args.ptr_row != nullptr) {
return fill_workspace(args.ptr_row, ElementOutput(args.reduction_identity), cosize(mRow_layout), stream, cuda_adapter);
}
return Status::kSuccess;
}
else
#endif
if constexpr (FinalReduction) {
else if constexpr (FinalReduction) {
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
size_t tile_counters_offset = product(ceil_div(make_shape(size<>(M),size<>(N),L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute);
tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -939,7 +937,7 @@ public:
for (int v = 0; v < size(frg_A); ++v) {
// Step1: swap
if (not (lane_m & m)) { // the first half of threads swap fragments from the first half of data to the second
swap(frg_A(v), frg_B(v));
cutlass::swap(frg_A(v), frg_B(v));
}
// Step2: shuffle
@@ -1023,9 +1021,7 @@ public:
}
else {
if (is_reduced_lane) {
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrRow), recast<ElementGmem>(filter(tCgBuf)));
copy_aligned(tCrRow, recast<ElementGmem>(tCgBuf));
}
}
sync_fn();
@@ -1054,9 +1050,7 @@ public:
}
else {
if (is_reduced_lane) {
// Filter so we don't issue redunant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrRow), filter(tCsBuf));
copy_aligned(tCrRow, tCsBuf);
}
}
sync_fn();
@@ -1296,7 +1290,6 @@ public:
else if constexpr (FinalReduction) {
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
size_t tile_counters_offset = product(ceil_div(make_shape(M,N,L), make_shape(tile_M, tile_N))) * tile_M * sizeof(ElementCompute);
tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -1348,19 +1341,18 @@ public:
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
#if !defined(CUTLASS_SKIP_REDUCTION_INIT)
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
if constexpr (IsAtomic) {
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
Layout mCol_layout = make_layout(make_shape(size<>(M),size<>(N),size<>(L)), args.dCol);
if (args.ptr_col != nullptr) {
return fill_workspace(args.ptr_col, ElementOutput(args.reduction_identity), cosize(mCol_layout), stream, cuda_adapter);
}
return Status::kSuccess;
}
else
#endif
if constexpr (FinalReduction) {
else if constexpr (FinalReduction) {
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
size_t tile_counters_offset = product(ceil_div(make_shape(M,N,L), make_shape(tile_M, tile_N))) * tile_M * sizeof(ElementCompute);
tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -1522,9 +1514,7 @@ public:
using ElementGmem = cute::conditional_t<FinalReduction, ElementCompute volatile, ElementCompute>;
Tensor tCgBuf = sm90_partition_for_epilogue<ReferenceSrc>(gBuf_nl(_,_,n,l), epi_tile, tiled_copy, thread_idx);
if (is_reduced_lane) {
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
copy_aligned(tCrCol, recast<ElementGmem>(tCgBuf));
}
sync_fn();
}
@@ -1542,9 +1532,7 @@ public:
// Dump warp reduction to smem workspace
Tensor tCsBuf = sm90_partition_for_epilogue<ReferenceSrc>(sBuf(_,_,get<1>(warp_mn)), epi_tile, tiled_copy, thread_idx);
if (is_reduced_lane) {
// Filter so we don't issue redunant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrCol), filter(tCsBuf));
copy_aligned(tCrCol, tCsBuf);
}
sync_fn();
include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp
-2
View File
@@ -300,7 +300,6 @@ struct Sm90VisitorImplBase {
tuple<Ops...> ops;
};
template <class... Ops>
struct Sm90VisitorImpl : Sm90VisitorImplBase<Ops...> {
@@ -658,7 +657,6 @@ struct Sm90SplitTreeVisitor : Sm90VisitorImpl<InputTree, AuxOutTrees..., OutputT
return ConsumerStoreCallbacks<decltype(callbacks_tuple)>(std::move(callbacks_tuple));
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
include/cutlass/epilogue/thread/activation.h
+52 -20
View File
@@ -258,6 +258,54 @@ struct LeakyReLU<Array<T, N> > {
}
};
// Y = min((X <= threshold ? 0 : X), upper_bound)
template <typename T>
struct ThresholdReLU {
static constexpr bool kIsHeavy = false;
struct Arguments {
T threshold = T(0);
T upper_bound = CUTLASS_STL_NAMESPACE::numeric_limits<T>::max();
};
CUTLASS_HOST_DEVICE
T operator()(T value, T threshold, T upper_bound) const {
minimum_with_nan_propagation<T> mn;
return mn((value <= threshold ? T(0) : value), upper_bound);
}
CUTLASS_HOST_DEVICE
T operator()(T value, Arguments const& args = Arguments()) const {
return operator()(value, args.threshold, args.upper_bound);
}
};
template <typename T, int N>
struct ThresholdReLU<Array<T,N>> {
static constexpr bool kIsHeavy = false;
using Arguments = typename ThresholdReLU<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T,N> operator()(Array<T,N> const& values, T threshold, T upper_bound) const {
ThresholdReLU<T> relu;
Array<T,N> retvals;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
retvals[i] = relu(values[i], threshold, upper_bound);
}
return retvals;
}
CUTLASS_HOST_DEVICE
Array<T,N> operator()(Array<T,N> const& values, Arguments const& args = Arguments()) const {
return operator()(values, args.threshold, args.upper_bound);
}
};
// Tanh operator
template <typename T>
struct Tanh {
@@ -311,26 +359,7 @@ struct Sigmoid {
};
template <typename T, int N>
struct Sigmoid<Array<T, N> > {
static const bool kIsHeavy = true;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Array<T, N> y;
Sigmoid<T> sigmoid_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
y[i] = sigmoid_op(value[i]);
}
return y;
}
};
template <int N>
struct Sigmoid<Array<half_t, N>> {
using T = half_t;
struct Sigmoid<Array<T, N>> {
static const bool kIsHeavy = true;
CUTLASS_HOST_DEVICE
@@ -450,6 +479,9 @@ struct HardSwish<Array<half_t, N> > {
}
};
template <typename T>
using ScaledHardSwish = Scale<HardSwish<T>>;
//
// GELU function definitions implemented as described by
// Hendrycks, D., and Gimpel, K. in
include/cutlass/epilogue/thread/linear_combination.h
+5 -1
View File
@@ -169,7 +169,7 @@ public:
/// Constructs the function object, possibly loading from pointers in host memory
CUTLASS_HOST_DEVICE
LinearCombination(Params const &params, int group_idx = 0) {
explicit LinearCombination(Params const &params, int group_idx) {
if (params.alpha_ptr_array != nullptr && params.alpha_ptr_array[group_idx] != nullptr) {
alpha_ = *(params.alpha_ptr_array[group_idx]);
}
@@ -190,6 +190,10 @@ public:
}
}
CUTLASS_HOST_DEVICE
explicit LinearCombination(const Params & params)
: LinearCombination(params, /* group_idx */ 0) { }
/// Returns true if source is needed
CUTLASS_HOST_DEVICE
bool is_source_needed() const {
include/cutlass/epilogue/threadblock/epilogue.h
+10 -4
View File
@@ -39,11 +39,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
@@ -478,6 +474,12 @@ public:
// Iterate over accumulator tile
//
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wcuda-compat"
// Turn off clangs warning about loop unroll argument using parens.
#endif
#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter)
{
@@ -531,6 +533,10 @@ public:
destination_iterator.store(output_fragment);
++destination_iterator;
}
#ifdef __clang__
#pragma clang diagnostic pop
#endif
}
};
include/cutlass/epilogue/threadblock/epilogue_base.h
-4
View File
@@ -43,11 +43,7 @@
#include <utility>
#endif
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/matrix_shape.h"
include/cutlass/epilogue/threadblock/epilogue_gemm_k_reduction.h
-4
View File
@@ -38,11 +38,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h
-4
View File
@@ -38,11 +38,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
include/cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h
+2 -2
View File
@@ -39,11 +39,11 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#if defined(__CUDACC_RTC__)
#include <cuda/std/utility>
#else
#include <assert.h>
#include <utility>
#endif
include/cutlass/epilogue/threadblock/epilogue_with_absmax.h
+2 -2
View File
@@ -50,11 +50,11 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#if defined(__CUDACC_RTC__)
#include <cuda/std/utility>
#else
#include <assert.h>
#include <utility>
#endif
include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h
+2 -2
View File
@@ -39,11 +39,11 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#if defined(__CUDACC_RTC__)
#include <cuda/std/utility>
#else
#include <assert.h>
#include <utility>
#endif
include/cutlass/epilogue/threadblock/epilogue_with_reduction.h
-4
View File
@@ -39,11 +39,7 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include <assert.h>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/array.h"
include/cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h
+22
View File
@@ -303,6 +303,12 @@ public:
// Pipeline Loop
//
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wcuda-compat"
// Turn off clang warning about loop unroll argument using parens.
#endif
#pragma unroll(IterationsUnroll ? kIterations : 1)
for (int iter_idx = 1; iter_idx < kIterations + 1; ++iter_idx) {
@@ -377,8 +383,19 @@ public:
callbacks.end_step(iter_idx-1);
}
#ifdef __clang__
#pragma clang diagnostic pop
#endif
} else {
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wcuda-compat"
// Turn off clang warning about loop unroll argument using parens.
#endif
#pragma unroll(IterationsUnroll ? kIterations : 1)
for (int iter_idx = 0; iter_idx < kIterations; ++iter_idx) {
@@ -459,6 +476,11 @@ public:
callbacks.end_step(iter_idx);
}
#ifdef __clang__
#pragma clang diagnostic pop
#endif
}
callbacks.end_epilogue();
include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp
+21 -3
View File
@@ -335,7 +335,8 @@ struct VisitorAuxLoad{
template<
class ThreadMap,
class Element,
class StrideMNL
class StrideMNL,
bool EnableNullptr = true // Fallback scalar broadcast for nullptr params
>
struct VisitorRowBroadcast {
@@ -399,6 +400,16 @@ struct VisitorRowBroadcast {
CUTLASS_DEVICE void
begin_epilogue() {
if constexpr (EnableNullptr) {
if (params_ptr->ptr_row == nullptr) {
auto tC_rRow_vec = recast<Array<Element, VecLength>>(coalesce(tC_rRow));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tC_rRow_vec); ++i) {
tC_rRow_vec[i].fill(params_ptr->null_default);
}
return;
}
}
clear(tC_rRow);
auto src_v = filter(tC_gRow);
auto coord_v = filter(tC_cRow);
@@ -406,7 +417,7 @@ struct VisitorRowBroadcast {
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);
cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const *)&src_v(i), guard);
}
}
@@ -464,7 +475,8 @@ struct VisitorRowBroadcast {
template<
class ThreadMap,
class Element,
class StrideMNL = Stride<_1,_0,_0>
class StrideMNL = Stride<_1,_0,_0>,
bool EnableNullptr = true // Fallback scalar broadcast for nullptr params
>
struct VisitorColBroadcast {
@@ -523,6 +535,12 @@ struct VisitorColBroadcast {
CUTLASS_DEVICE void
begin_epilogue() {
if constexpr (EnableNullptr) {
if (params_ptr->ptr_col == nullptr) {
fill(tC_rCol, params_ptr->null_default);
return;
}
}
clear(tC_rCol);
Tensor pred = make_tensor<bool>(shape(tC_gCol));
CUTLASS_PRAGMA_UNROLL
include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp
+1 -4
View File
@@ -519,10 +519,7 @@ struct VisitorRowReduction {
// 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));
}
copy(tRS_rSrc, tRS_sRows);
__syncthreads();
include/cutlass/epilogue/threadblock/output_tile_thread_map.h
+2 -2
View File
@@ -391,7 +391,7 @@ struct OutputTileOptimalThreadMap {
1>;
/// Initial offset function
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
static MatrixCoord initial_offset(int thread_idx) {
// int warp_idx = __shfl_sync(0xffffffff, thread_idx / kWarpSize, 0);
@@ -462,7 +462,7 @@ struct OutputTileOptimalThreadMap {
static int const kThreads = Threads;
/// Function to compute each thread's initial offset
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
static MatrixCoord initial_offset(int thread_idx) {
// int warp_idx = __shfl_sync(0xffffffff, thread_idx / kWarpSize, 0);
include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h
+11 -3
View File
@@ -212,15 +212,23 @@ public:
// When the optimization is enabled, small tiles require separate logic.
bool kN32_optimization = (WarpShape::kN * Detail::kLanesInQuad * Policy::kElementsPerAccess * sizeof_bits<Element>::value) % 1024 == 0;
if (kN32_optimization) {
int ptr_idx = ((warp_column_ * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
if (ptr_idx == 0) {
ptr = pointers_[0];
} else if (ptr_idx == 1) {
ptr = pointers_[1];
if constexpr (AccessType::kElements >= 2) {
ptr = pointers_[1];
}
} else if (ptr_idx == 2) {
ptr = pointers_[2];
if constexpr (AccessType::kElements >= 3) {
ptr = pointers_[2];
}
} else if (ptr_idx == 3) {
ptr = pointers_[3];
if constexpr (AccessType::kElements >= 4) {
ptr = pointers_[3];
}
}
}
include/cutlass/fast_math.h
+2 -7
View File
@@ -38,7 +38,7 @@
#include <cmath>
#include <type_traits>
#endif
#include <cuda/std/utility>
#include "cutlass/cutlass.h"
#include "cutlass/array.h"
#include "cutlass/uint128.h"
@@ -54,12 +54,7 @@ namespace cutlass {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
CUTLASS_HOST_DEVICE void swap(T &lhs, T &rhs) {
T tmp = lhs;
lhs = rhs;
rhs = tmp;
}
using ::cuda::std::swap;
/******************************************************************************
* Static math utilities
include/cutlass/float8.h
+2 -2
View File
@@ -1053,8 +1053,8 @@ float_e5m2_t::float_e5m2_t(float_e4m3_t x) {
/// datatype in runtime argument list.
///
/// Currently supported runtime datatypes compatible with type_erased_dynamic_float8_t:
/// QMMAFormat::E5M2
/// QMMAFormat::E4M3
/// MXF8F6F4Format::E5M2
/// MXF8F6F4Format::E4M3
///
///////////////////////////////////////////////////////////////
include/cutlass/floating_point_nvrtc.h
+6
View File
@@ -35,6 +35,12 @@
#pragma once
#include <cutlass/detail/helper_macros.hpp> // CUTLASS_HOST_DEVICE
#include <cutlass/platform/platform.h> // uint32_t
#if !defined(__CUDACC_RTC__)
#include <cstring> // std::memcpy
#endif
namespace cutlass {
///////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/functional.h
+1 -1
View File
@@ -50,7 +50,7 @@
#ifdef _MSC_VER
// Provides support for alternate operators such as 'and', 'or', ...
#include <iso646.h>
#include <ciso646>
#endif // _MSC_VER
namespace cutlass {
include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl
+17 -14
View File
@@ -35,6 +35,8 @@
#include "cutlass/pipeline/sm90_pipeline.hpp"
#include "cutlass/gemm/collective/collective_mma_decl.hpp"
#include "cutlass/gemm/collective/collective_builder_decl.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/tensor.hpp"
// SM90 Collective Builders should be used only starting CUDA 12.0
#if (__CUDACC_VER_MAJOR__ >= 12)
@@ -236,8 +238,9 @@ struct CollectiveBuilder<
GmmaMajorA, ElementAMma, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using SmemLayoutAtomB = decltype(detail::ss_smem_selector<
GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes;
static constexpr int Sm90ReducedSmemCapacityBytes =
detail::sm90_smem_capacity_bytes;
static constexpr int PipelineStages = detail::compute_stage_count_or_override<Sm90ReducedSmemCapacityBytes,
ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
@@ -343,7 +346,7 @@ public:
return t;
}
else {
return cute::stride(t);
return cute::stride(t);
}
}
@@ -415,15 +418,15 @@ public:
static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int PipelineStages = IsMixedInput ?
detail::compute_stage_count_or_override_single_affine_transformed_input<detail::sm90_smem_capacity_bytes,
RealElementA, RealElementB, ElementScale, ElementZero, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{}) :
detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes,
ElementAMma, ElementBMma, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{});
static constexpr int PipelineStages = IsMixedInput ?
detail::compute_stage_count_or_override_single_affine_transformed_input<detail::sm90_smem_capacity_bytes,
RealElementA, RealElementB, ElementScale, ElementZero, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{})
: detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes,
ElementAMma, ElementBMma, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{});
using DispatchPolicy = cute::conditional_t<IsMixedInput,
MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<PipelineStages, ClusterShape_MNK, KernelScheduleType>,
MainloopSm90TmaGmmaRmemAWarpSpecialized<PipelineStages, ClusterShape_MNK, KernelScheduleType>>;
MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<PipelineStages, ClusterShape_MNK, KernelScheduleType>
, MainloopSm90TmaGmmaRmemAWarpSpecialized<PipelineStages, ClusterShape_MNK, KernelScheduleType>>;
using SmemCopyAtomA = cute::conditional_t<SwapAB, void, Copy_Atom<cute::AutoVectorizingCopy, ElementA>>;
using SmemCopyAtomB = cute::conditional_t<SwapAB, Copy_Atom<cute::AutoVectorizingCopy, ElementB>, void>;
@@ -761,13 +764,13 @@ struct CollectiveBuilder<
static constexpr int NumLoadWarpGroups = cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecialized> ? 2 : 1;
using AlignmentTypeA = cute::uint_byte_t<static_cast<int>(sizeof(ElementA)) * AlignmentA>;
using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeA>, ElementA>;
using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeA>, ElementA>;
using GmemTiledCopyA = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomA, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentA, TagToStrideA_t<GmemLayoutATag>,
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeB>, ElementB>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeB>, ElementB>;
using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomB, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
@@ -867,13 +870,13 @@ struct CollectiveBuilder<
static constexpr int NumLoadWarpGroups = 1;
using AlignmentTypeA = cute::uint_byte_t<static_cast<int>(sizeof(ElementA)) * AlignmentA>;
using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeA>, ElementA>;
using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeA>, ElementA>;
using GmemTiledCopyA = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomA, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentA, TagToStrideA_t<GmemLayoutATag>,
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeB>, ElementB>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeB>, ElementB>;
using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomB, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
include/cutlass/gemm/collective/collective_builder_decl.hpp
+12
View File
@@ -54,6 +54,18 @@ struct StageCountAutoCarveout {
explicit StageCountAutoCarveout(cute::Int<carveout_bytes>) {}
};
namespace detail {
// Forward Declaration
template<class CollectiveEpilogue>
constexpr int
compute_carveout_from_epi();
} // namespace detail
template<class CollectiveEpilogue>
struct StageCountAutoCarveoutEpi : StageCountAutoCarveout<detail::compute_carveout_from_epi<CollectiveEpilogue>()> {};
using StageCountAuto = StageCountAutoCarveout<0>;
// Used to automatically let the builder pick the kernel schedule.
include/cutlass/gemm/collective/collective_mma.hpp
+2 -1
View File
@@ -41,9 +41,10 @@
#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input.hpp
+1370
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp
+1 -1
View File
@@ -374,7 +374,7 @@ struct CollectiveMma<
// Prepare the TMA loads for A and B
//
constexpr uint32_t cluster_shape_x = get<0>(DispatchPolicy::ClusterShape());
constexpr uint32_t cluster_shape_x = get<0>(typename DispatchPolicy::ClusterShape());
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
Tensor gA_mkl = get<0>(load_inputs);
include/cutlass/gemm/device/gemm_universal_adapter.h
+58 -19
View File
@@ -85,13 +85,40 @@ class GemmUniversalAdapter;
////////////////////////////// CUTLASS 3.x API /////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
namespace detail {
// Work-around for some DispatchPolicy types not having a Stages member.
// In that case, the Stages value is 0. Most code should static_assert
// that the number of stages is valid.
// Whether DispatchPolicy::Stages is valid.
// It should also be convertible to int, but if not, that will show up
// as a build error when GemmUniversalAdapter attempts to assign it to kStages.
template <class DispatchPolicy, class Enable = void>
struct has_Stages : cute::false_type {};
template <class DispatchPolicy>
struct has_Stages<DispatchPolicy, cute::void_t<decltype(DispatchPolicy::Stages)>> : cute::true_type {};
template<class DispatchPolicy>
constexpr int stages_member(DispatchPolicy) {
if constexpr (has_Stages<DispatchPolicy>::value) {
return DispatchPolicy::Stages;
}
else {
return 0;
}
}
} // namespace detail
template <class GemmKernel_>
class GemmUniversalAdapter<
GemmKernel_,
cute::enable_if_t<gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
cute::enable_if_t<gemm::detail::IsCutlass3GemmKernel<GetUnderlyingKernel_t<GemmKernel_>>::value>>
{
public:
using GemmKernel = GemmKernel_;
using GemmKernel = GetUnderlyingKernel_t<GemmKernel_>;
using TileShape = typename GemmKernel::TileShape;
using ElementA = typename GemmKernel::ElementA;
using ElementB = typename GemmKernel::ElementB;
@@ -158,7 +185,7 @@ public:
CUTE_STATIC_V(cute::tile_size<1>(typename CollectiveMainloop::TiledMma{})) / WarpsInMmaN,
CUTE_STATIC_V(cute::tile_size<2>(typename CollectiveMainloop::TiledMma{}))>;
static int constexpr kStages = CollectiveMainloop::DispatchPolicy::Stages;
static int constexpr kStages = detail::stages_member(typename CollectiveMainloop::DispatchPolicy{});
// Inspect TiledCopy for A and B to compute the alignment size
static int constexpr kAlignmentA = cutlass::detail::get_alignment_count_from_gmem_tiled_copy<
@@ -336,7 +363,7 @@ public:
}
/// Primary run() entry point API that is static allowing users to create and manage their own params.
/// Supplied params struct must be construct by calling GemmKernel::to_underling_arguments()
/// Supplied params struct must be construct by calling GemmKernel::to_underlying_arguments()
static Status
run(Params& params,
cudaStream_t stream = nullptr,
@@ -358,10 +385,10 @@ public:
[[maybe_unused]] constexpr bool is_static_1x1x1 =
cute::is_static_v<typename GemmKernel::DispatchPolicy::ClusterShape> and
cute::size(typename GemmKernel::DispatchPolicy::ClusterShape{}) == 1;
dim3 cluster(cute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
cute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
cute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{}));
void* kernel_params[] = {&params};
[[maybe_unused]] dim3 cluster(cute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
cute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
cute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{}));
[[maybe_unused]] void* kernel_params[] = {&params};
if constexpr (kEnableCudaHostAdapter) {
//
@@ -377,13 +404,23 @@ public:
#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST("GemmUniversal::run: Launching kernel with CUDA host adapter");
#endif
launch_result = cuda_adapter->launch(grid,
cluster,
block,
smem_size,
stream,
kernel_params,
0);
if constexpr (is_static_1x1x1) {
launch_result = cuda_adapter->launch(grid,
block,
smem_size,
stream,
kernel_params,
0);
}
else {
launch_result = cuda_adapter->launch(grid,
cluster,
block,
smem_size,
stream,
kernel_params,
0);
}
}
else {
CUTLASS_TRACE_HOST("GemmUniversal::run: kEnableCudaHostAdapter is true, but CUDA host adapter is null");
@@ -392,8 +429,10 @@ public:
}
else {
CUTLASS_ASSERT(cuda_adapter == nullptr);
void const* kernel = (void const*) device_kernel<GemmKernel>;
if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 90) {
[[maybe_unused]] void const* kernel = (void const*) device_kernel<GemmKernel>;
static constexpr bool kClusterLaunch = GemmKernel::ArchTag::kMinComputeCapability == 90
;
if constexpr (kClusterLaunch) {
if constexpr (is_static_1x1x1) {
#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST("GemmUniversal::run: Launching static 1x1x1 kernel");
@@ -526,11 +565,11 @@ public:
template <class GemmKernel_>
class GemmUniversalAdapter<
GemmKernel_,
cute::enable_if_t<not gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
cute::enable_if_t<not gemm::detail::IsCutlass3GemmKernel<GetUnderlyingKernel_t<GemmKernel_>>::value>>
{
public:
using GemmKernel = GemmKernel_;
using GemmKernel = GetUnderlyingKernel_t<GemmKernel_>;
static bool const kInternalTranspose =
!cutlass::epilogue::threadblock::detail::is_2x_evt_v<typename GemmKernel::Epilogue> && // 2.x EVT does not require internal transpose
include/cutlass/gemm/dispatch_policy.hpp
+4 -1
View File
@@ -105,7 +105,8 @@ struct KernelCpAsyncWarpSpecializedPingpong { };
struct KernelCpAsyncWarpSpecializedCooperative { };
struct KernelTma { };
struct KernelTmaWarpSpecialized { };
struct KernelTmaWarpSpecializedPingpong { };
struct KernelTmaWarpSpecializedPingpong {
};
struct KernelTmaWarpSpecializedCooperative {
};
@@ -247,6 +248,7 @@ struct MainloopSm90TmaGmmaRmemAWarpSpecialized {
"KernelSchedule must be one of the warp specialized policies");
};
template<
int Stages_,
class ClusterShape_ = Shape<_1,_1,_1>,
@@ -310,6 +312,7 @@ struct MainloopSm90TmaGmmaWarpSpecializedSparse {
using Schedule = KernelSchedule;
};
//////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm
include/cutlass/gemm/group_array_problem_shape.hpp
+1 -1
View File
@@ -69,7 +69,7 @@ struct GroupProblemShape {
CUTLASS_HOST_DEVICE
UnderlyingProblemShape const
get_host_problem_shape(int32_t group_idx) const {
return host_problem_shapes[group_idx];
return host_problem_shapes != nullptr ? host_problem_shapes[group_idx] : UnderlyingProblemShape{};
}
CUTLASS_HOST_DEVICE
include/cutlass/gemm/kernel/default_gemm_grouped_per_group_scale.h
+384
View File
@@ -0,0 +1,384 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 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 kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
the appropriate threadblock-scoped epilogue.
Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
accommodated by exchanging A and B operands and assuming transposed layouts. Partial
specializations here choose 'device::GemmTransposed' to implement this functionality.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/complex.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/numeric_types.h"
#include "cutlass/gemm/kernel/gemm_grouped_per_group_scale.h"
#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
#include "cutlass/gemm/kernel/default_gemm.h"
#include "cutlass/gemm/kernel/default_gemm_complex.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"
#include "cutlass/layout/permute.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_,
/// Element type for internal accumulation
typename ElementAccumulator,
/// 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 EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Whether the schedule of problems to visit has been precomputed
GroupScheduleMode GroupScheduleMode_ = GroupScheduleMode::kDeviceOnly,
/// Operation performed by GEMM
typename Operator = typename device::DefaultGemmConfiguration<
OperatorClass, ArchTag, ElementA_, ElementB_, ElementC_,
ElementAccumulator>::Operator,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
/// Permute result D
typename PermuteDLayout = layout::NoPermute,
///
typename Enable = void
>
struct DefaultGemmGroupedPerGroupScale;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Real-valued GEMM kernels
//
template <
/// Element type for A matrix operand
typename ElementA,
/// Layout type for A matrix operand
typename LayoutA,
/// 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,
/// 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,
/// Element type for internal accumulation
typename ElementAccumulator,
/// 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 EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Whether the schedule of problems to visit has been precomputed
GroupScheduleMode GroupScheduleMode_,
/// Operation performed by GEMM
typename Operator,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear,
/// Permute result D
typename PermuteDLayout
>
struct DefaultGemmGroupedPerGroupScale<
ElementA,
LayoutA,
ComplexTransform::kNone, // transform A
kAlignmentA,
ElementB,
LayoutB,
ComplexTransform::kNone, // transform B
kAlignmentB,
ElementC,
LayoutC,
ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
GroupScheduleMode_,
Operator,
SharedMemoryClear,
PermuteDLayout,
typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
> {
// If true, we must construct a 'transposed-and-exchanged' Mma operator.
static bool const kInternalTranspose = platform::is_same<LayoutC, layout::ColumnMajor>::value;
using MapArguments = kernel::detail::MapArguments<
ElementA,
LayoutA,
ComplexTransform::kNone,
kAlignmentA,
ElementB,
LayoutB,
ComplexTransform::kNone,
kAlignmentB,
LayoutC,
kInternalTranspose
>;
// Define the default GEMM kernel
using DefaultGemmKernel = typename kernel::DefaultGemm<
typename MapArguments::ElementA,
typename MapArguments::LayoutA,
MapArguments::kAlignmentA,
typename MapArguments::ElementB,
typename MapArguments::LayoutB,
MapArguments::kAlignmentB,
ElementC,
typename MapArguments::LayoutC,
ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
true,
Operator,
SharedMemoryClear,
false, /*GatherA*/
false, /*GatherB*/
false, /*ScatterD*/
PermuteDLayout
>::GemmKernel;
/// Define the kernel in terms of the default kernel
using GemmKernel = kernel::GemmGroupedPerGroupScale<
typename DefaultGemmKernel::Mma,
typename DefaultGemmKernel::Epilogue,
ThreadblockSwizzle,
GroupScheduleMode_,
kInternalTranspose
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Complex-valued GEMM kernels
//
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,
/// Element type for internal accumulation
typename ElementAccumulator,
/// 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 EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Whether the schedule of problems to visit has been precomputed
GroupScheduleMode GroupScheduleMode_,
/// Operation performed by GEMM
typename Operator,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear
>
struct DefaultGemmGroupedPerGroupScale<
ElementA,
LayoutA,
TransformA,
kAlignmentA,
ElementB,
LayoutB,
TransformB,
kAlignmentB,
ElementC,
LayoutC,
ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
GroupScheduleMode_,
Operator,
SharedMemoryClear,
layout::NoPermute, /*PermuteDLayout*/
typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
> {
// If true, we must construct a 'transposed-and-exchanged' Mma operator.
static bool const kInternalTranspose = platform::is_same<LayoutC, layout::ColumnMajor>::value;
using MapArguments = kernel::detail::MapArguments<
ElementA,
LayoutA,
TransformA,
kAlignmentA,
ElementB,
LayoutB,
TransformB,
kAlignmentB,
LayoutC,
kInternalTranspose
>;
using DefaultGemmKernel = typename kernel::DefaultGemmComplex<
typename MapArguments::ElementA,
typename MapArguments::LayoutA,
typename MapArguments::ElementB,
typename MapArguments::LayoutB,
ElementC,
typename MapArguments::LayoutC,
ElementAccumulator,
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
MapArguments::kTransformA,
MapArguments::kTransformB,
Operator,
false
>::GemmKernel;
/// Define the kernel in terms of the default kernel
using GemmKernel = kernel::GemmGroupedPerGroupScale<
typename DefaultGemmKernel::Mma,
typename DefaultGemmKernel::Epilogue,
ThreadblockSwizzle,
GroupScheduleMode_,
kInternalTranspose
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/ell_gemm.h
+3 -3
View File
@@ -691,7 +691,7 @@ struct EllGemm<Mma_, Epilogue_, ThreadblockSwizzle_, SplitKSerial, false> {
static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
constexpr bool is_double = (sizeof(Mma::IteratorA::Element) == 8);
constexpr bool is_double = (sizeof(typename Mma::IteratorA::Element) == 8);
constexpr bool is_multiple_alignment =
(kAlignmentA > 1) && (kAlignmentB > 1) && (kAlignmentC > 1);
const bool is_specialized_blocksize =
@@ -699,11 +699,11 @@ struct EllGemm<Mma_, Epilogue_, ThreadblockSwizzle_, SplitKSerial, false> {
&& params.ell_blocksize >= Mma::Shape::kK;
// Compute threadblock-scoped matrix multiply-add
if ((is_double || is_multiple_alignment) && is_specialized_blocksize) {
mma.operator()<false, true>(
mma.template operator()<false, true>(
gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator);
}
else {
mma.operator()<false, false>(
mma.template operator()<false, false>(
gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator);
}
}
include/cutlass/gemm/kernel/gemm_grouped_per_group_scale.h
+261
View File
@@ -0,0 +1,261 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 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 Problem visitor for grouped GEMMs
*/
#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/semaphore.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/trace.h"
#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
#include "cutlass/gemm/kernel/gemm_grouped_problem_visitor.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/kernel/gemm_grouped.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate
typename Epilogue_, ///! Epilogue
typename ThreadblockSwizzle_, ///! Threadblock swizzling function
GroupScheduleMode GroupScheduleMode_, ///! Type of scheduling to perform
bool Transposed = false
>
struct GemmGroupedPerGroupScale :
public GemmGrouped<Mma_, Epilogue_, ThreadblockSwizzle_, GroupScheduleMode_, Transposed> {
// Inherit constructors
using Base = GemmGrouped<Mma_, Epilogue_, ThreadblockSwizzle_, GroupScheduleMode_, Transposed>;
// Inherit type definitions
using typename Base::Mma;
using typename Base::Epilogue;
using typename Base::EpilogueOutputOp;
using typename Base::ThreadblockSwizzle;
using typename Base::Params;
using typename Base::SharedStorage;
// Explicitly inherit the kTransposed constant
static bool const kTransposed = Base::kTransposed;
/// Executes one GEMM
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage) {
//
// These types shadow the type-level definitions and support the ability to implement
// a 'transposed' GEMM that computes the transposed problems.
//
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;
//
// Problem visitor.
//
typename Base::ProblemVisitor problem_visitor(
params.problem_visitor,
shared_storage.problem_visitor,
blockIdx.x);
// Outer 'persistent' loop to iterate over tiles
while (problem_visitor.next_tile()) {
GemmCoord problem_size = problem_visitor.problem_size();
int32_t problem_idx = problem_visitor.problem_index();
int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx());
GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
cutlass::gemm::GemmCoord threadblock_offset(
int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM,
int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN,
0);
// Load element pointers. Exchange pointers and strides if working on the transpose
ElementA *ptr_A = reinterpret_cast<ElementA *>((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx]));
typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]);
ElementB *ptr_B = reinterpret_cast<ElementB *>((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx]));
typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]);
// Compute initial location in logical coordinates
cutlass::MatrixCoord tb_offset_A{
threadblock_offset.m(),
0,
};
cutlass::MatrixCoord tb_offset_B{
0,
threadblock_offset.n()
};
// Compute position within threadblock
int thread_idx = threadIdx.x;
// Construct iterators to A and B operands
typename Mma::IteratorA iterator_A(
LayoutA(ldm_A),
ptr_A,
{problem_size.m(), problem_size.k()},
thread_idx,
tb_offset_A);
typename Mma::IteratorB iterator_B(
LayoutB(ldm_B),
ptr_B,
{problem_size.k(), problem_size.n()},
thread_idx,
tb_offset_B);
typename Mma::FragmentC accumulators;
accumulators.clear();
// 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;
//
// Matrix multiply phase
//
// Construct thread-scoped matrix multiply
Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx);
// Compute threadblock-scoped matrix multiply-add
int gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK;
// Wait for all threads to finish their epilogue phases from the previous tile.
__syncthreads();
// Compute threadblock-scoped matrix multiply-add
mma(
gemm_k_iterations,
accumulators,
iterator_A,
iterator_B,
accumulators);
//
// Epilogue
//
ElementC *ptr_C = params.ptr_C[problem_idx];
ElementC *ptr_D = params.ptr_D[problem_idx];
LayoutC layout_C(params.ldc[problem_idx]);
LayoutC layout_D(params.ldd[problem_idx]);
typename Epilogue::OutputTileIterator::Params params_C(layout_C);
typename Epilogue::OutputTileIterator::Params params_D(layout_D);
// Tile iterator loading from source tensor.
typename Epilogue::OutputTileIterator iterator_C(
params_C,
ptr_C,
problem_size.mn(),
thread_idx,
threadblock_offset.mn()
);
// Tile iterator writing to destination tensor.
typename Epilogue::OutputTileIterator iterator_D(
params_D,
ptr_D,
problem_size.mn(),
thread_idx,
threadblock_offset.mn()
);
Epilogue epilogue(
shared_storage.kernel.epilogue,
thread_idx,
warp_idx,
lane_idx);
// The if branch is for the per-group scaling epilogue. The customized epilogue operator scales each gemm output by a scalar value.
// This branch is only enabled if EpilogueOutputOp is LinearCombination.
if constexpr (platform::is_same<EpilogueOutputOp,
::cutlass::epilogue::thread::LinearCombination<typename EpilogueOutputOp::ElementOutput,
EpilogueOutputOp::kCount, typename EpilogueOutputOp::ElementAccumulator,
typename EpilogueOutputOp::ElementCompute, EpilogueOutputOp::kScale,
EpilogueOutputOp::kRound>>::value)
{
EpilogueOutputOp output_op(params.output_op, problem_idx);
// Execute the epilogue operator to update the destination tensor.
epilogue(
output_op,
iterator_D,
accumulators,
iterator_C);
} else {
EpilogueOutputOp output_op(params.output_op);
// Execute the epilogue operator to update the destination tensor.
epilogue(
output_op,
iterator_D,
accumulators,
iterator_C);
}
// Next tile
problem_visitor.advance(gridDim.x);
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/gemm_grouped_problem_visitor.h
+1 -1
View File
@@ -68,7 +68,7 @@ struct GemmGroupedProblemSizeHelper {
CUTLASS_HOST_DEVICE
static void possibly_transpose_problem(cutlass::gemm::GemmCoord& problem) {
if (kTransposed) {
swap(problem.m(), problem.n());
cutlass::swap(problem.m(), problem.n());
}
}
include/cutlass/gemm/kernel/gemm_universal_with_visitor_streamk.h
+2 -2
View File
@@ -437,7 +437,7 @@ protected:
int m_begin = tile_work.tiled_coord.m() * Mma::Shape::kM;
int m_end = params.block_mapping.problem_size.m();
return Mma::IteratorA(
return typename Mma::IteratorA(
params.params_A,
ptr_A,
{ m_end, tile_work.k_end },
@@ -466,7 +466,7 @@ protected:
int n_begin = tile_work.tiled_coord.n() * Mma::Shape::kN;
int n_end = params.block_mapping.problem_size.n();
return Mma::IteratorB(
return typename Mma::IteratorB(
params.params_B,
ptr_B,
{ tile_work.k_end, n_end },
include/cutlass/gemm/kernel/grouped_problem_visitor.h
+2 -2
View File
@@ -66,10 +66,10 @@ struct BaseGroupedProblemVisitor {
int32_t problem_idx;
int32_t problem_start;
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
ProblemInfo() : problem_idx(kNoPrefetchEntry), problem_start(kNoPrefetchEntry) {}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
ProblemInfo(int32_t problem_idx_, int32_t problem_start_) :
problem_idx(problem_idx_), problem_start(problem_start_) {}
};
include/cutlass/gemm/kernel/params_universal_base.h
+1 -1
View File
@@ -182,7 +182,7 @@ struct UniversalParamsBase
CUTLASS_TRACE_HOST(" Initialize " << workspace_bytes << " workspace bytes");
cudaError_t result = cudaMemsetAsync(
semaphore,
static_cast<int *>(workspace),
0,
workspace_bytes,
stream);

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