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

Updates and Bug fixes to CUTLASS 3.3 (#1232)

Browse Source
This commit is contained in:
Pradeep Ramani
2023-12-05 06:50:49 -08:00
committed by GitHub
parent 4a1709e17e
commit e9e30c2304
31 changed files with 534 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files

295
include/cute/algorithm/copy.hpp
View File

@@ -32,6 +32,8 @@
#include <cute/config.hpp>
#include <cute/container/alignment.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_predicate.hpp>
@@ -44,30 +46,14 @@ namespace cute
// Accept mutable temporaries
//
template <class PrdTensor,
class SrcEngine, class SrcLayout,
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(pred, src, dst);
}
template <class... CopyArgs,
class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(copy_atom, pred, src, dst);
return copy(src, dst);
}
template <class VecType,
@@ -85,22 +71,48 @@ template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy(src, dst);
return copy_aligned(src, dst);
}
template <class... CopyArgs,
template <class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<CopyArgs...> const& copy_atom,
copy_if(PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(pred, src, dst);
}
template <class CopyPolicy,
class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(CopyPolicy const& copy_policy,
PrdTensor const& pred,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy_if(copy_policy, pred, src, dst);
}
template <class CopyPolicy,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(CopyPolicy const& copy_policy,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return copy(copy_atom, src, dst);
return copy(copy_policy, src, dst);
}
//
@@ -135,7 +147,7 @@ namespace detail {
// Trait that detects if atom's traits has a member function with(bool)
template<typename, typename Enable = void>
constexpr bool has_with_bool = false;
template<typename T>
constexpr bool has_with_bool<T, cute::void_t<decltype(declval<typename T::Traits>().with(declval<bool>()))>> = true;
@@ -157,15 +169,14 @@ copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
copy_atom.call(src, dst);
} else { // Loop over all but the first mode
constexpr int R = SrcLayout::rank;
auto src_v = group_modes<1,R>(src);
auto dst_v = group_modes<1,R>(dst);
Tensor src_v = group_modes<1,R>(src);
Tensor dst_v = group_modes<1,R>(dst);
CUTE_UNROLL
for (int i = 0; i < size<1>(src_v); ++i) {
// If copy traits can be transformed with a predicate value, do it, otherwise branch here
if constexpr (detail::has_with_bool<Copy_Atom<CopyArgs...>>) {
copy_atom.with(pred(i)).call(src_v(_,i), dst_v(_,i));
}
else {
} else {
if (pred(i)) {
copy_atom.call(src_v(_,i), dst_v(_,i));
}
@@ -186,23 +197,24 @@ void
copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
static_assert(sizeof_bits_v<VecType> >= 8 && sizeof_bits_v<VecType> % 8 == 0,
"Expected a vectorization type of at least a byte.");
using SrcType = typename SrcEngine::element_type;
using DstType = typename DstEngine::element_type;
if constexpr (sizeof(SrcType) == sizeof(DstType) && sizeof(VecType) > sizeof(DstType))
if constexpr (sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
sizeof_bits_v<VecType> > sizeof_bits_v<DstType>)
{
/* @pre is_aligned<N>(src.data()) &&
* is_aligned<N>(dst.data())
*/
// Preserve volatility of Src/Dst types.
using SrcVecType = conditional_t<is_volatile_v<SrcType>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<DstType>, VecType volatile, VecType >;
auto src_v = recast<SrcVecType>(src);
auto dst_v = recast<DstVecType>(dst);
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
#if 0
if (thread0()) {
print("copy_vec -- vectorizing copy from %3db to %3db\n", int(8*sizeof(SrcType)), int(8*sizeof(VecType)));
print(" "); print(layout(src)); print(" => "); print(layout(src_v)); print("\n");
print(" "); print(layout(dst)); print(" => "); print(layout(dst_v)); print("\n");
print("copy_vec<%db> -- vectorizing copy:\n", int(sizeof_bits_v<VecType>));
print(" "); print(src); print(" => "); print(src_v); print("\n");
print(" "); print(dst); print(" => "); print(dst_v); print("\n");
}
#endif
@@ -210,9 +222,9 @@ copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
} else {
#if 0
if (thread0()) {
print("copy_vec -- not vectorizing, copy with %3db and %3db\n", int(8*sizeof(SrcType)), int(8*sizeof(DstType)));
print(" "); print(layout(src)); print("\n");
print(" "); print(layout(dst)); print("\n");
print("copy_vec<%db> -- NOT vectorizing copy:\n", int(sizeof_bits_v<VecType>));
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
@@ -220,36 +232,6 @@ copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
}
}
//
// copy -- auto-vectorizing copy
//
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
constexpr int N = decltype(max_common_vector(src, dst))::value;
#if 0
if (thread0()) {
print("copy -- found a max_common_vector of %d\n", N);
print(" "); print(src.data()); print(" o "); print(layout(src)); print("\n");
print(" "); print(dst.data()); print(" o "); print(layout(dst)); print("\n");
}
#endif
if constexpr (N <= 1) {
return copy_if(TrivialPredTensor{}, src, dst);
} else {
constexpr int vec_bits = N * sizeof_bits<typename SrcEngine::value_type>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
return copy_vec<VecType>(src, dst);
}
}
//
// copy -- CopyAtom
//
@@ -266,23 +248,135 @@ copy(Copy_Atom<CopyArgs...> const& copy_atom,
return copy_if(copy_atom, TrivialPredTensor{}, src, dst);
}
template <class... CopyArgs,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<DefaultCopy, CopyArgs...> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(src, dst);
}
//////////////////////////////////////////
// Special Auto-Vectorizing Overloads
//////////////////////////////////////////
// Specialization for AutoVectorizingCopyAssumedAlignment<MaxVecBits>
template <int MaxVecBits, class... Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
constexpr int vec_elem = decltype(max_common_vector(src, dst))::value;
constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
// When layouts are static, accept vec_bits up to 128
// When layouts are dynamic, accept vec_bits up to MaxVecBits
constexpr int vec_bits = (is_static<SrcLayout>::value && is_static<DstLayout>::value) ?
cute::min(vec_elem * src_bits, 128) :
cute::min(vec_elem * src_bits, MaxVecBits);
#if 0
if (thread0()) {
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
if constexpr (vec_elem > 1 && vec_bits >= 8) {
return copy_vec<uint_bit_t<vec_bits>>(src, dst);
} else {
return copy_if(TrivialPredTensor{}, src, dst);
}
}
// Auto-vectorizing copy for static layouts
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopy{}, src, dst);
}
// Auto-vectorizing copy with assumed alignment of dynamic layout strides up to 128bit.
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
}
// Specializaton for Atom AutoVectorizingCopy
template <class... Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<AutoVectorizingCopy, Args...> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopy{}, src, dst);
}
// Specializaton for Atom AutoVectorizingCopyAssumedAlignment
template <int MaxVecBits, class... Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, Args...> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>{}, src, dst);
}
#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)
template <class... CT_Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...> const& atom, // Copy_Traits may or may not have the memory barrier in it already
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
static_assert(sizeof_bits<SrcType>::value == sizeof_bits<DstType>::value);
static_assert((is_gmem<SrcEngine>::value && is_smem<DstEngine>::value) ||
(is_smem<SrcEngine>::value && is_gmem<DstEngine>::value),
"Bulk Copy only supports gmem -> smem or smem -> gmem movement.");
// G2S or S2G dispatch
using BULK_COPY_OP = conditional_t<is_gmem<SrcEngine>::value,
SM90_BULK_COPY_G2S,
SM90_BULK_COPY_S2G>;
// Find the common subtensor of src and dst
auto tiler = max_common_layout(src, dst);
constexpr int vec_elem = decltype(size(tiler))::value;
constexpr int vec_bits = vec_elem * sizeof_bits_v<SrcType>;
static_assert(vec_bits >= 128, "Expected at least 128-bits for BLKCP");
// Construct a new concrete Atom of the vector size
using BulkAtom = Copy_Atom<Copy_Traits<BULK_COPY_OP, Int<vec_bits>, CT_Args...>, SrcType>;
auto bulk_atom = apply(atom.opargs_, [](auto const&... args) { return BulkAtom{args...}; });
#if 0
if (thread0()) {
print("copy blkcp -- found a max_common_layout of "); print(tiler); print("\n");
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
return copy(bulk_atom, logical_divide(src, tiler), logical_divide(dst, tiler));
}
// Backwards-compat. Throw out any extra Copy_Atom args.
template <class... CT_Args, class... CA_Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -292,36 +386,7 @@ copy(Copy_Atom<Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...>, CA_Args...> const&
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
static_assert(sizeof_bits<SrcType>::value == sizeof_bits<DstType>::value);
static_assert((is_gmem<SrcEngine>::value && is_smem<DstEngine>::value) ||
(is_smem<SrcEngine>::value && is_gmem<DstEngine>::value),
"Bulk Copy only supports gmem -> smem or smem -> gmem movement.");
// Do BulkCopy dispatch
using BULK_COPY_OP = conditional_t<is_gmem<SrcEngine>::value,
SM90_BULK_COPY_G2S,
SM90_BULK_COPY_S2G>;
constexpr int N = decltype(max_common_vector(src, dst))::value;
// Construct a new concrete Atom of the vector size
using N_BITS = Int<N*sizeof_bits<SrcType>::value>;
using COPY_ATOM = Copy_Atom<Copy_Traits<BULK_COPY_OP, N_BITS, CT_Args...>, SrcType>;
auto bulk_atom = apply(atom.opargs_, [&](auto const&... args) { return COPY_ATOM{args...}; });
// Tile the src and dst to the Atom
auto tiler = right_inverse(dst.layout()).compose(Int<N>{});
#if 0
if (thread0()) {
print("copy -- found a max_common_vector of %d\n", N);
print(" "); print(src.data()); print(" o "); print(layout(src)); print("\n");
print(" "); print(dst.data()); print(" o "); print(layout(dst)); print("\n");
}
#endif
return copy(bulk_atom, logical_divide(src, tiler), logical_divide(dst, tiler));
return copy(static_cast<Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...> const&>(atom), src, dst);
}
#endif // #if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)

