github-actions[bot]
122 lines
5.9 KiB
Plaintext
122 lines
5.9 KiB
Plaintext
#pragma once
|
|
|
|
#include <cute/tensor.hpp>
|
|
|
|
#include <cutlass/cutlass.h>
|
|
#include <cutlass/layout/layout.h>
|
|
#include <cutlass/numeric_types.h>
|
|
|
|
#include "mma_sm89_fp16.hpp"
|
|
#include "mma_traits_sm89_fp16.hpp"
|
|
|
|
using namespace cute;
|
|
|
|
template<int BYTES> struct BytesToType {};
|
|
template<> struct BytesToType<4> {
|
|
using Type = uint32_t;
|
|
static_assert(sizeof(Type) == 4);
|
|
};
|
|
template<> struct BytesToType<2> {
|
|
using Type = uint16_t;
|
|
static_assert(sizeof(Type) == 2);
|
|
};
|
|
|
|
template<int BM_, int BN_, int KStages_, int K_, int WARP_ROW_=2, int WARP_COL_=2, bool HasBias_=false, typename accum_t_=cutlass::half_t, typename out_t_=cutlass::bfloat16_t>
|
|
struct gemm_traits {
|
|
static constexpr int BLOCK_SIZE = 128;
|
|
static constexpr int K = K_;
|
|
static constexpr int BM = BM_;
|
|
static constexpr int BN = BN_;
|
|
static constexpr int BK = 128;
|
|
static constexpr int TILES_PER_BLOCK = BLOCK_SIZE / BK;
|
|
static constexpr int NUM_SFB_PER_STEP = BN / BLOCK_SIZE;
|
|
static constexpr int NTiles = K / BK;
|
|
static constexpr int KSF = K / BLOCK_SIZE;
|
|
static constexpr int KStages = KStages_;
|
|
static constexpr int WARP_ROW = WARP_ROW_;
|
|
static constexpr int WARP_COL = WARP_COL_;
|
|
static constexpr int NUM_WARPS = WARP_ROW * WARP_COL;
|
|
static constexpr int NUM_THREADS = NUM_WARPS * 32;
|
|
static constexpr int MMA_WARP_M = WARP_ROW * 16;
|
|
static constexpr int MMA_WARP_N = WARP_COL * 8;
|
|
static constexpr int MMA_WARP_K = 32;
|
|
using accum_t = accum_t_;
|
|
using out_t = out_t_;
|
|
using SwizzleLayoutO = std::conditional_t<
|
|
std::is_same_v<out_t_, cutlass::bfloat16_t>,
|
|
Swizzle<3, 3, 3>,
|
|
Swizzle<2, 4, 3>
|
|
>;
|
|
using SwizzleLayoutAB = Swizzle<2, 4, 3>;
|
|
using MMA_Atom_SM89 = std::conditional_t<
|
|
std::is_same_v<accum_t, cutlass::half_t>,
|
|
MMA_Atom<SM89_16x8x32_F16E4M3E4M3F16_TN>,
|
|
MMA_Atom<SM89_16x8x32_F32E4M3E4M3F32_TN>
|
|
>;
|
|
static constexpr int INPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(float_e4m3_t);
|
|
static constexpr int OUTPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(out_t_);
|
|
static constexpr int THREADS_PER_ROW = BK / INPUT_ELEMS_PER_COPY;
|
|
using GMEMLayout = Layout< Shape <Int<NUM_THREADS / THREADS_PER_ROW>, Int<THREADS_PER_ROW>>, Stride<Int<THREADS_PER_ROW>, _1>>;
|
|
using G2SCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, float_e4m3_t>;
|
|
using G2STiledCopy = decltype(
|
|
make_tiled_copy(
|
|
G2SCopyAtom{},
|
|
GMEMLayout{},
|
|
Layout<Shape<_1, Int<INPUT_ELEMS_PER_COPY>>>{}
|
|
)
|
|
);
|
|
using S2RCopyAtomA = Copy_Atom<SM75_U32x4_LDSM_N, float_e4m3_t>;
|
|
using S2RCopyAtomB = Copy_Atom<SM75_U32x2_LDSM_N, float_e4m3_t>;
|
|
using SmemLayoutAtom = decltype(composition(
|
|
Swizzle<2, 4, 3>{},
|
|
make_layout(make_shape(Int<8>{}, Int<BK>{}),
|
|
make_stride(Int<BK>{}, Int<1>{}))));
|
|
using SmemLayoutA = decltype(
|
|
tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BM>{}, Int<BK>{}, Int<KStages>{}))
|
|
);
|
|
using SmemLayoutB = decltype(
|
|
tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BN>{}, Int<BK>{}, Int<KStages>{}))
|
|
);
|
|
using MMATile = decltype(
|
|
make_tiled_mma(
|
|
MMA_Atom_SM89{},
|
|
Layout<Shape<Int<WARP_ROW>, Int<WARP_COL>, _1>>{},
|
|
Tile<Int<MMA_WARP_M>, Int<MMA_WARP_N>, Int<MMA_WARP_K>>{}
|
|
)
|
|
);
|
|
|
|
static constexpr int ELEMS_PER_TILE = MMA_WARP_M * MMA_WARP_N;
|
|
static constexpr int NUM_ELEMS_PER_WRITE = NUM_THREADS * sizeof(cute::uint128_t) / sizeof(out_t_);
|
|
static constexpr int OUT_PIPE = NUM_ELEMS_PER_WRITE / ELEMS_PER_TILE;
|
|
// using SmemLayoutC = Layout<Shape<Int<BM>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
|
|
|
|
using SmemLayoutC = decltype(
|
|
make_layout(
|
|
make_shape(Int<MMA_WARP_M>{}, Int<MMA_WARP_N*OUT_PIPE>{}),
|
|
make_stride(Int<MMA_WARP_N*OUT_PIPE>{}, Int<1>{})
|
|
)
|
|
);
|
|
static constexpr int THREADS_PER_ROW_WRITE = MMA_WARP_N * OUT_PIPE / OUTPUT_ELEMS_PER_COPY;
|
|
using R2SCopyAtomC = Copy_Atom<UniversalCopy<typename BytesToType<2*sizeof(out_t)>::Type>, out_t>;
|
|
using S2GCopyAtomC = Copy_Atom<UniversalCopy<cute::uint128_t>, out_t>;
|
|
using S2GCopyC = decltype(make_tiled_copy(S2GCopyAtomC{},
|
|
make_layout(make_shape(Int<NUM_THREADS / THREADS_PER_ROW_WRITE>{}, Int<THREADS_PER_ROW_WRITE>{}),
|
|
make_stride(Int<THREADS_PER_ROW_WRITE>{}, Int<1>{})),
|
|
make_layout(make_shape(Int<1>{}, Int<OUTPUT_ELEMS_PER_COPY>{}))));
|
|
|
|
using G2SBiasCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<float>, float>;
|
|
using G2SBiasCopy = decltype(make_tiled_copy(G2SBiasCopyAtom{}, make_layout(
|
|
make_shape(Int<1>{},Int<BN>{}), make_stride(Int<BN>{}, Int<1>{})),
|
|
make_layout(make_shape(Int<1>{},Int<1>{}), make_stride(Int<1>{}, Int<1>{}))));
|
|
using sfa_copy_vtype = float;
|
|
static constexpr int SFA_ELEMS_PER_COPY = sizeof(sfa_copy_vtype)/sizeof(float);
|
|
static constexpr int THREADS_SFA_COPY = BM * sizeof(float) / sizeof(sfa_copy_vtype);
|
|
// using G2SSFACopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, float>;
|
|
static constexpr bool HasBias = HasBias_;
|
|
using SmemLayoutBias = Layout<Shape<Int<1>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
|
|
using SmemLayoutSFA = Layout<Shape<Int<BM>, Int<KStages>>, Stride<Int<1>, Int<BM>>>;
|
|
using BiasThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_2, _0>, Stride<_0, _8>>>;
|
|
using SFAThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_0, _1>, Stride<_16, _0>>>;
|
|
static constexpr int SmemSize = cute::max(cute::cosize(SmemLayoutA{})+cute::cosize(SmemLayoutB{}), cute::cosize(SmemLayoutC{})*sizeof(out_t)) + cute::cosize(SmemLayoutBias{}) * sizeof(float) + cute::cosize(SmemLayoutSFA{})*sizeof(float);
|
|
};
|