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

3.6.0 update (#2005)

Browse Source 
* 3.6.0 update

* doc and swap stuff

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-12-25 01:34:40 -05:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
parent e1cd8c7866
commit 3d261a5974
258 changed files with 10863 additions and 3883 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/ampere/CMakeLists.txt
+1 -1
View File
@@ -28,7 +28,7 @@
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_ampere
cp_async.cu
cp_sync.cu
ldsm.cu
cooperative_gemm.cu
cooperative_copy.cu
test/unit/cute/ampere/cooperative_copy.cu
+8 -7
View File
@@ -46,6 +46,7 @@
#include <cute/swizzle.hpp> // cute::Swizzle
#include <cute/swizzle_layout.hpp> // cute::compose(cute::Swizzle)
#include <cute/numeric/numeric_types.hpp>
#include <cute/atom/copy_traits_sm80.hpp>
using namespace cute;
@@ -71,7 +72,7 @@ cooperative_copy_default_gs(T const* g_in, T* g_out, GMemLayout const& gmem_layo
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), gmem_layout);
Tensor s_tensor = make_tensor(make_smem_ptr(smem), smem_layout);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, s_tensor);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, s_tensor, AutoCopyAsync{});
cp_async_fence();
cp_async_wait<0>();
@@ -84,7 +85,7 @@ cooperative_copy_default_gs(T const* g_in, T* g_out, GMemLayout const& gmem_layo
}
__syncthreads();
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s_tensor, g_out_tensor);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s_tensor, g_out_tensor, AutoCopyAsync{});
}
// ss --> shared to shared
@@ -106,7 +107,7 @@ cooperative_copy_default_ss(T const* g_in, T* g_out, Layout1 const& layout1, Lay
Tensor s1_tensor = make_tensor(make_smem_ptr(smem1), layout2);
Tensor s2_tensor = make_tensor(make_smem_ptr(smem2), layout1);
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, g_in_tensor, s1_tensor);
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, g_in_tensor, s1_tensor, AutoCopyAsync{});
cp_async_fence();
cp_async_wait<0>();
@@ -119,10 +120,10 @@ cooperative_copy_default_ss(T const* g_in, T* g_out, Layout1 const& layout1, Lay
}
__syncthreads();
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s1_tensor, s2_tensor);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s1_tensor, s2_tensor, AutoCopyAsync{});
__syncthreads();
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, s2_tensor, g_out_tensor);
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, s2_tensor, g_out_tensor, AutoCopyAsync{});
}
// gg --> global to global
@@ -135,7 +136,7 @@ cooperative_copy_default_gg(T const* g_in, T* g_out, Layout1 const& layout1, Lay
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), layout1);
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), layout2);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, g_out_tensor);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, g_out_tensor, AutoCopyAsync{});
}
template <class Mode, int MaxVecBits, uint32_t ThreadBlockSize, class T, class Layout1, class Layout2>
@@ -252,7 +253,7 @@ typedef testing::Types<
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<128>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<64>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<32>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<16>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<16>>
> CooperativeCopyModeMaxVecBitsList;
TYPED_TEST_SUITE(SM80_CuTe_Ampere, CooperativeCopyModeMaxVecBitsList);
test/unit/cute/ampere/cooperative_gemm.cu
+297 -241
View File
@@ -40,406 +40,462 @@
using namespace cute;
TEST(SM80_CuTe_Ampere, CooperativeGemm1_Half_MMA) {
constexpr uint32_t thread_block_size = 128;
using value_type = cutlass::half_t;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_64, _64, _64>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm2_Double_MMA) {
constexpr uint32_t thread_block_size = 128;
using value_type = double;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_64, _64, _64>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,
Layout<Shape<_2,_2,_1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm3_Half_MMA_CustomSmemLayouts) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using value_type = cutlass::half_t;
constexpr uint32_t m = 128;
constexpr uint32_t n = 128;
constexpr uint32_t k = 128;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_128, _128, _128>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>, // 2x2x1 thread group
Tile<_32, _32, _16> // 32x32x16 MMA for LDSM, 1x2x1 value group`
>;
>{};
using smem_a_atom_layout_t = Layout<Shape<_64, _8>, Stride< _1,_64>>;
using smem_b_atom_layout_t = Layout<Shape< _8,_32>, Stride<_32, _1>>;
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
auto smem_a_atom_layout = Layout<Shape<_64, _8>, Stride< _1,_64>>{};
auto smem_b_atom_layout = Layout<Shape< _8,_32>, Stride<_32, _1>>{};
auto smem_c_atom_layout = make_layout(select<0,1>(shape_mnk));
test_cooperative_gemm_col_major_layout<smem_a_atom_layout_t,
smem_b_atom_layout_t,
smem_c_atom_layout_t,
m,
n,
k,
thread_block_size,
tiled_mma_t,
128,
test_cooperative_gemm_col_major_layout<thread_block_size,
max_vec_bits,
value_type,
value_type,
value_type>();
value_type>
(smem_a_atom_layout,
smem_b_atom_layout,
smem_c_atom_layout,
shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm4_Half_MMA_SwizzledSmemLayouts) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using value_type = cutlass::half_t;
constexpr uint32_t m = 128;
constexpr uint32_t n = 128;
constexpr uint32_t k = 128;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_128, _128, _128>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>, // 2x2x1 thread group
Tile<_32, _32, _16> // 32x32x16 MMA for LDSM, 1x2x1 value group`
>;
>{};
// RowMajor
using smem_rowmajor_atom_layout_t = decltype(
auto smem_a_atom_layout =
composition(Swizzle<3,3,3>{},
Layout<Shape < _8,_64>,
Stride<_64, _1>>{}));
Stride<_64, _1>>{});
// ColMajor
using smem_colmajor_atom_layout_t = decltype(
auto smem_b_atom_layout =
composition(Swizzle<3,3,3>{},
Layout<Shape <_64, _8>,
Stride< _1,_64>>{}));
using smem_a_atom_layout_t = smem_rowmajor_atom_layout_t;
using smem_b_atom_layout_t = smem_colmajor_atom_layout_t;
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{}));
Stride< _1,_64>>{});
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<k> {}), GenRowMajor{}));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n> {}, Int<k> {}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {}), GenRowMajor{}));
auto smem_c_atom_layout = make_layout(select<0, 1>(shape_mnk), GenRowMajor{});
using smem_a_atom_layout_t = smem_a_atom_layout_t;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk), GenRowMajor{});
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk), GenRowMajor{});
using smem_b_atom_layout_t = smem_b_atom_layout_t;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
auto smem_a_layout = tile_to_shape(
smem_a_atom_layout,
make_shape(shape<0>(gmem_a_layout), shape<1>(gmem_a_layout)));
using smem_c_atom_layout_t = smem_c_atom_layout_t;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
auto smem_b_layout = tile_to_shape(
smem_b_atom_layout,
make_shape(shape<0>(gmem_b_layout), shape<1>(gmem_b_layout)));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
SM75_U32x4_LDSM_N, // A
SM75_U16x8_LDSM_T, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
auto smem_c_layout = tile_to_shape(
smem_c_atom_layout,
make_shape(shape<0>(gmem_c_layout), shape<1>(gmem_c_layout)));
test_cooperative_gemm<thread_block_size,
max_vec_bits,
value_type,
value_type,
value_type>();
value_type>
(gmem_a_layout,
gmem_b_layout,
gmem_c_layout,
smem_a_layout,
smem_b_layout,
smem_c_layout,
tiled_mma,
cute::identity{}, // TransformLoadA
cute::identity{}, // TransformLoadB
cute::identity{}, // TransformLoadC
cute::identity{}, // TransformStoreC
SM75_U32x4_LDSM_N{}, // A
SM75_U16x8_LDSM_T{}, // B
AutoVectorizingCopyWithAssumedAlignment<128>{}); // C
}
TEST(SM80_CuTe_Ampere, CooperativeGemm5_Double_MMA_SwizzledSmemLayouts) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using value_type = double;
constexpr uint32_t m = 128;
constexpr uint32_t n = 64;
constexpr uint32_t k = 16;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_128, _64, _16>{};
auto tiled_mma =
TiledMMA<MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>, // Atom
Layout<Shape<_2, _2, _1>>, // Atom layout
Tile<Layout<Shape<_16, _2>, Stride<_2, _1>>, // 32x32x4 MMA with perm for load vectorization
Layout<Shape<_16, _2>, Stride<_2, _1>>,
Underscore>>;
Underscore>>{};
using smem_a_atom_layout_t = decltype(
auto smem_a_atom_layout =
composition(Swizzle<2,2,2>{},
Layout<Shape <_16, _4>,
Stride< _1,_16>>{})); // M, K
using smem_b_atom_layout_t = decltype(
Stride< _1,_16>>{}); // M, K
auto smem_b_atom_layout =
composition(Swizzle<2,2,2>{},
Layout<Shape <_16, _4>,
Stride< _1,_16>>{})); // N, K
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{}));
Stride< _1,_16>>{}); // N, K
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<k> {}), GenRowMajor{}));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n> {}, Int<k> {}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {}), GenRowMajor{}));
auto smem_c_atom_layout = make_layout(select<0, 1>(shape_mnk), GenRowMajor{});
using smem_a_atom_layout_t = smem_a_atom_layout_t;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
using smem_b_atom_layout_t = smem_b_atom_layout_t;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
using smem_c_atom_layout_t = smem_c_atom_layout_t;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk), GenRowMajor{});
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk), GenRowMajor{});
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<128>, // A
AutoVectorizingCopyWithAssumedAlignment<128>, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
auto smem_a_layout = tile_to_shape(
smem_a_atom_layout,
make_shape(shape<0>(gmem_a_layout), shape<1>(gmem_a_layout)));
auto smem_b_layout = tile_to_shape(
smem_b_atom_layout,
make_shape(shape<0>(gmem_b_layout), shape<1>(gmem_b_layout)));
auto smem_c_layout = tile_to_shape(
smem_c_atom_layout,
make_shape(shape<0>(gmem_c_layout), shape<1>(gmem_c_layout)));
test_cooperative_gemm<thread_block_size,
max_vec_bits,
value_type,
value_type,
value_type>();
value_type>
(gmem_a_layout,
gmem_b_layout,
gmem_c_layout,
smem_a_layout,
smem_b_layout,
smem_c_layout,
tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm6_MixedPrecisionFP16FP32_MMA) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using TA = cutlass::half_t;
using TB = cutlass::half_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_64, _64, _64>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32F16F16F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm7_MixedPrecisionBF16FP32_MMA) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using TA = cutlass::bfloat16_t;
using TB = cutlass::bfloat16_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_64, _64, _64>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32BF16BF16F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_MMA) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using TA = cutlass::tfloat32_t;
using TB = cutlass::tfloat32_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = Shape<_64, _64, _64>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm9_C64C64C64_MMA) {
TEST(SM80_CuTe_Ampere, CooperativeGemm9_C64C64C64_MMA_Dynamic) {
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using TA = cutlass::complex<double>;
using TB = cutlass::complex<double>;
using TC = cutlass::complex<double>;
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using tiled_mma_t =
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_8x8x4_C64C64C64C64_TN>,
Layout<Shape<_4, _4, _1>, Stride<_1, _4, _0>>,
Tile<Underscore, Underscore, Underscore>
>;
>{};
using ALayout = Layout<Shape<Int<13>,Int<35>>, Stride<Int<44>, Int<1> >>;
using BLayout = Layout<Shape< Int<7>, Int<35>>, Stride<Int<44>, Int<1> >>;
using CLayout = Layout<Shape<Int<13>, Int<7>>, Stride< Int<1>, Int<30>>>;
auto a_layout = make_layout(Shape<Int<13>,Int<35>>{}, make_stride(44, 1));
auto b_layout = make_layout(Shape< Int<7>, Int<35>>{}, make_stride(44, 1));
auto c_layout = make_layout(Shape<Int<13>, Int<7>>{}, make_stride(1, 30));
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA,
TB,
TC>();
TA, TB, TC>
(a_layout,
b_layout,
c_layout,
a_layout,
b_layout,
c_layout,
tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm9_C64C64C64_MMA) {
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using TA = cutlass::complex<double>;
using TB = cutlass::complex<double>;
using TC = cutlass::complex<double>;
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_8x8x4_C64C64C64C64_TN>,
Layout<Shape<_4, _4, _1>, Stride<_1, _4, _0>>,
Tile<Underscore, Underscore, Underscore>
>{};
auto a_layout = Layout<Shape<Int<13>,Int<35>>, Stride<Int<44>, Int<1> >>{};
auto b_layout = Layout<Shape< Int<7>, Int<35>>, Stride<Int<44>, Int<1> >>{};
auto c_layout = Layout<Shape<Int<13>, Int<7>>, Stride< Int<1>, Int<30>>>{};
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA, TB, TC>
(a_layout,
b_layout,
c_layout,
a_layout,
b_layout,
c_layout,
tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm10_F16F64F16_FMA) {
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using TA = cutlass::half_t;
using TB = double;
using TC = cutlass::half_t;
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using tiled_mma_t =
auto tiled_mma =
TiledMMA<
MMA_Atom<UniversalFMA<half_t, half_t, double, half_t>>,
Layout<Shape<_16, _16, _1>, Stride<_1, _16, _0>>,
Tile<Underscore, Underscore, Underscore>
>;
>{};
using ALayout = Layout<Shape<Int<64>,Int<64>>, Stride<Int<64>, Int< 1>>>;
using BLayout = Layout<Shape<Int<64>,Int<64>>, Stride<Int< 1>, Int<64>>>;
using CLayout = Layout<Shape<Int<64>,Int<64>>, Stride<Int< 1>, Int<64>>>;
auto a_layout = Layout<Shape<Int<64>,Int<64>>, Stride<Int<64>, Int< 1>>>{};
auto b_layout = Layout<Shape<Int<64>,Int<64>>, Stride<Int< 1>, Int<64>>>{};
auto c_layout = Layout<Shape<Int<64>,Int<64>>, Stride<Int< 1>, Int<64>>>{};
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA,
TB,
TC>();
TC>
(a_layout,
b_layout,
c_layout,
a_layout,
b_layout,
c_layout,
tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemmComposedStride) {
using T = cute::half_t;
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 16;
using T = cute::half_t;
using tiled_mma_t =
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<Underscore, Underscore, Underscore>
>;
>{};
using swizzle = cute::Swizzle<3, 3, 3>;
using offset = cute::_0;
using atom_tile_right = decltype(cute::make_layout(cute::Shape<cute::_8, cute::_64>{}, cute::LayoutRight{}));
using FP16AtomLayoutRight = decltype(cute::composition(swizzle{}, offset{}, atom_tile_right{}));
auto swizzle = cute::Swizzle<3, 3, 3>{};
auto offset = cute::_0{};
auto atom_tile_right = cute::make_layout(cute::Shape<cute::_8, cute::_64>{}, cute::LayoutRight{});
auto FP16AtomLayoutRight = cute::composition(swizzle, offset, atom_tile_right);
using shape = cute::Shape<cute::Int<128>, cute::Int<128>>;
using global_a_layout = decltype(cute::make_layout(shape{}, cute::LayoutRight{}));
using global_b_layout = decltype(cute::make_layout(shape{}, cute::LayoutLeft{}));
using global_c_layout = decltype(cute::make_layout(shape{}, cute::LayoutRight{}));
auto shape = cute::Shape<cute::Int<128>, cute::Int<128>>{};
auto global_a_layout = cute::make_layout(shape, cute::LayoutRight{});
auto global_b_layout = cute::make_layout(shape, cute::LayoutLeft{});
auto global_c_layout = cute::make_layout(shape, cute::LayoutRight{});
// This is for A row major, B col major according to CUTLASS default configs
using ALayout = decltype(cute::tile_to_shape(FP16AtomLayoutRight{}, global_a_layout{}));
using BLayout = decltype(cute::tile_to_shape(FP16AtomLayoutRight{}, global_b_layout{}));
using CLayout = global_c_layout;
auto a_layout = cute::tile_to_shape(FP16AtomLayoutRight, global_a_layout);
auto b_layout = cute::tile_to_shape(FP16AtomLayoutRight, global_b_layout);
auto c_layout = global_c_layout;
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
test_cooperative_gemm<thread_block_size,
MaxVecBits,
T,
T,
T>();
T, T, T>
(a_layout,
b_layout,
c_layout,
a_layout,
b_layout,
c_layout,
tiled_mma);
}
TEST(SM89_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_Transform) {
TEST(SM80_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_Transform) {
constexpr uint32_t thread_block_size = 64;
constexpr uint32_t max_vec_bits = 16;
using TA = cutlass::tfloat32_t;
using TB = cutlass::tfloat32_t;
using TC = float;
constexpr uint32_t m = 9;
constexpr uint32_t n = 9;
constexpr uint32_t k = 9;
constexpr uint32_t thread_block_size = 64;
using tiled_mma_t =
auto shape_mnk = Shape<C<9>, C<9>, C<9>>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
Layout<Shape<_1, _2, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 16, TA, TB, TC>(cute::negate{}, cute::negate{}, cute::negate{}, cute::negate{});
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>
(shape_mnk, tiled_mma, cute::negate{}, cute::negate{}, cute::negate{}, cute::negate{});
}
TEST(SM80_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_TransformPrecision) {
constexpr uint32_t thread_block_size = 64;
constexpr uint32_t max_vec_bits = 16;
using InputTA = cutlass::half_t;
using InputTB = cutlass::half_t;
using InputTC = cutlass::half_t;
using ComputeTA = cutlass::tfloat32_t;
using ComputeTB = cutlass::tfloat32_t;
using ComputeTC = float;
auto shape_mnk = Shape<C<9>, C<9>, C<9>>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
Layout<Shape<_1, _2, _1>>
>{};
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, InputTA, InputTB, InputTC>
(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_TransformPrecisionReg) {
constexpr uint32_t thread_block_size = 64;
constexpr uint32_t max_vec_bits = 16;
using InputTA = cutlass::half_t;
using InputTB = cutlass::half_t;
using InputTC = cutlass::half_t;
using ComputeTA = cutlass::tfloat32_t;
using ComputeTB = cutlass::tfloat32_t;
using ComputeTC = float;
auto shape_mnk = Shape<C<9>, C<9>, C<9>>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
Layout<Shape<_1, _2, _1>>
>{};
test_cooperative_gemm_col_major_layout_rmem_c<thread_block_size, max_vec_bits, InputTA, InputTB, InputTC>
(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm1_Half_MMA_Reg) {
using value_type = cutlass::half_t;
auto shape_mnk = Shape<_64, _64, _64>{};
constexpr uint32_t thread_block_size = 128;
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x8_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>
>{};
test_cooperative_gemm_col_major_layout_rmem_c<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm2_Double_MMA_Reg) {
constexpr uint32_t thread_block_size = 128;
using value_type = double;
auto shape_mnk = Shape<_64, _64, _64>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,
Layout<Shape<_2,_2,_1>>
>{};
test_cooperative_gemm_col_major_layout_rmem_c<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemm2_Double_MMA_Predicated_Reg) {
constexpr uint32_t thread_block_size = 128;
using value_type = double;
auto shape_mnk = Shape<C<62>, C<62>, C<62>>{};
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,
Layout<Shape<_2,_2,_1>>
>{};
test_cooperative_gemm_col_major_layout_rmem_c<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
test/unit/cute/ampere/cp_async.cu → test/unit/cute/ampere/cp_sync.cu
+1 -3
View File
@@ -69,14 +69,12 @@ test2(double const* g_in, double* g_out)
copy(g_tensor, s_tensor);
cp_async_fence();
cp_async_wait<0>();
__syncthreads();
g_out[threadIdx.x] = 2 * smem[threadIdx.x];
}
TEST(SM80_CuTe_Ampere, CpAsync)
TEST(SM80_CuTe_Ampere, CpSync)
{
constexpr int count = 32;
thrust::host_vector<double> h_in(count);
test/unit/cute/cooperative_gemm_common.hpp
+646 -308
View File
File diff suppressed because it is too large Load Diff
test/unit/cute/core/inverse_left.cpp
+1 -1
View File
@@ -104,7 +104,7 @@ TEST(CuTe_core, Inverse_left)
auto layout = Layout<Shape <_8, _4>,
Stride<_4, _1>>{};
test_left_inverse(filter(layout));
test_left_inverse(layout);
}
{
test/unit/cute/hopper/cooperative_gemm.cu
+35 -52
View File
@@ -44,91 +44,74 @@ using namespace cute;
#if USE_FP8
TEST(SM90_CuTe_Hopper, CooperativeGemmTilingF8) {
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 16;
using TA = uint8_t;
using TB = uint8_t;
using TC = uint32_t;
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 16;
using tiled_mma_t =
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x32_S32S8S8S32_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<_32, _32, _32>
>;
>{};
using swizzle = Swizzle<2, 4, 3>;
auto swizzle = Swizzle<2, 4, 3>{};
// This is for A row major, B col major according to CUTLASS default configs
using ALayout = decltype(composition(swizzle{}, Layout<Shape<_64, _64>, Stride<_64, _1>>{}));
using BLayout = decltype(composition(swizzle{}, Layout<Shape<_64, _64>, Stride<_1, _64>>{}));
auto a_layout = composition(swizzle, Layout<Shape<_64, _64>, Stride<_64, _1>>{});
auto b_layout = composition(swizzle, Layout<Shape<_64, _64>, Stride<_1, _64>>{});
auto c_layout = make_layout(Shape<_64, _64>{}, LayoutLeft{});
using CLayout = decltype(make_layout(Shape<_64, _64>{}, LayoutLeft{}));
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA,
TB,
TC>();
TA, TB, TC>
(a_layout,
b_layout,
c_layout,
a_layout,
b_layout,
c_layout,
tiled_mma);
}
#else
TEST(SM90_CuTe_Hopper, CooperativeGemmTilingF16) {
constexpr uint32_t thread_block_size = 64;
constexpr int max_vec_bits = 16;
using TA = half_t;
using TB = half_t;
using TC = half_t;
constexpr uint32_t thread_block_size = 64;
constexpr int MaxVecBits = 16;
using tiled_mma_t =
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _1, _1>, Stride<_1, _0, _0>>,
Tile<_32, _32, _32>
>;
using swizzle = Swizzle<3, 3, 3>;
>{};
// This is for A row major, B col major according to CUTLASS default configs
using ALayout = decltype(composition(swizzle{},
Layout<Shape<_64, _64>, Stride<_64, _1>>{}));
auto swizzle = Swizzle<3, 3, 3>{};
auto ALayout = composition(swizzle{}, Layout<Shape<_64, _64>, Stride<_64, _1>>{});
auto BLayout = composition(swizzle{}, Layout<Shape<_64, _64>, Stride<_1, _64>>{});
auto CLayout = make_layout(Shape<_64, _64>{}, LayoutLeft{});
using BLayout = decltype(composition(swizzle{},
Layout<Shape<_64, _64>, Stride<_1, _64>>{}));
using CLayout = decltype(make_layout(Shape<_64, _64>{}, LayoutLeft{}));
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
MaxVecBits,
test_cooperative_gemm<thread_block_size,
max_vec_bits,
TA,
TB,
TC>();
TC>
(ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
tiled_mma);
}
#endif
test/unit/cute/hopper/tma_load_testbed.hpp
+1 -1
View File
@@ -131,7 +131,7 @@ tma_test_device_cute(T const* g_in, T* g_out,
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = sizeof(ArrayEngine<T, size(sA)>);
constexpr int kTmaTransactionBytes = sizeof(ArrayEngine<T, CUTE_STATIC_V(size(filter_zeros(sA)))>);
if (threadIdx.x == 0)
{
test/unit/cute/hopper/tma_mcast_load_testbed.hpp
+1 -1
View File
@@ -146,7 +146,7 @@ tma_test_device_cute(T const* g_in, T* g_out, GmemLayout gmem_layout, SmemLayout
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = sizeof(ArrayEngine<T, size(sA)>);
constexpr int kTmaTransactionBytes = sizeof(ArrayEngine<T, CUTE_STATIC_V(size(filter_zeros(sA)))>);
if (elect_one_thr)
{
test/unit/cute/turing/cooperative_gemm.cu
+7 -9
View File
@@ -38,21 +38,19 @@
using namespace cute;
TEST(SM75_CuTe_Turing, CooperativeGemm1_MixedPrecisionFP16FP32_MMA) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using TA = cutlass::half_t;
using TB = cutlass::half_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = make_shape(_64{}, _64{}, _64{});
auto tiled_mma =
TiledMMA<
MMA_Atom<SM75_16x8x8_F32F16F16F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>(shape_mnk, tiled_mma);
}
test/unit/cute/volta/cooperative_gemm.cu
+138 -199
View File
@@ -40,105 +40,85 @@
using namespace cute;
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA) {
using value_type = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 32;
constexpr uint32_t k = 16;
constexpr uint32_t thread_block_size = 128;
using value_type = float;
using tiled_mma_t =
auto shape_mnk = make_shape(_64{}, _32{}, _16{});
auto tiled_mma =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_16, _8, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA_Predication) {
using value_type = float;
constexpr uint32_t m = 88;
constexpr uint32_t n = 20;
constexpr uint32_t k = 12;
constexpr uint32_t thread_block_size = 128;
using value_type = float;
using tiled_mma_t =
auto shape_mnk = make_shape(C<88>{}, C<20>{}, C<12>{});
auto tiled_mma =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_2, _64, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA_Predication2) {
using value_type = float;
constexpr uint32_t m = 88;
constexpr uint32_t n = 36;
constexpr uint32_t k = 24;
constexpr uint32_t thread_block_size = 128;
using value_type = float;
using tiled_mma_t =
auto shape_mnk = make_shape(C<88>{}, C<36>{}, C<24>{});
auto tiled_mma =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_4, _32, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA_Predication3) {
constexpr uint32_t thread_block_size = 128;
using value_type = float;
constexpr uint32_t m = 67;
constexpr uint32_t n = 13;
constexpr uint32_t k = 11;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = make_shape(C<67>{}, C<13>{}, C<11>{});
auto tiled_mma =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_1, _128, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm2_DoubleFMA) {
constexpr uint32_t thread_block_size = 128;
using value_type = double;
constexpr uint32_t m = 16;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
auto shape_mnk = make_shape(C<16>{}, C<32>{}, C<32>{});
auto tiled_mma =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_16, _8, _1>>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm3_Float_FMA_CustomPermutationMNK) {
using value_type = float;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 256;
using value_type = float;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_32{}, _32{}, _32{});
auto tiled_mma = TiledMMA<
MMA_Atom<
UniversalFMA<value_type, value_type, value_type, value_type>
>,
@@ -154,228 +134,188 @@ TEST(SM70_CuTe_Volta, CooperativeGemm3_Float_FMA_CustomPermutationMNK) {
>,
Underscore
>
>;
>{};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm4_Half_MMA) {
constexpr uint32_t thread_block_size = 128;
using value_type = cutlass::half_t;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_32{}, _32{}, _32{});
auto tiled_mma = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
>{};
using smem_a_atom_layout_t = typename tiled_mma_t::AtomLayoutB_TV;
using smem_b_atom_layout_t = typename tiled_mma_t::AtomLayoutA_TV;
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {})));
auto smem_a_atom_layout = typename decltype(tiled_mma)::AtomLayoutB_TV{};
auto smem_b_atom_layout = typename decltype(tiled_mma)::AtomLayoutA_TV{};
auto smem_c_atom_layout = make_layout(select<0, 1>(shape_mnk));
test_cooperative_gemm_col_major_layout<smem_a_atom_layout_t,
smem_b_atom_layout_t,
smem_c_atom_layout_t,
m,
n,
k,
thread_block_size,
tiled_mma_t,
value_type>();
test_cooperative_gemm_col_major_layout<thread_block_size,
value_type>
(smem_a_atom_layout,
smem_b_atom_layout,
smem_c_atom_layout,
shape_mnk,
tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm5_Half_MMA) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using value_type = cutlass::half_t;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_32{}, _32{}, _32{});
auto tiled_mma = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
>{};
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk));
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk));
using smem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using smem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using smem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
auto smem_a_layout = make_layout(select<0, 2>(shape_mnk));
auto smem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto smem_c_layout = make_layout(select<0, 1>(shape_mnk));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<128>, // A
AutoVectorizingCopyWithAssumedAlignment<128>, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
test_cooperative_gemm<thread_block_size,
max_vec_bits,
value_type,
value_type,
value_type>();
value_type>
(gmem_a_layout,
gmem_b_layout,
gmem_c_layout,
smem_a_layout,
smem_b_layout,
smem_c_layout,
tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm5_Half_MMA_Predicated) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 31;
constexpr uint32_t n = 27;
constexpr uint32_t k = 17;
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 16;
using value_type = cutlass::half_t;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(C<31>{}, C<27>{}, C<17>{});
auto tiled_mma = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
>{};
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk));
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk));
using smem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using smem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using smem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
auto smem_a_layout = make_layout(select<0, 2>(shape_mnk));
auto smem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto smem_c_layout = make_layout(select<0, 1>(shape_mnk));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<16>, // A
AutoVectorizingCopyWithAssumedAlignment<16>, // B
AutoVectorizingCopyWithAssumedAlignment<16>, // C
thread_block_size,
tiled_mma_t,
16,
test_cooperative_gemm<thread_block_size,
max_vec_bits,
value_type,
value_type,
value_type>();
value_type>
(gmem_a_layout,
gmem_b_layout,
gmem_c_layout,
smem_a_layout,
smem_b_layout,
smem_c_layout,
tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm6_Half_MAA_SwizzledSmemLayouts) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 128;
constexpr uint32_t n = 128;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 128;
using value_type = cutlass::half_t;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_128{}, _128{}, _64{});
auto tiled_mma = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
>{};
using smem_a_atom_layout_t = decltype(
composition(Swizzle<3,3,3>{},
Layout<Shape < _8,_64>,
Stride<_64, _1>>{}));
using smem_b_atom_layout_t = decltype(
composition(Swizzle<3,3,3>{},
Layout<Shape <_64, _8>,
Stride< _1,_64>>{}));
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{}));
auto smem_a_atom_layout = composition(Swizzle<3,3,3>{}, Layout<Shape < _8,_64>, Stride<_64, _1>>{});
auto smem_b_atom_layout = composition(Swizzle<3,3,3>{}, Layout<Shape <_64, _8>, Stride< _1,_64>>{});
auto smem_c_atom_layout = make_layout(select<0, 1>(shape_mnk), GenRowMajor{});
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<k> {}), GenRowMajor{}));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n> {}, Int<k> {}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {}), GenRowMajor{}));
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk), GenRowMajor{});
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenColMajor{});
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk), GenRowMajor{});
using smem_a_atom_layout_t = smem_a_atom_layout_t;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
auto smem_a_layout = tile_to_shape(
smem_a_atom_layout,
make_shape(shape<0>(gmem_a_layout), shape<1>(gmem_a_layout)));
// Transposed
using smem_b_atom_layout_t = smem_b_atom_layout_t;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
auto smem_b_layout = tile_to_shape(
smem_b_atom_layout,
make_shape(shape<0>(gmem_b_layout), shape<1>(gmem_b_layout)));
using smem_c_atom_layout_t = smem_c_atom_layout_t;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
auto smem_c_layout = tile_to_shape(
smem_c_atom_layout,
make_shape(shape<0>(gmem_c_layout), shape<1>(gmem_c_layout)));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<128>, // A
AutoVectorizingCopyWithAssumedAlignment<128>, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
test_cooperative_gemm<thread_block_size,
max_vec_bits,
value_type,
value_type,
value_type>();
value_type>
(gmem_a_layout,
gmem_b_layout,
gmem_c_layout,
smem_a_layout,
smem_b_layout,
smem_c_layout,
tiled_mma);
}
TEST(SM70_CuTe_Volta, CooperativeGemm7_TransformNegate_FMA) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 64;
using TA = float;
using TB = float;
using TC = double;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_32{}, _32{}, _32{});
auto tiled_mma = TiledMMA<
MMA_Atom<UniversalFMA<TC, TA, TB, TC>>,
Layout<Shape<_16, _8, _1>>
>;
>{};
auto aload = cute::negate {};
auto bload = cute::negate {};
auto cload = cute::negate {};
auto cstore = cute::negate {};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 64, TA, TB, TC>(
aload, bload, cload, cstore);
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>(
shape_mnk, tiled_mma, aload, bload, cload, cstore);
}
TEST(SM70_CuTe_Volta, CooperativeGemm7_TransformNegate_MMA) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using value_type = cutlass::half_t;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_32{}, _32{}, _32{});
auto tiled_mma = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
>{};
auto aload = cute::negate {};
auto bload = cute::negate {};
auto cload = cute::negate {};
auto cstore = cute::negate {};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>(
aload, bload, cload, cstore);
test_cooperative_gemm_col_major_layout<thread_block_size, value_type>(
shape_mnk, tiled_mma, aload, bload, cload, cstore);
}
template<class ConstantType>
@@ -398,26 +338,25 @@ struct convert_to {
};
TEST(SM70_CuTe_Volta, CooperativeGemm7_TransformCustomOp_FMA) {
constexpr uint32_t thread_block_size = 128;
constexpr uint32_t max_vec_bits = 64;
using TA = float;
using TB = float;
using TC = double;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
auto shape_mnk = make_shape(_32{}, _32{}, _32{});
auto tiled_mma = TiledMMA<
MMA_Atom<UniversalFMA<TC, TA, TB, TC>>,
Layout<Shape<_16, _8, _1>>
>;
>{};
auto aload = increment_by_x<float>{1.111f};
auto bload = convert_to<float, double> {};
auto cload = cute::negate {};
auto cstore = cute::negate {};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 64, TA, TB, TC>(
aload, bload, cload, cstore);
test_cooperative_gemm_col_major_layout<thread_block_size, max_vec_bits, TA, TB, TC>(
shape_mnk, tiled_mma, aload, bload, cload, cstore);
}
test/unit/cute/volta/vectorization_auto.cu
-1
View File
@@ -67,7 +67,6 @@ kernel(GmemTensor gC, RmemTiler tiler, CopyPolicy policy)
// NOTE: only 1 thread, this thread produce a block of 8x8 output. The fringe will not be touched.
//copy(rC, tCgC); // Enable auto-vectorization if static
//copy_vec<T>(rC, tCgC); // Disable auto-vectorization always
copy(policy, rC, tCgC); // Use a policy to establish vectorization assumptions
}