29
include/cute/arch/copy.hpp
View File

@@ -33,7 +33,7 @@
#include <cute/config.hpp>
#include <cute/arch/util.hpp>
#include <cute/numeric/uint128.hpp>
#include <cute/numeric/int.hpp>
namespace cute
{
@@ -48,7 +48,7 @@ struct UniversalCopy
using SRegisters = S[1];
using DRegisters = D[1];
template<class S_, class D_>
template <class S_, class D_>
CUTE_HOST_DEVICE static constexpr void
copy(S_ const& src,
D_ & dst)
@@ -57,25 +57,36 @@ struct UniversalCopy
}
// Accept mutable temporaries
template<class S_, class D_>
template <class S_, class D_>
CUTE_HOST_DEVICE static constexpr void
copy(S_ const& src,
D_ && dst)
{
copy(src, dst);
UniversalCopy<S,D>::copy(src, dst);
}
};
//
// Placeholder for the copy algorithm's default, auto-vectorizing behavior
// Placeholder for the copy algorithm's stronger auto-vectorizing behavior
// that assumes alignment of dynamic layouts up to MaxVecBits
//
struct DefaultCopy
template <int MaxVecBits = 128>
struct AutoVectorizingCopyWithAssumedAlignment
: UniversalCopy<uint_bit_t<MaxVecBits>>
{
using SRegisters = uint128_t[1];
using DRegisters = uint128_t[1];
static_assert(MaxVecBits == 8 || MaxVecBits == 16 || MaxVecBits == 32 || MaxVecBits == 64 || MaxVecBits == 128,
"Expected MaxVecBits to be 8 or 16 or 32 or 64 or 128 for alignment and performance.");
};
using AutoVectorizingCopy = DefaultCopy;
//
// Placeholder for the copy algorithm's default auto-vectorizing behavior
// that does not assume alignment of dynamic layouts
//
using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
// Alias
using DefaultCopy = AutoVectorizingCopy;
} // end namespace cute

