Skip to content

GitLab

"...pruning/git@developer.sourcefind.cn:OpenDAS/nni.git" did not exist on "ef4459e4d82a090580b4676a4e2d3f0960ff3803"
Unverified Commit fe70549f authored by FeiyangChen's avatar FeiyangChen Committed by GitHub
Browse files

[Feat] Support mma gemm with stride (#701) 



* gemm_with_stride sm89

* fix offset issue

* bug fix

* format

* sm80 support

* add sm90

* add testing

* format

* add static_assert for wgmma

* Enhance error message for inner_box_dim validation in LowerBulkCopy

* lint fix

---------
Co-authored-by: default avatarLeiWang1999 <leiwang1999@outlook.com>
parent 569b0127
Show whitespace changes
Inline Side-by-side
Showing with 556 additions and 251 deletions
src/op/bulk_copy.cc
View file @ fe70549f
...@@ -273,7 +273,9 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const { ...@@ -273,7 +273,9 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
<< "inner_box_dim: " << *inner_box_dim << " is not divisible by 256"; << "inner_box_dim: " << *inner_box_dim << " is not divisible by 256";
instruction_dim = 256; instruction_dim = 256;
} }
ICHECK((*inner_box_dim) % instruction_dim == 0); ICHECK((*inner_box_dim) % instruction_dim == 0)
<< "inner_box_dim: " << *inner_box_dim
<< " is not divisible by instruction_dim: " << instruction_dim;
desc.smem_box.Set(0, PrimExpr(instruction_dim)); desc.smem_box.Set(0, PrimExpr(instruction_dim));
int inner_box_dim_ = instruction_dim * shared_tensor->dtype.bytes(); int inner_box_dim_ = instruction_dim * shared_tensor->dtype.bytes();
......
src/op/gemm.cc
View file @ fe70549f
...@@ -47,14 +47,18 @@ Gemm::Gemm(Array<PrimExpr> args, BufferMap vmap) { ...@@ -47,14 +47,18 @@ Gemm::Gemm(Array<PrimExpr> args, BufferMap vmap) {
K = args[7].as<IntImm>().value()->value; K = args[7].as<IntImm>().value()->value;
policy = static_cast<GemmWarpPolicy>(args[8].as<IntImm>().value()->value); policy = static_cast<GemmWarpPolicy>(args[8].as<IntImm>().value()->value);
clear_accum = args[9].as<Bool>().value(); clear_accum = args[9].as<Bool>().value();
if (args.size() > 10) { stride_A = args[10].as<IntImm>().value()->value;
kPack = args[10].as<IntImm>().value()->value; stride_B = args[11].as<IntImm>().value()->value;
offset_A = args[12].as<IntImm>().value()->value;
offset_B = args[13].as<IntImm>().value()->value;
if (args.size() > 14) {
kPack = args[14].as<IntImm>().value()->value;
if (kPack != 1 && kPack != 2) { if (kPack != 1 && kPack != 2) {
ICHECK(false) << "kPack must be 1 or 2"; ICHECK(false) << "kPack must be 1 or 2";
} }
} }
if (args.size() > 11) { if (args.size() > 15) {
wg_wait = args[11].as<IntImm>().value()->value; wg_wait = args[15].as<IntImm>().value()->value;
} }
} }
...@@ -284,6 +288,19 @@ bool Gemm::CheckWGMMA() const { ...@@ -284,6 +288,19 @@ bool Gemm::CheckWGMMA() const {
} }
} }
static int GetArchInt(Target target) {
int arch_int = 0;
auto s = target->GetAttr<String>("arch");
ICHECK(s.defined());
const char *arch_str = s.value().c_str();
if (arch_str[0] == 's' && arch_str[1] == 'm' && arch_str[2] == '_') {
arch_int = atoi(&arch_str[3]);
} else {
arch_int = 0;
}
return arch_int;
}
Stmt Gemm::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const { Stmt Gemm::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
auto block_size = *as_const_int(T.thread_bounds->extent); auto block_size = *as_const_int(T.thread_bounds->extent);
GemmInst gemm_inst = GetGemmInst(block_size, T.target); GemmInst gemm_inst = GetGemmInst(block_size, T.target);
...@@ -301,6 +318,10 @@ Stmt Gemm::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const { ...@@ -301,6 +318,10 @@ Stmt Gemm::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
ss << warp_m << ", " << warp_n << ", "; ss << warp_m << ", " << warp_n << ", ";
ss << trans_A << ", " << trans_B; ss << trans_A << ", " << trans_B;
ss << ", " << clear_accum; ss << ", " << clear_accum;
if (TargetIsCuda(T.target) && (GetArchInt(T.target) >= 75)) {
ss << ", " << stride_A << ", " << stride_B;
ss << ", " << offset_A << ", " << offset_B;
}
if (TargetIsCDNA(T.target)) { if (TargetIsCDNA(T.target)) {
// for cdna gemm, we need to specify kPack // for cdna gemm, we need to specify kPack
ss << ", " << kPack; ss << ", " << kPack;
......
src/op/gemm.h
View file @ fe70549f
...@@ -45,6 +45,8 @@ private: ...@@ -45,6 +45,8 @@ private:
PrimExpr Aptr, Bptr, Cptr; PrimExpr Aptr, Bptr, Cptr;
bool trans_A, trans_B; bool trans_A, trans_B;
int M, N, K; int M, N, K;
int stride_A, stride_B;
int offset_A, offset_B;
bool clear_accum = false; bool clear_accum = false;
// k_pack please ref to bitblas/tl/mfma_macro_generator.py::k_pack // k_pack please ref to bitblas/tl/mfma_macro_generator.py::k_pack
// only will be enabled under cdna mfma instructions // only will be enabled under cdna mfma instructions
......
src/tl_templates/cuda/gemm_sm80.h
View file @ fe70549f
...@@ -73,135 +73,143 @@ template <int N, int num_warp_n, bool transpose> struct SelectCopy { ...@@ -73,135 +73,143 @@ template <int N, int num_warp_n, bool transpose> struct SelectCopy {
DefaultCopy>; DefaultCopy>;
}; };
template <int Bits, int N, int K, bool K_inner, int num_warp_n, template <int Bits, int N, int K, bool K_inner, int num_warp_n, int leading_dim,
typename Enable = void> typename Enable = void>
struct OperandTraits { struct OperandTraits {
// Primary template, use padded layout and default copy // Primary template, use padded layout and default copy
static constexpr int stride = K_inner ? K : N; static constexpr int stride = leading_dim;
static constexpr int padded = static constexpr int padded =
stride % (256 / Bits) == 0 ? stride + 128 / Bits : stride; stride % (256 / Bits) == 0 ? stride + 128 / Bits : stride;
using Layout = typename std::conditional< using Layout = typename std::conditional<
K_inner, Layout<Shape<Int<N>, Int<K>>, Shape<Int<padded>, _1>>, K_inner, Layout<Shape<Int<N>, Int<leading_dim>>, Shape<Int<padded>, _1>>,
Layout<Shape<Int<N>, Int<K>>, Shape<_1, Int<padded>>>>::type; Layout<Shape<Int<leading_dim>, Int<K>>, Shape<_1, Int<padded>>>>::type;
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, true, num_warp_n, struct OperandTraits<16, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 64 == 32>::type> { typename std::enable_if<leading_dim % 64 == 32>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{})); Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, true, num_warp_n, struct OperandTraits<16, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 64 == 0>::type> { typename std::enable_if<leading_dim % 64 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{})); Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, false, num_warp_n, struct OperandTraits<16, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 64 == 32>::type> { typename std::enable_if<leading_dim % 64 == 32>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{})); Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = typename SelectCopy<N, num_warp_n, false>::type; using Copy = typename SelectCopy<N, num_warp_n, false>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, false, num_warp_n, struct OperandTraits<16, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 64 == 0>::type> { typename std::enable_if<leading_dim % 64 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{})); Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = typename SelectCopy<N, num_warp_n, false>::type; using Copy = typename SelectCopy<N, num_warp_n, false>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, true, num_warp_n, struct OperandTraits<32, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 32 == 0>::type> { typename std::enable_if<leading_dim % 32 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{})); Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, true, num_warp_n, struct OperandTraits<32, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 32 == 16>::type> { typename std::enable_if<leading_dim % 32 == 16>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{})); Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, false, num_warp_n, struct OperandTraits<32, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 32 == 0>::type> { typename std::enable_if<leading_dim % 32 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 2, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{})); Swizzle<3, 2, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = UniversalCopy<tfloat32_t>; using Copy = UniversalCopy<tfloat32_t>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, false, num_warp_n, struct OperandTraits<32, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 32 == 16>::type> { typename std::enable_if<leading_dim % 32 == 16>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 3>{}, Layout<Shape<_16, _8>, Stride<_1, _16>>{})); Swizzle<2, 2, 3>{}, Layout<Shape<_16, _8>, Stride<_1, _16>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = UniversalCopy<tfloat32_t>; using Copy = UniversalCopy<tfloat32_t>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<8, N, K, true, num_warp_n, struct OperandTraits<8, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 128 == 64>::type> { typename std::enable_if<leading_dim % 128 == 64>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{})); Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<8, N, K, true, num_warp_n, struct OperandTraits<8, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 128 == 0>::type> { typename std::enable_if<leading_dim % 128 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{})); Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<64, N, K, true, num_warp_n, struct OperandTraits<64, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 16 == 0>::type> { typename std::enable_if<leading_dim % 16 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 0, 4>{}, Layout<Shape<_4, _16>, Stride<_16, _1>>{})); Swizzle<2, 0, 4>{}, Layout<Shape<_4, _16>, Stride<_16, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<64, N, K, false, num_warp_n, struct OperandTraits<64, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 16 == 0>::type> { typename std::enable_if<leading_dim % 16 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 2>{}, Layout<Shape<_16, _4>, Stride<_1, _16>>{})); Swizzle<2, 2, 2>{}, Layout<Shape<_16, _4>, Stride<_1, _16>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type_raw, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename B_type_raw, typename C_type_raw> int offset_b, typename A_type_raw, typename B_type_raw,
typename C_type_raw>
class GemmTensorOp { class GemmTensorOp {
public: public:
using A_type = using A_type =
...@@ -215,10 +223,10 @@ public: ...@@ -215,10 +223,10 @@ public:
using Instruction = using Instruction =
DispatchInstruction<A_type, B_type, C_type, num_warp_m, num_warp_n, N>; DispatchInstruction<A_type, B_type, C_type, num_warp_m, num_warp_n, N>;
using OperandATraits = using OperandATraits = OperandTraits<sizeof_bits<A_type>::value, M, K,
OperandTraits<sizeof_bits<A_type>::value, M, K, !trans_A, num_warp_m>; !trans_A, num_warp_m, lda>;
using OperandBTraits = using OperandBTraits =
OperandTraits<sizeof_bits<B_type>::value, N, K, trans_B, num_warp_n>; OperandTraits<sizeof_bits<B_type>::value, N, K, trans_B, num_warp_n, ldb>;
using SmemLayoutA = typename OperandATraits::Layout; using SmemLayoutA = typename OperandATraits::Layout;
using SmemLayoutB = typename OperandBTraits::Layout; using SmemLayoutB = typename OperandBTraits::Layout;
...@@ -244,12 +252,38 @@ public: ...@@ -244,12 +252,38 @@ public:
return layout; return layout;
} }
template <int offset, int NN, int KK, bool trans, int lddim, typename Engine0,
typename Layout0>
static CUTE_DEVICE auto get_region_tensor(Tensor<Engine0, Layout0> &sa) {
if constexpr (offset == 0) {
return composition(
sa,
Layout<Shape<Int<NN>, Int<KK>>,
Stride<_1, typename std::conditional<trans, Int<NN>,
Int<lddim>>::type>>{});
} else {
if constexpr (trans) {
static_assert(offset % KK == 0, "Offset must be a multiple of K");
constexpr int offset_n = offset / KK;
return flat_divide(sa, Shape<Int<NN>, Int<KK>>{})(_, _, _0{},
Int<offset_n>{});
} else {
static_assert(offset % NN == 0, "Offset must be a multiple of N");
constexpr int offset_n = offset / NN;
return flat_divide(sa, Shape<Int<NN>, Int<KK>>{})(_, _, Int<offset_n>{},
_0{});
}
}
}
static CUTE_DEVICE void body(A_type_raw *pA, B_type_raw *pB, C_type_raw *pC) { static CUTE_DEVICE void body(A_type_raw *pA, B_type_raw *pB, C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sA = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)), Tensor sA_all = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)),
SmemLayoutA{}); SmemLayoutA{});
Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)), Tensor sB_all = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)),
SmemLayoutB{}); SmemLayoutB{});
Tensor sA = get_region_tensor<offset_a, M, K, !trans_A, lda>(sA_all);
Tensor sB = get_region_tensor<offset_b, N, K, trans_B, ldb>(sB_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma); auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma);
...@@ -287,8 +321,9 @@ public: ...@@ -287,8 +321,9 @@ public:
static CUTE_DEVICE void body_rs(A_type_raw *pA, B_type_raw *pB, static CUTE_DEVICE void body_rs(A_type_raw *pA, B_type_raw *pB,
C_type_raw *pC) { C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)), Tensor sB_all = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)),
SmemLayoutB{}); SmemLayoutB{});
Tensor sB = get_region_tensor<offset_b, N, K, trans_B, ldb>(sB_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_B = make_tiled_copy_B(SmemCopyB{}, tiled_mma); auto tiled_copy_B = make_tiled_copy_B(SmemCopyB{}, tiled_mma);
...@@ -322,8 +357,9 @@ public: ...@@ -322,8 +357,9 @@ public:
static CUTE_DEVICE void body_sr(A_type_raw *pA, B_type_raw *pB, static CUTE_DEVICE void body_sr(A_type_raw *pA, B_type_raw *pB,
C_type_raw *pC) { C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sA = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)), Tensor sA_all = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)),
SmemLayoutA{}); SmemLayoutA{});
Tensor sA = get_region_tensor<offset_a, M, K, !trans_A, lda>(sA_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma); auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma);
...@@ -360,29 +396,32 @@ public: ...@@ -360,29 +396,32 @@ public:
namespace tl { namespace tl {
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) {
using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body(pA, pB, accum); MMA::body(pA, pB, accum);
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) {
using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_rs(pA, pB, accum); MMA::body_rs(pA, pB, accum);
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_sr(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_sr(A_type *pA, B_type *pB, C_type *accum) {
using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_sr(pA, pB, accum); MMA::body_sr(pA, pB, accum);
} }
......
src/tl_templates/cuda/gemm_sm89.h
View file @ fe70549f
...@@ -91,135 +91,143 @@ template <int N, int num_warp_n, bool transpose> struct SelectCopy { ...@@ -91,135 +91,143 @@ template <int N, int num_warp_n, bool transpose> struct SelectCopy {
DefaultCopy>; DefaultCopy>;
}; };
template <int Bits, int N, int K, bool K_inner, int num_warp_n, template <int Bits, int N, int K, bool K_inner, int num_warp_n, int leading_dim,
typename Enable = void> typename Enable = void>
struct OperandTraits { struct OperandTraits {
// Primary template, use padded layout and default copy // Primary template, use padded layout and default copy
static constexpr int stride = K_inner ? K : N; static constexpr int stride = leading_dim;
static constexpr int padded = static constexpr int padded =
stride % (256 / Bits) == 0 ? stride + 128 / Bits : stride; stride % (256 / Bits) == 0 ? stride + 128 / Bits : stride;
using Layout = typename std::conditional< using Layout = typename std::conditional<
K_inner, Layout<Shape<Int<N>, Int<K>>, Shape<Int<padded>, _1>>, K_inner, Layout<Shape<Int<N>, Int<leading_dim>>, Shape<Int<padded>, _1>>,
Layout<Shape<Int<N>, Int<K>>, Shape<_1, Int<padded>>>>::type; Layout<Shape<Int<leading_dim>, Int<K>>, Shape<_1, Int<padded>>>>::type;
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, true, num_warp_n, struct OperandTraits<16, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 64 == 32>::type> { typename std::enable_if<leading_dim % 64 == 32>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{})); Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, true, num_warp_n, struct OperandTraits<16, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 64 == 0>::type> { typename std::enable_if<leading_dim % 64 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{})); Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, false, num_warp_n, struct OperandTraits<16, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 64 == 32>::type> { typename std::enable_if<leading_dim % 64 == 32>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{})); Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = typename SelectCopy<N, num_warp_n, false>::type; using Copy = typename SelectCopy<N, num_warp_n, false>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, false, num_warp_n, struct OperandTraits<16, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 64 == 0>::type> { typename std::enable_if<leading_dim % 64 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{})); Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = typename SelectCopy<N, num_warp_n, false>::type; using Copy = typename SelectCopy<N, num_warp_n, false>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, true, num_warp_n, struct OperandTraits<32, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 32 == 0>::type> { typename std::enable_if<leading_dim % 32 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{})); Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, true, num_warp_n, struct OperandTraits<32, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 32 == 16>::type> { typename std::enable_if<leading_dim % 32 == 16>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{})); Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, false, num_warp_n, struct OperandTraits<32, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 32 == 0>::type> { typename std::enable_if<leading_dim % 32 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 2, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{})); Swizzle<3, 2, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = UniversalCopy<tfloat32_t>; using Copy = UniversalCopy<tfloat32_t>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, false, num_warp_n, struct OperandTraits<32, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 32 == 16>::type> { typename std::enable_if<leading_dim % 32 == 16>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 3>{}, Layout<Shape<_16, _8>, Stride<_1, _16>>{})); Swizzle<2, 2, 3>{}, Layout<Shape<_16, _8>, Stride<_1, _16>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = UniversalCopy<tfloat32_t>; using Copy = UniversalCopy<tfloat32_t>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<8, N, K, true, num_warp_n, struct OperandTraits<8, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 128 == 64>::type> { typename std::enable_if<leading_dim % 128 == 64>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{})); Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<8, N, K, true, num_warp_n, struct OperandTraits<8, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 128 == 0>::type> { typename std::enable_if<leading_dim % 128 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{})); Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<64, N, K, true, num_warp_n, struct OperandTraits<64, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 16 == 0>::type> { typename std::enable_if<leading_dim % 16 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 0, 4>{}, Layout<Shape<_4, _16>, Stride<_16, _1>>{})); Swizzle<2, 0, 4>{}, Layout<Shape<_4, _16>, Stride<_16, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<64, N, K, false, num_warp_n, struct OperandTraits<64, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 16 == 0>::type> { typename std::enable_if<leading_dim % 16 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 2>{}, Layout<Shape<_16, _4>, Stride<_1, _16>>{})); Swizzle<2, 2, 2>{}, Layout<Shape<_16, _4>, Stride<_1, _16>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type_raw, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename B_type_raw, typename C_type_raw> int offset_b, typename A_type_raw, typename B_type_raw,
typename C_type_raw>
class GemmTensorOp { class GemmTensorOp {
public: public:
using A_type = using A_type =
...@@ -233,10 +241,10 @@ public: ...@@ -233,10 +241,10 @@ public:
using Instruction = using Instruction =
DispatchInstruction<A_type, B_type, C_type, num_warp_m, num_warp_n, N>; DispatchInstruction<A_type, B_type, C_type, num_warp_m, num_warp_n, N>;
using OperandATraits = using OperandATraits = OperandTraits<sizeof_bits<A_type>::value, M, K,
OperandTraits<sizeof_bits<A_type>::value, M, K, !trans_A, num_warp_m>; !trans_A, num_warp_m, lda>;
using OperandBTraits = using OperandBTraits =
OperandTraits<sizeof_bits<B_type>::value, N, K, trans_B, num_warp_n>; OperandTraits<sizeof_bits<B_type>::value, N, K, trans_B, num_warp_n, ldb>;
using SmemLayoutA = typename OperandATraits::Layout; using SmemLayoutA = typename OperandATraits::Layout;
using SmemLayoutB = typename OperandBTraits::Layout; using SmemLayoutB = typename OperandBTraits::Layout;
...@@ -262,12 +270,44 @@ public: ...@@ -262,12 +270,44 @@ public:
return layout; return layout;
} }
template <int offset, int NN, int KK, bool trans, int lddim, typename Engine0,
typename Layout0>
static CUTE_DEVICE auto get_region_tensor(Tensor<Engine0, Layout0> &sa) {
if constexpr (offset == 0) {
return composition(
sa,
Layout<Shape<Int<NN>, Int<KK>>,
Stride<_1, typename std::conditional<trans, Int<NN>,
Int<lddim>>::type>>{});
} else {
if constexpr (trans) {
static_assert(offset % KK == 0, "Offset must be a multiple of K");
constexpr int offset_n = offset / KK;
return flat_divide(sa, Shape<Int<NN>, Int<KK>>{})(_, _, _0{},
Int<offset_n>{});
} else {
static_assert(offset % NN == 0, "Offset must be a multiple of N");
constexpr int offset_n = offset / NN;
return flat_divide(sa, Shape<Int<NN>, Int<KK>>{})(_, _, Int<offset_n>{},
_0{});
}
}
}
static CUTE_DEVICE void body(A_type_raw *pA, B_type_raw *pB, C_type_raw *pC) { static CUTE_DEVICE void body(A_type_raw *pA, B_type_raw *pB, C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sA = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)), Tensor sA_all = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)),
SmemLayoutA{}); SmemLayoutA{});
Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)), Tensor sB_all = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)),
SmemLayoutB{}); SmemLayoutB{});
// Tensor sA = composition(sA_all, Layout<Shape<Int<M>, Int<K>>,
// Stride<_1, typename std::conditional<trans_A, Int<lda>,
// Int<M>>::type>>{});
// Tensor sB = composition(sB_all, Layout<Shape<Int<N>, Int<K>>,
// Stride<_1, typename std::conditional<trans_B, Int<N>,
// Int<ldb>>::type>>{});
Tensor sA = get_region_tensor<offset_a, M, K, !trans_A, lda>(sA_all);
Tensor sB = get_region_tensor<offset_b, N, K, trans_B, ldb>(sB_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma); auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma);
...@@ -306,8 +346,11 @@ public: ...@@ -306,8 +346,11 @@ public:
static CUTE_DEVICE void body_rs(A_type_raw *pA, B_type_raw *pB, static CUTE_DEVICE void body_rs(A_type_raw *pA, B_type_raw *pB,
C_type_raw *pC) { C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)), Tensor sB_all = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)),
SmemLayoutB{}); SmemLayoutB{});
// Tensor sB = flat_divide(sB_all, Shape<Int<N>, Int<K>>{})(_, _, _0{},
// _0{});
Tensor sB = get_region_tensor<offset_b, N, K, trans_B, ldb>(sB_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_B = make_tiled_copy_B(SmemCopyB{}, tiled_mma); auto tiled_copy_B = make_tiled_copy_B(SmemCopyB{}, tiled_mma);
...@@ -342,8 +385,11 @@ public: ...@@ -342,8 +385,11 @@ public:
static CUTE_DEVICE void body_sr(A_type_raw *pA, B_type_raw *pB, static CUTE_DEVICE void body_sr(A_type_raw *pA, B_type_raw *pB,
C_type_raw *pC) { C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sA = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)), Tensor sA_all = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)),
SmemLayoutA{}); SmemLayoutA{});
// Tensor sA = flat_divide(sA_all, Shape<Int<M>, Int<K>>{})(_, _, _0{},
// _0{});
Tensor sA = get_region_tensor<offset_a, M, K, !trans_A, lda>(sA_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma); auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma);
...@@ -380,29 +426,32 @@ public: ...@@ -380,29 +426,32 @@ public:
namespace tl { namespace tl {
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) {
using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body(pA, pB, accum); MMA::body(pA, pB, accum);
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) {
using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_rs(pA, pB, accum); MMA::body_rs(pA, pB, accum);
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_sr(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_sr(A_type *pA, B_type *pB, C_type *accum) {
using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, using MMA = cute::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_sr(pA, pB, accum); MMA::body_sr(pA, pB, accum);
} }
......
src/tl_templates/cuda/gemm_sm90.h
View file @ fe70549f
...@@ -194,16 +194,16 @@ struct DispatchInstruction<half_t, half_t, float, num_warp_m, num_warp_n, N> { ...@@ -194,16 +194,16 @@ struct DispatchInstruction<half_t, half_t, float, num_warp_m, num_warp_n, N> {
}; };
#endif #endif
template <int Bits, int N, int K, bool K_inner, int num_warp_n, template <int Bits, int N, int K, bool K_inner, int num_warp_n, int leading_dim,
typename Enable = void> typename Enable = void>
struct OperandTraits { struct OperandTraits {
// Primary template, use padded layout and default copy // Primary template, use padded layout and default copy
static constexpr int stride = K_inner ? K : N; static constexpr int stride = leading_dim;
static constexpr int padded = static constexpr int padded =
stride % (256 / Bits) == 0 ? stride + 128 / Bits : stride; stride % (256 / Bits) == 0 ? stride + 128 / Bits : stride;
using Layout = typename std::conditional< using Layout = typename std::conditional<
K_inner, Layout<Shape<Int<N>, Int<K>>, Shape<Int<padded>, _1>>, K_inner, Layout<Shape<Int<N>, Int<leading_dim>>, Shape<Int<padded>, _1>>,
Layout<Shape<Int<N>, Int<K>>, Shape<_1, Int<padded>>>>::type; Layout<Shape<Int<leading_dim>, Int<K>>, Shape<_1, Int<padded>>>>::type;
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
...@@ -224,124 +224,132 @@ template <int N, int num_warp_n, bool transpose> struct SelectCopy { ...@@ -224,124 +224,132 @@ template <int N, int num_warp_n, bool transpose> struct SelectCopy {
DefaultCopy>; DefaultCopy>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, true, num_warp_n, struct OperandTraits<16, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 64 == 32>::type> { typename std::enable_if<leading_dim % 64 == 32>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{})); Swizzle<2, 3, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, true, num_warp_n, struct OperandTraits<16, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 64 == 0>::type> { typename std::enable_if<leading_dim % 64 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{})); Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, false, num_warp_n, struct OperandTraits<16, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 64 == 32>::type> { typename std::enable_if<leading_dim % 64 == 32>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{})); Swizzle<2, 3, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = typename SelectCopy<N, num_warp_n, false>::type; using Copy = typename SelectCopy<N, num_warp_n, false>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<16, N, K, false, num_warp_n, struct OperandTraits<16, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 64 == 0>::type> { typename std::enable_if<leading_dim % 64 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{})); Swizzle<3, 3, 3>{}, Layout<Shape<_64, _8>, Stride<_1, _64>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = typename SelectCopy<N, num_warp_n, false>::type; using Copy = typename SelectCopy<N, num_warp_n, false>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, true, num_warp_n, struct OperandTraits<32, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 32 == 0>::type> { typename std::enable_if<leading_dim % 32 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{})); Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, true, num_warp_n, struct OperandTraits<32, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 32 == 16>::type> { typename std::enable_if<leading_dim % 32 == 16>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{})); Swizzle<2, 2, 3>{}, Layout<Shape<_8, _16>, Stride<_16, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename SelectCopy<N, num_warp_n, true>::type; using Copy = typename SelectCopy<N, num_warp_n, true>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, false, num_warp_n, struct OperandTraits<32, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 32 == 0>::type> { typename std::enable_if<leading_dim % 32 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 2, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{})); Swizzle<3, 2, 3>{}, Layout<Shape<_32, _8>, Stride<_1, _32>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = UniversalCopy<tfloat32_t>; using Copy = UniversalCopy<tfloat32_t>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<32, N, K, false, num_warp_n, struct OperandTraits<32, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 32 == 16>::type> { typename std::enable_if<leading_dim % 32 == 16>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 3>{}, Layout<Shape<_16, _8>, Stride<_1, _16>>{})); Swizzle<2, 2, 3>{}, Layout<Shape<_16, _8>, Stride<_1, _16>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = UniversalCopy<tfloat32_t>; using Copy = UniversalCopy<tfloat32_t>;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<8, N, K, true, num_warp_n, struct OperandTraits<8, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 128 == 64>::type> { typename std::enable_if<leading_dim % 128 == 64>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{})); Swizzle<2, 4, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N, using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
SM75_U32x4_LDSM_N>::type; SM75_U32x4_LDSM_N>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<8, N, K, true, num_warp_n, struct OperandTraits<8, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 128 == 0>::type> { typename std::enable_if<leading_dim % 128 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{})); Swizzle<3, 4, 3>{}, Layout<Shape<_8, _128>, Stride<_128, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N, using Copy = typename std::conditional<N == 8 * num_warp_n, SM75_U32x2_LDSM_N,
SM75_U32x4_LDSM_N>::type; SM75_U32x4_LDSM_N>::type;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<64, N, K, true, num_warp_n, struct OperandTraits<64, N, K, true, num_warp_n, leading_dim,
typename std::enable_if<K % 16 == 0>::type> { typename std::enable_if<leading_dim % 16 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 0, 4>{}, Layout<Shape<_4, _16>, Stride<_16, _1>>{})); Swizzle<2, 0, 4>{}, Layout<Shape<_4, _16>, Stride<_16, _1>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{})); using Layout =
decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<leading_dim>>{}));
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int N, int K, int num_warp_n> template <int N, int K, int num_warp_n, int leading_dim>
struct OperandTraits<64, N, K, false, num_warp_n, struct OperandTraits<64, N, K, false, num_warp_n, leading_dim,
typename std::enable_if<N % 16 == 0>::type> { typename std::enable_if<leading_dim % 16 == 0>::type> {
using LayoutAtom = decltype(composition( using LayoutAtom = decltype(composition(
Swizzle<2, 2, 2>{}, Layout<Shape<_16, _4>, Stride<_1, _16>>{})); Swizzle<2, 2, 2>{}, Layout<Shape<_16, _4>, Stride<_1, _16>>{}));
using Layout = decltype(tile_to_shape(LayoutAtom{}, Shape<Int<N>, Int<K>>{}, using Layout = decltype(tile_to_shape(
Step<_2, _1>{})); LayoutAtom{}, Shape<Int<leading_dim>, Int<K>>{}, Step<_2, _1>{}));
using Copy = DefaultCopy; using Copy = DefaultCopy;
}; };
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type_raw, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename B_type_raw, typename C_type_raw> int offset_b, typename A_type_raw, typename B_type_raw,
typename C_type_raw>
class GemmTensorOp { class GemmTensorOp {
public: public:
using A_type = using A_type =
...@@ -355,10 +363,11 @@ public: ...@@ -355,10 +363,11 @@ public:
using Instruction = using Instruction =
DispatchInstruction<A_type, B_type, C_type, num_warp_m, num_warp_n, N>; DispatchInstruction<A_type, B_type, C_type, num_warp_m, num_warp_n, N>;
using OperandATraits = using OperandATraits = OperandTraits<sizeof_bits<A_type>::value, M, K,
OperandTraits<sizeof_bits<A_type>::value, M, K, !trans_A, num_warp_m>; !trans_A, num_warp_m, lda>;
using OperandBTraits = using OperandBTraits =
OperandTraits<sizeof_bits<B_type>::value, N, K, trans_B, num_warp_n>; OperandTraits<sizeof_bits<B_type>::value, N, K, trans_B, num_warp_n, ldb>;
using SmemLayoutA = typename OperandATraits::Layout; using SmemLayoutA = typename OperandATraits::Layout;
using SmemLayoutB = typename OperandBTraits::Layout; using SmemLayoutB = typename OperandBTraits::Layout;
using SmemCopyA = Copy_Atom<typename OperandATraits::Copy, A_type>; using SmemCopyA = Copy_Atom<typename OperandATraits::Copy, A_type>;
...@@ -383,12 +392,38 @@ public: ...@@ -383,12 +392,38 @@ public:
return layout; return layout;
} }
template <int offset, int NN, int KK, bool trans, int lddim, typename Engine0,
typename Layout0>
static CUTE_DEVICE auto get_region_tensor(Tensor<Engine0, Layout0> &sa) {
if constexpr (offset == 0) {
return composition(
sa,
Layout<Shape<Int<NN>, Int<KK>>,
Stride<_1, typename std::conditional<trans, Int<NN>,
Int<lddim>>::type>>{});
} else {
if constexpr (trans) {
static_assert(offset % KK == 0, "Offset must be a multiple of K");
constexpr int offset_n = offset / KK;
return flat_divide(sa, Shape<Int<NN>, Int<KK>>{})(_, _, _0{},
Int<offset_n>{});
} else {
static_assert(offset % NN == 0, "Offset must be a multiple of N");
constexpr int offset_n = offset / NN;
return flat_divide(sa, Shape<Int<NN>, Int<KK>>{})(_, _, Int<offset_n>{},
_0{});
}
}
}
static CUTE_DEVICE void body(A_type_raw *pA, B_type_raw *pB, C_type_raw *pC) { static CUTE_DEVICE void body(A_type_raw *pA, B_type_raw *pB, C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sA = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)), Tensor sA_all = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)),
SmemLayoutA{}); SmemLayoutA{});
Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)), Tensor sB_all = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)),
SmemLayoutB{}); SmemLayoutB{});
Tensor sA = get_region_tensor<offset_a, M, K, !trans_A, lda>(sA_all);
Tensor sB = get_region_tensor<offset_b, N, K, trans_B, ldb>(sB_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma); auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma);
...@@ -426,8 +461,9 @@ public: ...@@ -426,8 +461,9 @@ public:
static CUTE_DEVICE void body_rs(A_type_raw *pA, B_type_raw *pB, static CUTE_DEVICE void body_rs(A_type_raw *pA, B_type_raw *pB,
C_type_raw *pC) { C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)), Tensor sB_all = make_tensor(make_smem_ptr(reinterpret_cast<B_type *>(pB)),
SmemLayoutB{}); SmemLayoutB{});
Tensor sB = get_region_tensor<offset_b, N, K, trans_B, ldb>(sB_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_B = make_tiled_copy_B(SmemCopyB{}, tiled_mma); auto tiled_copy_B = make_tiled_copy_B(SmemCopyB{}, tiled_mma);
...@@ -461,8 +497,9 @@ public: ...@@ -461,8 +497,9 @@ public:
static CUTE_DEVICE void body_sr(A_type_raw *pA, B_type_raw *pB, static CUTE_DEVICE void body_sr(A_type_raw *pA, B_type_raw *pB,
C_type_raw *pC) { C_type_raw *pC) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
Tensor sA = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)), Tensor sA_all = make_tensor(make_smem_ptr(reinterpret_cast<A_type *>(pA)),
SmemLayoutA{}); SmemLayoutA{});
Tensor sA = get_region_tensor<offset_a, M, K, !trans_A, lda>(sA_all);
TileMma tiled_mma; TileMma tiled_mma;
auto thr_mma = tiled_mma.get_thread_slice(tid); auto thr_mma = tiled_mma.get_thread_slice(tid);
auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma); auto tiled_copy_A = make_tiled_copy_A(SmemCopyA{}, tiled_mma);
...@@ -503,67 +540,86 @@ namespace tl { ...@@ -503,67 +540,86 @@ namespace tl {
namespace tl_mma { namespace tl_mma {
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) {
using MMA = using MMA =
cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body(pA, pB, accum); MMA::body(pA, pB, accum);
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) {
using MMA = using MMA =
cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_rs(pA, pB, accum); MMA::body_rs(pA, pB, accum);
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum, typename A_type, typename B_type, bool trans_B, bool clear_accum, int lda, int ldb, int offset_a,
typename C_type> int offset_b, typename A_type, typename B_type, typename C_type>
CUTLASS_DEVICE void gemm_sr(A_type *pA, B_type *pB, C_type *accum) { CUTLASS_DEVICE void gemm_sr(A_type *pA, B_type *pB, C_type *accum) {
using MMA = using MMA =
cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A, cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
trans_B, clear_accum, A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_sr(pA, pB, accum); MMA::body_sr(pA, pB, accum);
} }
} // namespace tl_mma } // namespace tl_mma
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum = false, bool use_wgmma = true, bool trans_B, bool clear_accum = false, int lda = 0, int ldb = 0,
int offset_a = 0, int offset_b = 0, bool use_wgmma = true,
int wg_wait = 0, typename A_type, typename B_type, typename C_type> int wg_wait = 0, typename A_type, typename B_type, typename C_type>
TL_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) { TL_DEVICE void gemm_ss(A_type *pA, B_type *pB, C_type *accum) {
if constexpr (use_wgmma) { if constexpr (use_wgmma) {
static_assert((trans_A && lda == M) || (!trans_A && lda == K),
"Hopper wgmma doesn't support custom stride for A");
static_assert((trans_B && ldb == K) || (!trans_B && ldb == N),
"Hopper wgmma doesn't support custom stride for B");
static_assert(offset_a == 0 && offset_b == 0,
"offset_a and offset_b must be zero for wgmma");
using MMA = cute::tl_wgmma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, using MMA = cute::tl_wgmma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n,
trans_A, trans_B, clear_accum, trans_A, trans_B, clear_accum,
A_type, B_type, C_type>; A_type, B_type, C_type>;
MMA::body<wg_wait>(pA, pB, accum); MMA::body<wg_wait>(pA, pB, accum);
} else { } else {
using MMA = cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, using MMA =
trans_A, trans_B, clear_accum, cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body(pA, pB, accum); MMA::body(pA, pB, accum);
} }
} }
template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A, template <int M, int N, int K, int num_warp_m, int num_warp_n, bool trans_A,
bool trans_B, bool clear_accum = false, bool use_wgmma = true, bool trans_B, bool clear_accum = false, int lda = 0, int ldb = 0,
int offset_a = 0, int offset_b = 0, bool use_wgmma = true,
int wg_wait = 0, typename A_type, typename B_type, typename C_type> int wg_wait = 0, typename A_type, typename B_type, typename C_type>
TL_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) { TL_DEVICE void gemm_rs(A_type *pA, B_type *pB, C_type *accum) {
if constexpr (use_wgmma) { if constexpr (use_wgmma) {
static_assert((trans_A && lda == M) || (!trans_A && lda == K),
"Hopper wgmma doesn't support custom stride for A");
static_assert((trans_B && ldb == K) || (!trans_B && ldb == N),
"Hopper wgmma doesn't support custom stride for B");
static_assert(offset_a == 0 && offset_b == 0,
"offset_a and offset_b must be zero for wgmma");
using MMA = cute::tl_wgmma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, using MMA = cute::tl_wgmma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n,
trans_A, trans_B, clear_accum, trans_A, trans_B, clear_accum,
A_type, B_type, C_type>; A_type, B_type, C_type>;
MMA::body_rs<wg_wait>(pA, pB, accum); MMA::body_rs<wg_wait>(pA, pB, accum);
} else { } else {
using MMA = cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, using MMA =
trans_A, trans_B, clear_accum, cute::tl_mma::GemmTensorOp<M, N, K, num_warp_m, num_warp_n, trans_A,
A_type, B_type, C_type>; trans_B, clear_accum, lda, ldb, offset_a,
offset_b, A_type, B_type, C_type>;
MMA::body_rs(pA, pB, accum); MMA::body_rs(pA, pB, accum);
} }
} }
......
testing/python/kernel/test_tilelang_kernel_gemm_with_stride.py 0 → 100644
View file @ fe70549f
import tilelang.testing
import tilelang
import tilelang.language as T
import torch
def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
@T.prim_func
def main(
A: T.Tensor((M, K), dtype),
B: T.Tensor((K, N), dtype),
C: T.Tensor((M, N), dtype),
):
# Initialize Kernel Context
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
A_shared = T.alloc_shared((block_M, block_K * 2), dtype, scope="shared")
B_shared = T.alloc_shared((block_K, block_N * 2), dtype, scope="shared")
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
# Clear local accumulation
T.clear(C_local)
T.clear(B_shared)
T.clear(A_shared)
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=0):
# Copy tile of A
# T.copy(A[by * block_M, ko * block_K], A_shared)
for i, k in T.Parallel(block_M, block_K):
A_shared[i, k + block_K] = A[by * block_M + i, ko * block_K + k]
# Copy tile of B
# T.copy(B[ko * block_K, bx * block_N], B_shared)
for i, k in T.Parallel(block_K, block_N):
B_shared[i, k] = B[ko * block_K + i, bx * block_N + k]
# Perform a tile-level GEMM on the shared buffers
# Currently we dispatch to the cute/hip on Nvidia/AMD GPUs
T.gemm(A_shared[:, block_K:], B_shared[0:block_K, 0:block_N], C_local)
# Copy result back to global memory
T.copy(C_local, C[by * block_M, bx * block_N])
return main
def run_gemm_with_stride_ss(M: int, N: int, K: int, block_M: int, block_N: int, block_K: int):
# 1. Define the kernel (matmul) and compile/lower it into an executable module
func = matmul(M, N, K, block_M, block_N, block_K)
# 2. Compile the kernel into a torch function
# out_idx specifies the index of the output buffer in the argument list
# if out_idx is specified, the tensor will be created during runtime
# target currently can be "cuda" or "hip" or "cpu".
jit_kernel = tilelang.compile(
func,
out_idx=[2],
target="cuda",
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
})
# Create random input tensors on the GPU
a = torch.randn(M, K, device="cuda", dtype=torch.float16)
b = torch.randn(K, N, device="cuda", dtype=torch.float16)
# Run the kernel through the Profiler
c = jit_kernel(a, b)
print(c)
# Reference multiplication using PyTorch
ref_c = a @ b
# Validate correctness
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
print("Kernel output matches PyTorch reference.")
@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(7, 5)
def test_tilelang_kernel_gemm_with_stride():
run_gemm_with_stride_ss(128, 128, 64, 32, 32, 32)
if __name__ == "__main__":
tilelang.testing.main()
tilelang/language/gemm.py
View file @ fe70549f
...@@ -69,10 +69,32 @@ def gemm( ...@@ -69,10 +69,32 @@ def gemm(
else: else:
raise ValueError(f"Unsupported argument type: {type(object)} for buffer {object}") raise ValueError(f"Unsupported argument type: {type(object)} for buffer {object}")
def retrieve_stride(object: Union[tir.Buffer, tir.BufferRegion]) -> List[int]:
if isinstance(object, tir.Buffer):
strides = []
stride = 1
for s in reversed(object.shape):
strides.insert(0, stride)
stride *= s
return strides
elif isinstance(object, tir.BufferRegion):
buffer, _ = object.buffer, object.region
strides = []
stride = 1
for s in reversed(buffer.shape):
strides.insert(0, stride)
stride *= s
return strides
else:
raise ValueError(f"Unsupported argument type: {type(object)} for buffer {object}")
A_shape = retrieve_shape(A) A_shape = retrieve_shape(A)
B_shape = retrieve_shape(B) B_shape = retrieve_shape(B)
C_shape = retrieve_shape(C) C_shape = retrieve_shape(C)
A_stride = retrieve_stride(A)
B_stride = retrieve_stride(B)
assert len(C_shape) == 2, "current only support C as a 2D tensor" assert len(C_shape) == 2, "current only support C as a 2D tensor"
assert len(A_shape) >= 2, "current only support A as a 2D or higher-order tensor" assert len(A_shape) >= 2, "current only support A as a 2D or higher-order tensor"
assert len(B_shape) >= 2, "current only support B as a 2D or higher-order tensor" assert len(B_shape) >= 2, "current only support B as a 2D or higher-order tensor"
...@@ -90,6 +112,9 @@ def gemm( ...@@ -90,6 +112,9 @@ def gemm(
K_B = B_shape[-1] if transpose_B else B_shape[-2] K_B = B_shape[-1] if transpose_B else B_shape[-2]
assert K == K_B, f"T.gemm K shape check failed: K_A = {K}, K_B = {K_B}" assert K == K_B, f"T.gemm K shape check failed: K_A = {K}, K_B = {K_B}"
stride_a = A_stride[-2]
stride_b = B_stride[-2]
def retrieve_ptr(object: Union[tir.Buffer, tir.BufferRegion], def retrieve_ptr(object: Union[tir.Buffer, tir.BufferRegion],
access_type: str = "r") -> tir.PrimExpr: access_type: str = "r") -> tir.PrimExpr:
if isinstance(object, tir.Buffer): if isinstance(object, tir.Buffer):
...@@ -105,12 +130,33 @@ def gemm( ...@@ -105,12 +130,33 @@ def gemm(
strides.insert(0, stride) strides.insert(0, stride)
stride *= s stride *= s
offset = 0 offset = 0
for i in range(len(indices)): # not offset the last two dimension
for i in range(len(indices) - 2):
offset += indices[i] * strides[i] offset += indices[i] * strides[i]
return buffer.access_ptr(access_mask=access_type, offset=offset) return buffer.access_ptr(access_mask=access_type, offset=offset)
else: else:
raise ValueError(f"Unsupported argument type: {type(object)} for buffer {object}") raise ValueError(f"Unsupported argument type: {type(object)} for buffer {object}")
def retrieve_offset(object: Union[tir.Buffer, tir.BufferRegion]) -> tir.PrimExpr:
"""Retrieve the offset of the buffer or buffer region."""
if isinstance(object, tir.Buffer):
return [0] * len(object.shape)
elif isinstance(object, tir.BufferRegion):
_, region = object.buffer, object.region
indices = []
for r in region:
indices.append(r.min)
return indices
else:
raise ValueError(f"Unsupported argument type: {type(object)} for buffer {object}")
A_offset = retrieve_offset(A)
B_offset = retrieve_offset(B)
assert A_offset[-2] == 0, "The offset of the first dimension of A must be 0"
assert B_offset[-2] == 0, "The offset of the first dimension of B must be 0"
offset_a = A_offset[-1]
offset_b = B_offset[-1]
Aptr = retrieve_ptr(A, "r") Aptr = retrieve_ptr(A, "r")
Bptr = retrieve_ptr(B, "r") Bptr = retrieve_ptr(B, "r")
Cptr = retrieve_ptr(C, "rw") Cptr = retrieve_ptr(C, "rw")
...@@ -127,6 +173,10 @@ def gemm( ...@@ -127,6 +173,10 @@ def gemm(
K, K,
policy, policy,
clear_accum, clear_accum,
stride_a,
stride_b,
offset_a,
offset_b,
k_pack, k_pack,
wg_wait, wg_wait,
) )
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment