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

v4.5 dev update. (#3153)

Browse Source
This commit is contained in:
Junkai-Wu
2026-04-07 12:16:05 -04:00
committed by GitHub
parent 418d38a5de
commit a221da7ccf
265 changed files with 4913 additions and 1478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/copy_sm100_tma.hpp
+101
View File
@@ -663,4 +663,105 @@ struct SM100_TMA_2SM_LOAD_IM2COL_MULTICAST
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM100_TMA_LOAD_2D_GATHER4
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1_i0, int32_t const& crd1_i1, int32_t const& crd1_i2, int32_t const& crd1_i3)
{
#if defined(CUTE_ARCH_TMA_SM100_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1_i0), "r"(crd1_i1), "r"(crd1_i2), "r"(crd1_i3), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM100_ENABLED.");
#endif
}
struct PREFETCH
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr,
int32_t const& crd0, int32_t const& crd1_i0, int32_t const& crd1_i1, int32_t const& crd1_i2, int32_t const& crd1_i3)
{
#if defined(CUTE_ARCH_TMA_SM100_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
asm volatile (
"cp.async.bulk.prefetch.tensor.2d.L2.global.tile::gather4"
" [%0, {%1, %2, %3, %4, %5}];"
:
: "l"(gmem_int_desc),
"r"(crd0), "r"(crd1_i0), "r"(crd1_i1), "r"(crd1_i2), "r"(crd1_i3)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM100_ENABLED.");
#endif
}
};
};
struct SM100_TMA_LOAD_MULTICAST_2D_GATHER4
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1_i0, int32_t const& crd1_i1, int32_t const& crd1_i2, int32_t const& crd1_i3)
{
#if defined(CUTE_ARCH_TMA_SM100_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%4, %5, %6, %7, %8}], [%2], %3, %9;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"h"(multicast_mask),
"r"(crd0), "r"(crd1_i0), "r"(crd1_i1), "r"(crd1_i2), "r"(crd1_i3), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM100_ENABLED.");
#endif
}
using PREFETCH = SM100_TMA_LOAD_2D_GATHER4::PREFETCH;
};
struct SM100_TMA_STORE_2D_SCATTER4
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, void const* smem_ptr,
int32_t const& crd0, int32_t const& crd1_i0, int32_t const& crd1_i1, int32_t const& crd1_i2, int32_t const& crd1_i3)
{
#if defined(CUTE_ARCH_TMA_SM100_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.global.shared::cta.tile::scatter4.bulk_group [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr),
"r"(crd0), "r"(crd1_i0), "r"(crd1_i1), "r"(crd1_i2), "r"(crd1_i3)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM100_ENABLED.");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // end namespace cute
include/cute/atom/copy_traits_sm100_tma.hpp
+333 -4
View File
@@ -239,6 +239,331 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
}
};
//////////////////////////////////////////////////////////////////////////////
////////////////////////// TMA_LOAD_GATHER ///////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
// Utility for unpacking TMA_LOAD arguments into a CopyOp
template <class CopyOp, class... Args>
struct TMA_LOAD_GATHER_Unpack
{
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, "SM100_TMA_LOAD_2D_GATHER4 requires the destination be shared memory.");
auto [src_crd, src_idx] = unzip_tensor(src);
auto src_coord = src_crd.data().coord_;
static_assert(rank(src_coord) == 2, "SM100_TMA_LOAD_2D_GATHER4 requires 2D tensors");
Tensor idx = filter(src_idx);
static_assert(size(idx) == 4, "SM100_TMA_LOAD_2D_GATHER4 requires 4 indices");
auto coord = make_tuple(get<0>(src_coord), idx(0), idx(1), idx(2), idx(3));
void* dst_ptr = cute::raw_pointer_cast(dst.data());
#if 0
auto [c0,c1,c2,c3,c4] = coord;
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>{},
coord, make_seq<5>{});
}
};
struct SM100_TMA_LOAD_2D_GATHER4_OP : SM100_TMA_LOAD_2D_GATHER4 {};
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM100_TMA_LOAD_2D_GATHER4, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_LOAD arguments
TmaDescriptor tma_desc_;
using AuxParams = AuxParams_;
AuxParams aux_params_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return &tma_desc_;
}
// Construct an executable SM90_TMA_LOAD with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_LOAD_2D_GATHER4_OP, NumBitsPerTMA>
with(
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm100 const& cache_hint = TMA::CacheHintSm100::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)};
}
// Construct an executable SM90_TMA_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_LOAD_2D_GATHER4_OP, NumBitsPerTMA>
with(
TmaDescriptor const* new_tma_desc,
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm100 const& cache_hint = TMA::CacheHintSm100::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)};
}
// Generate the TMA coord tensor
template <class GShape>
CUTE_HOST_DEVICE constexpr
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM100_TMA_LOAD_2D_GATHER4 before calling .with()
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) = delete;
};
template <class NumBitsPerTMA>
struct Copy_Traits<SM100_TMA_LOAD_2D_GATHER4_OP, NumBitsPerTMA>
: TMA_LOAD_GATHER_Unpack<SM100_TMA_LOAD_2D_GATHER4_OP, NumBitsPerTMA>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
tuple<
TmaDescriptor const*,
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) {}
};
template <class NumBitsPerTMA, class... Args>
struct Copy_Traits<SM100_TMA_LOAD_2D_GATHER4_OP::PREFETCH, NumBitsPerTMA, Args...>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
tuple<TmaDescriptor const*> const opargs_;
// Construct with any other Traits' TMA Desc
template <class... CopyArgs>
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_crd, src_idx] = unzip_tensor(src);
auto src_coord = src_crd.data().coord_;
static_assert(rank(src_coord) == 2, "SM100_TMA_LOAD_2D_GATHER4 requires 2D tensors");
Tensor idx = filter(src_idx);
static_assert(size(idx) == 4, "SM100_TMA_LOAD_2D_GATHER4 requires 4 indices");
auto coord = make_tuple(get<0>(src_coord), idx(0), idx(1), idx(2), idx(3));
return detail::explode_tuple(detail::CallCOPY<SM100_TMA_LOAD_2D_GATHER4_OP::PREFETCH>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
coord, make_seq<5>{});
}
};
struct SM100_TMA_LOAD_MULTICAST_2D_GATHER4_OP : SM100_TMA_LOAD_MULTICAST_2D_GATHER4 {};
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM100_TMA_LOAD_MULTICAST_2D_GATHER4, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_LOAD arguments
TmaDescriptor tma_desc_;
using AuxParams = AuxParams_;
AuxParams aux_params_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return &tma_desc_;
}
// Construct an executable SM90_TMA_LOAD with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_LOAD_MULTICAST_2D_GATHER4_OP, NumBitsPerTMA>
with(
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm100 const& cache_hint = TMA::CacheHintSm100::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {&tma_desc_, &tma_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)};
}
// Construct an executable SM90_TMA_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_LOAD_MULTICAST_2D_GATHER4_OP, NumBitsPerTMA>
with(
TmaDescriptor const* new_tma_desc,
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm100 const& cache_hint = TMA::CacheHintSm100::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {new_tma_desc, &tma_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)};
}
// Generate the TMA coord tensor
template <class GShape>
CUTE_HOST_DEVICE constexpr
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM100_TMA_LOAD_2D_GATHER4 before calling .with()
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) = delete;
};
template <class NumBitsPerTMA>
struct Copy_Traits<SM100_TMA_LOAD_MULTICAST_2D_GATHER4_OP, NumBitsPerTMA>
: TMA_LOAD_GATHER_Unpack<SM100_TMA_LOAD_MULTICAST_2D_GATHER4_OP, NumBitsPerTMA>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
tuple<
TmaDescriptor const*,
uint64_t*, // smem mbarrier
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 cache)
: opargs_(desc, mbar, mask, cache) {}
};
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM100_TMA_STORE_2D_SCATTER4, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
TmaDescriptor tma_desc_;
using AuxParams = AuxParams_;
AuxParams aux_params_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return &tma_desc_;
}
// Generate the TMA coord tensor
template <class GShape>
CUTE_HOST_DEVICE constexpr
auto
get_tma_tensor(GShape const& g_shape) const {
static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
return make_coord_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
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 [dsc_crd, dsc_idx] = unzip_tensor(dst);
auto dsc_coord = dsc_crd.data().coord_;
static_assert(rank(dsc_coord) == 2, "SM100_TMA_STORE_2D_SCATTER4 requires 2D tensors");
Tensor idx = filter(dsc_idx);
static_assert(size(idx) == 4, "SM100_TMA_STORE_2D_SCATTER4 requires 4 indices");
auto coord = make_tuple(get<0>(dsc_coord), idx(0), idx(1), idx(2), idx(3));
void* src_ptr = cute::raw_pointer_cast(src.data());
#if 0
auto [c0,c1,c2,c3,c4] = coord;
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<SM100_TMA_STORE_2D_SCATTER4>{},
make_tuple(&traits.tma_desc_), seq<0>{},
make_tuple(src_ptr), seq<0>{},
coord, make_seq<5>{});
}
};
////////////////////////////////////
// Make TMA
///////////////////////////////////
@@ -428,12 +753,14 @@ make_tma_atom_A_sm100(CopyOp const& copy_op,
// The size of the multicasting
auto num_multicast = [&](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_MULTICAST>) {
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_LOAD_MULTICAST_2D_GATHER4>) {
return size<2>(cluster_shape); // VMNK: Use only the N-CTAs in the Multicast
} else
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD> ||
is_same_v<CopyOp, SM90_TMA_STORE> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD>) {
is_same_v<CopyOp, SM100_TMA_2SM_LOAD> ||
is_same_v<CopyOp, SM100_TMA_LOAD_2D_GATHER4>) {
return Int<1>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
@@ -479,12 +806,14 @@ make_tma_atom_B_sm100(CopyOp const& copy_op,
// The size of the multicasting
auto num_multicast = [&](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_MULTICAST>) {
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_LOAD_MULTICAST_2D_GATHER4>) {
return size<1>(cluster_shape); // VMNK: Use only the M-CTAs in the Multicast
} else
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD> ||
is_same_v<CopyOp, SM90_TMA_STORE> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD>) {
is_same_v<CopyOp, SM100_TMA_2SM_LOAD> ||
is_same_v<CopyOp, SM100_TMA_LOAD_2D_GATHER4>) {
return Int<1>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
include/cute/atom/copy_traits_sm90_tma.hpp
+35 -19
View File
@@ -42,6 +42,7 @@
#include <cute/numeric/integral_ratio.hpp>
#include <cute/arch/copy_sm100_tma.hpp>
#include <cutlass/cuda_host_adapter.hpp>
namespace cute
@@ -1149,13 +1150,27 @@ make_tma_copy_atom(CopyOp,
auto smem_layout = get_nonswizzle_portion(slayout);
auto tma_gbasis = detail::construct_tma_gbasis<TmaInternalType>(gtensor, smem_layout, cta_v_map);
auto tma_gbasis_tuple = conditional_return<is_same_v<CopyOp, SM100_TMA_LOAD_2D_GATHER4>
||is_same_v<CopyOp, SM100_TMA_LOAD_MULTICAST_2D_GATHER4>
||is_same_v<CopyOp, SM100_TMA_STORE_2D_SCATTER4>>(
[](auto tma_gbasis) constexpr {
static_assert(rank_v<decltype(tma_gbasis)> == 2, "TMA Gather/Scatter only supports 2D tensors");
auto tma_gbasis_g4 = tma_gbasis.compose(make_identity_layout(make_shape(shape<0>(tma_gbasis), _1{})));
auto tma_gbasis_g4_size = size(tma_gbasis_g4) * _4{};
return make_tuple(tma_gbasis_g4, tma_gbasis_g4_size);
},
[](auto tma_gbasis) constexpr {
auto tma_gbasis_size = size(tma_gbasis);
return make_tuple(tma_gbasis, tma_gbasis_size);
})(tma_gbasis);
auto _tma_gbasis = get<0>(tma_gbasis_tuple);
auto _tma_gbasis_size = get<1>(tma_gbasis_tuple);
//
// Construct the TMA Desc and the strides of the TMA Tensor
//
auto [tma_desc, aux_params] = detail::make_tma_copy_desc<TmaInternalType>(gtensor,
tma_gbasis,
_tma_gbasis,
smem_swizzle,
num_multicast);
@@ -1163,7 +1178,7 @@ make_tma_copy_atom(CopyOp,
// Construct the Copy_Traits
//
constexpr int num_bits_per_tma = size(tma_gbasis) * sizeof_bits_v<TmaInternalType>;
constexpr int num_bits_per_tma = _tma_gbasis_size * sizeof_bits_v<TmaInternalType>;
using Traits = Copy_Traits<CopyOp, cute::C<num_bits_per_tma>, decltype(aux_params)>;
using Atom = Copy_Atom<Traits, typename GEngine::value_type>;
@@ -1397,17 +1412,16 @@ template <class... Args,
class CtaCoord,
class TShape, class TStride,
class SEngine, class SLayout,
class GEngine, class GLayout>
class... GTensors,
__CUTE_REQUIRES(conjunction_v<is_tensor<GTensors>...>)>
CUTE_DEVICE
auto
tma_partition(Copy_Atom<Args...> const& copy_atom,
CtaCoord const& cta_coord,
Layout<TShape,TStride> const& cta_layout, // T: CTA coord -> logical multicast id
Tensor<SEngine,SLayout> const& stensor, // SMEM Tensor (TMATile, Rest...)
Tensor<GEngine,GLayout> const& gtensor) // GMEM Tensor (TMATile, Rest...)
GTensors const&... gtensors) // GMEM Tensor (TMATile, Rest...)
{
CUTE_STATIC_ASSERT_V(size<0>(stensor) == size<0>(gtensor));
// Invert the smem to get the largest contiguous vector in the smem layout
Layout inv_smem_layout = right_inverse(get_nonswizzle_portion(layout<0>(stensor)));
// Scale that up to cover all of the smem_coords
@@ -1417,22 +1431,24 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
Layout tma_layout_v = make_layout(Int<Copy_Atom<Args...>::NumValSrc>{});
auto layout_V = make_tile(logical_divide(layout_v, tma_layout_v));
// Append with _ until we cover all Rest... modes
auto glayout_V = append<GLayout::rank>(layout_V, _);
auto slayout_V = append<SLayout::rank>(layout_V, _);
// Transform tile mode and coalesce
Tensor gtensor_v = coalesce(gtensor.compose(glayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
Tensor stensor_v = coalesce(stensor.compose(slayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
// Offset inside the TMA-mode for the multicast
auto multicast_offset = cta_layout(cta_coord) * (size(tma_layout_v) / cosize(cta_layout));
auto multicast_coord = make_coord(make_coord(multicast_offset, Int<0>{}));
auto gcoord = append<GLayout::rank>(multicast_coord, Int<0>{});
auto scoord = append<SLayout::rank>(multicast_coord, Int<0>{});
Tensor gresult = domain_offset(gcoord, gtensor_v);
Tensor sresult = domain_offset(scoord, stensor_v);
// Existing convention is to return stensor last
return cute::transform(make_tuple(gtensors..., stensor), [&](auto && tensor) {
auto R = rank(tensor);
CUTE_STATIC_ASSERT_V(size<0>(stensor) == size<0>(tensor));
return cute::make_tuple(gresult, sresult);
// Append with _ until we cover all Rest... modes
auto tlayout_V = append<R>(layout_V, _);
// Transform tile mode and coalesce
Tensor tensor_v = coalesce(tensor.compose(tlayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
// Offset inside the TMA-mode for the multicast
auto coord = append<R>(multicast_coord, Int<0>{});
return domain_offset(coord, tensor_v);
});
}
// Explicit defaults for cta_coord and cta_layout
include/cute/tensor_zip.hpp
+70 -2
View File
@@ -72,9 +72,12 @@ struct ZipIterator
template <class... Index>
CUTE_HOST_DEVICE constexpr
ZipIterator operator+(cute::tuple<Index...> const& idxs) const {
auto operator+(cute::tuple<Index...> const& idxs) const {
static_assert(sizeof...(Index) == sizeof...(Iters), "Expect same number of offsets as iterators.");
return cute::transform(iters_, idxs, [](auto&& iter, auto&& idx) { return iter + idx; });
return cute::transform_apply(iters_, idxs,
[](auto&& iter, auto&& idx) { return iter + idx; },
[](auto... iter) { return ZipIterator<decltype(iter)...>(iter...); }
);
}
template <class... Index>
@@ -149,6 +152,13 @@ struct ZipLayout
template <class... Layouts>
struct is_layout<ZipLayout<Layouts...>> : true_type {};
template <class Layout>
struct is_zip_layout : false_type {};
template <class... Layouts>
struct is_zip_layout<ZipLayout<Layouts...>> : true_type {};
//
// make_zip_tensor and unzip_tensor
//
@@ -191,6 +201,23 @@ size(ZipLayout<Layouts...> const& layouts)
return size<Is...>(get<0>(layouts.layouts_));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
get(ZipLayout<Layouts...> const& layouts)
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return get<Is...>(t); }));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
layout(ZipLayout<Layouts...> const& layouts)
{
return get<Is...>(layouts);
}
//
// Manipulation
//
@@ -243,4 +270,45 @@ slice_and_offset(Coord const& c, ZipLayout<Layouts...> const& layouts)
return cute::make_tuple(ZipLayout(get<0>(result)), get<1>(result));
}
template <int B, int E, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
group(ZipLayout<Layouts...> const& layouts)
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return group<B,E>(t); }));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
shape(ZipLayout<Layouts...> const& layouts) {
return shape<Is...>(get<0>(layouts.layouts_));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
coshape(ZipLayout<Layouts...> const& layouts) {
return cute::transform(layouts.layouts_, [&](auto t){ return coshape<Is...>(t); });
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
cosize(ZipLayout<Layouts...> const& layouts)
{
return size(coshape<Is...>(layouts));
}
template <class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
nullspace(ZipLayout<Layouts...> const& layouts) {
return cute::fold(layouts.layouts_, make_layout(size(layouts)),
[](auto null, auto layout) {
return composition(null, nullspace(composition(layout, null)));
});
}
} // end namespace cute
include/cutlass/conv/collective/builders/sm100_common.inl
+12
View File
@@ -186,6 +186,18 @@ sm100_make_tiled_mma() {
}
}
template <class ArchTag, int KernelSmemCarveout>
constexpr int sm100_reduced_smem_capacity_bytes() {
if constexpr (cute::is_same_v<ArchTag, arch::Sm100>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else if constexpr (cute::is_same_v<ArchTag, arch::Sm103>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else {
static_assert(cutlass::detail::dependent_false<ArchTag>, "Invalid ArchTag, only Sm10x are supported.");
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::conv::collective::detail
include/cutlass/conv/collective/builders/sm100_umma_builder.inl
+9 -4
View File
@@ -44,6 +44,7 @@ using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
conv::Operator ConvOp,
class ElementA,
class GmemLayoutA,
@@ -58,7 +59,7 @@ template <
class KernelScheduleType
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ConvOp,
ElementA,
@@ -73,6 +74,9 @@ struct CollectiveBuilder<
StageCountType,
KernelScheduleType,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
|| cute::is_same_v<ArchTag, arch::Sm103>
) &&
(cute::is_same_v<KernelScheduleType, KernelImplicitTmaWarpSpecialized1SmSm100> ||
cute::is_same_v<KernelScheduleType, KernelImplicitTmaWarpSpecialized2SmSm100> ||
cute::is_same_v<KernelScheduleType, KernelStridedDgradTmaWs1SmSm100> ||
@@ -191,12 +195,12 @@ private:
CLCResponseStorage +
TmemBasePtrsStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
static constexpr int PipelineStages = detail::compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma, ElementBMma, SmemTileShape>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma, ElementBMma, SmemTileShape>(StageCountType{});
constexpr static int NumSpatialDimensions = detail::gmem_layout_tags_to_spatial_dims<GmemLayoutA, GmemLayoutB>();
@@ -206,7 +210,8 @@ private:
NumSpatialDimensions,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>;
ClusterShape_MNK,
ArchTag>;
public:
using CollectiveOp = cutlass::conv::collective::CollectiveConv<
include/cutlass/conv/collective/sm100_implicit_gemm_umma_warpspecialized.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int NumSpatialDims,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShapeMNKL_, // (MmaAtomShapeM, MmaAtomShapeN, TileK, optional: TileL)
class ElementA_,
@@ -79,7 +80,8 @@ struct CollectiveConv<
NumSpatialDims,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShapeMNKL_,
ElementA_,
ElementB_,
@@ -96,7 +98,8 @@ struct CollectiveConv<
NumSpatialDims,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = decltype(cute::take<0,3>(TileShapeMNKL_{})); // (MmaAtomShapeM, MmaAtomShapeN, TileK)
using ElementA = ElementA_;
using ElementB = ElementB_;
include/cutlass/conv/device/conv_universal_adapter.hpp
+2 -1
View File
@@ -351,7 +351,8 @@ public:
}
else {
if constexpr (ConvKernel::ArchTag::kMinComputeCapability == 100 ||
ConvKernel::ArchTag::kMinComputeCapability == 101) {
ConvKernel::ArchTag::kMinComputeCapability == 101
) {
launch_result = ClusterLauncher::launch_with_fallback_cluster(
grid,
cluster,
include/cutlass/conv/dispatch_policy.hpp
+3 -2
View File
@@ -116,14 +116,15 @@ template<
int NumSpatialDimensions_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = cute::Shape<cute::C<1>,cute::C<1>,cute::C<1>>
class ClusterShape_ = cute::Shape<cute::C<1>,cute::C<1>,cute::C<1>>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedImplicitGemm {
static constexpr int Stages = Stages_;
static constexpr int NumSpatialDimensions = NumSpatialDimensions_;
static constexpr Operator ConvOp = ConvOp_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelScheduleImplicitTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
static_assert(NumSpatialDimensions >= 1);
include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp
+1 -2
View File
@@ -188,7 +188,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
@@ -413,7 +412,7 @@ public:
// Kernel level shared memory storage
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// In a warp specialized kernel, collectives expose data movement and compute operations separately
CollectiveMainloop collective_mainloop(params.mainloop, cluster_shape, cta_rank_in_cluster);
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
include/cutlass/detail/layout.hpp
+3
View File
@@ -310,6 +310,9 @@ constexpr bool is_tma_copy_engine() {
|| cute::is_base_of_v<cute::SM90_TMA_STORE_IM2COL, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_2SM_LOAD, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_2SM_LOAD_MULTICAST, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_LOAD_2D_GATHER4, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_LOAD_MULTICAST_2D_GATHER4, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM100_TMA_STORE_2D_SCATTER4, GmemTiledCopy>
) {
return true;
}
include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_tma_warpspecialized.hpp
+3 -1
View File
@@ -138,10 +138,12 @@ private:
constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideD, StrideD>;
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2;
constexpr static uint32_t NumTmaDescriptorsPerSm = IsGroupedGemmKernel ? (NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2) : 1;
using SmemLayoutC = decltype(tile_to_shape(
SmemLayoutAtomC{},
include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp
+5 -5
View File
@@ -181,11 +181,6 @@ private:
// TMA store delay only benefits with loop unrolling
constexpr static bool DelayTmaStore = DelayTmaStore_ and UnrollEpiLoop;
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2;
struct CollectiveStorageWithC {
alignas(SmemAlignmentC) ArrayEngine<SmemElementC, cosize_v<SmemLayoutC>> smem_C;
alignas(SmemAlignmentD) ArrayEngine<SmemElementD, cosize_v<SmemLayoutD>> smem_D;
@@ -241,6 +236,11 @@ public:
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideD, StrideD>;
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = IsGroupedGemmKernel ? (NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2) : 1;
// Host side epilogue arguments
struct Arguments {
typename FusionCallbacks::Arguments thread{};
include/cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl
+6 -5
View File
@@ -185,8 +185,7 @@ struct CollectiveBuilder<
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute,
NumComplexComponents, NumComputeMtxs
@@ -209,7 +208,8 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>,
AccumulatorCopyAtom,
ArchTag>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedFastF32<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
@@ -219,8 +219,9 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>
>;
AccumulatorCopyAtom,
ArchTag>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl
+6 -5
View File
@@ -253,8 +253,7 @@ struct CollectiveBuilder<
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute,
/*Cmplx=*/ 1, /*Mtxs=*/ NumComputeMtxs
@@ -276,7 +275,8 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>,
AccumulatorCopyAtom,
ArchTag>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedFastF32<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
@@ -286,8 +286,9 @@ struct CollectiveBuilder<
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>
>;
AccumulatorCopyAtom,
ArchTag>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
include/cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl
+11 -5
View File
@@ -93,6 +93,7 @@ sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync(StageCountAu
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
@@ -106,7 +107,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassBlockScaledTensorOp,
ElementPairA,
GmemLayoutATag,
@@ -119,7 +120,11 @@ struct CollectiveBuilder<
ClusterShape_MNK, // Static cluster shape (_1, _1, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag> >
cute::enable_if_t<
cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100, BuilderScheduleTag> &&
(cute::is_same_v<ArchTag, arch::Sm100>
)
>
>
{
using ElementSFA = typename detail::blockscaled::blockscaled_type<BuilderScheduleTag, ElementPairA>::sf_type;
@@ -238,12 +243,12 @@ struct CollectiveBuilder<
CLCPipelineStorage +
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<SmemShapeA_M, BlockTileB_N, SmemShapeA_K>;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled_mixed_tma_cpasync<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
static_assert(PipelineStages > 0, "Smem usage is too high. Can't create any SMEM buffers for A, B, SFA, and SFB.");
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
@@ -251,7 +256,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
cute::tuple<ElementA, ElementSF>,
StridePairA,
include/cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl
+3 -3
View File
@@ -278,8 +278,7 @@ struct CollectiveBuilder<
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Calculate SMEM capacity based on ArchTag
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, 0>();
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockscaled_sparse<
ReducedSmemCapacityBytes,
@@ -297,7 +296,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>;
ClusterShape_MNK,
ArchTag>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
+7 -5
View File
@@ -255,8 +255,7 @@ struct CollectiveBuilder<
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
@@ -271,20 +270,23 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockScaled<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl
+5 -4
View File
@@ -375,8 +375,7 @@ struct CollectiveBuilder<
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopABPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
@@ -417,12 +416,14 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockwiseScaling<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>>;
ClusterShape_MNK,
ArchTag>>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
include/cutlass/gemm/collective/builders/sm100_common.inl
+12
View File
@@ -1000,6 +1000,18 @@ struct TrivialBlockscaledMma<
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, Instr, BuilderScheduleTag>());
};
template <class ArchTag, int KernelSmemCarveout>
constexpr int sm100_reduced_smem_capacity_bytes() {
if constexpr (cute::is_same_v<ArchTag, arch::Sm100>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else if constexpr (cute::is_same_v<ArchTag, arch::Sm103>) {
return cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
}
else {
static_assert(cutlass::detail::dependent_false<ArchTag>, "Invalid ArchTag, only Sm10x are supported.");
}
}
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl
+3 -3
View File
@@ -142,8 +142,7 @@ struct CollectiveBuilder<
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage;
@@ -156,7 +155,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
include/cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl
+5 -4
View File
@@ -210,8 +210,7 @@ struct CollectiveBuilder<
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
@@ -230,7 +229,8 @@ struct CollectiveBuilder<
AccumulatorPipelineStageCount,
TransformationStageCount,
ClusterShape_MNK,
AccumulatorCopyAtom
AccumulatorCopyAtom,
ArchTag
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32<
PipelineStages,
@@ -238,7 +238,8 @@ struct CollectiveBuilder<
AccumulatorPipelineStageCount,
TransformationStageCount,
ClusterShape_MNK,
AccumulatorCopyAtom
AccumulatorCopyAtom,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl
+3 -3
View File
@@ -307,8 +307,7 @@ struct CollectiveBuilder<
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
static constexpr int ScaleGranularityK = get_ScaleGranularityK<LayoutScale>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_mixed_input<
@@ -325,7 +324,8 @@ struct CollectiveBuilder<
Transform2MmaPipelineStageCount,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
include/cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl
+9 -6
View File
@@ -41,6 +41,7 @@ namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ArchTag,
class ElementA,
class GmemLayoutATag,
int AlignmentA,
@@ -54,7 +55,7 @@ template <
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
ArchTag,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATag,
@@ -67,7 +68,9 @@ struct CollectiveBuilder<
ClusterShape_MNK, // Static cluster shape (_1, _1, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmSm100, BuilderScheduleTag> >
cute::enable_if_t<cute::is_same_v<KernelMixedTmaCpAsyncWarpSpecialized1SmSm100, BuilderScheduleTag> &&
(cute::is_same_v<ArchTag, arch::Sm100>
)>
>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
@@ -135,20 +138,20 @@ struct CollectiveBuilder<
CLCPipelineStorage +
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<SmemShapeA_M, BlockTileB_N, SmemShapeA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
cutlass::gemm::MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
include/cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl
+5 -5
View File
@@ -132,9 +132,7 @@ struct CollectiveBuilder<
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Use complex type to calculate SMEM stage count
@@ -150,13 +148,15 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedPlanarComplex<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl
+3 -3
View File
@@ -381,8 +381,7 @@ struct CollectiveBuilder<
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Calculate SMEM capacity based on ArchTag
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, 0>();
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_sparse<
ReducedSmemCapacityBytes,
@@ -399,7 +398,8 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
ClusterShape_MNK,
ArchTag>,
TileShape_MNK,
ElementA,
LayoutPairAE,
include/cutlass/gemm/collective/builders/sm100_umma_builder.inl
+7 -6
View File
@@ -288,9 +288,7 @@ struct CollectiveBuilder<
IsArrayOfPointersGemm
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
@@ -307,20 +305,23 @@ struct CollectiveBuilder<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
ClusterShape_MNK,
ArchTag
>
>;
include/cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl
+5 -5
View File
@@ -498,9 +498,7 @@ struct CollectiveBuilder<
TensorMapStorage +
TmaPrefetchStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ReducedSmemCapacityBytes = detail::sm100_reduced_smem_capacity_bytes<ArchTag, KernelSmemCarveout>();
using SmemTileShape = cute::Shape<Int<MMA_M>, Int<MMA_N/cute::size(AtomThrID{})>, _128>; // SmemAllocTypes are uint8_t. We always allocate 128bytes
static constexpr auto PipelineStages = cutlass::gemm::collective::detail::sm103_compute_stage_count_or_override_blockscaled<
ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
@@ -512,7 +510,8 @@ struct CollectiveBuilder<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK,
PrefetchType
PrefetchType,
ArchTag
>,
cutlass::gemm::MainloopSm103TmaUmmaWarpSpecializedBlockScaled<
get<0>(PipelineStages),
@@ -520,7 +519,8 @@ struct CollectiveBuilder<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK,
PrefetchType
PrefetchType,
ArchTag
>
>;
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp
+10 -7
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
@@ -143,11 +146,6 @@ struct CollectiveMma<
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementPairA = ElementPairA_;
using ElementPairB = ElementPairB_;
using ElementAMma = typename TiledMma::ValTypeA;
@@ -264,6 +262,11 @@ struct CollectiveMma<
static constexpr bool IsF8F6F4 = detail::is_sm100_mma_f8f6f4<TiledMma, ElementA, ElementB>();
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = IsGroupedGemmKernel ? (SchedulerPipelineStageCount + Stages + 2) : 1;
using TmaInternalElementA = cute::conditional_t<IsF8F6F4, ElementAMma, ElementA>;
using TmaInternalElementB = cute::conditional_t<IsF8F6F4, ElementBMma, ElementB>;
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp
+5 -2
View File
@@ -64,6 +64,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -84,7 +85,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -110,7 +112,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp
+5 -2
View File
@@ -63,6 +63,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -83,7 +84,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -111,7 +113,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
// TileShape refers to MmaTileShape to adapt for runtime cluster
using TileShape = TileShape_;
using TiledMma_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
Layout<Shape<_1,_1,_1>>,
include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp
+5 -2
View File
@@ -63,6 +63,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementPairA_,
@@ -83,7 +84,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementPairA_,
LayoutPairA_,
@@ -109,7 +111,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using TiledMMA_SF = TiledMMA<MMA_Atom<typename TiledMma::MMA_ScaleFactor>,
Layout<Shape<_1,_1,_1>>,
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized.hpp
+10 -7
View File
@@ -63,6 +63,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -83,7 +84,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -109,7 +111,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
@@ -123,11 +126,6 @@ struct CollectiveMma<
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
@@ -227,6 +225,11 @@ struct CollectiveMma<
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = IsGroupedGemmKernel ? (SchedulerPipelineStageCount + Stages + 2) : 1;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp
+5 -2
View File
@@ -60,6 +60,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -80,7 +81,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StridePairA_,
@@ -106,7 +108,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp
+5 -2
View File
@@ -69,6 +69,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -93,7 +94,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
float,
StrideA_,
@@ -124,7 +126,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp
+5 -2
View File
@@ -68,6 +68,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -92,7 +93,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -160,7 +162,8 @@ public:
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
int TransformationPipelineStageCount_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -86,7 +87,8 @@ struct CollectiveMma<
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -144,7 +146,8 @@ public:
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp
+12 -7
View File
@@ -64,6 +64,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape,
class TileShape_, // Static cluster shape or dynamic (int, int, _1)
class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -84,7 +85,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -116,7 +118,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
@@ -128,15 +131,17 @@ struct CollectiveMma<
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = IsGroupedGemmKernel ? (SchedulerPipelineStageCount + Stages + 2) : 1;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -85,7 +86,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -111,7 +113,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp
+5 -2
View File
@@ -60,6 +60,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -80,7 +81,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -107,7 +109,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
// TileShape refers to MmaTileShape to adapt for runtime cluster shape
using TileShape = TileShape_;
include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -111,7 +113,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
// TileShape refers to MmaTileShape to adapt for runtime cluster
using TileShape = TileShape_;
include/cutlass/gemm/collective/sm100_mma_warpspecialized.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp
+5 -2
View File
@@ -62,6 +62,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -82,7 +83,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StridePairA_,
@@ -108,7 +110,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using ElementA = ElementA_;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp
+5 -2
View File
@@ -82,6 +82,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -106,7 +107,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
float,
StrideA_,
@@ -137,7 +139,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp
+5 -2
View File
@@ -66,6 +66,7 @@ template <
int NumBandsToCompute_,
int ScalingFactor_,
int AccPromotionInterval_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape,
class TileShape_,
@@ -90,7 +91,8 @@ struct CollectiveMma<
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -156,7 +158,8 @@ public:
ScalingFactor_,
AccPromotionInterval_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp
+5 -2
View File
@@ -78,6 +78,7 @@ template <
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
int TransformationPipelineStageCount_,
class ArchTag_,
class AccumulatorCopyAtom_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -99,7 +100,8 @@ struct CollectiveMma<
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>,
AccumulatorCopyAtom_,
ArchTag_>,
TileShape_,
complex<float>,
StrideA_,
@@ -155,7 +157,8 @@ public:
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>;
AccumulatorCopyAtom_,
ArchTag_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp
+5 -2
View File
@@ -65,6 +65,7 @@ template <
int Transform2MmaPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ArchTag_,
class ClusterShape,
class TileShape_,
class ElementAOptionalTuple_,
@@ -86,7 +87,8 @@ struct CollectiveMma<
Transform2MmaPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementAOptionalTuple_,
StridePairA_,
@@ -115,7 +117,8 @@ public:
Transform2MmaPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
using TiledMma = TiledMma_;
using KernelSchedule = typename DispatchPolicy::Schedule;
include/cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp
+5 -2
View File
@@ -70,6 +70,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape,
class TileShape_, // Static cluster shape or dynamic (int, int, _1)
class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -90,7 +91,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
StrideA_,
@@ -122,7 +124,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
include/cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp
+5 -2
View File
@@ -61,6 +61,7 @@ template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
@@ -81,7 +82,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
ClusterShape,
ArchTag_>,
TileShape_,
ElementA_,
LayoutPairAE_,
@@ -107,7 +109,8 @@ struct CollectiveMma<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
ClusterShape,
ArchTag_>;
using TileShape = TileShape_;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
include/cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp
+5 -2
View File
@@ -66,6 +66,7 @@ template <
int LoadSFPipelineStageCount,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, int)
cutlass::sm103::detail::KernelPrefetchType PrefetchType,
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -89,7 +90,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>,
PrefetchType,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -117,7 +119,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>;
PrefetchType,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
include/cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp
+5 -2
View File
@@ -66,6 +66,7 @@ template <
int LoadSFPipelineStageCount,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ArchTag_,
class ClusterShape, // Static cluster shape or dynamic (int, int, int)
cutlass::sm103::detail::KernelPrefetchType PrefetchType,
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
@@ -89,7 +90,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>,
PrefetchType,
ArchTag_>,
TileShape_,
ElementPairA_,
StridePairA_,
@@ -117,7 +119,8 @@ struct CollectiveMma<
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape,
PrefetchType>;
PrefetchType,
ArchTag_>;
using TileShape = TileShape_;
// Due to an MSVC bug, we can't use decltype(make_tiled_mma()) interface.
include/cutlass/gemm/dispatch_policy.hpp
+66 -44
View File
@@ -976,12 +976,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100UmmaCpAsyncWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -989,12 +990,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100UmmaMixedTmaCpAsyncWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelMixedTmaCpAsyncWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1003,12 +1005,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100UmmaMixedTmaCpAsyncWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelMixedTmaCpAsyncWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1018,12 +1021,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1033,12 +1037,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedBlockwiseScaling {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedMmaTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1048,12 +1053,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedBlockwiseScaling {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelPtrArrayTmaWarpSpecializedMmaTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1063,12 +1069,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1077,13 +1084,14 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedSparse {
constexpr static int Stages = Stages_;
constexpr static int MetadataS2TStages = 4;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelSparseTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1092,13 +1100,14 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedBlockScaledSparse {
constexpr static int Stages = Stages_;
constexpr static int MetadataS2TStages = 4;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelSparseTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1132,7 +1141,8 @@ template<
class ClusterShape_ = Shape<_1,_1,_1>,
// The TMEM_LOAD atom to be used for loading local accumulator
// from TMEM to registers
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
@@ -1143,7 +1153,7 @@ struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::FastF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1163,7 +1173,8 @@ template<
// Transformation <-> MMA
int TransformationPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_;
@@ -1171,7 +1182,7 @@ struct MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1191,7 +1202,8 @@ template<
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
// ClusterShape for the kernel
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
@@ -1199,7 +1211,7 @@ struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
constexpr static int Transform2MmaPipelineStageCount = Transform2MmaPipelineStageCount_;
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::MixedInput;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedMixedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1214,12 +1226,13 @@ template<
int SchedulerPipelineStageCount_,
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100TmaUmmaWarpSpecializedPlanarComplex {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
@@ -1229,12 +1242,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1244,12 +1258,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100RCGroupGemmTmaUmmaWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1259,12 +1274,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100RCGroupGemmTmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1274,12 +1290,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1291,12 +1308,13 @@ template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
class ClusterShape_ = Shape<_1,_1,_1>,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
@@ -1331,7 +1349,8 @@ template<
class ClusterShape_ = Shape<_1,_1,_1>,
// The TMEM_LOAD atom to be used for loading local accumulator
// from TMEM to registers
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_32dp32b32x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
@@ -1342,7 +1361,7 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::FastF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelPtrArrayTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1362,7 +1381,8 @@ template<
// Transformation <-> MMA
int TransformationPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x,
class ArchTag_ = arch::Sm100
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_;
@@ -1370,7 +1390,7 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using ArchTag = ArchTag_;
using Schedule = KernelPtrArrayTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1385,13 +1405,14 @@ template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch,
class ArchTag_ = arch::Sm103
>
struct MainloopSm103TmaUmmaWarpSpecializedBlockScaled {
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm103;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
@@ -1407,13 +1428,14 @@ template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch,
class ArchTag_ = arch::Sm103
>
struct MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled {
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm103;
using ArchTag = ArchTag_;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+1 -1
View File
@@ -144,7 +144,7 @@ public:
using TileSchedulerParams = typename TileScheduler::Params;
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
static constexpr bool IsTensorMapUpdateAsync = not IsSchedDynamicPersistent;
static constexpr bool IsTensorMapUpdateAsync = IsGroupedGemmKernel && not IsSchedDynamicPersistent;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
static constexpr uint32_t MinTensorMapWorkspaceAlignment = 64;
include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp
+11 -1
View File
@@ -44,8 +44,18 @@ public StaticPersistentTileScheduler<
public:
using BaseScheduler = StaticPersistentTileScheduler<StaticPersistentTileScheduler100>;
public:
using BaseScheduler::StaticPersistentTileScheduler;
using Params = PersistentTileSchedulerSm90Params;
// Explicit forwarding constructors replacing inheriting-constructor syntax
// (`using BaseScheduler::StaticPersistentTileScheduler;`) which newer CUDA
// host compilers reject in dependent-base contexts: the injected-class-name
// resolves to a type rather than a constructor name.
CUTLASS_HOST_DEVICE
StaticPersistentTileScheduler100() = default;
CUTLASS_DEVICE explicit
StaticPersistentTileScheduler100(Params const& params_)
: BaseScheduler(params_) {}
using RasterOrder = typename Params::RasterOrder;
using RasterOrderOptions = typename Params::RasterOrderOptions;
struct CLCResponse { uint32_t data[4] = {0}; };
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+4 -2
View File
@@ -242,8 +242,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
// Calculate workspace pointers
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+4 -2
View File
@@ -245,8 +245,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
// Calculate workspace pointers
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp
+4 -2
View File
@@ -193,8 +193,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp
+4 -2
View File
@@ -206,8 +206,10 @@ public:
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " << max_active_clusters);
}
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters};
KernelHardwareInfo hw_info = args.hw_info;
hw_info.sm_count = sm_count;
hw_info.max_active_clusters = max_active_clusters;
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
include/cutlass/kernel_hardware_info.h
+19 -9
View File
@@ -81,21 +81,26 @@ struct KernelHardwareInfo {
}
// Query maximum number of active clusters that could co-exist on the target device
// based on kernel properties such as cluster dims and threadblock dims
// based on kernel properties such as cluster dims and threadblock dims.
// When a green context stream is provided, the occupancy query is scoped to the
// green context's SM partition, returning the max active clusters for that partition.
static inline int
query_device_max_active_clusters(
dim3 cluster_dims,
uint32_t threads_per_block,
void const* kernel_ptr) {
void const* kernel_ptr,
cudaStream_t stream = nullptr) {
int max_active_clusters = 0;
#if defined(CUTLASS_SM90_CLUSTER_LAUNCH_ENABLED)
ClusterLauncher::LaunchConfig cluster_launch_config = ClusterLauncher::make_cluster_launch_config(
cluster_dims /* minimum grid dim */, cluster_dims, {threads_per_block, 1, 1});
cluster_dims /* minimum grid dim */, cluster_dims, {threads_per_block, 1, 1},
0 /* smem_size */, stream /* green ctx stream or nullptr */);
// Given the kernel function and launch configuration, return the maximum number of clusters that could co-exist on the target device.
// When stream is a green context stream, this returns the max active clusters for that partition.
cudaError_t result = cudaOccupancyMaxActiveClusters(&max_active_clusters, kernel_ptr, &cluster_launch_config.launch_config);
if (result != cudaSuccess) {
CUTLASS_TRACE_HOST(
" cudaGetDevice() returned error "
" cudaOccupancyMaxActiveClusters() returned error "
<< cudaGetErrorString(result));
return 0;
}
@@ -108,26 +113,31 @@ struct KernelHardwareInfo {
#endif
}
// Simpler version of the above query function that fetches relevant information from the Kernel
// Simpler version of the above query function that fetches relevant information from the Kernel.
// When a green context stream is provided, the occupancy query is scoped to that partition.
template <typename Kernel>
static inline int
query_device_max_active_clusters() {
query_device_max_active_clusters(cudaStream_t stream = nullptr) {
dim3 cluster_dims(cute::size<0>(typename Kernel::ClusterShape{}),
cute::size<1>(typename Kernel::ClusterShape{}),
cute::size<2>(typename Kernel::ClusterShape{}));
uint32_t threads_per_block = Kernel::MaxThreadsPerBlock;
void const* kernel_ptr = (void*)(device_kernel<Kernel>);
return query_device_max_active_clusters(cluster_dims, threads_per_block, kernel_ptr);
return query_device_max_active_clusters(cluster_dims, threads_per_block, kernel_ptr, stream);
}
// Create a KernelHardwareInfo by querying device properties.
// When a green context stream is provided, max_active_clusters is queried
// against that stream's green context partition instead of the full device.
template <typename Kernel>
static inline KernelHardwareInfo
make_kernel_hardware_info(int const device_id = 0, int sm_count = 0, int max_active_clusters = 0) {
make_kernel_hardware_info(int const device_id = 0, int sm_count = 0, int max_active_clusters = 0,
cudaStream_t stream = nullptr) {
if (sm_count == 0) {
sm_count = query_device_multiprocessor_count(device_id);
}
if (max_active_clusters == 0) {
max_active_clusters = query_device_max_active_clusters<Kernel>();
max_active_clusters = query_device_max_active_clusters<Kernel>(stream);
}
return {device_id, sm_count, max_active_clusters};
}
include/cutlass/version.h
+2 -2
View File
@@ -35,8 +35,8 @@
#include <string>
#define CUTLASS_MAJOR 4
#define CUTLASS_MINOR 4
#define CUTLASS_PATCH 2
#define CUTLASS_MINOR 5
#define CUTLASS_PATCH 0
#ifdef CUTLASS_VERSIONS_GENERATED
#include "cutlass/version_extended.h"