3.6.0 update (#2005)
* 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
co-authored by
yuzhai
Haicheng Wu
yuzhai
Haicheng Wu
parent
e1cd8c7866
commit
3d261a5974
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user