16
include/cute/atom/copy_atom.hpp
View File

@@ -353,7 +353,7 @@ struct ThrCopy
template <class STensor>
CUTE_HOST_DEVICE
auto
partition_S(STensor&& stensor) {
partition_S(STensor&& stensor) const {
//static_assert(sizeof(typename remove_cvref_t<STensor>::value_type) == sizeof(typename TiledCopy::ValType),
// "Expected ValType for tiling SrcTensor.");
auto thr_tensor = make_tensor(std::forward<STensor>(stensor).data(), TiledCopy::tidfrg_S(stensor.layout()));
@@ -363,7 +363,7 @@ struct ThrCopy
template <class DTensor>
CUTE_HOST_DEVICE
auto
partition_D(DTensor&& dtensor) {
partition_D(DTensor&& dtensor) const {
//static_assert(sizeof(typename remove_cvref_t<DTensor>::value_type) == sizeof(typename TiledCopy::ValType),
// "Expected ValType for tiling DstTensor.");
auto thr_tensor = make_tensor(std::forward<DTensor>(dtensor).data(), TiledCopy::tidfrg_D(dtensor.layout()));
@@ -479,10 +479,10 @@ make_tiled_copy_C_atom(Copy_Atom<Args...> const& copy_atom,
return make_tiled_copy_impl(copy_atom, layout_tv, tiler);
}
/** Produce a TiledCopy from logical thread and values layouts.
* The thread and value layouts map coordinates to thr_idx and val_idx.
/** Produce a TiledCopy from logical thread and values layouts.
* The thread and value layouts map coordinates to thr_idx and val_idx.
* The product of these layouts is taken to produce the TV layout and the Tiler.
* Useful when threads and values need very specific mappings onto coordinates
* Useful when threads and values need very specific mappings onto coordinates
* in the target tensors.
*/
template <class... Args,
@@ -510,16 +510,16 @@ make_tiled_copy(Copy_Atom<Args...> const& copy_atom,
return make_tiled_copy_impl(copy_atom, layout_tv, product_each(shape(layout_mn)));
}
/** Produce a TiledCopy from thread and value offset maps.
/** Produce a TiledCopy from thread and value offset maps.
* The TV Layout maps threads and values to the codomain of the data_layout.
* It is verified that the intended codomain is valid within data_layout.
* It is verified that the intended codomain is valid within data_layout.
* Useful when threads and values don't care about owning specific coordinates, but
* care more about the vector-width and offsets between them.
*/
template <class... Args, class AtomTVLayout, class DataLayout>
CUTE_HOST_DEVICE constexpr
auto
make_cotiled_copy(Copy_Atom<Args...> const& copy_atom,
make_cotiled_copy(Copy_Atom<Args...> const& copy_atom,
AtomTVLayout const& atom_tv_layout, // atom (thr,val) -> data addr
DataLayout const& data_layout) // coord -> data addr The target layout
{

49
include/cute/atom/copy_traits.hpp
View File

@@ -59,7 +59,7 @@ namespace cute
template <class CopyOperation, class... CopyOpArgs>
struct Copy_Traits
{
static_assert(sizeof(CopyOperation) == 0, "Copy_Traits not implemented for this Copy_Operation.");
static_assert(dependent_false<CopyOperation>, "Copy_Traits not implemented for this CopyOperation.");
};
template <class S, class D>
@@ -77,8 +77,8 @@ struct Copy_Traits<UniversalCopy<S,D>>
using RefLayout = SrcLayout;
};
template <>
struct Copy_Traits<DefaultCopy>
template <int MaxVecBits>
struct Copy_Traits<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
@@ -108,23 +108,24 @@ copy_explode(PtrS&& s, int_sequence<Is...>,
} // end namespace detail
//
// Generic copy_unpack for any Copy_Traits
// Generic copy_unpack for common argument-based Copy_Traits
//
template <class Operation, class... Args,
class TS, class SLayout,
class TD, class DLayout>
template <class CopyOp, class... Args,
class SEngine, class SLayout,
class DEngine, class DLayout>
CUTE_HOST_DEVICE constexpr
void
copy_unpack(Copy_Traits<Operation, Args...> const&,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
copy_unpack(Copy_Traits<CopyOp,Args...> const&,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> & dst)
{
// Specializations can generalize on these checks
//static_assert(is_smem<TS>::value, "Expected smem for this Copy_Traits<Operation>");
//static_assert(is_rmem<TD>::value, "Expected rmem for this Copy_Traits<Operation>");
//static_assert(is_smem<TS>::value, "Expected smem for this Copy_Traits<CopyOp>");
//static_assert(is_rmem<TD>::value, "Expected rmem for this Copy_Traits<CopyOp>");
using RegistersSrc = typename Operation::SRegisters;
using RegistersDst = typename Operation::DRegisters;
using RegistersSrc = typename CopyOp::SRegisters;
using RegistersDst = typename CopyOp::DRegisters;
using RegTypeSrc = typename remove_extent<RegistersSrc>::type;
using RegTypeDst = typename remove_extent<RegistersDst>::type;
constexpr int RegNumSrc = extent<RegistersSrc>::value;
@@ -134,26 +135,26 @@ copy_unpack(Copy_Traits<Operation, Args...> const&,
Tensor rD = recast<RegTypeDst>(dst);
CUTE_STATIC_ASSERT_V(size(rS) == Int<RegNumSrc>{},
"In CopyAtom, src layout doesn't vectorize into registers. This src layout is incompatible with this tiled copy.");
"Copy_Traits: src failed to vectorize into registers. Layout is incompatible with this CopyOp.");
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumDst>{},
"In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
"Copy_Traits: dst failed to vectorize into registers. Layout is incompatible with this CopyOp.");
detail::copy_explode<Operation>(rS, make_int_sequence<RegNumSrc>{},
rD, make_int_sequence<RegNumDst>{});
detail::copy_explode<CopyOp>(rS, make_int_sequence<RegNumSrc>{},
rD, make_int_sequence<RegNumDst>{});
}
//
// Accept mutable temporaries
//
template <class Operation, class... Args,
class TS, class SLayout,
class TD, class DLayout>
template <class CopyOp, class... Args,
class SEngine, class SLayout,
class DEngine, class DLayout>
CUTE_HOST_DEVICE constexpr
void
copy_unpack(Copy_Traits<Operation, Args...> const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> && dst)
copy_unpack(Copy_Traits<CopyOp,Args...> const& traits,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> && dst)
{
copy_unpack(traits, src, dst);
}

17
include/cute/layout.hpp
View File

@@ -1184,11 +1184,10 @@ left_inverse(Underscore const& _)
//
/* Return a layout that points to the maximum number of contiguous elements
* that logically correspond in the layouts of @a a and @a b. This is,
* the elements that could reasonably be "vectorized" in the layouts.
* that logically correspond in the layouts of @a a and @a b.
*
* @returns Layout R
* @post For all 0 <= i < size(R), a(R(i)) == i && b(R(i)) == i
* @post For all 0 <= i < size(R), a(R(i)) == i and b(R(i)) == i
*/
template <class ShapeA, class StrideA,
class ShapeB, class StrideB>
@@ -1200,8 +1199,7 @@ max_common_layout(Layout<ShapeA,StrideA> const& a,
Layout inv_b = right_inverse(b);
Layout common = coalesce(composition(a, inv_b));
// NOTE: If one of the layouts is dynamic, we can't prove alignment+vectorization is valid
// We assume dynamic shapes/strides obey alignment requirements (i.e. are large and multiples of the vector)
// Keep only the static identity component of the common layout
if constexpr (is_static<decltype(shape<0>(common))>::value &&
is_constant<1, decltype(stride<0>(common))>::value) {
// Truncate to the size of the contiguous vector (static stride-1 mode)
@@ -1212,11 +1210,11 @@ max_common_layout(Layout<ShapeA,StrideA> const& a,
}
/* Return Int<N> such that N is the maximum number of contiguous elements
* that logically correspond in the layouts of @a a and @a b. This is,
* the number of elements that could reasonably be "vectorized" in the layouts.
* that logically correspond in the layouts of @a a and @a b.
*
* @returns Int<N> with N >= 1
* @post For all 0 <= n < N, a(b[n]) == n (NOTE: Problems with negative strides/coords in this post-condition)
* @post For all 0 <= n < N, a(b.get_1d_coord(n)) == n
* (NOTE: Problems with negative strides/coords in this post-condition)
*/
template <class ShapeA, class StrideA,
class ShapeB, class StrideB>
@@ -1227,8 +1225,7 @@ max_common_vector(Layout<ShapeA,StrideA> const& a,
{
Layout common = coalesce(composition(a, right_inverse(b)));
// NOTE: If one of the layouts is dynamic, we can't prove alignment+vectorization is valid
// We assume dynamic shapes/strides obey alignment requirements (i.e. are large and multiples of the vector)
// Keep only the static identity component of the common layout
if constexpr (is_static<decltype(shape<0>(common))>::value &&
is_constant<1, decltype(stride<0>(common))>::value) {
// Truncate to the size of the contiguous vector (static stride-1 mode)

24
include/cute/swizzle_layout.hpp
View File

@@ -449,6 +449,30 @@ recast_layout(Swizzle<B,M,S> const& swizzle)
// Other operations
//
template <int B, int M, int S, class Offset, class LayoutB, class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
max_common_layout(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& a,
Layout<Shape,Stride> const& b)
{
auto common = max_common_layout(a.layout_b(), b);
auto base = Int<(1 << M)>{};
if constexpr (base < size(common)) {
return common.compose(base); // Truncate common to size base
} else {
return common;
}
}
template <class Shape, class Stride, int B, int M, int S, class Offset, class LayoutB>
CUTE_HOST_DEVICE constexpr
auto
max_common_layout(Layout<Shape,Stride> const& a,
ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& b)
{
return max_common_layout(b, a);
}
template <int B, int M, int S, class Offset, class LayoutB, class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto

45
include/cute/tensor.hpp
View File

@@ -674,7 +674,7 @@ recast(Tensor&& tensor)
// max_common_vector
//
/* Return Int<N> such that N is the maximum number of continguous elements
/* Return Int<N> such that N is the maximum number of contiguous elements
* that logically correspond in the tensors of @a a and @a b. This is,
* the number of elements that could reasonably be vectorized into a single load/store.
*
@@ -682,6 +682,9 @@ recast(Tensor&& tensor)
*
* A return value of Int<0> indicates that no such conclusion can be made and no
* vectorization should be attempted.
*
* Note that the return value does NOT include alignment concerns such as the pointer value and
* the divisbility of dynamic strides.
*/
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -713,6 +716,46 @@ max_common_vector(Tensor<SrcEngine,SrcLayout> const& a,
CUTE_GCC_UNREACHABLE;
}
/* Return a layout that points to the maximum number of contiguous elements
* that logically correspond in the tensors of @a a and @a b. This is,
* the elements that could reasonably be "vectorized" into a single load/store.
*
* @returns Layout R such that composition(a.layout(), R) and composition(b.layout(), R)
* are both identity Layouts.
*
* Note that the returned layout does NOT include alignment concerns such as the pointer value and
* the divisbility of dynamic strides.
*/
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE constexpr
auto
max_common_layout(Tensor<SrcEngine,SrcLayout> const& a,
Tensor<DstEngine,DstLayout> const& b)
{
using SrcType = typename Tensor<SrcEngine,SrcLayout>::value_type;
using DstType = typename Tensor<DstEngine,DstLayout>::value_type;
using SrcRef = typename Tensor<SrcEngine,SrcLayout>::reference;
using DstRef = typename Tensor<SrcEngine,SrcLayout>::reference;
// Determine if vectorization candidates at all
if constexpr (// Should be the same value_types, else the copy is also performing a cast
sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
// The types should be trivially copyable so that vectorization is valid
is_trivially_copyable<SrcType>::value &&
is_trivially_copyable<DstType>::value &&
// Should be load/storing real data, rather than implicit iterators or such
is_reference<SrcRef>::value &&
is_reference<DstRef>::value)
{
return max_common_layout(a.layout(), b.layout());
} else {
return Layout<_1,_0>{};
}
CUTE_GCC_UNREACHABLE;
}
//
// Key algebraic operations -- Divide and Product
//

2
include/cutlass/cluster_launch.hpp
View File

@@ -42,8 +42,8 @@
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#else
#include <cstdio>
#include <type_traits>
#include <cstdio>
#endif
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)))

4
include/cutlass/epilogue/collective/builders/sm90_builder.inl
View File

@@ -154,7 +154,7 @@ sm90_get_smem_store_op_for_accumulator() {
}
else {
// auto-vectorizing store
return DefaultCopy{};
return AutoVectorizingCopyWithAssumedAlignment{};
}
}
@@ -175,7 +175,7 @@ sm90_get_smem_load_op_for_source() {
}
else {
// auto-vectorizing load
return DefaultCopy{};
return AutoVectorizingCopyWithAssumedAlignment{};
}
}

6
include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
View File

@@ -438,7 +438,7 @@ struct Sm90ReLUAuxStore {
using VecType = uint_bit_t<V>;
Tensor tC_rAux_vec = recast<VecType>(tC_rAux);
Tensor tC_gAux_vec = recast<VecType>(tC_gAux);
Tensor tC_cAux_vec = tC_cAux.compose(make_layout(Int<size(tC_rAux_vec)>{}, Int<V>{}));
Tensor tC_cAux_vec = tC_cAux.compose(make_layout(Int<size(tC_rAux_vec)>{}, Int<V>{})); // only works if vector is logically sequential
auto predicate_fn = [&] (auto&&... coords) { return elem_less(tC_cAux_vec(coords...), residue_mn); };
copy_if(FunctionPredTensor(predicate_fn), tC_rAux_vec, tC_gAux_vec);
}
@@ -662,7 +662,7 @@ struct Sm90AuxLoad<
}
if (elem_less(repeat_like(residue_mn, _0{}), residue_mn)) { // (partially) in-bounds CTA tile
copy(tC_gAux, tC_rAux);
copy_aligned(tC_gAux, tC_rAux);
}
}
}
@@ -677,7 +677,7 @@ struct Sm90AuxLoad<
}
if (elem_less(repeat_like(residue_mn, _0{}), residue_mn)) {
copy(tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
copy_aligned(tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
}
}
}

6
include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
View File

@@ -641,8 +641,9 @@ struct Sm90RowBroadcast {
if (epi_m == 0) { // Assumes M-major subtile loop
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
int bcast_pipe_index = (load_iteration / EpiTiles) % Stages;
copy(filter(tCsRow(_,_,_,epi_m,epi_n,bcast_pipe_index)), filter(tCrRow));
copy_aligned(filter(tCsRow(_,_,_,epi_m,epi_n,bcast_pipe_index)), filter(tCrRow));
}
}
@@ -774,7 +775,8 @@ struct Sm90ColBroadcast {
}
// Filter so we don't issue redundant copies over stride-0 modes
copy(filter(tCgCol), filter(tCrCol));
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCgCol), filter(tCrCol));
}
template <typename ElementAccumulator, int FragmentSize>

