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

Updates for CUTLASS 3.5.0 (#1468)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/atom/copy_atom.hpp
+1
View File
@@ -756,6 +756,7 @@ print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and
////////////////////////////////////////////////////////////////////////////////////////////////////
#include <cute/atom/copy_traits_sm50.hpp>
#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.hpp
+3 -55
View File
@@ -92,59 +92,6 @@ struct Copy_Traits<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>>
using RefLayout = SrcLayout;
};
namespace detail {
// Utility for exploding pointers, arrays, or tensors into Operation::copy
template <class Operation,
class PtrSrc, int... Is,
class PtrDst, int... Id>
CUTE_HOST_DEVICE constexpr
void
copy_explode_index(PtrSrc&& s, int_sequence<Is...>,
PtrDst&& d, int_sequence<Id...>)
{
return Operation::copy(s[Is]..., d[Id]...);
}
// Utility for exploding tuples into ::copy
template <class Operation,
class TupleArg, int... I>
CUTE_HOST_DEVICE constexpr
void
copy_explode(TupleArg&& t, int_sequence<I...>)
{
return Operation::copy(get<I>(static_cast<TupleArg&&>(t))...);
}
template <class Operation,
class TupleSrc, int... Is,
class TupleDst, int... Id>
CUTE_HOST_DEVICE constexpr
void
copy_explode(TupleSrc&& s, int_sequence<Is...>,
TupleDst&& d, int_sequence<Id...>)
{
return Operation::copy(get<Is>(static_cast<TupleSrc&&>(s))...,
get<Id>(static_cast<TupleDst&&>(d))...);
}
template <class Operation,
class TupleAux, int... Ia,
class TupleSrc, int... Is,
class TupleDst, int... Id>
CUTE_HOST_DEVICE constexpr
void
copy_explode(TupleAux&& a, int_sequence<Ia...>,
TupleSrc&& s, int_sequence<Is...>,
TupleDst&& d, int_sequence<Id...>)
{
return Operation::copy(get<Ia>(static_cast<TupleAux&&>(a))...,
get<Is>(static_cast<TupleSrc&&>(s))...,
get<Id>(static_cast<TupleDst&&>(d))...);
}
} // end namespace detail
//
// Generic copy_unpack for common argument-based Copy_Traits
//
@@ -177,8 +124,9 @@ copy_unpack(Copy_Traits<CopyOp,Args...> const&,
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumDst>{},
"Copy_Traits: dst failed to vectorize into registers. Layout is incompatible with this CopyOp.");
detail::copy_explode_index<CopyOp>(rS, make_int_sequence<RegNumSrc>{},
rD, make_int_sequence<RegNumDst>{});
detail::explode(detail::CallCOPY<CopyOp>{},
rS, make_int_sequence<RegNumSrc>{},
rD, make_int_sequence<RegNumDst>{});
}
//
include/cute/atom/copy_traits_sm50.hpp
+58
View File
@@ -0,0 +1,58 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/arch/copy_sm50.hpp>
#include <cute/atom/copy_traits.hpp>
#include <cute/layout.hpp>
namespace cute
{
template <>
struct Copy_Traits<SM50_Shuffle_U32_2x2Trans>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_32>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape <_32,_64>,
Stride<_64, _1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape <Shape < _2, _16>,Shape <_32, _2>>,
Stride<Stride<_32, _128>,Stride< _1, _64>>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
};
} // end namespace cute
include/cute/atom/copy_traits_sm90_im2col.hpp
+158 -109
View File
@@ -73,14 +73,16 @@ struct TMA_LOAD_IM2COL_Unpack
CUTE_STATIC_ASSERT_V(rank<1>(src_coord_offset) == rank<3>(src_coord_offset));
if constexpr (detail::is_prefetch<CopyOp>) {
return detail::copy_explode<CopyOp>(traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord_cwhdn_offset_srt, tuple_seq<decltype(src_coord_cwhdn_offset_srt)>{});
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord_cwhdn_offset_srt, tuple_seq<decltype(src_coord_cwhdn_offset_srt)>{});
} else {
static_assert(is_smem<TD>::value, "SM90_TMA_LOAD_IM2COL requires the destination be shared memory.");
void* dst_ptr = cute::raw_pointer_cast(dst.data());
return detail::copy_explode<CopyOp>(traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
make_tuple(dst_ptr), seq<0>{},
src_coord_cwhdn_offset_srt, tuple_seq<decltype(src_coord_cwhdn_offset_srt)>{});
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
make_tuple(dst_ptr), seq<0>{},
src_coord_cwhdn_offset_srt, tuple_seq<decltype(src_coord_cwhdn_offset_srt)>{});
}
}
};
@@ -349,8 +351,9 @@ struct Copy_Traits<SM90_TMA_STORE_IM2COL, NumBitsPerTMA, TMATensor>
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = flatten(take<0,3>(dst(Int<0>{})));
return detail::copy_explode<SM90_TMA_STORE_IM2COL>(make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_STORE_IM2COL>{},
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
@@ -537,6 +540,133 @@ make_im2col_tma_copy_desc(
return cute::make_tuple(tma_desc, tma_tensor);
}
template <class CopyOp,
class GEngine, class GLayout,
class SLayout,
class VShape, class VStride,
class LowerCornerStride,
class UpperCornerStride,
class LowerPaddingStride,
class UpperPaddingStride,
class TraversalStride,
class LowerSRTStride,
class DilationStride>
CUTE_HOST_RTC
auto
make_tma_atom_im2col(CopyOp,
Tensor<GEngine,GLayout> const& gtensor, // Full GMEM Tensor: ((w, h, d, n), c)
SLayout const& slayout, // CTA Tile of SMEM, potentially swizzled
int32_t const& num_multicast, // The number of CTAs involved in multicasting
Layout<VShape,VStride> const& cta_v_map, // V: CTA val idx -> gmem mode
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt) // dilation
{
//
// TMA parameter checking
//
CUTE_STATIC_ASSERT_V(product_each(shape(slayout)) == product_each(shape(cta_v_map)),
"TMA requires CTA_Tile and SLayout top-level shape equivalence.");
//
// TMA slayout manipulation
//
// Invert the smem to get the largest contiguous vector in the smem layout
auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout));
// trunc_smem_idx -> trunc_smem_coord
// Map from smem idx to a gmem mode
auto sidx_to_gmode = coalesce(composition(cta_v_map, inv_smem_layout));
#if 0
print("g_layout : "); print(gtensor.layout()); print("\n");
print("s_layout : "); print(slayout); print("\n");
print("cta_t_map : "); print(cta_t_map); print("\n");
print("cta_v_map : "); print(cta_v_map); print("\n");
print("inv_smem : "); print(inv_smem_layout); print("\n");
print("sidx_to_gmode : "); print(sidx_to_gmode); print("\n");
#endif
//
// TMA gtensor manipulation
//
// Generate a TupleBasis for the gtensor
auto glayout_basis = make_identity_layout(product_each(shape(gtensor)));
// Tile the modes of gtensor with the truncated cta_v_map o inv_smem_layout_trunc
auto tma_layout_full = flatten(composition(glayout_basis, sidx_to_gmode));
// Truncate any incompatibilities -- no starting in the middle of gmodes
auto smem_rank = find_if(stride(tma_layout_full), [](auto e) {
[[maybe_unused]] auto v = basis_value(e);
return not is_constant<1,decltype(v)>{};
});
static_assert(smem_rank >= 2, "IM2COL expects at least 2 modes of the smem to vectorize with gmem.");
// IM2COL uses a maximum of 2 modes
constexpr int smem_tma_rank = cute::min(int(smem_rank), 2);
// Keep only the static-1 basis modes into gmem
auto tma_layout_trunc = take<0,smem_tma_rank>(tma_layout_full);
// Split according to the portion each multicast CTA will be responsible for
auto tma_layout_vt = logical_divide(tma_layout_trunc, shape_div(size(tma_layout_trunc), num_multicast));
#if 0
print("glayout_basis : "); print(glayout_basis); print("\n");
print("tma_layout_full : "); print(tma_layout_full); print("\n");
print("tma_layout_trunc: "); print(tma_layout_trunc); print("\n");
print("tma_layout_vt : "); print(tma_layout_vt); print("\n");
#endif
auto range_c = size<0,0>(tma_layout_vt);
auto range_whdn = size<0,1>(tma_layout_vt);
Tensor gtensor_cwhdn = make_tensor(gtensor.data(),
flatten(make_layout(basis_get(stride<0,0>(tma_layout_vt), gtensor.layout()),
basis_get(stride<0,1>(tma_layout_vt), gtensor.layout()))));
auto [tma_desc, tma_tensor] = make_im2col_tma_copy_desc(
gtensor_cwhdn,
range_c,
range_whdn,
detail::get_swizzle_portion(slayout),
tma_layout_vt,
lower_corner_whd,
upper_corner_whd,
lower_padding_whd,
upper_padding_whd,
stride_whd,
lower_srt,
stride_srt);
//
// Construct the Copy_Traits
//
using T = typename GEngine::value_type;
constexpr int num_bits_per_tma = decltype(size(tma_layout_trunc))::value * sizeof(T) * 8;
using Traits = Copy_Traits<CopyOp, cute::C<num_bits_per_tma>, decltype(tma_tensor)>;
using Atom = Copy_Atom<Traits, typename GEngine::value_type>;
#if 0
print("num_bits : "); print(num_bits_per_tma); print("\n");
#endif
Traits tma_traits{tma_desc, tma_tensor};
// Return the Copy_Atom
return Atom{tma_traits};
}
/// Make a TiledCopy for im2col TMA load.
///
/// @param copy_op The copy implementation: either
@@ -584,99 +714,12 @@ make_tma_copy_im2col(CopyOp const& copy_op,
// TMA parameter checking
//
CUTE_STATIC_ASSERT_V(product_each(shape(slayout)) == product_each(shape(cta_v_map)),
"TMA requires CTA_Tile and SLayout top-level shape equivalence.");
CUTE_STATIC_ASSERT_V(size(slayout) % cosize(cta_t_map) == Int<0>{},
"Number of active CTAs in TMA must divide domain size of slayout.");
//
// TMA slayout manipulation
//
// Invert the smem to get the largest contiguous vector in the smem layout
auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout));
// trunc_smem_idx -> trunc_smem_coord
// Map from smem idx to a gmem mode
auto sidx_to_gmode = coalesce(composition(cta_v_map, inv_smem_layout));
#if 0
print("g_layout : "); print(gtensor.layout()); print("\n");
print("s_layout : "); print(slayout); print("\n");
print("cta_t_map : "); print(cta_t_map); print("\n");
print("cta_v_map : "); print(cta_v_map); print("\n");
print("inv_smem : "); print(inv_smem_layout); print("\n");
print("sidx_to_gmode : "); print(sidx_to_gmode); print("\n");
#endif
//
// TMA gtensor manipulation
//
// Generate a TupleBasis for the gtensor
auto glayout_basis = make_identity_layout(product_each(shape(gtensor)));
// Tile the modes of gtensor with the truncated cta_v_map o inv_smem_layout_trunc
auto tma_layout_full = flatten(composition(glayout_basis, sidx_to_gmode));
// Truncate any incompatibilities -- no starting in the middle of gmodes
auto smem_rank = find_if(stride(tma_layout_full), [](auto e) {
[[maybe_unused]] auto v = basis_value(e);
return not is_constant<1,decltype(v)>{};
});
static_assert(smem_rank >= 2, "IM2COL expects at least 2 modes of the smem to vectorize with gmem.");
// IM2COL uses a maximum of 2 modes
constexpr int smem_tma_rank = cute::min(int(smem_rank), 2);
// Keep only the static-1 basis modes into gmem
auto tma_layout_trunc = take<0,smem_tma_rank>(tma_layout_full);
// Split according to the portion each multicast CTA will be responsible for
auto tma_layout_vt = logical_divide(tma_layout_trunc, shape_div(size(tma_layout_trunc), cosize(cta_t_map)));
#if 0
print("glayout_basis : "); print(glayout_basis); print("\n");
print("tma_layout_full : "); print(tma_layout_full); print("\n");
print("tma_layout_trunc: "); print(tma_layout_trunc); print("\n");
print("tma_layout_vt : "); print(tma_layout_vt); print("\n");
#endif
auto range_c = size<0,0>(tma_layout_vt);
auto range_whdn = size<0,1>(tma_layout_vt);
Tensor gtensor_cwhdn = make_tensor(gtensor.data(),
flatten(make_layout(basis_get(stride<0,0>(tma_layout_vt), gtensor.layout()),
basis_get(stride<0,1>(tma_layout_vt), gtensor.layout()))));
auto [tma_desc, tma_tensor] = make_im2col_tma_copy_desc(
gtensor_cwhdn,
range_c,
range_whdn,
detail::get_swizzle_portion(slayout),
tma_layout_vt,
lower_corner_whd,
upper_corner_whd,
lower_padding_whd,
upper_padding_whd,
stride_whd,
lower_srt,
stride_srt);
//
// Construct the Copy_Traits
//
using T = typename GEngine::value_type;
constexpr int num_bits_per_tma = decltype(size<0>(tma_layout_vt))::value * sizeof(T) * 8;
using Traits = Copy_Traits<CopyOp, cute::C<num_bits_per_tma>, decltype(tma_tensor)>;
#if 0
print("num_bits : "); print(NumBitsPerTMA{}); print("\n");
#endif
Traits tma_traits{tma_desc, tma_tensor};
Copy_Atom atom = make_tma_atom_im2col(copy_op, gtensor, slayout, cosize(cta_t_map), cta_v_map,
lower_corner_whd, upper_corner_whd, lower_padding_whd,
upper_padding_whd, stride_whd, lower_srt, stride_srt);
//
// Construct the TiledCopy
@@ -684,25 +727,31 @@ make_tma_copy_im2col(CopyOp const& copy_op,
auto cta_tiler = product_each(shape(cta_v_map));
// (CTA V, CTA T) -> smem_coord
auto layout_vt = composition(inv_smem_layout, make_layout(shape(tma_layout_vt)));
auto num_elems_per_tma = size<1>(typename decltype(atom)::RefLayout{}) / static_value<sizeof_bits<typename GEngine::value_type>>();
// smem idx -> smem coord
auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout));
// CTA V -> smem_coord
auto layout_v = composition(inv_smem_layout, num_elems_per_tma);
// Scale that up to cover all of the smem_coords
//
// The smem vector might not cover all of the tile,
// so multiply it up to cover the entire tile.
// "T" here (the parallel index) is a CTA index.
auto layout_VT = tile_to_shape(layout_vt, make_shape(size(cta_v_map)/size<1>(layout_vt), size<1>(layout_vt)));
// Flip it and change the domain of the T from logical thr to thr_idx
auto layout_TV = make_layout(composition(layout<1>(layout_VT), cta_t_map), layout<0>(layout_VT));
auto layout_V = tile_to_shape(make_layout(layout_v), size(cta_v_map));
// CTA T -> smem idx
auto layout_t = make_layout(cosize(cta_t_map), shape_div(num_elems_per_tma, cosize(cta_t_map)));
// CTA TID -> smem coord
auto layout_T = composition(inv_smem_layout, composition(layout_t, cta_t_map));
// Combine with the T mapping
[[maybe_unused]] auto layout_TV = make_layout(layout_T, layout_V);
#if 0
print("cta_tiler : "); print(cta_tiler); print("\n");
print("layout_VT : "); print(layout_VT); print("\n");
print("layout_v : "); print(layout_v); print("\n");
print("layout_V : "); print(layout_V); print("\n");
print("layout_t : "); print(layout_t); print("\n");
print("layout_T : "); print(layout_T); print("\n");
print("layout_TV : "); print(layout_TV); print("\n");
#endif
using T = typename GEngine::value_type;
return TiledCopy<Copy_Atom<Traits,T>, decltype(layout_TV), decltype(cta_tiler)>{tma_traits};
return TiledCopy<decltype(atom), decltype(layout_TV), decltype(cta_tiler)>{atom};
}
/// Make a TiledCopy for im2col TMA with no offsets.
include/cute/atom/copy_traits_sm90_tma.hpp
+40 -31
View File
@@ -69,8 +69,9 @@ struct TMA_LOAD_Unpack
{
auto src_coord = src.data().coord_;
if constexpr (detail::is_prefetch<CopyOp>) {
return detail::copy_explode<CopyOp>(traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord, tuple_seq<decltype(src_coord)>{});
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord, tuple_seq<decltype(src_coord)>{});
} else {
static_assert(is_smem<TD>::value, "SM90_TMA_LOAD requires the destination be shared memory.");
void* dst_ptr = cute::raw_pointer_cast(dst.data());
@@ -81,9 +82,10 @@ struct TMA_LOAD_Unpack
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::copy_explode<CopyOp>(traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
make_tuple(dst_ptr), seq<0>{},
src_coord, tuple_seq<decltype(src_coord)>{});
return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
make_tuple(dst_ptr), seq<0>{},
src_coord, tuple_seq<decltype(src_coord)>{});
}
}
};
@@ -337,8 +339,9 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
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::copy_explode<SM90_TMA_STORE>(make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_STORE>{},
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
@@ -1278,7 +1281,7 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
// Factor out the single-instrucion portion
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<rank_v<decltype(gtensor)>>(layout_V, _);
auto slayout_V = append<rank_v<decltype(stensor)>>(layout_V, _);
@@ -1288,39 +1291,45 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
#if 0
if (thread0()) {
print("gtensor : "); print(gtensor); print("\n");
print("stensor : "); print(stensor); print("\n");
print("cta_coord : "); print(cta_coord); print("\n");
print("cta_layout : "); print(cta_layout); print("\n");
print("gtensor : "); print(gtensor); print("\n");
print("stensor : "); print(stensor); print("\n");
print("layout_V : "); print(layout_V); print("\n");
print("gtensor_v : "); print(gtensor_v); print("\n");
print("stensor_v : "); print(stensor_v); print("\n");
}
#endif
// Restride the cta-into-tma-instr layout
Layout tma_layout_t = composition(make_layout(Int<1>{}, shape_div(size(tma_layout_v), cosize(cta_layout))), cta_layout);
auto tma_layout_tv = make_tile(make_tile(make_layout(tma_layout_t, tma_layout_v), _));
// 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 scoord = append<SLayout::rank>(multicast_coord, Int<0>{});
auto gcoord = append<GLayout::rank>(multicast_coord, Int<0>{});
// Append with _ until we cover all Rest... modes
auto gtma_layout_tv = append<rank_v<decltype(gtensor)>>(tma_layout_tv, _);
auto stma_layout_tv = append<rank_v<decltype(stensor)>>(tma_layout_tv, _);
Tensor gresult = domain_offset(gcoord, gtensor_v);
Tensor sresult = domain_offset(scoord, stensor_v);
// Transform TMA mode
Tensor gtensor_tv = gtensor_v.compose(gtma_layout_tv); // (((Thr,Frg),TMA_Iter), Rest...)
Tensor stensor_tv = stensor_v.compose(stma_layout_tv); // (((Thr,Frg),TMA_Iter), Rest...)
return cute::make_tuple(gresult, sresult);
}
#if 0
if (thread0()) {
print("tma_layout_tv : "); print(tma_layout_tv); print("\n");
print("gtensor_tv : "); print(gtensor_tv); print("\n");
print("stensor_tv : "); print(stensor_tv); print("\n");
// TMA Multicast Masks Calculation
template <int Mode, class CtaLayout, class CtaCoord>
CUTE_HOST_DEVICE constexpr
auto
create_tma_multicast_mask(CtaLayout const& cta_layout_vmnk,
CtaCoord const& cta_coord_vmnk)
{
auto cta_coord_slicer = replace<Mode>(cta_coord_vmnk, _);
auto [cta_layout, elected_cta] = slice_and_offset(cta_coord_slicer, cta_layout_vmnk);
// Get the instruction code
uint16_t mcast_mask = 0;
for (int i = 0; i < size(cta_layout); ++i) {
mcast_mask |= uint16_t(1) << cta_layout(i);
}
#endif
auto c = make_coord(make_coord(make_coord(cta_coord, _), _));
auto c_s = append<rank_v<decltype(stensor_tv)>>(c, _);
auto c_g = append<rank_v<decltype(gtensor_tv)>>(c, _);
return cute::make_tuple(group_modes<0,2>(gtensor_tv(c_g)), group_modes<0,2>(stensor_tv(c_s)));
// Shift by the instruction's elected block rank (dynamic)
mcast_mask <<= elected_cta;
return mcast_mask;
}
} // end namespace cute
include/cute/atom/mma_atom.hpp
+1
View File
@@ -715,6 +715,7 @@ print(MMA_Atom<MMA_Traits<Args...>> const&)
using Atom = MMA_Atom<MMA_Traits<Args...>>;
print("MMA_Atom\n");
print(" ThrID: "); print(typename Atom::ThrID{}); print("\n");
print(" Shape_MNK: "); print(typename Atom::Shape_MNK{}); print("\n");
print(" LayoutA_TV: "); print(typename Atom::LayoutA_TV{}); print("\n");
print(" LayoutB_TV: "); print(typename Atom::LayoutB_TV{}); print("\n");
print(" LayoutC_TV: "); print(typename Atom::LayoutC_TV{}); print("\n");
include/cute/atom/mma_traits.hpp
+15 -15
View File
@@ -149,17 +149,17 @@ mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
//CUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
if constexpr (detail::supports_output_scaling<MMATraits>::value) {
detail::explode_with_d_scaling(MMA_Op::fma,
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{},
traits.accumulate_);
detail::explode(MMA_Op::fma,
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{},
&(traits.accumulate_), seq<0>{});
}
else {
detail::explode(MMA_Op::fma,
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
}
}
else {
@@ -169,19 +169,19 @@ mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumD>{});
CUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
if constexpr (detail::supports_output_scaling<MMATraits>::value) {
detail::explode_with_d_scaling(MMA_Op::fma,
detail::explode(MMA_Op::fma,
rD, make_int_sequence<RegNumD>{},
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{},
traits.accumulate_);
&(traits.accumulate_), seq<0>{});
}
else {
detail::explode(MMA_Op::fma,
rD, make_int_sequence<RegNumD>{},
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
rD, make_int_sequence<RegNumD>{},
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
}
}
}
@@ -198,7 +198,7 @@ template <class MMA_Op, class... MMA_Args,
CUTE_HOST_DEVICE constexpr
void
mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
Tensor<TD, DLayout> && D,
Tensor<TD, DLayout> && D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
include/cute/atom/mma_traits_sm90_gmma.hpp
+1 -1
View File
@@ -208,7 +208,7 @@ make_gmma_desc(Tensor<TEngine,TLayout> const& tensor)
// Start address (4LSB not included)
uint32_t start_address = cast_smem_ptr_to_uint(raw_pointer_cast(u128_tensor.data()));
desc.bitfield.start_address_ = start_address >> 4;
desc.bitfield.start_address_ = static_cast<uint16_t>(start_address >> 4);
constexpr uint8_t base_offset = 0;
desc.bitfield.base_offset_ = base_offset;