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

CUTLASS 3.8 Release (#2059)

Browse Source 
* CUTLASS 3.8 Release

* update

* Update README.md

* Revert "Update README.md"

This reverts commit b353e36fe83e0815f99b44e46c0c95494c44726b.

* update

* update

---------

Co-authored-by: Haicheng Wu <57973641+hwu36@users.noreply.github.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
mihir-awatramani
2025-01-25 02:44:06 -05:00
committed by GitHub
co-authored by Haicheng Wu Haicheng Wu
parent 9eb01fa0b0
commit 389e493055
290 changed files with 91222 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/atom/copy_atom.hpp
+19
View File
@@ -751,14 +751,33 @@ print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and
#include <cute/atom/copy_traits_sm75.hpp>
#include <cute/atom/copy_traits_sm80.hpp>
#include <cute/atom/copy_traits_sm90.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
// Config
#if (__CUDACC_VER_MAJOR__ >= 12)
# define CUTE_COPY_ATOM_TMA_SM90_ENABLED
# define CUTE_COPY_ATOM_TMA_SM100_ENABLED
#endif
#if (!defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED))
# define CUTE_COPY_ATOM_TMA_SM90_ENABLED
#endif
#if (!defined(CUTE_COPY_ATOM_TMA_SM100_ENABLED))
# define CUTE_COPY_ATOM_TMA_SM100_ENABLED
#endif
#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)
#include <cute/atom/copy_traits_sm90_tma.hpp>
#endif
#if defined(CUTE_COPY_ATOM_TMA_SM100_ENABLED)
#include <cute/atom/copy_traits_sm100_tma.hpp>
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cute/atom/copy_traits_sm100.hpp
+3797
View File
File diff suppressed because it is too large Load Diff
include/cute/atom/copy_traits_sm100_im2col.hpp
+488
View File
@@ -0,0 +1,488 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
/*! \file
\brief im2col make_tma_copy
*/
#include "cute/arch/copy_sm90.hpp"
#include "cute/arch/copy_sm90_desc.hpp"
#include "cute/atom/copy_traits_sm90_im2col.hpp"
#include "cute/tensor.hpp"
namespace cute {
struct SM100_TMA_2SM_LOAD_IM2COL_OP : SM100_TMA_2SM_LOAD_IM2COL {};
/// @brief Non-executable specialization of Copy_Traits for SM100
/// im2col TMA load, with TMA descriptor but no barrier.
///
/// Use `.with(memory_barrier)` to construct an executable version.
template <class NumBitsPerTMA, class TMATensor>
struct Copy_Traits<SM100_TMA_2SM_LOAD_IM2COL, NumBitsPerTMA, TMATensor>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
Im2ColTmaDescriptor tma_desc_;
TMATensor tma_tensor_;
CUTE_HOST_DEVICE constexpr
Im2ColTmaDescriptor const*
get_tma_descriptor() const
{
return &tma_desc_;
}
template <class GShape>
CUTE_HOST_DEVICE constexpr
TMATensor const
get_tma_tensor(GShape const&) const
{
return tma_tensor_;
}
/// @brief Get an executable specialization.
///
/// Copy_Traits specializations with SM100_TMA_2SM_LOAD_IM2COL are not
/// directly executable. Instead, call this "with" member function
/// to get an executable specialization. "Executable" means that
/// @c copy_unpack works.
///
/// @param tma_mbar Memory barrier for synchronization
///
/// @param multicast_mask Multicast mask (unused; only exists
/// for consistency with the actual multicast Copy_Traits
/// specialization)
///
/// @return Executable specialization of @c Copy_Traits
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_2SM_LOAD_IM2COL_OP, NumBitsPerTMA>
with(uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const
{
return {{}, {&tma_desc_, &tma_mbar}};
}
// Copy_Traits specializations with SM100_TMA_2SM_LOAD_IM2COL
// are not directly executable. Instead, call .with
// to get an executable specialization.
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;
};
/// TMA load, with TMA descriptor and barrier.
template <class NumBitsPerTMA>
struct Copy_Traits<SM100_TMA_2SM_LOAD_IM2COL_OP, NumBitsPerTMA>
: TMA_LOAD_IM2COL_Unpack<SM100_TMA_2SM_LOAD_IM2COL_OP>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_2SM_LOAD_IM2COL arguments
tuple<
Im2ColTmaDescriptor const*,
uint64_t* // smem mbarrier
> const opargs_;
};
//////////////////////////////////////////////////////////////////////////////
///////////////////////////// TMA_LOAD_MULTICAST /////////////////////////////
//////////////////////////////////////////////////////////////////////////////
struct SM100_TMA_2SM_LOAD_IM2COL_MULTICAST_OP : SM100_TMA_2SM_LOAD_IM2COL_MULTICAST {};
/// @brief Non-executable specialization of Copy_Traits for SM100
/// im2col TMA load, with TMA descriptor but no barrier or multicast
/// mask.
///
/// Use `.with(memory_barrier)` to construct an executable version.
template <class NumBitsPerTMA, class TMATensor>
struct Copy_Traits<SM100_TMA_2SM_LOAD_IM2COL_MULTICAST, NumBitsPerTMA, TMATensor>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
Im2ColTmaDescriptor tma_desc_;
TMATensor tma_tensor_;
CUTE_HOST_DEVICE constexpr
Im2ColTmaDescriptor const*
get_tma_descriptor() const
{
return &tma_desc_;
}
template <class GShape>
CUTE_HOST_DEVICE constexpr
TMATensor const
get_tma_tensor(GShape const&) const
{
return tma_tensor_;
}
/// @brief Get an executable specialization.
///
/// Copy_Traits specializations with SM100_TMA_2SM_LOAD_IM2COL_MULTICAST
/// are not directly executable. Instead, call this "with" member
/// function to get an executable specialization. "Executable"
/// means that @c copy_unpack works.
///
/// @param tma_mbar Memory barrier for synchronization
///
/// @param multicast_mask Multicast mask (defaults to a single CTA)
///
/// @return Executable specialization of @c Copy_Traits
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_2SM_LOAD_IM2COL_MULTICAST_OP, NumBitsPerTMA>
with(uint64_t& tma_mbar, uint16_t const& multicast_mask) const
{
return {{}, {&tma_desc_, &tma_mbar, multicast_mask}};
}
// Copy_Traits specializations with SM100_TMA_LOAD_IM2COL_MULTICAST
// are not directly executable. Instead, call .with to get an
// executable specialization.
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;
};
/// @brief Executable specialization of Copy_Traits for SM100 multicast
/// im2col TMA load, with TMA descriptor, barrier, and multicast mask.
template <class NumBitsPerTMA>
struct Copy_Traits<SM100_TMA_2SM_LOAD_IM2COL_MULTICAST_OP, NumBitsPerTMA>
: TMA_LOAD_IM2COL_Unpack<SM100_TMA_2SM_LOAD_IM2COL_MULTICAST_OP>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit.
using SrcLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2, NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_2SM_LOAD_IM2COL_MULTICAST arguments
tuple<
Im2ColTmaDescriptor const*,
uint64_t*, // smem mbarrier
uint16_t // multicast mask
> const opargs_;
};
////////////////////////////////////
// Make TMA
///////////////////////////////////
#if !defined(__CUDACC_RTC__)
/** Make a CuTe CTA-collective TiledCopy for a TMA operation.
*
* @param CopyOp The target copy operation: SM100_TMA_2SM_LOAD
* @param gtensor The GMEM Tensor to be involved in the TMA.
* @param slayout The SMEM Layout to be involved in the TMA.
* @param cluster_tile The Cluster-local tile that each Cluster will be tiling GMEM with.
* This is often the cluster_tile_shape that is used to tile the GMEM:
* local_tile(gtensor, cluster_tile_shape, cluster_coord)
* -> Cluster-local tile of GMEM
* @param mma The TiledMMA that defines the Cluster-Tile to Block-Tile partitioning.
*
* This code attempts to maximize the TMA box size. It does this by tracing
* the SMEM "vector" -- the inverse of the smem layout -- to find the largest
* contiguous array of smem that can be written to/from global memory given
* the constraints that the TMA instruction imposes.
*
* This is accomplished by assigning "basis" strides to the GMEM to track which
* modes of SMEM map to which modes of GMEM, then reordering the modes of GMEM according
* to the SMEM vector, and then using those GMEM/SMEM modes to fill in the desc.
*
* Examples:
*/
template <class CopyOp,
class GEngine, class GLayout,
class SLayout,
class Cluster_Tile,
class... Args,
class LowerCornerStride,
class UpperCornerStride,
class LowerPaddingStride,
class UpperPaddingStride,
class TraversalStride,
class LowerSRTStride,
class DilationStride>
CUTE_HOST
auto
make_im2col_tma_copy_A_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (M,K,...)
SLayout const& slayout, // (MMA, MMA_M, MMA_K)
Cluster_Tile const& cluster_tile, // (TILE_M,TILE_N,TILE_K)
TiledMMA<Args...> const& mma,
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd,
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt,
TMA::DescriptorAuxParams const& aux_params = {})
{
constexpr int R = GLayout::rank;
// Keep only MK modes from MNK
auto cluster_tile_shape = append<R>(make_shape(get<0>(cluster_tile), get<2>(cluster_tile)), Int<1>{});
auto cluster_layout = make_identity_layout(cluster_tile_shape);
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_A(cluster_layout))(_, repeat<R>(_));
auto cta_t_vmnk_strides = [](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL_MULTICAST>) {
return Stride<_0,_0,_1,_0>{}; // VMNK: Use only the N-CTAs in the Multicast
} else
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL>) {
return Stride<_0,_0,_0,_0>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
auto cta_t_shape = shape(mma.get_thr_layout_vmnk());
// cta rank -> logical cta idx
auto cta_t_map = make_layout(cta_t_shape, compact_col_major(cta_t_shape, cta_t_vmnk_strides));
return detail::make_tma_copy_im2col(copy_op, gtensor, slayout,
cta_t_map, cta_v_tile,
lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd, stride_whd,
lower_srt, stride_srt, aux_params);
}
template <class CopyOp,
class GEngine, class GLayout,
class SLayout,
class Cluster_Tile,
class... Args,
class LowerCornerStride,
class UpperCornerStride,
class LowerPaddingStride,
class UpperPaddingStride,
class TraversalStride,
class LowerSRTStride,
class DilationStride>
CUTE_HOST
auto
make_im2col_tma_copy_B_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (N,K,...)
SLayout const& slayout, // (MMA, MMA_N, MMA_K)
Cluster_Tile const& cluster_tile, // (TILE_M,TILE_N,TILE_K)
TiledMMA<Args...> const& mma,
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd,
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt,
TMA::DescriptorAuxParams const& aux_params = {})
{
constexpr int R = GLayout::rank;
// Keep only NK modes from MNK
auto cluster_tile_shape = append<R>(make_shape(get<1>(cluster_tile), get<2>(cluster_tile)), Int<1>{});
auto cluster_layout = make_identity_layout(cluster_tile_shape);
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_B(cluster_layout))(_, repeat<R>(_));
auto cta_t_vmnk_strides = [](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL_MULTICAST>) {
return Stride<_0,_1,_0,_0>{}; // VMNK: Use only the M-CTAs in the Multicast
} else
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL>) {
return Stride<_0,_0,_0,_0>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
auto cta_t_shape = shape(mma.get_thr_layout_vmnk());
// cta rank -> logical cta idx
auto cta_t_map = make_layout(cta_t_shape, compact_col_major(cta_t_shape, cta_t_vmnk_strides));
return detail::make_tma_copy_im2col(copy_op, gtensor, slayout,
cta_t_map, cta_v_tile,
lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd, stride_whd,
lower_srt, stride_srt, aux_params);
}
/////////////////////////////////////
// Experimental Make Im2col TMA Atom
/////////////////////////////////////
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class MMA_Tiler,
class... Args,
class ClusterShapeVMNK,
class LowerCornerStride,
class UpperCornerStride,
class LowerPaddingStride,
class UpperPaddingStride,
class TraversalStride,
class LowerSRTStride,
class DilationStride>
CUTE_HOST
auto
make_im2col_tma_atom_A_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (M, K, ...)
SLayout const& slayout, // (MMA, MMA_M, MMA_K, ...)
MMA_Tiler const& mma_tiler, // (TILE_M, TILE_N, TILE_K, ...)
TiledMMA<Args...> const& mma,
ClusterShapeVMNK const& cluster_shape, // (CTA_V, CTA_M, CTA_N, CTA_K)
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd,
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt,
TMA::DescriptorAuxParams const& aux_params = {})
{
constexpr int R = GLayout::rank;
// Keep only MK modes from MNK
auto cluster_tile_shape = append<R>(make_shape(get<0>(mma_tiler), get<2>(mma_tiler)), Int<1>{});
auto cluster_layout = make_identity_layout(cluster_tile_shape);
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_A(cluster_layout))(_, repeat<R>(_));
// The size of the multicasting
auto num_multicast = [&](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL_MULTICAST>) {
return size<2>(cluster_shape); // VMNK: Use only the N-CTAs in the Multicast
} else
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL> ||
is_same_v<CopyOp, SM90_TMA_STORE_IM2COL> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL>) {
return Int<1>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
return detail::make_tma_atom_im2col(copy_op, gtensor, slayout, num_multicast, cta_v_tile,
lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd,
stride_whd, lower_srt, stride_srt, aux_params);
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class MMA_Tiler,
class... Args,
class ClusterShapeVMNK,
class LowerCornerStride,
class UpperCornerStride,
class LowerPaddingStride,
class UpperPaddingStride,
class TraversalStride,
class LowerSRTStride,
class DilationStride>
CUTE_HOST
auto
make_im2col_tma_atom_B_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (N, K, ...)
SLayout const& slayout, // (MMA, MMA_N, MMA_K, ...)
MMA_Tiler const& mma_tiler, // (TILE_M, TILE_N, TILE_K, ...)
TiledMMA<Args...> const& mma,
ClusterShapeVMNK const& cluster_shape, // (CTA_V, CTA_M, CTA_N, CTA_K)
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd,
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt,
TMA::DescriptorAuxParams const& aux_params = {})
{
constexpr int R = GLayout::rank;
// Keep only NK modes from MNK
auto cluster_tile_shape = append<R>(make_shape(get<1>(mma_tiler), get<2>(mma_tiler)), Int<1>{});
auto cluster_layout = make_identity_layout(cluster_tile_shape);
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_B(cluster_layout))(_, repeat<R>(_));
// The size of the multicasting
auto num_multicast = [&](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL_MULTICAST>) {
return size<1>(cluster_shape); // VMNK: Use only the M-CTAs in the Multicast
} else
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL> ||
is_same_v<CopyOp, SM90_TMA_STORE_IM2COL> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_IM2COL>) {
return Int<1>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
return detail::make_tma_atom_im2col(copy_op, gtensor, slayout, num_multicast, cta_v_tile,
lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd,
stride_whd, lower_srt, stride_srt, aux_params);
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace cute
include/cute/atom/copy_traits_sm100_tma.hpp
+487
View File
@@ -0,0 +1,487 @@
/***************************************************************************************************
* Copyright (c) 2021 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#if !defined(__CUDACC_RTC__)
#include <cuda.h>
#endif
#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cute/arch/copy_sm100_tma.hpp>
#include <cute/atom/copy_traits.hpp>
namespace cute
{
//////////////////////////////////////////////////////////////////////////////
////////////////////////////// TMA_LOAD ////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
struct SM100_TMA_2SM_LOAD_OP : SM100_TMA_2SM_LOAD {};
// The non-executable SM100_TMA_2SM_LOAD with tma_desc and no tma_mbar
// Use .with(tma_mbar) to construct an executable version
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM100_TMA_2SM_LOAD, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_2SM_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 SM100_TMA_2SM_LOAD with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_2SM_LOAD_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 SM100_TMA_2SM_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_2SM_LOAD_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)}};
}
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_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM100_TMA_2SM_LOAD 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;
};
// The executable SM100_TMA_2SM_LOAD with tma_desc and tma_mbar
template <class NumBitsPerTMA>
struct Copy_Traits<SM100_TMA_2SM_LOAD_OP, NumBitsPerTMA>
: TMA_LOAD_Unpack<SM100_TMA_2SM_LOAD_OP, NumBitsPerTMA>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_2SM_LOAD arguments
tuple<
TmaDescriptor const*,
uint64_t*, // smem mbarrier
uint64_t // cache hint
> const opargs_;
};
//////////////////////////////////////////////////////////////////////////////
///////////////////////////// TMA_LOAD_MULTICAST /////////////////////////////
//////////////////////////////////////////////////////////////////////////////
struct SM100_TMA_2SM_LOAD_MULTICAST_OP : SM100_TMA_2SM_LOAD_MULTICAST {};
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_2SM_LOAD_MULTICAST_OP 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 SM100_TMA_2SM_LOAD_MULTICAST_OP with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
with(
uint64_t& tma_load_mbar,
uint16_t const& multicast_mask,
TMA::CacheHintSm100 const& cache_hint = TMA::CacheHintSm100::EVICT_NORMAL) const {
return {{}, {&tma_desc_, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)}};
}
// Construct an executable SM100_TMA_2SM_LOAD_MULTICAST_OP with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
with(
TmaDescriptor const* new_tma_desc,
uint64_t& tma_load_mbar,
uint16_t const& multicast_mask,
TMA::CacheHintSm100 const& cache_hint = TMA::CacheHintSm100::EVICT_NORMAL) const {
return {{}, {new_tma_desc, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)}};
}
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_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
}
// Don't try to execute a copy with SM100_TMA_2SM_LOAD_MULTICAST_OP 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_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
: TMA_LOAD_Unpack<SM100_TMA_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
{
using ThrID = Layout<_2>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_2,NumBitsPerTMA>, Stride<NumBitsPerTMA,_1>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM100_TMA_2SM_LOAD_MULTICAST_OP arguments
tuple<
TmaDescriptor const*,
uint64_t*, // smem mbarrier
uint16_t, // multicast mask
uint64_t // cache hint
> const opargs_;
};
////////////////////////////////////
// Make TMA
///////////////////////////////////
#if !defined(__CUDACC_RTC__)
/** Make a CuTe CTA-collective TiledCopy for a TMA operation.
*
* @param CopyOp The target copy operation: SM100_TMA_2SM_LOAD
* @param gtensor The GMEM Tensor to be involved in the TMA.
* @param slayout The SMEM Layout to be involved in the TMA.
* @param cluster_tile The Cluster-local tile that each Cluster will be tiling GMEM with.
* This is often the cluster_tile_shape that is used to tile the GMEM:
* local_tile(gtensor, cluster_tile_shape, cluster_coord)
* -> Cluster-local tile of GMEM
* @param mma The TiledMMA that defines the Cluster-Tile to Block-Tile partitioning.
*
* This code attempts to maximize the TMA box size. It does this by tracing
* the SMEM "vector" -- the inverse of the smem layout -- to find the largest
* contiguous array of smem that can be written to/from global memory given
* the constraints that the TMA instruction imposes.
*
* This is accomplished by assigning "basis" strides to the GMEM to track which
* modes of SMEM map to which modes of GMEM, then reordering the modes of GMEM according
* to the SMEM vector, and then using those GMEM/SMEM modes to fill in the desc.
*
* Examples:
*/
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class Cluster_Tiler,
class... Args>
CUTE_HOST
auto
make_tma_copy_A_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (M, K, ...)
SLayout const& slayout, // (MMA, MMA_M, MMA_K, ...)
Cluster_Tiler const& cluster_tiler, // (TILER_M, TILER_N, TILER_K, ...)
TiledMMA<Args...> const& mma)
{
// Keep only MK modes from MNK
auto cluster_tiler_mk = remove<1>(cluster_tiler);
// cluster tile coord -> gtensor coord
auto g_tile = make_identity_layout(shape(gtensor)).compose(cluster_tiler_mk); // (TILE_M, TILE_K, ...)
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_A(g_tile))(_, repeat<rank(g_tile)>(_)); // (MMA, MMA_M, MMA_K, ...)
auto cta_t_vmnk_strides = [](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_MULTICAST>) {
return Stride<_0,_0,_1,_0>{}; // 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>) {
return Stride<_0,_0,_0,_0>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
auto cta_t_shape = shape(mma.get_thr_layout_vmnk());
// cta rank -> logical cta idx
auto cta_t_map = coalesce(make_layout(cta_t_shape, compact_col_major(cta_t_shape, cta_t_vmnk_strides)));
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
return detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_map, cta_v_tile);
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class Cluster_Tiler,
class... Args>
CUTE_HOST
auto
make_tma_copy_B_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (N, K, ...)
SLayout const& slayout, // (MMA, MMA_N, MMA_K, ...)
Cluster_Tiler const& cluster_tiler, // (TILE_M, TILE_N, TILE_K, ...)
TiledMMA<Args...> const& mma)
{
// Keep only NK modes from MNK
auto cluster_tiler_nk = remove<0>(cluster_tiler);
// cluster tile coord -> gtensor coord
auto g_tile = make_identity_layout(shape(gtensor)).compose(cluster_tiler_nk); // (TILE_N, TILE_K, ...)
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_B(g_tile))(_, repeat<rank(g_tile)>(_)); // (MMA, MMA_N, MMA_K, ...)
auto cta_t_vmnk_strides = [](){
if constexpr (is_same_v<CopyOp, SM90_TMA_LOAD_MULTICAST> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD_MULTICAST>) {
return Stride<_0,_1,_0,_0>{}; // 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>) {
return Stride<_0,_0,_0,_0>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
auto cta_t_shape = shape(mma.get_thr_layout_vmnk());
// cta rank -> logical cta idx
auto cta_t_map = coalesce(make_layout(cta_t_shape, compact_col_major(cta_t_shape, cta_t_vmnk_strides)));
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
return detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_map, cta_v_tile);
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class Cluster_Tiler,
class... Args>
CUTE_HOST
auto
make_tma_copy_C_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (M, N, ...)
SLayout const& slayout, // (MMA, MMA_M, MMA_N, ...)
Cluster_Tiler const& cluster_tiler, // (TILE_M, TILE_N, TILE_K, ...)
TiledMMA<Args...> const& mma)
{
// Keep only MN modes from MNK
auto cluster_tiler_mn = remove<2>(cluster_tiler);
// cluster tile coord -> gtensor coord
auto g_tile = make_identity_layout(shape(gtensor)).compose(cluster_tiler_mn); // (TILE_M, TILE_N, ...)
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_C(g_tile))(_, repeat<rank(g_tile)>(_)); // (MMA, MMA_M, MMA_N, ...)
static_assert(is_same_v<CopyOp, SM90_TMA_LOAD> ||
is_same_v<CopyOp, SM90_TMA_STORE> ||
is_same_v<CopyOp, SM100_TMA_2SM_LOAD>,
"Unsupported TMA Op, expected a non-multicast TMA");
// No multicast, so only 1 CTA involved
auto cta_t_map = Layout<_1,_0>{};
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
return detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_map, cta_v_tile);
}
////////////////////////////////////
// Experimental Make TMA Atom
///////////////////////////////////
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class MMA_Tiler,
class... Args,
class ClusterShapeVMNK>
CUTE_HOST
auto
make_tma_atom_A_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (M, K, ...)
SLayout const& slayout, // (MMA, MMA_M, MMA_K, ...)
MMA_Tiler const& mma_tiler, // (TILE_M, TILE_N, TILE_K, ...)
TiledMMA<Args...> const& mma,
ClusterShapeVMNK const& cluster_shape) // (CTA_V, CTA_M, CTA_N, CTA_K)
{
// Keep only MK modes from MNK
auto mma_tiler_mk = remove<1>(mma_tiler);
// cluster tile coord -> gtensor coord
auto g_tile = make_identity_layout(shape(gtensor)).compose(mma_tiler_mk); // (TILE_M, TILE_K, ...)
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_A(g_tile))(_, repeat<rank(g_tile)>(_)); // (MMA, MMA_M, MMA_K, ...)
#if 0
print("(tma_a) slayout: "); print(slayout); print("\n");
print("(tma_a) mma_tiler_nk: "); print(mma_tiler_nk); print("\n");
print("(tma_a) g_tile: "); print(g_tile); print("\n");
print("(tma_a) mma_tiler: "); print(mma_tiler); print("\n");
print("(tma_a) cta_v_tile: "); print(cta_v_tile); print("\n");
#endif
// 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>) {
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>) {
return Int<1>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
return detail::make_tma_copy_atom<TmaType>(copy_op, gtensor, slayout, num_multicast, cta_v_tile);
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class MMA_Tiler,
class... Args,
class ClusterShapeVMNK>
CUTE_HOST
auto
make_tma_atom_B_sm100(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor, // (N, K, ...)
SLayout const& slayout, // (MMA, MMA_N, MMA_K, ...)
MMA_Tiler const& mma_tiler, // (TILE_M, TILE_N, TILE_K, ...)
TiledMMA<Args...> const& mma,
ClusterShapeVMNK const& cluster_shape) // (CTA_V, CTA_M, CTA_N, CTA_K)
{
// Keep only NK modes from MNK
auto mma_tiler_nk = remove<0>(mma_tiler);
// cluster tile coord -> gtensor coord
auto g_tile = make_identity_layout(shape(gtensor)).compose(mma_tiler_nk); // (TILE_N, TILE_K, ...)
// cta val idx -> gmem mode
auto cta_v_tile = layout<1>(mma.thrfrg_B(g_tile))(_, repeat<rank(g_tile)>(_)); // (MMA, MMA_N, MMA_K, ...)
#if 0
print("(tma_b) slayout: "); print(slayout); print("\n");
print("(tma_b) mma_tiler_nk: "); print(mma_tiler_nk); print("\n");
print("(tma_b) g_tile: "); print(g_tile); print("\n");
print("(tma_b) mma_tiler: "); print(mma_tiler); print("\n");
print("(tma_b) cta_v_tile: "); print(cta_v_tile); print("\n");
#endif
// 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>) {
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>) {
return Int<1>{}; // VMNK: Use no CTAs in Non-Multicast
} else {
static_assert(dependent_false<CopyOp>, "Unsupported TMA");
}
}();
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
return detail::make_tma_copy_atom<TmaType>(copy_op, gtensor, slayout, num_multicast, cta_v_tile);
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace cute
include/cute/atom/copy_traits_sm90_tma_swizzle.hpp
+23
View File
@@ -56,6 +56,13 @@ get_tma_swizzle_bits(Swizzle<B,M,S>)
case 0: return TMA::SmemSwizzleBits::DISABLE;
}
} else
if constexpr (M == 5 || M == 6) {
static_assert(B == 2, "Expected B = 2 when M == 5 or 6. Unsupported layout swizzle.");
// S-condition as well?
return TMA::SmemSwizzleBits::B128;
} else
{
static_assert(M < 0, "Unsupported layout swizzle.");
}
@@ -78,9 +85,25 @@ get_tma_swizzle_base(Swizzle<B,M,S>)
static_assert(S == 3, "Expected S = 3 when M == 4. Unsupported layout swizzle.");
return TMA::SmemSwizzleBase::SWIZZLE_BASE_16B;
}
else if constexpr (M == 5) {
static_assert(B == 2, "Expected B = 2 when M == 5. Unsupported layout swizzle.");
static_assert(S == 2, "Expected S = 2 when M == 5. Unsupported layout swizzle.");
return TMA::SmemSwizzleBase::SWIZZLE_BASE_32B;
} else if constexpr (M == 6) {
static_assert(B == 2, "Expected B = 2 when M == 5. Unsupported layout swizzle.");
return TMA::SmemSwizzleBase::SWIZZLE_BASE_64B;
}
#if 1
else {
static_assert(4 <= M && M <= 6, "Expected 128b=16B=(2^4)B to 512b=64B=(2^6)B base swizzle.");
}
#else
else {
static_assert(M == 4, "Expected 128b=16B=(2^4)B base swizzle.");
}
#endif
}
template <class Layout>
include/cute/atom/mma_atom.hpp
+9
View File
@@ -154,6 +154,10 @@ struct MMA_Atom<MMA_Traits<MMAOperation, Args...>>
if constexpr (has_dereference<FrgTypeA>::value) {
// If the intended FrgTypeA is a view (of the current tensor), forward the whole
static_assert(is_same<ValTypeA, typename remove_cvref_t<ATensor>::value_type>::value
|| (sizeof_bits_v<typename remove_cvref_t<ATensor>::value_type> == 8 &&
(sizeof_bits_v<ValTypeA> == 8 || sizeof_bits_v<ValTypeA> == 6 || sizeof_bits_v<ValTypeA> == 4))
, "Expecting ValTypeA type");
return make_tensor<FrgTypeA>(static_cast<ATensor&&>(atensor));
} else {
@@ -176,6 +180,10 @@ struct MMA_Atom<MMA_Traits<MMAOperation, Args...>>
if constexpr (has_dereference<FrgTypeB>::value) {
// If the intended FrgTypeB is a view (of the current tensor), forward the whole
static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value
|| (sizeof_bits_v<typename remove_cvref_t<BTensor>::value_type> == 8 &&
(sizeof_bits_v<ValTypeB> == 8 || sizeof_bits_v<ValTypeB> == 6 || sizeof_bits_v<ValTypeB> == 4))
, "Expecting ValTypeB type");
return make_tensor<FrgTypeB>(static_cast<BTensor&&>(btensor));
} else {
@@ -1109,4 +1117,5 @@ print_svg(TiledMMA<Args...> const &mma) {
#include <cute/atom/mma_traits_sm80.hpp>
#include <cute/atom/mma_traits_sm90.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cute/atom/mma_traits_sm100.hpp
+2425
View File
File diff suppressed because it is too large Load Diff
include/cute/atom/partitioner.hpp
+109
View File
@@ -0,0 +1,109 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#else
#include <type_traits>
#endif
#include <cute/config.hpp>
#include <cute/tensor.hpp>
namespace cute {
//
// A generic tiling of thread-value layouts
//
template <class Layout_TV_, // (tid,vid) -> coord [Need not be 2D...]
class Tiler_MN_> // coord space
struct TV_Tiler
{
using Tiler_MN = Tiler_MN_;
using TiledLayout_TV = Layout_TV_;
// Tile a tensor or a layout from shape
// (M,N,...)
// to shape
// ((ThrV,FrgV),(RestM,RestN,...))
// where
// ThrV: The threads local to a tile.
// FrgV: The values local to a tile.
// RestM: The values tiled in M.
// RestN: The values tiled in N.
template <class Tensor>
CUTE_HOST_DEVICE constexpr static
auto
apply(Tensor&& tensor)
{
// If Layout_TV and Tiler_MN were composable in general, then this won't be needed!
// ((thr_id,val_id),(RestM,RestN,...))
return zipped_divide(tensor, Tiler_MN{}).compose(TiledLayout_TV{}, _);
}
template <class SliceCoord>
struct TV_Partitioner
{
SliceCoord coord_;
template <class TargetTensor>
CUTE_HOST_DEVICE
auto
partition(TargetTensor&& target) {
Tensor thr_tensor = make_tensor(static_cast<TargetTensor&&>(target).data(), apply(target.layout()));
return thr_tensor(coord_, repeat<rank_v<TargetTensor>>(_));
}
};
template <class SliceCoord>
CUTE_HOST_DEVICE static
auto
get_slice(SliceCoord const& coord)
{
return TV_Partitioner<SliceCoord>{coord};
}
};
template <class Layout_TV,
class Tiler_MN>
CUTE_HOST_DEVICE
auto
make_tiler_impl(Layout_TV const&,
Tiler_MN const&)
{
return TV_Tiler<Layout_TV, Tiler_MN>{};
}
}