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

Updates for 3.4 release. (#1305)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2024-01-16 13:42:51 -05:00
committed by GitHub
parent acba5beee5
commit 2f589ffa76
166 changed files with 5996 additions and 4702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/atom/copy_atom.hpp
+20 -15
View File
@@ -65,9 +65,9 @@ struct Copy_Atom<Copy_Traits<Args...>, CopyInternalType>
using ValType = CopyInternalType;
using ValLayoutSrc = decltype(upcast<sizeof_bits<ValType>::value>(BitLayoutSrc{}));
using ValLayoutDst = decltype(upcast<sizeof_bits<ValType>::value>(BitLayoutDst{}));
using ValLayoutRef = decltype(upcast<sizeof_bits<ValType>::value>(BitLayoutRef{}));
using ValLayoutSrc = decltype(recast_layout<uint1_t, ValType>(BitLayoutSrc{}));
using ValLayoutDst = decltype(recast_layout<uint1_t, ValType>(BitLayoutDst{}));
using ValLayoutRef = decltype(recast_layout<uint1_t, ValType>(BitLayoutRef{}));
CUTE_STATIC_ASSERT_V(size<0>(ValLayoutSrc{}) == size(ThrID{}), "CopyOperation is not valid for Src of ValType.");
CUTE_STATIC_ASSERT_V(size<0>(ValLayoutDst{}) == size(ThrID{}), "CopyOperation is not valid for Dst of ValType.");
@@ -479,20 +479,24 @@ make_tiled_copy(Copy_Atom<Args...> const& copy_atom,
ThrLayout const& thr_layout = {}, // (m,n) -> thr_idx
ValLayout const& val_layout = {}) // (m,n) -> val_idx
{
constexpr int R = cute::max(rank_v<ThrLayout>, rank_v<ValLayout>);
auto thr_layout_mn = append<R>(thr_layout, Layout<_1>{});
auto val_layout_mn = append<R>(val_layout, Layout<_1>{});
// Take the raked_products to compute the Layout_MN
auto layout_mn = raked_product(thr_layout_mn, val_layout_mn);
// (M,N) -> (thr_idx, val_idx)
auto layout_mn = raked_product(thr_layout, val_layout);
// (thr_idx, val_idx) -> (M,N)
auto layout_tv = right_inverse(layout_mn).with_shape(make_shape(size(thr_layout), size(val_layout)));
// print("thr_layout: "); print(thr_layout_mn); print("\n");
// print("val_layout: "); print(val_layout_mn); print("\n");
// print("layout_mn : "); print(layout_mn); print("\n");
// print("layout_tv : "); print(layout_tv); print("\n");
// Tiler for extracting relevant elements
// (M,N) -> tensor coord
auto tiler = product_each(shape(layout_mn));
return make_tiled_copy_impl(copy_atom, layout_tv, product_each(shape(layout_mn)));
#if 0
print("thr_layout: "); print(thr_layout); print("\n");
print("val_layout: "); print(val_layout); print("\n");
print("layout_mn : "); print(layout_mn); print("\n");
print("layout_tv : "); print(layout_tv); print("\n");
print("tiler : "); print(tiler); print("\n");
#endif
return make_tiled_copy_impl(copy_atom, layout_tv, tiler);
}
/** Produce a TiledCopy from thread and value offset maps.
@@ -622,7 +626,7 @@ print(Copy_Atom<Copy_Traits<Args...>, T> const&)
print(" ValLayoutSrc: "); print(typename Atom::ValLayoutSrc{}); print("\n");
print(" ValLayoutDst: "); print(typename Atom::ValLayoutDst{}); print("\n");
print(" ValLayoutRef: "); print(typename Atom::ValLayoutRef{}); print("\n");
print(" ValueType: %db\n", int(sizeof_bits<typename Atom::ValType>::value));
print(" ValueType: "); print(sizeof_bits<typename Atom::ValType>::value); print("b\n");
}
template <class Atom, class... Args>
@@ -755,6 +759,7 @@ 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>
// Config
#if (__CUDACC_VER_MAJOR__ >= 12)
# define CUTE_COPY_ATOM_TMA_SM90_ENABLED
include/cute/atom/copy_traits_sm90_tma.hpp
+13 -8
View File
@@ -673,15 +673,14 @@ fill_tma_gmem_shape_stride(Tensor<GEngine,GLayout> const& gtensor, /
// Trivial contribution of this gmem mode to this tma mode
auto ej = unwrap(get<i>(tma_gbasis_stride));
gmem_prob_shape[i] = basis_get(ej, gmem_shape);
gmem_prob_stride[i] = basis_get(ej, gmem_stride) * sizeof_bits_v<TmaInternalType> / 8;
gmem_prob_stride[i] = basis_get(ej, gmem_stride);
} else {
// Apply a recurrence to each gmem mode that contributes to this tma mode
for_each(get<i>(tma_gbasis_stride), [&](auto ej) {
// Problem shape
uint64_t shape_j = basis_get(ej, gmem_shape);
// Problem stride (in bytes)
uint64_t stride_j = basis_get(ej, gmem_stride) * sizeof_bits_v<TmaInternalType> / 8;
uint64_t stride_j = basis_get(ej, gmem_stride);
uint64_t old_stride = gmem_prob_stride[i];
gmem_prob_stride[i] = gcd(gmem_prob_stride[i], stride_j);
@@ -764,8 +763,14 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The origin
assert(gmem_prob_shape[4] <= (uint64_t(1) << 32)); // Size must be max 2^32
// TMA descriptor does not store the zeroth stride and assumes it is 1 (TmaInternalType element).
assert(gmem_prob_stride[0] == sizeof(TmaInternalType) && "Majorness of smem doesn't match majorness of gmem");
assert(gmem_prob_stride[0] == 1 && "Majorness of smem doesn't match majorness of gmem");
// convert strides to byte strides
for(uint64_t& stride : gmem_prob_stride) {
stride = (stride * sizeof_bits_v<TmaInternalType>) / 8;
}
// Assert the byte strides. Tma Descriptor uses byte strides
assert((gmem_prob_stride[1]) < (uint64_t(1) << 40)); // Stride must be max 2^40
assert((gmem_prob_stride[1] & 0b1111) == 0); // Stride must be multiple of 16B (128b)
assert((gmem_prob_stride[2]) < (uint64_t(1) << 40)); // Stride must be max 2^40
@@ -866,8 +871,8 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The origin
}
#endif // (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
auto recast_ratio = cute::ratio(Int<sizeof_bits<typename GEngine::value_type>::value>{},
Int<sizeof_bits< TmaInternalType>::value>{});
auto recast_ratio = cute::trait_ratio(sizeof_bits<typename GEngine::value_type>{},
sizeof_bits< TmaInternalType>{});
auto gbasis = make_basis_like(shape(gtensor));
@@ -943,7 +948,7 @@ make_tma_copy_atom(CopyOp,
// Construct the Copy_Traits
//
constexpr int num_bits_per_tma = decltype(size(tma_gbasis))::value * sizeof_bits_v<TmaInternalType>;
constexpr int num_bits_per_tma = size(tma_gbasis) * sizeof_bits<TmaInternalType>::value;
using Traits = Copy_Traits<CopyOp, cute::C<num_bits_per_tma>, decltype(aux_params)>;
using Atom = Copy_Atom<Traits, typename GEngine::value_type>;
@@ -985,7 +990,7 @@ make_tma_copy_tiled(CopyOp const& copy_op,
[[maybe_unused]] auto cta_tiler = product_each(shape(cta_v_map));
auto num_elems_per_tma = size<1>(typename decltype(atom)::RefLayout{}) / Int<sizeof_bits_v<typename GEngine::value_type>>{};
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));
include/cute/atom/mma_atom.hpp
+4 -4
View File
@@ -55,10 +55,10 @@ struct MMA_Atom<MMA_Traits<Args...>>
using Traits = MMA_Traits<Args...>;
// Element value types from the MMA_Traits
using ValTypeD = typename Traits::ElementDVal;
using ValTypeA = typename Traits::ElementAVal;
using ValTypeB = typename Traits::ElementBVal;
using ValTypeC = typename Traits::ElementCVal;
using ValTypeD = typename Traits::ValTypeD;
using ValTypeA = typename Traits::ValTypeA;
using ValTypeB = typename Traits::ValTypeB;
using ValTypeC = typename Traits::ValTypeC;
// Thr-Val layouts from the MMA_Traits
using Shape_MNK = typename Traits::Shape_MNK;
include/cute/atom/mma_traits.hpp
+17 -17
View File
@@ -50,14 +50,14 @@ struct supports_output_scaling<X, void_t<decltype(declval<X>().accumulate_)>> {
/**
* concept MMA_Traits
* {
* using ElementDVal = // Logical A-value type
* using ElementAVal = // Logical B-value type
* using ElementBVal = // Logical C-value type
* using ElementCVal = // Logical D-value type (NOTE: Not used? Assumed == ElementDVal)
* using ValTypeD = // Logical A-value type
* using ValTypeA = // Logical B-value type
* using ValTypeB = // Logical C-value type
* using ValTypeC = // Logical D-value type (NOTE: Not used? Assumed == ValTypeD)
*
* using ElementAFrg = // A-type consumed by MMA (if ommitted, same as ElementAVal)
* using ElementBFrg = // B_type consumed by MMA (if ommitted, same as ElementBVal)
* using ElementCFrg = // C_type consumed by MMA (if ommitted, same as ElementCVal)
* using FrgTypeA = // A-type consumed by MMA (if ommitted, same as ValTypeA)
* using FrgTypeB = // B_type consumed by MMA (if ommitted, same as ValTypeB)
* using FrgTypeC = // C_type consumed by MMA (if ommitted, same as ValTypeC)
*
* using Shape_MNK = // Logical MxNxK shape of the MMA
*
@@ -78,10 +78,10 @@ struct MMA_Traits
template <class D, class A, class B, class C>
struct MMA_Traits<UniversalFMA<D,A,B,C>>
{
using ElementDVal = D;
using ElementAVal = A;
using ElementBVal = B;
using ElementCVal = C;
using ValTypeD = D;
using ValTypeA = A;
using ValTypeB = B;
using ValTypeC = C;
// Logical shape of the MMA
using Shape_MNK = Shape<_1,_1,_1>;
@@ -209,19 +209,19 @@ mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
namespace detail {
template <class X, class = void>
struct FrgTypeA_or_Default { using type = typename X::ElementAVal; };
struct FrgTypeA_or_Default { using type = typename X::ValTypeA; };
template <class X>
struct FrgTypeA_or_Default<X,void_t<typename X::ElementAFrg>> { using type = typename X::ElementAFrg; };
struct FrgTypeA_or_Default<X,void_t<typename X::FrgTypeA>> { using type = typename X::FrgTypeA; };
template <class X, class = void>
struct FrgTypeB_or_Default { using type = typename X::ElementBVal; };
struct FrgTypeB_or_Default { using type = typename X::ValTypeB; };
template <class X>
struct FrgTypeB_or_Default<X,void_t<typename X::ElementBFrg>> { using type = typename X::ElementBFrg; };
struct FrgTypeB_or_Default<X,void_t<typename X::FrgTypeB>> { using type = typename X::FrgTypeB; };
template <class X, class = void>
struct FrgTypeC_or_Default { using type = typename X::ElementCVal; };
struct FrgTypeC_or_Default { using type = typename X::ValTypeC; };
template <class X>
struct FrgTypeC_or_Default<X,void_t<typename X::ElementCFrg>> { using type = typename X::ElementCFrg; };
struct FrgTypeC_or_Default<X,void_t<typename X::FrgTypeC>> { using type = typename X::FrgTypeC; };
} // end namespace detail
include/cute/atom/mma_traits_sm61.hpp
+8 -8
View File
@@ -41,10 +41,10 @@ namespace cute
template <>
struct MMA_Traits<SM61_DP4A>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_1,_1,_4>;
using ThrID = Layout<_1>;
@@ -58,10 +58,10 @@ struct MMA_Traits<SM61_DP4A>
template <>
struct MMA_Traits<SM61_DP2A>
{
using ElementDVal = int32_t;
using ElementAVal = int16_t;
using ElementBVal = int16_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int16_t;
using ValTypeB = int16_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_1,_1,_2>;
using ThrID = Layout<_1>;
include/cute/atom/mma_traits_sm70.hpp
+32 -32
View File
@@ -63,10 +63,10 @@ using SM70_8x8_32b = Layout<Shape <Shape <_2, _2,_2>,Shape <_2,_2, _2>>,
template <>
struct MMA_Traits<SM70_8x8x4_F16F16F16F16_TN>
{
using ElementDVal = half_t;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = half_t;
using ValTypeD = half_t;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -80,10 +80,10 @@ struct MMA_Traits<SM70_8x8x4_F16F16F16F16_TN>
template <>
struct MMA_Traits<SM70_8x8x4_F16F16F16F16_NT>
{
using ElementDVal = half_t;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = half_t;
using ValTypeD = half_t;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -97,10 +97,10 @@ struct MMA_Traits<SM70_8x8x4_F16F16F16F16_NT>
template <>
struct MMA_Traits<SM70_8x8x4_F16F16F16F16_NN>
{
using ElementDVal = half_t;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = half_t;
using ValTypeD = half_t;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -114,10 +114,10 @@ struct MMA_Traits<SM70_8x8x4_F16F16F16F16_NN>
template <>
struct MMA_Traits<SM70_8x8x4_F16F16F16F16_TT>
{
using ElementDVal = half_t;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = half_t;
using ValTypeD = half_t;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -131,10 +131,10 @@ struct MMA_Traits<SM70_8x8x4_F16F16F16F16_TT>
template <>
struct MMA_Traits<SM70_8x8x4_F32F16F16F32_TN>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -148,10 +148,10 @@ struct MMA_Traits<SM70_8x8x4_F32F16F16F32_TN>
template <>
struct MMA_Traits<SM70_8x8x4_F32F16F16F32_NT>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -165,10 +165,10 @@ struct MMA_Traits<SM70_8x8x4_F32F16F16F32_NT>
template <>
struct MMA_Traits<SM70_8x8x4_F32F16F16F32_NN>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
@@ -182,10 +182,10 @@ struct MMA_Traits<SM70_8x8x4_F32F16F16F32_NN>
template <>
struct MMA_Traits<SM70_8x8x4_F32F16F16F32_TT>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
include/cute/atom/mma_traits_sm75.hpp
+8 -8
View File
@@ -41,10 +41,10 @@ namespace cute
template <>
struct MMA_Traits<SM75_16x8x8_F32F16F16F32_TN>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
using Shape_MNK = Shape<_16,_8,_8>;
using ThrID = Layout<_32>;
@@ -61,10 +61,10 @@ struct MMA_Traits<SM75_16x8x8_F32F16F16F32_TN>
template <>
struct MMA_Traits<SM75_8x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_8,_8,_16>;
using ThrID = Layout<_32>;
include/cute/atom/mma_traits_sm80.hpp
+96 -96
View File
@@ -66,10 +66,10 @@ using SM80_16x8_Row = Layout<Shape <Shape < _4,_8>,Shape < _2,_2>>,
template <>
struct MMA_Traits<SM80_16x8x8_F16F16F16F16_TN>
{
using ElementDVal = half_t;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = half_t;
using ValTypeD = half_t;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_16,_8,_8>;
using ThrID = Layout<_32>;
@@ -81,10 +81,10 @@ struct MMA_Traits<SM80_16x8x8_F16F16F16F16_TN>
template <>
struct MMA_Traits<SM80_16x8x16_F16F16F16F16_TN>
{
using ElementDVal = half_t;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = half_t;
using ValTypeD = half_t;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_16,_8,_16>;
using ThrID = Layout<_32>;
@@ -103,20 +103,20 @@ template <>
struct MMA_Traits<SM80_16x8x8_F32F16F16F32_TN>
: MMA_Traits<SM80_16x8x8_F16F16F16F16_TN>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
};
template <>
struct MMA_Traits<SM80_16x8x16_F32F16F16F32_TN>
: MMA_Traits<SM80_16x8x16_F16F16F16F16_TN>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = half_t;
using ValTypeB = half_t;
using ValTypeC = float;
};
///////////////////////////////////////////////////////////////////////////////
@@ -127,20 +127,20 @@ template <>
struct MMA_Traits<SM80_16x8x8_F32BF16BF16F32_TN>
: MMA_Traits<SM80_16x8x8_F16F16F16F16_TN>
{
using ElementDVal = float;
using ElementAVal = bfloat16_t;
using ElementBVal = bfloat16_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = bfloat16_t;
using ValTypeB = bfloat16_t;
using ValTypeC = float;
};
template <>
struct MMA_Traits<SM80_16x8x16_F32BF16BF16F32_TN>
: MMA_Traits<SM80_16x8x16_F16F16F16F16_TN>
{
using ElementDVal = float;
using ElementAVal = bfloat16_t;
using ElementBVal = bfloat16_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = bfloat16_t;
using ValTypeB = bfloat16_t;
using ValTypeC = float;
};
///////////////////////////////////////////////////////////////////////////////
@@ -150,10 +150,10 @@ struct MMA_Traits<SM80_16x8x16_F32BF16BF16F32_TN>
template <>
struct MMA_Traits<SM80_16x8x4_F32TF32TF32F32_TN>
{
using ElementDVal = float;
using ElementAVal = cutlass::tfloat32_t;
using ElementBVal = cutlass::tfloat32_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = cutlass::tfloat32_t;
using ValTypeB = cutlass::tfloat32_t;
using ValTypeC = float;
using Shape_MNK = Shape<_16,_8,_4>;
using ThrID = Layout<_32>;
@@ -166,10 +166,10 @@ struct MMA_Traits<SM80_16x8x4_F32TF32TF32F32_TN>
template <>
struct MMA_Traits<SM80_16x8x8_F32TF32TF32F32_TN>
{
using ElementDVal = float;
using ElementAVal = cutlass::tfloat32_t;
using ElementBVal = cutlass::tfloat32_t;
using ElementCVal = float;
using ValTypeD = float;
using ValTypeA = cutlass::tfloat32_t;
using ValTypeB = cutlass::tfloat32_t;
using ValTypeC = float;
using Shape_MNK = Shape<_16,_8,_8>;
using ThrID = Layout<_32>;
@@ -187,10 +187,10 @@ struct MMA_Traits<SM80_16x8x8_F32TF32TF32F32_TN>
template <>
struct MMA_Traits<SM80_8x8x4_F64F64F64F64_TN>
{
using ElementDVal = double;
using ElementAVal = double;
using ElementBVal = double;
using ElementCVal = double;
using ValTypeD = double;
using ValTypeA = double;
using ValTypeB = double;
using ValTypeC = double;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = Layout<_32>;
@@ -204,10 +204,10 @@ template <>
struct MMA_Traits<SM80_8x8x4_C64C64C64C64_TN>
: MMA_Traits<SM80_8x8x4_F64F64F64F64_TN>
{
using ElementDVal = complex<double>;
using ElementAVal = complex<double>;
using ElementBVal = complex<double>;
using ElementCVal = complex<double>;
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
using ValTypeB = complex<double>;
using ValTypeC = complex<double>;
};
// Custom complex fp64 MMA composed of 3 fp64 MMAs -- same layouts
@@ -215,10 +215,10 @@ template <>
struct MMA_Traits<SM80_8x8x4_GC64C64C64GC64_TN>
: MMA_Traits<SM80_8x8x4_F64F64F64F64_TN>
{
using ElementDVal = typename SM80_8x8x4_GC64C64C64GC64_TN::GaussComplex;
using ElementAVal = complex<double>;
using ElementBVal = complex<double>;
using ElementCVal = typename SM80_8x8x4_GC64C64C64GC64_TN::GaussComplex;
using ValTypeD = typename SM80_8x8x4_GC64C64C64GC64_TN::GaussComplex;
using ValTypeA = complex<double>;
using ValTypeB = complex<double>;
using ValTypeC = typename SM80_8x8x4_GC64C64C64GC64_TN::GaussComplex;
};
///////////////////////////////////////////////////////////////////////////////
@@ -228,10 +228,10 @@ struct MMA_Traits<SM80_8x8x4_GC64C64C64GC64_TN>
template <>
struct MMA_Traits<SM80_8x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_8,_8,_16>;
using ThrID = Layout<_32>;
@@ -247,10 +247,10 @@ struct MMA_Traits<SM80_8x8x16_S32S8S8S32_TN_SATURATE>
template <>
struct MMA_Traits<SM80_16x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16,_8,_16>;
using ThrID = Layout<_32>;
@@ -267,10 +267,10 @@ struct MMA_Traits<SM80_16x8x16_S32S8S8S32_TN_SATURATE>
template <>
struct MMA_Traits<SM80_16x8x32_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16,_8,_32>;
using ThrID = Layout<_32>;
@@ -293,10 +293,10 @@ template <>
struct MMA_Traits<SM80_8x8x16_S32S8U8S32_TN>
: MMA_Traits<SM80_8x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = uint8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = uint8_t;
using ValTypeC = int32_t;
};
template <>
@@ -307,10 +307,10 @@ template <>
struct MMA_Traits<SM80_16x8x16_S32S8U8S32_TN>
: MMA_Traits<SM80_16x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = uint8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = uint8_t;
using ValTypeC = int32_t;
};
template <>
@@ -321,10 +321,10 @@ template <>
struct MMA_Traits<SM80_16x8x32_S32S8U8S32_TN>
: MMA_Traits<SM80_16x8x32_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = int8_t;
using ElementBVal = uint8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = int8_t;
using ValTypeB = uint8_t;
using ValTypeC = int32_t;
};
template <>
@@ -339,10 +339,10 @@ template <>
struct MMA_Traits<SM80_8x8x16_S32U8S8S32_TN>
: MMA_Traits<SM80_8x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = uint8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = uint8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
};
template <>
@@ -353,10 +353,10 @@ template <>
struct MMA_Traits<SM80_16x8x16_S32U8S8S32_TN>
: MMA_Traits<SM80_16x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = uint8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = uint8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
};
template <>
@@ -367,10 +367,10 @@ template <>
struct MMA_Traits<SM80_16x8x32_S32U8S8S32_TN>
: MMA_Traits<SM80_16x8x32_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = uint8_t;
using ElementBVal = int8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = uint8_t;
using ValTypeB = int8_t;
using ValTypeC = int32_t;
};
template <>
@@ -385,10 +385,10 @@ template <>
struct MMA_Traits<SM80_8x8x16_S32U8U8S32_TN>
: MMA_Traits<SM80_8x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = uint8_t;
using ElementBVal = uint8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = uint8_t;
using ValTypeB = uint8_t;
using ValTypeC = int32_t;
};
template <>
@@ -399,10 +399,10 @@ template <>
struct MMA_Traits<SM80_16x8x16_S32U8U8S32_TN>
: MMA_Traits<SM80_16x8x16_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = uint8_t;
using ElementBVal = uint8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = uint8_t;
using ValTypeB = uint8_t;
using ValTypeC = int32_t;
};
template <>
@@ -413,10 +413,10 @@ template <>
struct MMA_Traits<SM80_16x8x32_S32U8U8S32_TN>
: MMA_Traits<SM80_16x8x32_S32S8S8S32_TN>
{
using ElementDVal = int32_t;
using ElementAVal = uint8_t;
using ElementBVal = uint8_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = uint8_t;
using ValTypeB = uint8_t;
using ValTypeC = int32_t;
};
template <>
@@ -430,10 +430,10 @@ struct MMA_Traits<SM80_16x8x32_S32U8U8S32_TN_SATURATE>
template <>
struct MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC>
{
using ElementDVal = int32_t;
using ElementAVal = cute::uint1b_t;
using ElementBVal = cute::uint1b_t;
using ElementCVal = int32_t;
using ValTypeD = int32_t;
using ValTypeA = cute::uint1b_t;
using ValTypeB = cute::uint1b_t;
using ValTypeC = int32_t;
using Shape_MNK = Shape<_16,_8,_256>;
using ThrID = Layout<_32>;
include/cute/atom/mma_traits_sm90.hpp
+24 -24
View File
@@ -44,10 +44,10 @@ namespace cute {
template <>
struct MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
{
using ElementDVal = double;
using ElementAVal = double;
using ElementBVal = double;
using ElementCVal = double;
using ValTypeD = double;
using ValTypeA = double;
using ValTypeB = double;
using ValTypeC = double;
using Shape_MNK = Shape<_16,_8,_4>;
using ThrID = Layout<_32>;
@@ -62,10 +62,10 @@ struct MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
template <>
struct MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
{
using ElementDVal = double;
using ElementAVal = double;
using ElementBVal = double;
using ElementCVal = double;
using ValTypeD = double;
using ValTypeA = double;
using ValTypeB = double;
using ValTypeC = double;
using Shape_MNK = Shape<_16,_8,_8>;
using ThrID = Layout<_32>;
@@ -80,10 +80,10 @@ struct MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
template <>
struct MMA_Traits<SM90_16x8x16_F64F64F64F64_TN>
{
using ElementDVal = double;
using ElementAVal = double;
using ElementBVal = double;
using ElementCVal = double;
using ValTypeD = double;
using ValTypeA = double;
using ValTypeB = double;
using ValTypeC = double;
using Shape_MNK = Shape<_16,_8,_16>;
using ThrID = Layout<_32>;
@@ -103,30 +103,30 @@ template <>
struct MMA_Traits<SM90_16x8x4_C64C64C64C64_TN>
: MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
{
using ElementDVal = complex<double>;
using ElementAVal = complex<double>;
using ElementBVal = complex<double>;
using ElementCVal = complex<double>;
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
using ValTypeB = complex<double>;
using ValTypeC = complex<double>;
};
template <>
struct MMA_Traits<SM90_16x8x8_C64C64C64C64_TN>
: MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
{
using ElementDVal = complex<double>;
using ElementAVal = complex<double>;
using ElementBVal = complex<double>;
using ElementCVal = complex<double>;
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
using ValTypeB = complex<double>;
using ValTypeC = complex<double>;
};
template <>
struct MMA_Traits<SM90_16x8x16_C64C64C64C64_TN>
: MMA_Traits<SM90_16x8x16_F64F64F64F64_TN>
{
using ElementDVal = complex<double>;
using ElementAVal = complex<double>;
using ElementBVal = complex<double>;
using ElementCVal = complex<double>;
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
using ValTypeB = complex<double>;
using ValTypeC = complex<double>;
};
} // end namespace cute
include/cute/atom/mma_traits_sm90_gmma.hpp
+1408 -1408
View File
File diff suppressed because it is too large Load Diff