leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 29 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