8
include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp
View File

@@ -915,8 +915,9 @@ public:
using ElementGmem = conditional_t<FinalReduction, ElementCompute volatile, ElementCompute>;
Tensor tCgBuf = sm90_partition_for_epilogue<ReferenceSrc>(gBuf_nl(_,_,n,l), epi_tile, tiled_copy, thread_idx);
if (is_reduced_lane) {
// Filter so we don't issue redunant copies over stride-0 modes
copy(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
}
sync_fn();
}
@@ -934,7 +935,8 @@ public:
Tensor tCsBuf = sm90_partition_for_epilogue<ReferenceSrc>(sBuf(_,_,get<1>(warp_mn)), epi_tile, tiled_copy, thread_idx);
if (is_reduced_lane) {
// Filter so we don't issue redunant copies over stride-0 modes
copy(filter(tCrCol), filter(tCsBuf));
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrCol), filter(tCsBuf));
}
sync_fn();

2
include/cutlass/fast_math.h
View File

@@ -231,7 +231,7 @@ int ceil_div(int a, int b) {
* log2_up/down codes?
*/
template <typename value_t>
CUTLASS_HOST_DEVICE int clz(value_t x) {
CUTLASS_HOST_DEVICE value_t clz(value_t x) {
for (int i = 31; i >= 0; --i) {
if ((1 << i) & x)
return value_t(31 - i);

4
include/cutlass/float8.h
View File

@@ -1,4 +1,4 @@
/**************************************************************************************************
/***************************************************************************************************
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
@@ -217,7 +217,7 @@ struct alignas(1) float8_base {
// Extract the bits in the FP32 type
uint8_t sign = uint8_t((s >> 24 & 0x80));
int32_t exp = int32_t((s >> FP32_NUM_MANTISSA_BITS) & 0xff) - FP32_EXPONENT_BIAS;
uint32_t mantissa = s & 0x7fffff;
int mantissa = s & 0x7fffff;
uint8_t u = 0;
uint8_t const kF8_NaN = 0x7f;

1
include/cutlass/gemm/device/gemm_universal_base.h
View File

@@ -375,6 +375,7 @@ public:
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Static initializers
/////////////////////////////////////////////////////////////////////////////////////////////////

2
include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp
View File

@@ -347,7 +347,7 @@ public:
// The number of tiles for which reduction is required is either:
// (a) the total number of output tiles (in the case of split-K)
// (b) the number of stream-K tiles
// To calcualte the total number of output tiles in the split-K case, we
// To calculate the total number of output tiles in the split-K case, we
// note that, in the split-K case, the units_per_problem_ member of Params will be
// the total number of output tiles.
auto reduction_tiles = params.splits_ > 1 ? params.units_per_problem_ : params.sk_tiles_;