diff --git a/example/ck_tile/18_flatmm/mixed_prec/a8w4_moe_flatmm.cpp b/example/ck_tile/18_flatmm/mixed_prec/a8w4_moe_flatmm.cpp index 248d46b5f3..b2eb6f010e 100644 --- a/example/ck_tile/18_flatmm/mixed_prec/a8w4_moe_flatmm.cpp +++ b/example/ck_tile/18_flatmm/mixed_prec/a8w4_moe_flatmm.cpp @@ -98,8 +98,8 @@ float a8w4_moe_gemm(const MoeFlatmmHostArgs& args, const ck_tile::stream_config& static_assert(sizeof(ComputeDataType) >= sizeof(BDataType), "mixed_prec_flatmm requires ADataType is a wider type than BDataType"); - using GemmPipelineProblem = ck_tile::GemmPipelineProblem; @@ -133,9 +133,9 @@ float a8w4_moe_gemm(const MoeFlatmmHostArgs& args, const ck_tile::stream_config& constexpr int BlockedXDLN_PerWarp = 2; // determined by scale shuffle pattern using GemmEpilogue = ck_tile::CShuffleEpilogue< - ck_tile::CShuffleEpilogueProblem; @@ -968,7 +969,7 @@ struct MoeFlatmmKernel a_scale_gather_block_tile, // weight scale with granularityK = 32 b_scale_block_window, // weight scale with granularityK = 32 num_loop, - kargs.k_padded_zeros, + // kargs.k_padded_zeros, smem_ptr_ping, smem_ptr_pong); } diff --git a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index 76d191a40c..71b245f3cb 100644 --- a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -513,8 +513,6 @@ struct UniversalFlatmmPipelineAgBgCrPolicy using BlockFlatmmPolicy = BlockFlatmmASmemBSmemCRegV1CustomPolicy< typename Problem::ADataType, - // BlockGemmASmemBSmemCRegV1CustomPolicy +// struct F8xMXF4FlatmmPipelineProblem : FlatmmPipelineProblem +// { +// using BlockGemmShape = BlockGemmShape_; +// +// using QuantType = BDataType_; +// +// static constexpr index_t flatNPerWarp = BlockGemmShape::flatNPerWarp; +// +// static constexpr int MXF4ScaleGranularityK = 32; +// +// static constexpr int ContinuousKPerThread = 32; // it's fixed for fp4 +// static constexpr int ContinuousScaleMPerThread = 2; // it's fixed for fp8 +// static constexpr int ContinuousScaleNPerThread = 2; // it's fixed for fp4 +// static constexpr int ContinuousScaleKPerThread = 2; // it's fixed for fp4 +// static constexpr index_t flatKPerWarp = 64 * ContinuousKPerThread; +// }; +// +// template +// struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 +// : FlatmmPipelineAGmemBGmemCRegV1 +// { +// using Underlying = FlatmmPipelineAGmemBGmemCRegV1; +// +// using ADataType = remove_cvref_t; +// using BDataType = remove_cvref_t; +// using CDataType = remove_cvref_t; +// using BlockGemmShape = remove_cvref_t; // TileFlatmmShape +// +// using ComputeType = ADataType; +// static_assert(sizeof(ADataType) >= sizeof(BDataType)); +// +// using ALayout = remove_cvref_t; +// using BLayout = remove_cvref_t; +// using CLayout = remove_cvref_t; +// +// using BlockFlatmm = +// remove_cvref_t())>; +// +// static constexpr auto config = +// BlockFlatmm::BlockPolicy::template GetWarpGemmMWarpNWarp(); +// +// using WG = remove_cvref_t())>; +// +// static constexpr index_t DsWritePreIssue = 3; // default 2, ds write at MIter - 2 +// static constexpr index_t DsReadPreload = 2; // default 2, preload 2 ds read +// +// static constexpr index_t BlockSize = Problem::kBlockSize; +// static constexpr index_t WaveSize = get_warp_size(); +// +// static constexpr index_t kMPerBlock = BlockGemmShape::kM; +// static constexpr index_t kNPerBlock = BlockGemmShape::kN; +// static constexpr index_t kKPerBlock = BlockGemmShape::kK; +// +// static constexpr index_t flatKPerWarp = Problem::flatKPerWarp; +// static constexpr index_t flatNPerWarp = Problem::flatNPerWarp; +// +// static constexpr index_t GetVectorSizeA() { return Problem::VectorSizeA; } +// static constexpr index_t GetVectorSizeB() { return 32; /* fixed for fp4 shuffle layout*/ } +// static constexpr index_t GetVectorSizeC() { return Problem::VectorSizeC; } +// +// static constexpr bool kPadM = Problem::kPadM; +// static constexpr bool kPadN = Problem::kPadN; +// static constexpr bool kPadK = Problem::kPadK; +// +// static constexpr index_t kLdsAlignmentInBytes = 16; +// static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; +// static constexpr bool UsePersistentKernel = Problem::Traits::UsePersistentKernel; +// +// static constexpr auto I0 = number<0>(); +// static constexpr auto I1 = number<1>(); +// static constexpr auto I2 = number<2>(); +// static constexpr auto idxM = I0; +// static constexpr auto idxN = I1; +// static constexpr auto idxK = I2; +// using BlockTile = remove_cvref_t; +// using BlockWarps = remove_cvref_t; +// using WarpTile = remove_cvref_t; +// +// static constexpr index_t MWarp = config.template at<1>(); +// static constexpr index_t NWarp = config.template at<2>(); +// +// static constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WG::kM); +// static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN); +// static constexpr index_t KIterPerWarp = kKPerBlock / WG::kK; // 256 / 128 = 2 +// +// static constexpr index_t KFlatPerBlockPerIter = flatKPerWarp; +// static constexpr index_t NFlatPerBlockPerIter = flatNPerWarp; +// +// static constexpr index_t MPerBlockPerIter = kMPerBlock / MIterPerWarp; +// static constexpr index_t KPerBlockPerIter = kKPerBlock / KIterPerWarp; +// +// static constexpr int MXFP4PackedSize = 2; +// static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType); +// static constexpr index_t BK1 = Problem::VectorLoadSize / sizeof(BDataType) * MXFP4PackedSize; +// static constexpr index_t m_preload = (MIterPerWarp * KIterPerWarp >= DsReadPreload) +// ? DsReadPreload +// : MIterPerWarp * KIterPerWarp; +// +// static constexpr int ContinuousKPerThread = Problem::ContinuousKPerThread; +// static constexpr int ContinuousScaleMPerThread = Problem::ContinuousScaleMPerThread; +// static constexpr int ContinuousScaleNPerThread = Problem::ContinuousScaleNPerThread; +// static constexpr int ContinuousScaleKPerThread = Problem::ContinuousScaleKPerThread; +// +// // static constexpr int ScaleKPerWarp = +// // ContinuousScaleNPerThread * ContinuousScaleKPerThread * get_warp_size(); +// static constexpr int ScaleKFlatPerWarp = +// ContinuousScaleNPerThread * ContinuousScaleKPerThread * get_warp_size(); +// +// static constexpr int XDLK_PerThread = +// WarpTile::at(I2) / (get_warp_size() / WarpTile::at(I1)); // 8 +// +// static constexpr int XDL_PerWeightK = 1; // 1 for tile block = 128 load 128 per instrution +// static constexpr int XDL_PerScaleK = XDL_PerWeightK * ContinuousScaleKPerThread; // 4 +// static constexpr int XDL_PerScaleN = ContinuousScaleNPerThread; // 2 +// static_assert(XDL_PerScaleK % XDL_PerWeightK == 0); +// static_assert(KIterPerWarp % XDL_PerScaleK == 0); +// static_assert(NIterPerWarp % XDL_PerScaleN == 0); +// +// static constexpr int MXFP4KPerWarp = KIterPerWarp / XDL_PerWeightK; +// static constexpr int ScaleKPerWarp = KIterPerWarp / XDL_PerScaleK; +// static constexpr int ScaleNPerWarp = NIterPerWarp / XDL_PerScaleN; +// +// static constexpr int MXFP4K_PerScaleK = MXFP4KPerWarp / ScaleKPerWarp; +// +// static constexpr bool HasHotLoop = Problem::HasHotLoop; +// static constexpr auto TailNum = Problem::TailNum; +// +// #ifdef __gfx942__ +// static constexpr index_t mfma_per_wg = 2; +// #else +// static constexpr index_t mfma_per_wg = 1; +// #endif +// static constexpr index_t dsread_per_wg = +// WG::kM * WG::kK * sizeof(ADataType) / WaveSize / Problem::VectorLoadSize; +// static_assert((WG::kM * WG::kK * sizeof(ADataType) / WaveSize) % Problem::VectorLoadSize == 0); +// +// static constexpr index_t dsread_num_perK = dsread_per_wg * MIterPerWarp; +// static constexpr index_t dswrite_num_perK = dsread_num_perK / (MWarp * NWarp); +// static constexpr index_t dswrite_rep = (dswrite_num_perK + MIterPerWarp - 1) / MIterPerWarp; +// static constexpr index_t Aload_num_perK = dswrite_num_perK; +// static constexpr index_t Aload_rep = dswrite_rep; +// static constexpr index_t Bload_num_perK = kNPerBlock * WG::kK / NWarp / BK1 / WaveSize; +// static constexpr index_t ScaleBload_K1 = ContinuousScaleNPerThread * ContinuousScaleKPerThread; +// static constexpr index_t ScaleBload_num = +// kNPerBlock * kKPerBlock / NWarp / 32 / ScaleBload_K1 / +// WaveSize; // BlockN * BlockK / NWarp / ScalePerK / ScaleB_K1 / wavesize +// static constexpr index_t Bload_total_num = +// Bload_num_perK * KIterPerWarp + ScaleBload_num + 0X3f0; +// static constexpr index_t KPerScaleLoad = KIterPerWarp / ScaleBload_num; +// static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2; +// static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter; +// +// static constexpr index_t mfma_perM_perK = NIterPerWarp * mfma_per_wg; +// static constexpr index_t dswrite_mIter = (DsWritePreIssue - 1) % MIterPerWarp; +// static constexpr index_t dswrite_kIter = (DsWritePreIssue - 1) / MIterPerWarp; +// +// // For the basic gemm pipelien DoubleSmemBuffer set to be false naturally. +// static constexpr bool DoubleSmemBuffer = false; +// +// CK_TILE_HOST_DEVICE static constexpr auto +// SchedulerPerM(index_t dsread_perM, index_t dswrite_perM, index_t load_perM) +// { +// #if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS +// // GFX950 use BUFFER_LOAD_LDS to fill lds_buffer_A. +// // There is no separate DS_WRITE instruction at all. +// dswrite_perM = 0; +// #endif +// // Init inst order +// index_t max_data_inst = dsread_perM > load_perM +// ? (dsread_perM > dswrite_perM ? dsread_perM : dswrite_perM) +// : (load_perM > dswrite_perM ? load_perM : dswrite_perM); +// index_t sum_data_inst = dsread_perM + load_perM + dswrite_perM; +// index_t round_data_inst = (sum_data_inst + mfma_perM_perK - 1) / mfma_perM_perK; +// +// index_t inst_order[NIterPerWarp * 10]; +// _Pragma("unroll") for(int idx = 0; idx < NIterPerWarp * 10; idx++) { inst_order[idx] = 0; } +// +// index_t index = 0; +// _Pragma("unroll") for(int j = 0; j < max_data_inst; j++) +// { +// if(dswrite_perM > j) +// { +// inst_order[index] = 1; +// index++; +// } +// if(load_perM > j) +// { +// inst_order[index] = 2; +// index++; +// } +// if(dsread_perM > j) +// { +// inst_order[index] = 3; +// index++; +// } +// } +// +// // Schedule IGLP +// _Pragma("unroll") for(int j = 0; j < mfma_perM_perK; j++) +// { +// index_t inst_idx = 0; +// if(j == 0) +// ; +// else if(j == 1) +// inst_idx = mfma_perM_perK == 2 ? 1 : mfma_perM_perK - 2; +// else if(j == 2) +// inst_idx = mfma_perM_perK - 1; +// else +// inst_idx = mfma_perM_perK - j; +// +// __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA +// +// _Pragma("unroll") for(int r = 0; r < round_data_inst; r++) +// { +// if(r % 2 == 0) +// { +// if(inst_order[inst_idx + r * mfma_perM_perK] == 1) +// { +// __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write +// } +// if(inst_order[inst_idx + r * mfma_perM_perK] == 2) +// { +// __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read +// } +// if(inst_order[inst_idx + r * mfma_perM_perK] == 3) +// { +// __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read +// } +// } +// else +// { +// if(inst_order[(r + 1) * mfma_perM_perK - 1 - inst_idx] == 1) +// { +// __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write +// } +// if(inst_order[(r + 1) * mfma_perM_perK - 1 - inst_idx] == 2) +// { +// __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read +// } +// if(inst_order[(r + 1) * mfma_perM_perK - 1 - inst_idx] == 3) +// { +// __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read +// } +// } +// } +// } +// } +// CK_TILE_HOST_DEVICE static constexpr auto HotLoopScheduler() +// { +// // Keypoint of pipeline optimize is workload balance in time +// // instruction schedule example(128X256X256, 1X4, 16X16X128): +// // Iter MNK MFMA ds_read ds_write A_load b_load +// // -1 M6N0: 57 - 8 - - +// // -1 M6N1: 58 1 - - - +// // -1 M6N2: 59 - - 7 - +// // -1 M6N3: 60 2 - - - +// // -1 M7N0: 61 - - - - +// // -1 M7N1: 62 3 - - - +// // -1 M7N2: 63 - - 8 - +// // -1 M7N3: 64 4 - - - +// // 0 M0N0K0: 1 - - - 1 +// // 0 M0N1: 2 5 - - - +// // 0 M0N2: 3 - - - 2 +// // 0 M0N3: 4 6 - - - +// // 0 M1N0: 5 - - - 3 +// // 0 M1N1: 6 7 - - - +// // 0 M1N2: 7 - - - 4 +// // 0 M1N3: 8 8 - - - +// // 0 M2N0: 9 - - - 5 +// // 0 M2N1: 10 9 - - - +// // 0 M2N2: 11 - - - 6 +// // 0 M2N3: 12 10 - - - +// // 0 M3N0: 13 - 1 - 7 +// // 0 M3N1: 14 11 - - - +// // 0 M3N2: 15 - - - 8 +// // 0 M3N3: 16 12 - - - +// // 0 M4N0: 17 - 2 - - +// // 0 M4N1: 18 13 - - - +// // 0 M4N2: 19 - - 1 - +// // 0 M4N3: 20 14 - - - +// // 0 M5N0: 21 - 3 - - +// // 0 M5N1: 22 15 - - - +// // 0 M5N2: 23 - - 2 - +// // 0 M5N3: 24 16 - - - +// // 0 M6N0: 25 - 4 - - +// // 0 M6N1: 26 17 - - - +// // 0 M6N2: 27 - - 3 - +// // 0 M6N3: 28 18 - - - +// // 0 M7N0: 29 - - - - +// // 0 M7N1: 30 19 - - - +// // 0 M7N2: 31 - - 4 - +// // 0 M7N3: 32 20 - - - +// // 0 M0N0K1: 33 - - - 9 +// // 0 M0N1: 34 21 - - - +// // 0 M0N2: 35 - - - 10 +// // 0 M0N3: 36 22 - - - +// // 0 M1N0: 37 - - - 11 +// // 0 M1N1: 38 23 - - - +// // 0 M1N2: 39 - - - 12 +// // 0 M1N3: 40 24 - - - +// // 0 M2N0: 41 - - - 13 +// // 0 M2N1: 42 25 - - - +// // 0 M2N2: 43 - - - 14 +// // 0 M2N3: 44 26 - - - +// // 0 M3N0: 45 - 5 - 15 +// // 0 M3N1: 46 27 - - - +// // 0 M3N2: 47 - - - 16 +// // 0 M3N3: 48 28 - - - +// // 0 M4N0: 49 - 6 - - +// // 0 M4N1: 50 29 - - - +// // 0 M4N2: 51 - - 5 - +// // 0 M4N3: 52 30 - - - +// // 0 M5N0: 53 - 7 - - +// // 0 M5N1: 54 31 - - - +// // 0 M5N2: 55 - - 6 - +// // 0 M5N3: 56 32 - - - +// // 0 M6N0: 57 - 8 - - +// // 0 M6N1: 58 1 - - - +// // 0 M6N2: 59 - - 7 - +// // 0 M6N3: 60 2 - - - +// // 0 M7N0: 61 - - - - +// // 0 M7N1: 62 3 - - - +// // 0 M7N2: 63 - - 8 - +// // 0 M7N3: 64 4 - - - +// +// _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) +// { +// _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) +// { +// index_t dsread_perM = 0; +// index_t dswrite_perM = 0; +// index_t load_perM = 0; +// +// // Calculate ds_read number per M +// dsread_perM = dsread_per_wg; +// +// // Calculate buffer_load number per M +// if(mIter < HalfMIter) +// { +// load_perM = +// ((Aload_num_perK - (MIterPerWarp - 1 - mIter) * Aload_rep) > 0 ? Aload_rep +// : 0) + +// ((Bload_num_perK - (HalfMIter - 1 - mIter) * Bload_rep) > 0 ? Bload_rep +// : 0); +// } +// else +// { +// load_perM = (Aload_num_perK - (MIterPerWarp - 1 - mIter) * Aload_rep) > 0 +// ? Aload_rep +// : 0; +// } +// if((kIter % KPerScaleLoad == 0) && (mIter == 0)) +// { +// load_perM = load_perM + 1; +// } +// SchedulerPerM(dsread_perM, dswrite_perM, load_perM); +// } +// } +// // Add Aload when Aload data > needed +// if(Aload_num_perK == 0) +// __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read +// __builtin_amdgcn_sched_barrier(0); +// } +// +// CK_TILE_HOST_DEVICE static constexpr auto Last2ndHotLoopScheduler() +// { +// _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) +// { +// _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) +// { +// index_t dsread_perM = 0; +// index_t dswrite_perM = 0; +// index_t load_perM = 0; +// +// // Calculate ds_read number per M +// dsread_perM = dsread_per_wg; +// +// // Calculate buffer_load number per M +// if(mIter < HalfMIter) +// { +// load_perM = +// ((Bload_num_perK - (HalfMIter - 1 - mIter) * Bload_rep) > 0 ? Bload_rep +// : 0); +// } +// SchedulerPerM(dsread_perM, dswrite_perM, load_perM); +// } +// } +// __builtin_amdgcn_sched_barrier(0); +// } +// +// CK_TILE_HOST_DEVICE static constexpr auto LastHotLoopScheduler() +// { +// _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) +// { +// _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) +// { +// index_t dsread_perM = 0; +// index_t dswrite_perM = 0; +// index_t load_perM = 0; +// +// // Calculate ds_read number per M +// if((kIter * MIterPerWarp + mIter) < (KIterPerWarp * MIterPerWarp - m_preload)) +// dsread_perM = dsread_per_wg; +// +// SchedulerPerM(dsread_perM, dswrite_perM, load_perM); +// } +// } +// // __builtin_amdgcn_sched_barrier(0); +// } +// +// CK_TILE_HOST_DEVICE static constexpr auto GetADramTileDistribution() +// { +// return PipelinePolicy::template MakeADramTileDistribution(); +// } +// +// CK_TILE_HOST_DEVICE static constexpr auto GetAScaleDramTileDistribution() +// { +// return PipelinePolicy::template MakeAScaleDramTileDistribution(); +// } +// +// template +// CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindowTmp a_copy_dram_window_, +// const AElementFunction& a_element_func, +// const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, +// const DequantAWindow& scale_a_window, +// const DequantBFlatWindow& scale_b_flat_window, +// const index_t num_loop, +// const index_t k_padded_zeros, +// void* p_smem_ping, +// void* p_smem_pong) const +// { +// #ifndef __gfx950__ +// static_assert(false, "Only gfx950 is supported for MXFP4 flatmm pipeline now."); +// #endif +// static_assert( +// std::is_same_v>, +// "wrong!"); +// +// static_assert(kMPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<0>{}], +// "wrong!"); +// static_assert(kKPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<1>{}], +// "wrong!"); +// +// constexpr auto MIter_2nd_last = (MIterPerWarp >= 2) ? MIterPerWarp - 2 : MIterPerWarp - 1; +// const index_t iMWarp = get_warp_id() / NWarp; +// +// using CWarpDstr = typename WG::CWarpDstr; +// using CWarpTensor = typename WG::CWarpTensor; +// +// constexpr auto c_warp_y_lengths = +// to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); +// constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; +// +// __builtin_amdgcn_sched_barrier(0); +// +// auto a_copy_dram_window = replace_bottom_tensor_view( +// PipelinePolicy::template TransformF8xF4_ATensorView( +// a_copy_dram_window_.get_bottom_tensor_view()), +// a_copy_dram_window_); +// +// // ======= a scale related start ====== +// auto scale_a_dram_window = make_tile_window( +// scale_a_window.get_bottom_tensor_view(), +// make_tuple(number{}, number{}), // TODO actually 64 / 16 = 4 && 128 / 32 =4 +// scale_a_window.get_window_origin(), +// PipelinePolicy::template MakeAScaleDramTileDistribution()); +// // ping pong buffer for scale A +// statically_indexed_array< +// statically_indexed_array, +// MIterPerWarp> +// scale_a_dram_windows; +// statically_indexed_array +// scale_a_tile_tensor; +// // ======= a scale related end ====== +// +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// scale_a_dram_windows(mIter)(kIter) = scale_a_dram_window; +// move_tile_window(scale_a_dram_windows(mIter)(kIter), +// {mIter * kMPerBlock, kIter * kKPerBlock / 32}); +// }); +// }); +// +// // A tile in LDS +// ADataType* p_a_lds_ping = static_cast(p_smem_ping); +// ADataType* p_a_lds_pong = static_cast(p_smem_pong); +// +// constexpr auto write_a_lds_block_desc = +// PipelinePolicy::template MakeF8xF4_WriteALdsBlockDescriptor(); +// constexpr auto read_a_lds_block_desc = +// PipelinePolicy::template MakeF8xF4_ReadALdsBlockDescriptor(); +// +// auto write_a_lds_block_ping = +// make_tensor_view(p_a_lds_ping, write_a_lds_block_desc); +// auto write_a_lds_block_pong = +// make_tensor_view(p_a_lds_pong, write_a_lds_block_desc); +// auto read_a_lds_block_ping = +// make_tensor_view(p_a_lds_ping, read_a_lds_block_desc); +// auto read_a_lds_block_pong = +// make_tensor_view(p_a_lds_pong, read_a_lds_block_desc); +// +// auto a_copy_lds_window_ping = +// make_tile_window(write_a_lds_block_ping, +// make_tuple(number{}, number{}), +// {0, 0}, +// PipelinePolicy::template MakeADramTileDistribution()); +// auto a_copy_lds_window_pong = +// make_tile_window(write_a_lds_block_pong, +// make_tuple(number{}, number{}), +// {0, 0}, +// PipelinePolicy::template MakeADramTileDistribution()); +// +// // ping-pong window for A LDS +// auto a_warp_window_ping_tmp = +// make_tile_window(read_a_lds_block_ping, +// make_tuple(number{}, number{}), +// {iMWarp * WG::kM, 0}, +// PipelinePolicy::template MakeF8xF4_ALDS_TileDistribution()); +// auto a_warp_window_pong_tmp = +// make_tile_window(read_a_lds_block_pong, +// make_tuple(number{}, number{}), +// {iMWarp * WG::kM, 0}, +// PipelinePolicy::template MakeF8xF4_ALDS_TileDistribution()); +// +// statically_indexed_array< +// statically_indexed_array, +// MIterPerWarp> +// a_warp_windows_ping; +// +// statically_indexed_array< +// statically_indexed_array, +// MIterPerWarp> +// a_warp_windows_pong; +// +// auto A_Lds_Stride = 8; +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// a_warp_windows_ping(mIter)(kIter) = a_warp_window_ping_tmp; +// a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp; +// +// auto weight_k_idx = kIter / number{}; +// auto weight_k_rank = kIter % number{}; +// move_tile_window( +// a_warp_windows_ping(mIter)(kIter), +// {mIter * MPerBlockPerIter, +// weight_k_rank * A_Lds_Stride + weight_k_idx * XDL_PerWeightK * WG::kK}); +// move_tile_window( +// a_warp_windows_pong(mIter)(kIter), +// {mIter * MPerBlockPerIter, +// weight_k_rank * A_Lds_Stride + weight_k_idx * XDL_PerWeightK * WG::kK}); +// }); +// }); +// +// // Block GEMM +// auto block_flatmm = BlockFlatmm(); +// // Acc register tile +// auto c_block_tile = block_flatmm.MakeCBlockTile(); +// +// // B flat DRAM window for load +// auto b_flat_distribution = +// PipelinePolicy::template MakeFp4BFlatDramTileDistribution(); +// auto scale_b_flat_distribution = +// PipelinePolicy::template MakeFp4ScaleBFlatDramTileDistribution(); +// +// auto b_flat_dram_window = make_tile_window( +// b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views +// make_tuple(number{}, number{}), +// b_flat_dram_block_window_tmp.get_window_origin(), +// b_flat_distribution); +// +// auto scale_b_flat_dram_window = make_tile_window( +// scale_b_flat_window.get_bottom_tensor_view(), // from kernel gemm_pad_views +// make_tuple(number{}, number{}), +// scale_b_flat_window.get_window_origin(), +// scale_b_flat_distribution); +// +// using MXFP4_Buffer = decltype(load_tile(b_flat_dram_window)); +// // use v4i32 as the data type between basicblock to avoid unpack and repack operation. +// // using V4UInt_Buffer = thread_buffer; +// // union UnionB +// // { +// // V4UInt_Buffer u = 0; +// // MXFP4_Buffer mxfp4; +// // } ub; +// +// // pingpong buffer for B +// statically_indexed_array< +// statically_indexed_array, +// NIterPerWarp> +// b_flat_dram_windows; +// statically_indexed_array, +// NIterPerWarp> +// b_warp_tensor_ping; +// statically_indexed_array, +// NIterPerWarp> +// b_warp_tensor_pong; +// +// statically_indexed_array< +// statically_indexed_array, +// ScaleNPerWarp> +// scale_b_flat_dram_windows; +// statically_indexed_array< +// statically_indexed_array, +// ScaleNPerWarp> +// scale_b_warp_tensor_ping; +// statically_indexed_array< +// statically_indexed_array, +// ScaleNPerWarp> +// scale_b_warp_tensor_pong; +// +// using ABlockTile = decltype(load_tile(a_copy_dram_window)); +// ABlockTile a_block_tile; +// +// enum +// { +// PrefillBeforeGemm = 1, +// PrefillAfterGemm = 2, +// PrefillAlways = PrefillBeforeGemm | PrefillAfterGemm, +// }; +// #if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS +// auto prefill_lds_a_stage1 = +// [&]([[maybe_unused]] auto lds_tile_a, auto dram_tile_a, auto prefill_location) { +// // global -> lds +// if constexpr(prefill_location & PrefillAfterGemm) +// async_load_tile(lds_tile_a, dram_tile_a); +// }; +// auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { +// // async_load_fence(); +// // __builtin_amdgcn_s_waitcnt(0x03fc); +// // data has been stored in lds, no need more operation. +// static_assert(std::is_same_v, +// "buffer_load_lds don't support element func fot A before mfma"); +// }; +// #else +// auto prefill_lds_a_stage1 = +// [&]([[maybe_unused]] auto lds_tile_a, auto dram_tile_a, auto prefill_location) { +// // global -> vgpr +// if constexpr(prefill_location & PrefillBeforeGemm) +// a_block_tile = load_tile(dram_tile_a); +// }; +// auto prefill_lds_a_stage2 = [&]([[maybe_unused]] auto lds_tile_a) { +// // vgpr -> lds +// auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); +// store_tile(lds_tile_a, a_block_tile_transformed); +// }; +// #endif +// +// // HEAD +// // Prefetch A0 +// prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window, number{}); +// +// // move A window to next k +// move_tile_window(a_copy_dram_window, {0, kKPerBlock}); +// +// // prefetch B +// static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { +// static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { +// if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) +// { +// auto scale_n_iter = nIter / number{}; +// auto scale_k_iter = kIter / number{}; +// +// scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = +// scale_b_flat_dram_window; +// move_tile_window( +// scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), +// {scale_n_iter * NFlatPerBlockPerIter, scale_k_iter * ScaleKFlatPerWarp}); +// scale_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = +// load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); +// } +// auto packed_n_idx = nIter / number{}; +// auto packed_n_rank = nIter % number{}; +// +// b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; +// move_tile_window(b_flat_dram_windows(nIter)(kIter), +// {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + +// packed_n_rank, +// kIter * KFlatPerBlockPerIter}); +// +// b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); +// }); +// }); +// // move B window to next flat K +// move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); +// move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); +// +// prefill_lds_a_stage2(a_copy_lds_window_ping); +// +// __builtin_amdgcn_sched_barrier(0); +// +// // Prefetch A1 +// prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window, number{}); +// // move A window to next k +// move_tile_window(a_copy_dram_window, {0, kKPerBlock}); +// +// // initialize C +// tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); +// +// __builtin_amdgcn_s_waitcnt(Bload_total_num); +// block_sync_lds(); +// +// // preload A00,A10... from lds +// statically_indexed_array{})(number<0>{}))), +// m_preload> +// a_warp_tensor; +// +// static_for<0, m_preload, 1>{}([&](auto loadIter) { +// constexpr auto mIter = loadIter % MIterPerWarp; +// constexpr auto kIter = loadIter / MIterPerWarp; +// a_warp_tensor(loadIter) = +// load_tile(a_warp_windows_ping(number{})(number{})); +// scale_a_tile_tensor(loadIter) = +// load_tile(scale_a_dram_windows(number{})(number{})); +// }); +// +// __builtin_amdgcn_sched_barrier(0); +// +// // statically_indexed_array dequant_B_n; +// +// //TODO: maybe need quant scale(change fp32 scale into fp8e8m0) +// +// auto convert_ascale_e8m0 = [&](auto a_scale_fp32) { +// return bit_cast(uint8x4_t{ 0, 0, 0, bit_cast(static_cast(a_scale_fp32)) }); +// }; +// +// // MAIN LOOP +// index_t iCounter = (num_loop - 1) / 2; +// while(iCounter > 0) +// { +// // prefetch B(2i+1) +// static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { +// static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { +// if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) +// { +// auto scale_n_iter = nIter / number{}; +// auto scale_k_iter = kIter / number{}; +// +// scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = +// scale_b_flat_dram_window; +// +// move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), +// {scale_n_iter * NFlatPerBlockPerIter, +// scale_k_iter * ScaleKFlatPerWarp}); +// +// scale_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = +// load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); +// } +// +// auto packed_n_idx = nIter / number{}; +// auto packed_n_rank = nIter % number{}; +// +// b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; +// +// move_tile_window( +// b_flat_dram_windows(nIter)(kIter), +// {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + +// packed_n_rank, +// kIter * KFlatPerBlockPerIter}); +// +// b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); +// }); +// }); +// +// // Prefill A(2i+1) +// prefill_lds_a_stage2(a_copy_lds_window_pong); +// +// // Prefetch A(2i+2) +// prefill_lds_a_stage1( +// a_copy_lds_window_ping, a_copy_dram_window, number{}); +// // GEMM 2i +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; +// static_for<0, NIterPerWarp / ContinuousScaleNPerThread, 1>{}([&](auto nIter) { +// static_for<0, ContinuousScaleNPerThread, 1>{}([&](auto inxdl) { +// // read C warp tensor from C block tensor +// CWarpTensor c_warp_tensor; +// +// c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); +// +// // warp GEMM +// WG{}.template operator()<0, 0>(c_warp_tensor, +// a_warp_tensor(number{}), +// b_warp_tensor_ping(nIter * ContinuousScaleNPerThread + inxdl)(kIter), +// convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), +// convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( +// kIter / number{}).get_thread_buffer()[0])); +// +// // write C warp tensor into C block tensor +// c_block_tile.set_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), +// c_warp_tensor.get_thread_buffer()); +// }); +// }); +// // preload next A from lds +// if constexpr((kIter * MIterPerWarp + mIter) < +// (KIterPerWarp * MIterPerWarp - m_preload)) +// { +// constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; +// constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); +// a_warp_tensor(number{}) = +// load_tile(a_warp_windows_ping(number{})(number{})); +// scale_a_tile_tensor(number{}) = +// load_tile(scale_a_dram_windows(number{})(number{})); +// } +// +// // barrier +// if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) +// { +// __builtin_amdgcn_s_waitcnt(Bload_total_num); +// block_sync_lds(); +// } +// }); +// }); +// prefill_lds_a_stage1( +// a_copy_lds_window_ping, a_copy_dram_window, number{}); +// +// // move A window to next k +// move_tile_window(a_copy_dram_window, {0, kKPerBlock}); +// +// // move B window to next flat K +// move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); +// move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); +// +// static_for<0, m_preload, 1>{}([&](auto loadIter) { +// constexpr auto mIter = loadIter % MIterPerWarp; +// constexpr auto kIter = loadIter / MIterPerWarp; +// a_warp_tensor(loadIter) = +// load_tile(a_warp_windows_pong(number{})(number{})); +// }); +// // HotLoopScheduler(); +// +// // Next K +// +// // prefetch B(2i+2) +// static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { +// static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { +// if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) +// { +// auto scale_n_iter = nIter / number{}; +// auto scale_k_iter = kIter / number{}; +// +// scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = +// scale_b_flat_dram_window; +// +// move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), +// {scale_n_iter * NFlatPerBlockPerIter, +// scale_k_iter * ScaleKFlatPerWarp}); +// +// scale_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = +// load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); +// } +// +// auto packed_n_idx = nIter / number{}; +// auto packed_n_rank = nIter % number{}; +// +// b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; +// move_tile_window( +// b_flat_dram_windows(nIter)(kIter), +// {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + +// packed_n_rank, +// kIter * KFlatPerBlockPerIter}); +// +// b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); +// }); +// }); +// +// // Prefill A(2i+2) +// prefill_lds_a_stage2(a_copy_lds_window_ping); +// +// // Prefetch A(2i+3) +// prefill_lds_a_stage1( +// a_copy_lds_window_pong, a_copy_dram_window, number{}); +// +// // GEMM 2i+1 +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; +// static_for<0, NIterPerWarp / ContinuousScaleNPerThread, 1>{}([&](auto nIter) { +// static_for<0, ContinuousScaleNPerThread, 1>{}([&](auto inxdl) { +// // read C warp tensor from C block tensor +// CWarpTensor c_warp_tensor; +// +// c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); +// +// // warp GEMM +// WG{}.template operator()<0, 0>(c_warp_tensor, +// a_warp_tensor(number{}), +// b_warp_tensor_ping(nIter * ContinuousScaleNPerThread + inxdl)(kIter), +// convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), +// convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( +// kIter / number{}).get_thread_buffer()[0])); +// +// // write C warp tensor into C block tensor +// c_block_tile.set_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), +// c_warp_tensor.get_thread_buffer()); +// }); +// }); +// // preload next A from lds +// if constexpr((kIter * MIterPerWarp + mIter) < +// (KIterPerWarp * MIterPerWarp - m_preload)) +// { +// constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; +// constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); +// a_warp_tensor(number{}) = +// load_tile(a_warp_windows_pong(number{})(number{})); +// scale_a_tile_tensor(number{}) = +// load_tile(scale_a_dram_windows(number{})(number{})); +// } +// +// // barrier +// if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) +// { +// __builtin_amdgcn_s_waitcnt(Bload_total_num); +// block_sync_lds(); +// } +// }); +// }); +// prefill_lds_a_stage1( +// a_copy_lds_window_pong, a_copy_dram_window, number{}); +// +// // move A window to next k +// move_tile_window(a_copy_dram_window, {0, kKPerBlock}); +// // move B window to next flat K +// move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); +// move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); +// +// static_for<0, m_preload, 1>{}([&](auto loadIter) { +// constexpr auto mIter = loadIter % MIterPerWarp; +// constexpr auto kIter = loadIter / MIterPerWarp; +// a_warp_tensor(loadIter) = +// load_tile(a_warp_windows_ping(number{})(number{})); +// }); +// // HotLoopScheduler(); +// +// iCounter--; +// } +// +// // TAIL +// if constexpr(TailNum == TailNumber::Even) +// { +// // prefetch B(loopK) +// const int b_k_off = b_flat_dram_window.get_tile_distribution().calculate_index()[I1] / +// ContinuousKPerThread / WG::kN * ContinuousKPerThread; +// static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { +// static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { +// if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) +// { +// auto scale_n_iter = nIter / number{}; +// auto scale_k_iter = kIter / number{}; +// +// scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = +// scale_b_flat_dram_window; +// +// move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), +// {scale_n_iter * NFlatPerBlockPerIter, +// scale_k_iter * ScaleKFlatPerWarp}); +// +// scale_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = +// load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); +// } +// }); +// +// const int b_k_off_inter = kIter * kKPerBlock / MXFP4KPerWarp + b_k_off; +// if(b_k_off_inter < kKPerBlock - k_padded_zeros) +// { +// static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { +// auto packed_n_idx = nIter / number{}; +// auto packed_n_rank = nIter % number{}; +// +// b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; +// +// move_tile_window( +// b_flat_dram_windows(nIter)(kIter), +// {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + +// packed_n_rank, +// kIter * KFlatPerBlockPerIter}); +// +// b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); +// }); +// } +// }); +// +// // Prefill A(loopK) +// prefill_lds_a_stage2(a_copy_lds_window_pong); +// +// // GEMM loopK-1 +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; +// static_for<0, NIterPerWarp / ContinuousScaleNPerThread, 1>{}([&](auto nIter) { +// static_for<0, ContinuousScaleNPerThread, 1>{}([&](auto inxdl) { +// // read C warp tensor from C block tensor +// CWarpTensor c_warp_tensor; +// +// c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); +// +// // warp GEMM +// WG{}.template operator()<0, 0>(c_warp_tensor, +// a_warp_tensor(number{}), +// b_warp_tensor_ping(nIter * ContinuousScaleNPerThread + inxdl)(kIter), +// convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), +// convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( +// kIter / number{}).get_thread_buffer()[0])); +// +// // write C warp tensor into C block tensor +// c_block_tile.set_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), +// c_warp_tensor.get_thread_buffer()); +// }); +// }); +// // preload next A from lds +// if constexpr((kIter * MIterPerWarp + mIter) < +// (KIterPerWarp * MIterPerWarp - m_preload)) +// { +// constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; +// constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); +// a_warp_tensor(number{}) = +// load_tile(a_warp_windows_ping(number{})(number{})); +// scale_a_tile_tensor(number{}) = +// load_tile(scale_a_dram_windows(number{})(number{})); +// } +// +// // barrier +// if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) +// { +// __builtin_amdgcn_s_waitcnt(Bload_total_num); +// block_sync_lds(); +// } +// }); +// }); +// +// static_for<0, m_preload, 1>{}([&](auto loadIter) { +// constexpr auto mIter = loadIter % MIterPerWarp; +// constexpr auto kIter = loadIter / MIterPerWarp; +// a_warp_tensor(loadIter) = +// load_tile(a_warp_windows_pong(number{})(number{})); +// }); +// +// __builtin_amdgcn_sched_barrier(0); +// // Last2ndHotLoopScheduler(); +// +// // GEMM loopK +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// if(kIter * WG::kK < kKPerBlock - k_padded_zeros) +// { +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; +// static_for<0, NIterPerWarp / ContinuousScaleNPerThread, 1>{}([&](auto nIter) { +// static_for<0, ContinuousScaleNPerThread, 1>{}([&](auto inxdl) { +// // read C warp tensor from C block tensor +// CWarpTensor c_warp_tensor; +// +// c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); +// +// // warp GEMM +// WG{}.template operator()<0, 0>(c_warp_tensor, +// a_warp_tensor(number{}), +// b_warp_tensor_ping(nIter * ContinuousScaleNPerThread + inxdl)(kIter), +// convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), +// convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( +// kIter / number{}).get_thread_buffer()[0])); +// +// // write C warp tensor into C block tensor +// c_block_tile.set_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), +// c_warp_tensor.get_thread_buffer()); +// }); +// }); +// if constexpr((kIter * MIterPerWarp + mIter) < +// (KIterPerWarp * MIterPerWarp - m_preload)) +// { +// constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; +// constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); +// a_warp_tensor(number{}) = +// load_tile(a_warp_windows_pong(number{})(number{})); +// scale_a_tile_tensor(number{}) = +// load_tile(scale_a_dram_windows(number{})(number{})); +// } +// // barrier +// // if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) +// // { +// // block_sync_lds(); +// // } +// }); +// } +// }); +// // LastHotLoopScheduler(); +// } +// else if constexpr(TailNum == TailNumber::Odd) +// { +// // GEMM loopK +// static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { +// static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { +// constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; +// static_for<0, NIterPerWarp / ContinuousScaleNPerThread, 1>{}([&](auto nIter) { +// static_for<0, ContinuousScaleNPerThread, 1>{}([&](auto inxdl) { +// // read C warp tensor from C block tensor +// CWarpTensor c_warp_tensor; +// +// c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); +// +// // warp GEMM +// WG{}.template operator()<0, 0>(c_warp_tensor, +// a_warp_tensor(number{}), +// b_warp_tensor_ping(nIter * ContinuousScaleNPerThread + inxdl)(kIter), +// convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), +// convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( +// kIter / number{}).get_thread_buffer()[0])); +// +// // write C warp tensor into C block tensor +// c_block_tile.set_y_sliced_thread_data( +// merge_sequences(sequence{}, c_warp_y_index_zeros), +// merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), +// c_warp_tensor.get_thread_buffer()); +// }); +// }); +// // preload next A from lds +// if constexpr((kIter * MIterPerWarp + mIter) < +// (KIterPerWarp * MIterPerWarp - m_preload)) +// { +// constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; +// constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); +// a_warp_tensor(number{}) = +// load_tile(a_warp_windows_ping(number{})(number{})); +// scale_a_tile_tensor(number{}) = +// load_tile(scale_a_dram_windows(number{})(number{})); +// } +// +// // barrier +// if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) +// { +// __builtin_amdgcn_s_waitcnt(Bload_total_num); +// block_sync_lds(); +// } +// }); +// }); +// // LastHotLoopScheduler(); +// } +// +// return c_block_tile; +// } +// +// template +// CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, +// const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, +// const DequantAWindow& scale_a_window, +// const DequantBFlatWindow& scale_b_flat_window, +// const index_t num_loop, +// const index_t k_padded_zeros, +// void* p_smem_ping, +// void* p_smem_pong) const +// { +// return operator()(a_dram_block_window_tmp, +// identity{}, +// b_flat_dram_block_window_tmp, +// scale_a_window, +// scale_b_flat_window, +// num_loop, +// k_padded_zeros, +// p_smem_ping, +// p_smem_pong); +// } +// +// template +// CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, +// const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, +// const DequantAWindow& scale_a_window, +// const DequantBFlatWindow& scale_b_flat_window, +// const index_t num_loop, +// void* p_smem_ping, +// void* p_smem_pong) const +// { +// return operator()(a_dram_block_window_tmp, +// identity{}, +// b_flat_dram_block_window_tmp, +// scale_a_window, +// scale_b_flat_window, +// num_loop, +// 0, +// p_smem_ping, +// p_smem_pong); +// } +// }; + template struct F8xMXF4FlatmmPipelineProblem : FlatmmPipelineProblem + ADataType_, + CDataType_, + BlockGemmShape_, + Traits_, + Scheduler_, + HasHotLoop_, + TailNum_, + ComputeDataType_> { using BlockGemmShape = BlockGemmShape_; - using QuantType = BDataType_; + // using QuantType = BDataType_; static constexpr index_t flatNPerWarp = BlockGemmShape::flatNPerWarp; - static constexpr int MXF4ScaleGranularityK = 32; + static constexpr int ScaleGranularityK = 32; - static constexpr int ContinuousKPerThread = 32; // it's fixed for fp4 - static constexpr int ContinuousScaleMPerThread = 2; // it's fixed for fp8 - static constexpr int ContinuousScaleNPerThread = 2; // it's fixed for fp4 - static constexpr int ContinuousScaleKPerThread = 2; // it's fixed for fp4 - static constexpr index_t flatKPerWarp = 64 * ContinuousKPerThread; + static constexpr int ContinuousKPerThread = 32; // it's fixed for fp4 + static constexpr int MXdlPack = 1; // it's fixed for fp4 + static constexpr int NXdlPack = 1; // it's fixed for fp4 + static constexpr int KXdlPack = 4; + static constexpr int ContinuousScaleNPerThread = 1; // it's fixed for fp4 + static constexpr int ContinuousScaleKPerThread = 4; // it's fixed for fp4 + // static constexpr index_t flatKPerWarp = BlockGemmShape::flatKPerWarp * KXdlPack; + static constexpr index_t flatKPerWarp = get_warp_size() * ContinuousKPerThread; }; template -struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 - : FlatmmPipelineAGmemBGmemCRegV1 +struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1 { using Underlying = FlatmmPipelineAGmemBGmemCRegV1; using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using BDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; // TileFlatmmShape @@ -1317,7 +2547,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 using WG = remove_cvref_t())>; static constexpr index_t DsWritePreIssue = 3; // default 2, ds write at MIter - 2 - static constexpr index_t DsReadPreload = 2; // default 2, preload 2 ds read + static constexpr index_t DsReadPreload = 4; // default 4 for MXFP4 (MXdlPack * KXdlPack) static constexpr index_t BlockSize = Problem::kBlockSize; static constexpr index_t WaveSize = get_warp_size(); @@ -1329,17 +2559,17 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 static constexpr index_t flatKPerWarp = Problem::flatKPerWarp; static constexpr index_t flatNPerWarp = Problem::flatNPerWarp; - static constexpr index_t GetVectorSizeA() { return Problem::VectorSizeA; } - static constexpr index_t GetVectorSizeB() { return 32; /* fixed for fp4 shuffle layout*/ } + static constexpr index_t GetVectorSizeA() { return 32; } /* fixed for fp4 shuffle layout*/ + static constexpr index_t GetVectorSizeB() { return 32; } /* fixed for fp4 shuffle layout*/ static constexpr index_t GetVectorSizeC() { return Problem::VectorSizeC; } static constexpr bool kPadM = Problem::kPadM; static constexpr bool kPadN = Problem::kPadN; static constexpr bool kPadK = Problem::kPadK; - static constexpr index_t kLdsAlignmentInBytes = 16; - static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; - static constexpr bool UsePersistentKernel = Problem::Traits::UsePersistentKernel; + // static constexpr index_t kLdsAlignmentInBytes = 16; + static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + static constexpr bool UsePersistentKernel = Problem::Traits::UsePersistentKernel; static constexpr auto I0 = number<0>(); static constexpr auto I1 = number<1>(); @@ -1356,7 +2586,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 static constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WG::kM); static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN); - static constexpr index_t KIterPerWarp = kKPerBlock / WG::kK; // 256 / 128 = 2 + static constexpr index_t KIterPerWarp = kKPerBlock / WG::kK; static constexpr index_t KFlatPerBlockPerIter = flatKPerWarp; static constexpr index_t NFlatPerBlockPerIter = flatNPerWarp; @@ -1364,66 +2594,45 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 static constexpr index_t MPerBlockPerIter = kMPerBlock / MIterPerWarp; static constexpr index_t KPerBlockPerIter = kKPerBlock / KIterPerWarp; - static constexpr int MXFP4PackedSize = 2; - static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType); - static constexpr index_t BK1 = Problem::VectorLoadSize / sizeof(BDataType) * MXFP4PackedSize; + static constexpr index_t APackedSize = numeric_traits::PackedSize; + static constexpr index_t BPackedSize = numeric_traits::PackedSize; + + static constexpr index_t MXdlPack = Problem::MXdlPack; + static constexpr index_t NXdlPack = Problem::NXdlPack; + static constexpr index_t KXdlPack = Problem::KXdlPack; + static constexpr index_t ScaleGranularityK = Problem::ScaleGranularityK; + + static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType) * APackedSize; + static constexpr index_t BK1 = Problem::VectorLoadSize / sizeof(BDataType) * BPackedSize; + static constexpr index_t m_preload = (MIterPerWarp * KIterPerWarp >= DsReadPreload) ? DsReadPreload : MIterPerWarp * KIterPerWarp; - static constexpr int ContinuousKPerThread = Problem::ContinuousKPerThread; - static constexpr int ContinuousScaleMPerThread = Problem::ContinuousScaleMPerThread; - static constexpr int ContinuousScaleNPerThread = Problem::ContinuousScaleNPerThread; - static constexpr int ContinuousScaleKPerThread = Problem::ContinuousScaleKPerThread; - - // static constexpr int ScaleKPerWarp = - // ContinuousScaleNPerThread * ContinuousScaleKPerThread * get_warp_size(); - static constexpr int ScaleKFlatPerWarp = - ContinuousScaleNPerThread * ContinuousScaleKPerThread * get_warp_size(); - - static constexpr int XDLK_PerThread = - WarpTile::at(I2) / (get_warp_size() / WarpTile::at(I1)); // 8 - - static constexpr int XDL_PerWeightK = 1; // 1 for tile block = 128 load 128 per instrution - static constexpr int XDL_PerScaleK = XDL_PerWeightK * ContinuousScaleKPerThread; // 4 - static constexpr int XDL_PerScaleN = ContinuousScaleNPerThread; // 2 - static_assert(XDL_PerScaleK % XDL_PerWeightK == 0); - static_assert(KIterPerWarp % XDL_PerScaleK == 0); - static_assert(NIterPerWarp % XDL_PerScaleN == 0); - - static constexpr int MXFP4KPerWarp = KIterPerWarp / XDL_PerWeightK; - static constexpr int ScaleKPerWarp = KIterPerWarp / XDL_PerScaleK; - static constexpr int ScaleNPerWarp = NIterPerWarp / XDL_PerScaleN; - - static constexpr int MXFP4K_PerScaleK = MXFP4KPerWarp / ScaleKPerWarp; - static constexpr bool HasHotLoop = Problem::HasHotLoop; static constexpr auto TailNum = Problem::TailNum; -#ifdef __gfx942__ - static constexpr index_t mfma_per_wg = 2; -#else - static constexpr index_t mfma_per_wg = 1; -#endif - static constexpr index_t dsread_per_wg = - WG::kM * WG::kK * sizeof(ADataType) / WaveSize / Problem::VectorLoadSize; - static_assert((WG::kM * WG::kK * sizeof(ADataType) / WaveSize) % Problem::VectorLoadSize == 0); + static constexpr index_t mfma_per_wg = 1; // 950 only + + static constexpr index_t dsread_per_wg = WG::kM * WG::kK / AK1 / WaveSize; + static_assert((WG::kM * WG::kK) % (AK1 * WaveSize) == 0); static constexpr index_t dsread_num_perK = dsread_per_wg * MIterPerWarp; - static constexpr index_t dswrite_num_perK = dsread_num_perK / (MWarp * NWarp); + static constexpr index_t dswrite_num_perK = dsread_num_perK / NWarp; static constexpr index_t dswrite_rep = (dswrite_num_perK + MIterPerWarp - 1) / MIterPerWarp; static constexpr index_t Aload_num_perK = dswrite_num_perK; static constexpr index_t Aload_rep = dswrite_rep; + static constexpr index_t Bload_num_perK = kNPerBlock * WG::kK / NWarp / BK1 / WaveSize; - static constexpr index_t ScaleBload_K1 = ContinuousScaleNPerThread * ContinuousScaleKPerThread; + static constexpr index_t Bload_num = Bload_num_perK * KIterPerWarp; static constexpr index_t ScaleBload_num = - kNPerBlock * kKPerBlock / NWarp / 32 / ScaleBload_K1 / - WaveSize; // BlockN * BlockK / NWarp / ScalePerK / ScaleB_K1 / wavesize - static constexpr index_t Bload_total_num = - Bload_num_perK * KIterPerWarp + ScaleBload_num + 0X3f0; - static constexpr index_t KPerScaleLoad = KIterPerWarp / ScaleBload_num; - static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2; - static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter; + kNPerBlock * kKPerBlock / NWarp / ScaleGranularityK / NXdlPack / KXdlPack / WaveSize; + static constexpr index_t ScaleAload_num = + kMPerBlock * kKPerBlock / MWarp / ScaleGranularityK / MXdlPack / KXdlPack / WaveSize; + + // static constexpr index_t KPerScaleLoad = KIterPerWarp / ScaleBload_num; + static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2; + static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter; static constexpr index_t mfma_perM_perK = NIterPerWarp * mfma_per_wg; static constexpr index_t dswrite_mIter = (DsWritePreIssue - 1) % MIterPerWarp; @@ -1435,11 +2644,6 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 CK_TILE_HOST_DEVICE static constexpr auto SchedulerPerM(index_t dsread_perM, index_t dswrite_perM, index_t load_perM) { -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS - // GFX950 use BUFFER_LOAD_LDS to fill lds_buffer_A. - // There is no separate DS_WRITE instruction at all. - dswrite_perM = 0; -#endif // Init inst order index_t max_data_inst = dsread_perM > load_perM ? (dsread_perM > dswrite_perM ? dsread_perM : dswrite_perM) @@ -1491,7 +2695,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 { if(inst_order[inst_idx + r * mfma_perM_perK] == 1) { - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + // __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write } if(inst_order[inst_idx + r * mfma_perM_perK] == 2) { @@ -1506,7 +2710,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 { if(inst_order[(r + 1) * mfma_perM_perK - 1 - inst_idx] == 1) { - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + // __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write } if(inst_order[(r + 1) * mfma_perM_perK - 1 - inst_idx] == 2) { @@ -1520,6 +2724,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 } } } + CK_TILE_HOST_DEVICE static constexpr auto HotLoopScheduler() { // Keypoint of pipeline optimize is workload balance in time @@ -1609,6 +2814,32 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 // Calculate ds_read number per M dsread_perM = dsread_per_wg; + // Calculate ds_write number per M + if(mIter == 0) + { + dswrite_perM = + (dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep) > 0 + ? dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep + : 0; + } + else if(mIter >= MIterPerWarp - DsWritePreIssue + 1) + { + dswrite_perM = 0; + } + else + { + dswrite_perM = (dswrite_num_perK - + (MIterPerWarp - DsWritePreIssue - mIter) * dswrite_rep) > 0 + ? dswrite_rep + : 0; + } + // Add ds write when ds write data > needed + if(dswrite_num_perK == 0 && kIter == (KIterPerWarp - 1 - dswrite_kIter)) + { + if(mIter == MIterPerWarp - 1 - dswrite_mIter) + dswrite_perM = 1; + } + // Calculate buffer_load number per M if(mIter < HalfMIter) { @@ -1624,10 +2855,10 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 ? Aload_rep : 0; } - if((kIter % KPerScaleLoad == 0) && (mIter == 0)) - { - load_perM = load_perM + 1; - } + // if((kIter % KPerScaleLoad == 0) && (mIter == 0)) + // { + // load_perM = load_perM + 1; + // } SchedulerPerM(dsread_perM, dswrite_perM, load_perM); } } @@ -1650,6 +2881,32 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 // Calculate ds_read number per M dsread_perM = dsread_per_wg; + // Calculate ds_write number per M + if(mIter == 0) + { + dswrite_perM = + (dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep) > 0 + ? dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep + : 0; + } + else if(mIter >= MIterPerWarp - DsWritePreIssue + 1) + { + dswrite_perM = 0; + } + else + { + dswrite_perM = (dswrite_num_perK - + (MIterPerWarp - DsWritePreIssue - mIter) * dswrite_rep) > 0 + ? dswrite_rep + : 0; + } + // Add ds write when ds write data > needed + if(dswrite_num_perK == 0 && kIter == (KIterPerWarp - 1 - dswrite_kIter)) + { + if(mIter == MIterPerWarp - 1 - dswrite_mIter) + dswrite_perM = 1; + } + // Calculate buffer_load number per M if(mIter < HalfMIter) { @@ -1690,23 +2947,20 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 CK_TILE_HOST_DEVICE static constexpr auto GetAScaleDramTileDistribution() { - return PipelinePolicy::template MakeAScaleDramTileDistribution(); + return PipelinePolicy::template MakeMXFP4_ScaleA_FlatDramTileDistribution(); } template - CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindowTmp a_copy_dram_window_, - const AElementFunction& a_element_func, - const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, - const DequantAWindow& scale_a_window, - const DequantBFlatWindow& scale_b_flat_window, - const index_t num_loop, - const index_t k_padded_zeros, - void* p_smem_ping, - void* p_smem_pong) const + typename ScaleADramBlockWindowTmp, + typename ScaleBDramBlockWindowTmp> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_copy_dram_window_tmp, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + const ScaleADramBlockWindowTmp& scale_a_window, + const ScaleBDramBlockWindowTmp& scale_b_window, + index_t num_loop, + void* __restrict__ p_smem_ping, + void* __restrict__ p_smem_pong) const { #ifndef __gfx950__ static_assert(false, "Only gfx950 is supported for MXFP4 flatmm pipeline now."); @@ -1720,8 +2974,8 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 static_assert(kKPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<1>{}], "wrong!"); - constexpr auto MIter_2nd_last = (MIterPerWarp >= 2) ? MIterPerWarp - 2 : MIterPerWarp - 1; - const index_t iMWarp = get_warp_id() / NWarp; + // constexpr auto MIter_2nd_last = max(0, MIterPerWarp - 2); + static_assert(NWarp == 4); using CWarpDstr = typename WG::CWarpDstr; using CWarpTensor = typename WG::CWarpTensor; @@ -1730,105 +2984,43 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + auto a_dram_window = + make_tile_window(PipelinePolicy::template MakeMXFP4_AAsyncLoadDramDescriptor( + a_copy_dram_window_tmp.get_bottom_tensor_view()), + a_copy_dram_window_tmp.get_window_lengths(), + a_copy_dram_window_tmp.get_window_origin(), + PipelinePolicy::template MakeADramTileDistribution()); + __builtin_amdgcn_sched_barrier(0); - auto a_copy_dram_window = replace_bottom_tensor_view( - PipelinePolicy::template TransformF8xF4_ATensorView( - a_copy_dram_window_.get_bottom_tensor_view()), - a_copy_dram_window_); - - // ======= a scale related start ====== - auto scale_a_dram_window = make_tile_window( - scale_a_window.get_bottom_tensor_view(), - make_tuple(number{}, number{}), // TODO actually 64 / 16 = 4 && 128 / 32 =4 - scale_a_window.get_window_origin(), - PipelinePolicy::template MakeAScaleDramTileDistribution()); - // ping pong buffer for scale A - statically_indexed_array< - statically_indexed_array, - MIterPerWarp> - scale_a_dram_windows; - statically_indexed_array - scale_a_tile_tensor; - // ======= a scale related end ====== - - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - scale_a_dram_windows(mIter)(kIter) = scale_a_dram_window; - move_tile_window(scale_a_dram_windows(mIter)(kIter), - {mIter * kMPerBlock, kIter * kKPerBlock / 32}); - }); - }); - // A tile in LDS ADataType* p_a_lds_ping = static_cast(p_smem_ping); ADataType* p_a_lds_pong = static_cast(p_smem_pong); - constexpr auto write_a_lds_block_desc = - PipelinePolicy::template MakeF8xF4_WriteALdsBlockDescriptor(); - constexpr auto read_a_lds_block_desc = - PipelinePolicy::template MakeF8xF4_ReadALdsBlockDescriptor(); + constexpr auto a_lds_block_desc = + PipelinePolicy::template MakeMXFP4_ALdsBlockDescriptor(); - auto write_a_lds_block_ping = - make_tensor_view(p_a_lds_ping, write_a_lds_block_desc); - auto write_a_lds_block_pong = - make_tensor_view(p_a_lds_pong, write_a_lds_block_desc); - auto read_a_lds_block_ping = - make_tensor_view(p_a_lds_ping, read_a_lds_block_desc); - auto read_a_lds_block_pong = - make_tensor_view(p_a_lds_pong, read_a_lds_block_desc); + auto a_lds_block_ping = + make_tensor_view(p_a_lds_ping, a_lds_block_desc); + auto a_lds_block_pong = + make_tensor_view(p_a_lds_pong, a_lds_block_desc); - auto a_copy_lds_window_ping = - make_tile_window(write_a_lds_block_ping, - make_tuple(number{}, number{}), - {0, 0}, - PipelinePolicy::template MakeADramTileDistribution()); - auto a_copy_lds_window_pong = - make_tile_window(write_a_lds_block_pong, - make_tuple(number{}, number{}), - {0, 0}, - PipelinePolicy::template MakeADramTileDistribution()); + auto a_store_lds_window_ping = make_tile_window( + a_lds_block_ping, make_tuple(number{}, number{}), {0, 0}); + auto a_store_lds_window_pong = make_tile_window( + a_lds_block_pong, make_tuple(number{}, number{}), {0, 0}); // ping-pong window for A LDS - auto a_warp_window_ping_tmp = - make_tile_window(read_a_lds_block_ping, + auto a_warp_window_ping = + make_tile_window(a_lds_block_ping, make_tuple(number{}, number{}), - {iMWarp * WG::kM, 0}, - PipelinePolicy::template MakeF8xF4_ALDS_TileDistribution()); - auto a_warp_window_pong_tmp = - make_tile_window(read_a_lds_block_pong, + {0, 0}, + PipelinePolicy::template MakeMXF4_ALDS_TileDistribution()); + auto a_warp_window_pong = + make_tile_window(a_lds_block_pong, make_tuple(number{}, number{}), - {iMWarp * WG::kM, 0}, - PipelinePolicy::template MakeF8xF4_ALDS_TileDistribution()); - - statically_indexed_array< - statically_indexed_array, - MIterPerWarp> - a_warp_windows_ping; - - statically_indexed_array< - statically_indexed_array, - MIterPerWarp> - a_warp_windows_pong; - - auto A_Lds_Stride = 8; - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - a_warp_windows_ping(mIter)(kIter) = a_warp_window_ping_tmp; - a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp; - - auto weight_k_idx = kIter / number{}; - auto weight_k_rank = kIter % number{}; - move_tile_window( - a_warp_windows_ping(mIter)(kIter), - {mIter * MPerBlockPerIter, - weight_k_rank * A_Lds_Stride + weight_k_idx * XDL_PerWeightK * WG::kK}); - move_tile_window( - a_warp_windows_pong(mIter)(kIter), - {mIter * MPerBlockPerIter, - weight_k_rank * A_Lds_Stride + weight_k_idx * XDL_PerWeightK * WG::kK}); - }); - }); + {0, 0}, + PipelinePolicy::template MakeMXF4_ALDS_TileDistribution()); // Block GEMM auto block_flatmm = BlockFlatmm(); @@ -1836,680 +3028,611 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 auto c_block_tile = block_flatmm.MakeCBlockTile(); // B flat DRAM window for load - auto b_flat_distribution = - PipelinePolicy::template MakeFp4BFlatDramTileDistribution(); - auto scale_b_flat_distribution = - PipelinePolicy::template MakeFp4ScaleBFlatDramTileDistribution(); - - auto b_flat_dram_window = make_tile_window( - b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views - make_tuple(number{}, number{}), - b_flat_dram_block_window_tmp.get_window_origin(), - b_flat_distribution); - - auto scale_b_flat_dram_window = make_tile_window( - scale_b_flat_window.get_bottom_tensor_view(), // from kernel gemm_pad_views - make_tuple(number{}, number{}), - scale_b_flat_window.get_window_origin(), - scale_b_flat_distribution); - - using MXFP4_Buffer = decltype(load_tile(b_flat_dram_window)); - // use v4i32 as the data type between basicblock to avoid unpack and repack operation. - using V4UInt_Buffer = thread_buffer; - union UnionB - { - V4UInt_Buffer u = 0; - MXFP4_Buffer mxfp4; - } ub; // pingpong buffer for B + auto b_flat_dram_windows = generate_tuple( + [&](auto nIter) { + constexpr auto packed_n_idx = nIter / number{}; + constexpr auto packed_n_rank = nIter % number{}; + auto window_i = make_tile_window( + b_flat_dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_flat_dram_block_window_tmp.get_window_origin(), + PipelinePolicy::template MakeMXFP4_BFlatDramTileDistribution()); + move_tile_window( + window_i, + {number{}, + number<0>{}}); + return window_i; + }, + number{}); statically_indexed_array< - statically_indexed_array, + statically_indexed_array, NIterPerWarp> - b_flat_dram_windows; - statically_indexed_array, - NIterPerWarp> - b_warp_tensor_ping; - statically_indexed_array, - NIterPerWarp> - b_warp_tensor_pong; + b_warp_tensor_ping, b_warp_tensor_pong; - statically_indexed_array< - statically_indexed_array, - ScaleNPerWarp> - scale_b_flat_dram_windows; - statically_indexed_array< - statically_indexed_array, - ScaleNPerWarp> - scale_b_warp_tensor_ping; - statically_indexed_array< - statically_indexed_array, - ScaleNPerWarp> - scale_b_warp_tensor_pong; + // pingpong buffer for Scale A and Scale B + auto scale_a_dram_window = make_tile_window( + scale_a_window.get_bottom_tensor_view(), + make_tuple(number{}, number<64 / WG::kM>{}), + scale_a_window.get_window_origin(), + PipelinePolicy::template MakeMXFP4_ScaleA_FlatDramTileDistribution()); - using ABlockTile = decltype(load_tile(a_copy_dram_window)); - ABlockTile a_block_tile; + auto scale_b_dram_window = make_tile_window( + scale_b_window.get_bottom_tensor_view(), + make_tuple(number{}, number<64 / WG::kN>{}), + scale_b_window.get_window_origin(), + PipelinePolicy::template MakeMXFP4_ScaleB_DramTileDistribution()); - enum - { - PrefillBeforeGemm = 1, - PrefillAfterGemm = 2, - PrefillAlways = PrefillBeforeGemm | PrefillAfterGemm, + // ping pong buffer for scale A + statically_indexed_array< + statically_indexed_array, + MIterPerWarp / MXdlPack> + scale_a_dram_windows; + statically_indexed_array, + MIterPerWarp / MXdlPack> + scale_a_tile_tensor_ping; + statically_indexed_array, + MIterPerWarp / MXdlPack> + scale_a_tile_tensor_pong; + + // ping pong buffer for scale B + statically_indexed_array< + statically_indexed_array, + NIterPerWarp / NXdlPack> + scale_b_dram_windows; + statically_indexed_array, + NIterPerWarp / NXdlPack> + scale_b_tile_tensor_ping; + statically_indexed_array, + NIterPerWarp / NXdlPack> + scale_b_tile_tensor_pong; + + auto async_load_tile_ = [](auto lds, auto dram) { + async_load_tile(lds, dram, number<-1>{}, true_type{}, false_type{}); }; -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS - auto prefill_lds_a_stage1 = - [&]([[maybe_unused]] auto lds_tile_a, auto dram_tile_a, auto prefill_location) { - // global -> lds - if constexpr(prefill_location & PrefillAfterGemm) - async_load_tile(lds_tile_a, dram_tile_a); - }; - auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { - // async_load_fence(); - // __builtin_amdgcn_s_waitcnt(0x03fc); - // data has been stored in lds, no need more operation. - static_assert(std::is_same_v, - "buffer_load_lds don't support element func fot A before mfma"); - }; -#else - auto prefill_lds_a_stage1 = - [&]([[maybe_unused]] auto lds_tile_a, auto dram_tile_a, auto prefill_location) { - // global -> vgpr - if constexpr(prefill_location & PrefillBeforeGemm) - a_block_tile = load_tile(dram_tile_a); - }; - auto prefill_lds_a_stage2 = [&]([[maybe_unused]] auto lds_tile_a) { - // vgpr -> lds - auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(lds_tile_a, a_block_tile_transformed); - }; -#endif // HEAD // Prefetch A0 - prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window, number{}); - - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); + async_load_tile_(a_store_lds_window_ping, a_dram_window); + move_tile_window(a_dram_window, {0, kKPerBlock}); // prefetch B static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { - if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) - { - auto scale_n_iter = nIter / number{}; - auto scale_k_iter = kIter / number{}; + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + b_warp_tensor_ping(nIter)(kIter) = load_tile_with_offset( + b_flat_dram_windows(nIter), number{}); + }); + // move B window to next flat K + move_tile_window(b_flat_dram_windows(nIter), {0, KIterPerWarp * KFlatPerBlockPerIter}); + }); - scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = - scale_b_flat_dram_window; - move_tile_window( - scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), - {scale_n_iter * NFlatPerBlockPerIter, scale_k_iter * ScaleKFlatPerWarp}); - scale_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = - load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); - } - auto packed_n_idx = nIter / number{}; - auto packed_n_rank = nIter % number{}; + // prefetch Scale A + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; + move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), + {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - move_tile_window(b_flat_dram_windows(nIter)(kIter), - {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + - packed_n_rank, - kIter * KFlatPerBlockPerIter}); - - ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); - b_warp_tensor_ping(nIter)(kIter) = ub.u; + scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) = + load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); }); }); - // move B window to next flat K - move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); - move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); + // move Scale A window to next K + move_tile_window(scale_a_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); - prefill_lds_a_stage2(a_copy_lds_window_ping); + // prefetch Scale B + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; + move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), + {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); + scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) = + load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); + }); + }); + // move Scale B window to next K + move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); __builtin_amdgcn_sched_barrier(0); // Prefetch A1 - prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window, number{}); - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - + if constexpr(HasHotLoop || TailNum == TailNumber::Even) + { + async_load_tile_(a_store_lds_window_pong, a_dram_window); + move_tile_window(a_dram_window, {0, kKPerBlock}); + } // initialize C - tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + clear_tile(c_block_tile); - __builtin_amdgcn_s_waitcnt(Bload_total_num); - block_sync_lds(); + statically_indexed_array a_warp_tensor; // preload A00,A10... from lds - statically_indexed_array{})(number<0>{}))), - m_preload> - a_warp_tensor; - + s_waitcnt_barrier(); static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MIterPerWarp; - constexpr auto kIter = loadIter / MIterPerWarp; - a_warp_tensor(loadIter) = - load_tile(a_warp_windows_ping(number{})(number{})); - scale_a_tile_tensor(loadIter) = - load_tile(scale_a_dram_windows(number{})(number{})); - }); + constexpr auto mIter = loadIter % MXdlPack; + constexpr auto kIter = loadIter / MXdlPack; + a_warp_tensor(loadIter) = load_tile_with_offset( + a_warp_window_ping, tuple, number>{}); + }); __builtin_amdgcn_sched_barrier(0); - // statically_indexed_array dequant_B_n; - - //TODO: maybe need quant scale(change fp32 scale into fp8e8m0) - - // auto convert_ascale_e8m0 = [&](auto a_scale_fp32) { - // return bit_cast(uint8x4_t{ 0, 0, 0, bit_cast(static_cast(a_scale_fp32)) }); - // }; - // MAIN LOOP - index_t iCounter = (num_loop - 1) / 2; - while(iCounter > 0) - { + auto main_body_implx2 = [&]() mutable { // prefetch B(2i+1) - static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) - { - auto scale_n_iter = nIter / number{}; - auto scale_k_iter = kIter / number{}; - - scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = - scale_b_flat_dram_window; - - move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), - {scale_n_iter * NFlatPerBlockPerIter, - scale_k_iter * ScaleKFlatPerWarp}); - - scale_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = - load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); - } - - auto packed_n_idx = nIter / number{}; - auto packed_n_rank = nIter % number{}; - - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - - move_tile_window( - b_flat_dram_windows(nIter)(kIter), - {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + - packed_n_rank, - kIter * KFlatPerBlockPerIter}); - - ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); - b_warp_tensor_pong(nIter)(kIter) = ub.u; + b_warp_tensor_pong(nIter)(kIter) = load_tile_with_offset( + b_flat_dram_windows(nIter), number{}); + if constexpr(kIter == KIterPerWarp - 1) + move_tile_window(b_flat_dram_windows(nIter), + {0, BlockGemmShape::flatKPerBlock}); }); }); - // Prefill A(2i+1) - prefill_lds_a_stage2(a_copy_lds_window_pong); + // prefetch Scale A and Scale B (2i+1) + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; + move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), + {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); + + scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) = + load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); + }); + }); + + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; + move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), + {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); + + scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) = + load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); + }); + }); + + // GEMM 2i + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; + constexpr auto m_iter = mIter_pack * MXdlPack + imxdl; + constexpr auto k_iter = kIter_pack * KXdlPack + ikxdl; + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + constexpr auto n_iter = nIter_pack * NXdlPack + inxdl; + + // read C warp tensor from C block tensor + CWarpTensor c_warp_tensor; + c_warp_tensor.get_thread_buffer() = + c_block_tile.get_y_sliced_thread_data( + merge_sequences(sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + // warp GEMM + WG{}.template + operator()( + c_warp_tensor, + a_warp_tensor(number{}), + b_warp_tensor_ping(nIter_pack * number{} + inxdl)( + kIter_pack * number{} + ikxdl), + scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) + .get_thread_buffer()[0], + scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) + .get_thread_buffer()[0]); + + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + constexpr auto addr = + m_iter % 2 + k_iter * 2 + m_iter / 2 * 4 + m_preload; + if constexpr(addr < (KIterPerWarp * MIterPerWarp) && + (nIter_pack == NIterPerWarp / NXdlPack - 1)) + { + constexpr auto AmIter = addr % 2 + addr / 4 * 2; + constexpr auto AkIter = addr / 2 % 2; + a_warp_tensor(number{}) = load_tile_with_offset( + a_warp_window_ping, + tuple, number>{}); + } + }); + }); + }); + }); + }); + // barrier as ds_load A(2i) and buffer_load_lds A(2i + 1) finished + s_waitcnt< // vmcnt + Bload_num + ScaleAload_num + ScaleBload_num>(); + block_sync_lds(); // Prefetch A(2i+2) - prefill_lds_a_stage1( - a_copy_lds_window_ping, a_copy_dram_window, number{}); - // GEMM 2i - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - // read C warp tensor from C block tensor - CWarpTensor c_warp_tensor; - - c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - - // warp GEMM - WG{}.template operator()<0, 0>(c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter)(kIter / number{}), - convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), - convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( - kIter / number{}).get_thread_buffer()[0])); - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); - }); - // preload next A from lds - if constexpr((kIter * MIterPerWarp + mIter) < - (KIterPerWarp * MIterPerWarp - m_preload)) - { - constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; - constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); - a_warp_tensor(number{}) = - load_tile(a_warp_windows_ping(number{})(number{})); - scale_a_tile_tensor(number{}) = - load_tile(scale_a_dram_windows(number{})(number{})); - } - - // barrier - if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) - { - __builtin_amdgcn_s_waitcnt(Bload_total_num); - block_sync_lds(); - } - }); - }); - prefill_lds_a_stage1( - a_copy_lds_window_ping, a_copy_dram_window, number{}); - - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); + async_load_tile_(a_store_lds_window_ping, a_dram_window); + move_tile_window(a_dram_window, {0, kKPerBlock}); // move B window to next flat K - move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); - move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); + move_tile_window(scale_a_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); + move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); + // preload A(2i+1) static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MIterPerWarp; - constexpr auto kIter = loadIter / MIterPerWarp; - a_warp_tensor(loadIter) = - load_tile(a_warp_windows_pong(number{})(number{})); + constexpr auto mIter = loadIter % MXdlPack; + constexpr auto kIter = loadIter / MXdlPack; + a_warp_tensor(loadIter) = load_tile_with_offset( + a_warp_window_pong, tuple, number>{}); }); - // HotLoopScheduler(); + HotLoopScheduler(); - // Next K + ////////////////////////////// Next K ////////////////////////////// // prefetch B(2i+2) - static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) - { - auto scale_n_iter = nIter / number{}; - auto scale_k_iter = kIter / number{}; - - scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = - scale_b_flat_dram_window; - - move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), - {scale_n_iter * NFlatPerBlockPerIter, - scale_k_iter * ScaleKFlatPerWarp}); - - scale_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = - load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); - } - - auto packed_n_idx = nIter / number{}; - auto packed_n_rank = nIter % number{}; - - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - move_tile_window( - b_flat_dram_windows(nIter)(kIter), - {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + - packed_n_rank, - kIter * KFlatPerBlockPerIter}); - - ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); - b_warp_tensor_ping(nIter)(kIter) = ub.u; + b_warp_tensor_ping(nIter)(kIter) = load_tile_with_offset( + b_flat_dram_windows(nIter), number{}); + if constexpr(kIter == KIterPerWarp - 1) + move_tile_window(b_flat_dram_windows(nIter), + {0, BlockGemmShape::flatKPerBlock}); }); }); - // Prefill A(2i+2) - prefill_lds_a_stage2(a_copy_lds_window_ping); + // prefetch Scale A and Scale B (2i+2) + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; + move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), + {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); - // Prefetch A(2i+3) - prefill_lds_a_stage1( - a_copy_lds_window_pong, a_copy_dram_window, number{}); + scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) = + load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); + }); + }); + + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; + move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), + {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); + + scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) = + load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); + }); + }); // GEMM 2i+1 - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - // read C warp tensor from C block tensor - CWarpTensor c_warp_tensor; - c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + // read C warp tensor from C block tensor + CWarpTensor c_warp_tensor; + c_warp_tensor.get_thread_buffer() = + c_block_tile.get_y_sliced_thread_data( + merge_sequences( + sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - // if constexpr(mIter == 0) - // dequant_mxfp4( - // b_warp_tensor_pong(nIter)(kIter / number{}), - // scale_b_warp_tensor_pong(nIter / number{})( - // kIter / number{}), - // nIter, - // kIter); + // warp GEMM + WG{}.template + operator()( + c_warp_tensor, + a_warp_tensor(number{}), + b_warp_tensor_pong(nIter_pack * number{} + inxdl)( + kIter_pack * number{} + ikxdl), + scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) + .get_thread_buffer()[0], // scale A + scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) + .get_thread_buffer()[0]); // scale B - // warp GEMM - WG{}.template operator()<0, 0>(c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter)(kIter / number{}), - convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), - convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( - kIter / number{}).get_thread_buffer()[0])); - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + + (kIter_pack * KXdlPack + ikxdl) * 2 + + (mIter_pack * MXdlPack + imxdl) / 2 * 4 + + m_preload; + if constexpr(addr < (KIterPerWarp * MIterPerWarp) && + (nIter_pack == NIterPerWarp / NXdlPack - 1)) + { + constexpr auto AmIter = addr % 2 + addr / 4 * 2; + constexpr auto AkIter = addr / 2 % 2; + a_warp_tensor(number{}) = load_tile_with_offset( + a_warp_window_pong, + tuple, number>{}); + } + }); + }); }); - // preload next A from lds - if constexpr((kIter * MIterPerWarp + mIter) < - (KIterPerWarp * MIterPerWarp - m_preload)) - { - constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; - constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); - a_warp_tensor(number{}) = - load_tile(a_warp_windows_pong(number{})(number{})); - scale_a_tile_tensor(number{}) = - load_tile(scale_a_dram_windows(number{})(number{})); - } - - // barrier - if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) - { - __builtin_amdgcn_s_waitcnt(Bload_total_num); - block_sync_lds(); - } }); }); - prefill_lds_a_stage1( - a_copy_lds_window_pong, a_copy_dram_window, number{}); + // barrier as ds_load A(2i + 1) and buffer_load_lds A(2i + 2) finished + s_waitcnt< // vmcnt + Bload_num + ScaleAload_num + ScaleBload_num>(); + block_sync_lds(); - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); + // Prefetch A(2i+3) + async_load_tile_(a_store_lds_window_pong, a_dram_window); + move_tile_window(a_dram_window, {0, kKPerBlock}); // move B window to next flat K - move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); - move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); + move_tile_window(scale_a_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); + move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); + // preload A(2i+2) static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MIterPerWarp; - constexpr auto kIter = loadIter / MIterPerWarp; - a_warp_tensor(loadIter) = - load_tile(a_warp_windows_ping(number{})(number{})); + constexpr auto mIter = loadIter % MXdlPack; + constexpr auto kIter = loadIter / MXdlPack; + a_warp_tensor(loadIter) = load_tile_with_offset( + a_warp_window_ping, tuple, number>{}); }); - // HotLoopScheduler(); + HotLoopScheduler(); + }; - iCounter--; + if constexpr(HasHotLoop) + { + index_t iCounter = (num_loop - 1) / 2; + do + { + main_body_implx2(); + iCounter--; + } while(iCounter > 0); } // TAIL if constexpr(TailNum == TailNumber::Even) { // prefetch B(loopK) - const int b_k_off = b_flat_dram_window.get_tile_distribution().calculate_index()[I1] / - ContinuousKPerThread / WG::kN * ContinuousKPerThread; - static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) - { - auto scale_n_iter = nIter / number{}; - auto scale_k_iter = kIter / number{}; - - scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = - scale_b_flat_dram_window; - - move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), - {scale_n_iter * NFlatPerBlockPerIter, - scale_k_iter * ScaleKFlatPerWarp}); - - scale_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = - load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); - } + b_warp_tensor_pong(nIter)(kIter) = load_tile_with_offset( + b_flat_dram_windows(nIter), + make_tuple(number<0>{}, number{})); }); - - const int b_k_off_inter = kIter * kKPerBlock / MXFP4KPerWarp + b_k_off; - if(b_k_off_inter < kKPerBlock - k_padded_zeros) - { - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - auto packed_n_idx = nIter / number{}; - auto packed_n_rank = nIter % number{}; - - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - - move_tile_window( - b_flat_dram_windows(nIter)(kIter), - {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + - packed_n_rank, - kIter * KFlatPerBlockPerIter}); - - ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); - b_warp_tensor_pong(nIter)(kIter) = ub.u; - }); - } }); - // Prefill A(loopK) - prefill_lds_a_stage2(a_copy_lds_window_pong); + // prefetch Scale A and Scale B (2i+1) + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; + move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), + {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); + + scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) = + load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); + }); + }); + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; + move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), + {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); + + scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) = + load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); + }); + }); // GEMM loopK-1 - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - // read C warp tensor from C block tensor - CWarpTensor c_warp_tensor; + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + // read C warp tensor from C block tensor + CWarpTensor c_warp_tensor; + c_warp_tensor.get_thread_buffer() = + c_block_tile.get_y_sliced_thread_data( + merge_sequences( + sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + // warp GEMM + WG{}.template + operator()( + c_warp_tensor, + a_warp_tensor(number{}), + b_warp_tensor_ping(nIter_pack * number{} + inxdl)( + kIter_pack * number{} + ikxdl), + scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) + .get_thread_buffer()[0], // scale A + scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) + .get_thread_buffer()[0]); // scale B - // if constexpr(mIter == 0) - // dequant_mxfp4( - // b_warp_tensor_ping(nIter)(kIter / number{}), - // scale_b_warp_tensor_ping(nIter / number{})( - // kIter / number{}), - // nIter, - // kIter); - - // warp GEMM - WG{}.template operator()<0, 0>(c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter)(kIter / number{}), - convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), - convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( - kIter / number{}).get_thread_buffer()[0])); - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + + (kIter_pack * KXdlPack + ikxdl) * 2 + + (mIter_pack * MXdlPack + imxdl) / 2 * 4 + + m_preload; + if constexpr(addr < (KIterPerWarp * MIterPerWarp) && + (nIter_pack == NIterPerWarp / NXdlPack - 1)) + { + constexpr auto AmIter = addr % 2 + addr / 4 * 2; + constexpr auto AkIter = addr / 2 % 2; + a_warp_tensor(number{}) = load_tile_with_offset( + a_warp_window_ping, + tuple, number>{}); + } + }); + }); }); - // preload next A from lds - if constexpr((kIter * MIterPerWarp + mIter) < - (KIterPerWarp * MIterPerWarp - m_preload)) - { - constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; - constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); - a_warp_tensor(number{}) = - load_tile(a_warp_windows_ping(number{})(number{})); - scale_a_tile_tensor(number{}) = - load_tile(scale_a_dram_windows(number{})(number{})); - } - - // barrier - if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) - { - __builtin_amdgcn_s_waitcnt(Bload_total_num); - block_sync_lds(); - } }); }); + // barrier as ds_load A(2i) and buffer_load_lds A(2i + 1) finished + s_waitcnt< // vmcnt + Bload_num + ScaleAload_num + ScaleBload_num>(); + block_sync_lds(); + // preload A(2i+1) static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MIterPerWarp; - constexpr auto kIter = loadIter / MIterPerWarp; - a_warp_tensor(loadIter) = - load_tile(a_warp_windows_pong(number{})(number{})); + constexpr auto mIter = loadIter % MXdlPack; + constexpr auto kIter = loadIter / MXdlPack; + a_warp_tensor(loadIter) = load_tile_with_offset( + a_warp_window_pong, tuple, number>{}); }); - __builtin_amdgcn_sched_barrier(0); - // Last2ndHotLoopScheduler(); + Last2ndHotLoopScheduler(); // GEMM loopK - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - if(kIter * WG::kK < kKPerBlock - k_padded_zeros) - { - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - // read C warp tensor from C block tensor - CWarpTensor c_warp_tensor; + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + // read C warp tensor from C block tensor + CWarpTensor c_warp_tensor; + c_warp_tensor.get_thread_buffer() = + c_block_tile.get_y_sliced_thread_data( + merge_sequences( + sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - c_warp_tensor.get_thread_buffer() = - c_block_tile.get_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + // warp GEMM + WG{}.template + operator()( + c_warp_tensor, + a_warp_tensor(number{}), + b_warp_tensor_pong(nIter_pack * number{} + inxdl)( + kIter_pack * number{} + ikxdl), + scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) + .get_thread_buffer()[0], // scale A + scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) + .get_thread_buffer()[0]); // scale B - // if constexpr(mIter == 0) - // dequant_mxfp4( - // b_warp_tensor_pong(nIter)(kIter / number{}), - // scale_b_warp_tensor_pong(nIter / number{})( - // kIter / number{}), - // nIter, - // kIter); - - // warp GEMM - // WG{}(c_warp_tensor, - // a_warp_tensor(number{}), - // dequant_B_n[nIter]); - WG{}.template operator()<0, 0>(c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter)(kIter / number{}), - convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), - convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( - kIter / number{}).get_thread_buffer()[0])); - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + + (kIter_pack * KXdlPack + ikxdl) * 2 + + (mIter_pack * MXdlPack + imxdl) / 2 * 4 + + m_preload; + if constexpr(addr < (KIterPerWarp * MIterPerWarp) && + (nIter_pack == NIterPerWarp / NXdlPack - 1)) + { + constexpr auto AmIter = addr % 2 + addr / 4 * 2; + constexpr auto AkIter = addr / 2 % 2; + a_warp_tensor(number{}) = load_tile_with_offset( + a_warp_window_pong, + tuple, number>{}); + } + }); }); - if constexpr((kIter * MIterPerWarp + mIter) < - (KIterPerWarp * MIterPerWarp - m_preload)) - { - constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; - constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); - a_warp_tensor(number{}) = - load_tile(a_warp_windows_pong(number{})(number{})); - scale_a_tile_tensor(number{}) = - load_tile(scale_a_dram_windows(number{})(number{})); - } - // barrier - // if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) - // { - // block_sync_lds(); - // } }); - } + }); }); - // LastHotLoopScheduler(); + LastHotLoopScheduler(); } else if constexpr(TailNum == TailNumber::Odd) { // GEMM loopK - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - // read C warp tensor from C block tensor - CWarpTensor c_warp_tensor; + static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { + static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { + static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + // read C warp tensor from C block tensor + CWarpTensor c_warp_tensor; + c_warp_tensor.get_thread_buffer() = + c_block_tile.get_y_sliced_thread_data( + merge_sequences( + sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + // warp GEMM + WG{}.template + operator()( + c_warp_tensor, + a_warp_tensor(number{}), + b_warp_tensor_ping(nIter_pack * number{} + inxdl)( + kIter_pack * number{} + ikxdl), + scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) + .get_thread_buffer()[0], // scale A + scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) + .get_thread_buffer()[0]); // scale B - // if constexpr(mIter == 0) - // dequant_mxfp4( - // b_warp_tensor_ping(nIter)(kIter / number{}), - // scale_b_warp_tensor_ping(nIter / number{})( - // kIter / number{}), - // nIter, - // kIter); - // warp GEMM - // WG{}(c_warp_tensor, a_warp_tensor(number{}), dequant_B_n[nIter]); - WG{}.template operator()<0, 0>(c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter)(kIter / number{}), - convert_ascale_e8m0(scale_a_tile_tensor(number{}).get_thread_buffer()[0]), - convert_ascale_e8m0(scale_b_warp_tensor_ping(nIter / number{})( - kIter / number{}).get_thread_buffer()[0])); - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + + (kIter_pack * KXdlPack + ikxdl) * 2 + + (mIter_pack * MXdlPack + imxdl) / 2 * 4 + + m_preload; + if constexpr(addr < (KIterPerWarp * MIterPerWarp) && + (nIter_pack == NIterPerWarp / NXdlPack - 1)) + { + constexpr auto AmIter = addr % 2 + addr / 4 * 2; + constexpr auto AkIter = addr / 2 % 2; + a_warp_tensor(number{}) = load_tile_with_offset( + a_warp_window_ping, + tuple, number>{}); + } + }); + }); }); - // preload next A from lds - if constexpr((kIter * MIterPerWarp + mIter) < - (KIterPerWarp * MIterPerWarp - m_preload)) - { - constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; - constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); - a_warp_tensor(number{}) = - load_tile(a_warp_windows_ping(number{})(number{})); - scale_a_tile_tensor(number{}) = - load_tile(scale_a_dram_windows(number{})(number{})); - } - - // barrier - if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) - { - __builtin_amdgcn_s_waitcnt(Bload_total_num); - block_sync_lds(); - } }); }); - // LastHotLoopScheduler(); + LastHotLoopScheduler(); + } + else + { + static_assert(false, "Wrong TailNum"); } - return c_block_tile; } - - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, - const DequantAWindow& scale_a_window, - const DequantBFlatWindow& scale_b_flat_window, - const index_t num_loop, - const index_t k_padded_zeros, - void* p_smem_ping, - void* p_smem_pong) const - { - return operator()(a_dram_block_window_tmp, - identity{}, - b_flat_dram_block_window_tmp, - scale_a_window, - scale_b_flat_window, - num_loop, - k_padded_zeros, - p_smem_ping, - p_smem_pong); - } - - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, - const DequantAWindow& scale_a_window, - const DequantBFlatWindow& scale_b_flat_window, - const index_t num_loop, - void* p_smem_ping, - void* p_smem_pong) const - { - return operator()(a_dram_block_window_tmp, - identity{}, - b_flat_dram_block_window_tmp, - scale_a_window, - scale_b_flat_window, - num_loop, - 0, - p_smem_ping, - p_smem_pong); - } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index a4b5cd9d10..6a8bb1ba6c 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp" +#include "ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp" namespace ck_tile { @@ -236,179 +237,428 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy } }; -struct F8xMXF4FlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy +struct F8xMXF4FlatmmPipelineAgBgCrPolicy : MXF4FlatmmPipelineAgBgCrPolicy { static constexpr auto I0 = number<0>{}; static constexpr auto I1 = number<1>{}; static constexpr auto I2 = number<2>{}; - static constexpr index_t KBPerLoad = 32; - static constexpr index_t N_Pack = 2; // it's fixed for fp4 - static constexpr index_t K_Pack = 2; // it's fixed for fp4 + static constexpr index_t kDramLoadPackBytes = 128; - template - CK_TILE_HOST_DEVICE static constexpr auto - TransformF8xF4_ATensorView(const NativeADramTensorView& a_dram_view) + static constexpr int MXdlPack = 1; + static constexpr int NXdlPack = 1; + static constexpr int KXdlPack = 4; + + template + static inline constexpr auto wg_attr_num_access = + std::is_same_v, pk_fp4_t> + ? WGAttrNumAccessEnum::Single + : WGAttrNumAccessEnum::Double; + + template + CK_TILE_HOST_DEVICE static constexpr auto GetBlockFlatmm() { -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS - constexpr int DynamicTileOffsetFlag = 0; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + static_assert( + sizeof(ADataType) * numeric_traits::PackedSize == + sizeof(BDataType) * numeric_traits::PackedSize, + "sizeof(ADataType) / APackedSize must be equal to sizeof(BDataType) / BPackedSize!"); + using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; + using WarpTile = typename Problem::BlockGemmShape::WarpTile; + using WarpGemm = WarpGemmDispatcher< // + ADataType, + BDataType, + typename Problem::CDataType, + WarpTile::at(I0), + WarpTile::at(I1), + WarpTile::at(I2), + Problem::TransposeC, + false, + false, + wg_attr_num_access>; + using BlockFlatmmPolicy = BlockFlatmmASmemBSmemCRegV1CustomPolicy< // + ADataType, + BDataType, + typename Problem::CDataType, + BlockWarps, + WarpGemm>; + return BlockFlatmmASmemBSmemCRegV1{}; + } + template + CK_TILE_DEVICE static constexpr auto + MakeMXFP4_AAsyncLoadDramDescriptor(const TensorView& naive_view) + { + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1); - static_assert(MPerXdl == 16 && NPerXdl == 16); + static_assert(std::is_same_v); - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackA(); + const auto& naive_desc = naive_view.get_tensor_descriptor(); + constexpr auto ndims = remove_cvref_t::get_num_of_dimension(); + static_assert(ndims == 2, "only support 2D tensor"); + const auto rows = naive_desc.get_length(number<0>{}); + const auto cols = naive_desc.get_length(number<1>{}); - constexpr int ContiguousThreadsCntInDS_READ_16B = 4; + constexpr index_t APackedSize = numeric_traits::PackedSize; + constexpr index_t K2 = GetSmemPackA() * APackedSize; // f4=32; f8=16 + constexpr index_t K1 = kDramLoadPackBytes * APackedSize / K2; // 8 + const index_t K0 = cols / (K1 * K2); + const auto col_lens = make_tuple(K0, number{}, number{}); - // implement swizzle pattern on global side - // because we can't adjust the ds_write pattern of BUFFER_LOAD_LDS. - auto swizzle_a_dram_view_1 = transform_tensor_view( - a_dram_view, - make_tuple( - // M-dim is not affected by swizzle pattern - make_unmerge_transform( - make_tuple(number{}, number{})), - // K-dim is the swizzle dimension - make_unmerge_transform(make_tuple(number{}, - number{}, - number{}))), - make_tuple(sequence<0>{}, sequence<1>{}), - make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{})); + constexpr index_t M1 = 4; // so that we can use imm offset to load lds + const index_t M0 = rows / M1; + const auto row_lens = make_tuple(M0, number{}); - auto swizzle_a_dram_view_2 = transform_tensor_view( - swizzle_a_dram_view_1, - make_tuple(make_pass_through_transform(number{}), - make_xor_transform(make_tuple(number{}, - number{})), - make_pass_through_transform(number{}), - make_pass_through_transform(number{})), + const auto desc_0 = + make_naive_tensor_descriptor_packed(container_concat(row_lens, col_lens)); + const auto desc_1 = transform_tensor_descriptor( + desc_0, + make_tuple(make_pass_through_transform(M0), + make_xor_transform(make_tuple(number{}, number{})), + make_pass_through_transform(K0), + make_pass_through_transform(number{})), make_tuple(sequence<0>{}, sequence<1, 3>{}, sequence<2>{}, sequence<4>{}), make_tuple(sequence<0>{}, sequence<1, 3>{}, sequence<2>{}, sequence<4>{})); - - return transform_tensor_view( - swizzle_a_dram_view_2, - make_tuple( - make_merge_transform_v3_division_mod( - make_tuple(number{}, number{})), - make_merge_transform_v3_division_mod(make_tuple(number{}, - number{}, - number{}))), + const auto desc = transform_tensor_descriptor( // + desc_1, + make_tuple(make_merge_transform_v3_division_mod(row_lens), + make_merge_transform_v3_division_mod(col_lens)), make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{}), make_tuple(sequence<0>{}, sequence<1>{})); -#else - return a_dram_view; -#endif + // printf("A async load dram desc %d x %d: \n", desc.get_length(I0), desc.get_length(I1)); + + return tensor_view, + TensorView::DstInMemOp>{naive_view.buf_, desc}; } template - CK_TILE_HOST_DEVICE static constexpr auto MakeF8xF4_ReadALdsBlockDescriptor() + CK_TILE_DEVICE static constexpr auto MakeADramTileDistribution() { + + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; + static_assert(std::is_same_v); + + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t APackedSize = numeric_traits::PackedSize; + + constexpr index_t K2 = GetSmemPackA() * APackedSize; // f4=32; f8=16 + constexpr index_t K1 = kDramLoadPackBytes * APackedSize / K2; // 8 + constexpr index_t K0 = KPerBlock / (K1 * K2); // KPerBlock/256 + + constexpr index_t M2 = get_warp_size() / K1; // 8 + constexpr index_t M1 = BlockSize / get_warp_size(); // 4 + constexpr index_t M0 = MPerBlock / (M2 * M1); + static_assert(M0 * M1 * M2 == MPerBlock, "M0, M1, M2 must cover whole MPerBlock!"); + static_assert(K0 * K1 * K2 == KPerBlock, "K0, K1, K2 must cover whole KPerBlock!"); + + return make_static_tile_distribution( + tile_distribution_encoding< // + sequence<1>, + tuple, sequence>, // ?,4,8 1,8,32 or 2,8,16 + tuple, sequence<1, 2>>, // M1 M2,K1 + tuple, sequence<2, 1>>, + sequence<1, 2, 2>, // M0,K0,K2 + sequence<0, 0, 2>>{}); + } + + template + CK_TILE_DEVICE static constexpr auto MakeMXFP4_ALdsBlockDescriptor() + { + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1); - static_assert(MPerXdl == 16 && NPerXdl == 16); + static_assert(std::is_same_v); /*reduce transform layers,compare with old ck*/ - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackA(); + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t APackedSize = numeric_traits::PackedSize; + constexpr index_t K2 = GetSmemPackA() * APackedSize; // f4=32; f8=16 + constexpr index_t K1 = kDramLoadPackBytes * APackedSize / K2; // 8 + constexpr index_t K0 = KPerBlock / (K1 * K2); // KPerBlock/256 + static_assert(K0 * K1 * K2 == KPerBlock, "K0, K1, K2 must cover whole KPerBlock!"); - constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, number{}, number{}), - make_tuple(number{}, number{}, number<1>{}), - number{}, + constexpr index_t M3 = 4; // so that we can use imm offset to load lds + constexpr index_t M2 = get_warp_size() / K1 / M3; // 2 + constexpr index_t M1 = MPerXdl / (M2 * M3); // 2 + constexpr index_t M0 = MPerBlock / (M1 * M2 * M3); // MPerBlock/16 + static_assert(M0 * M1 * M2 * M3 == MPerBlock, "M0, M1, M2, M3 must cover whole MPerBlock!"); + + constexpr index_t Pad = 4 * K2; // 4 * 32 + + constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor( // + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}, + number{}, + number{}), + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}, + number{}, + number<1>{}), + number{}, number<1>{}); - constexpr int ContiguousThreadsCntInDS_READ_16B = 4; - - constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor( + constexpr auto a_lds_block_desc_1 = transform_tensor_descriptor( a_lds_block_desc_0, - make_tuple(make_xor_transform(make_tuple(number{}, - number{})), - make_pass_through_transform(number{})), - make_tuple(sequence<1, 0>{}, sequence<2>{}), - make_tuple(sequence<1, 0>{}, sequence<2>{})); - + make_tuple(make_pass_through_transform(M0), + make_pass_through_transform(M1), + make_pass_through_transform(K0), + make_pass_through_transform(M2), + make_xor_transform(make_tuple(number{}, number{})), + make_pass_through_transform(number{})), + make_tuple(sequence<0>{}, + sequence<1>{}, + sequence<2>{}, + sequence<3>{}, + sequence<4, 5>{}, + sequence<6>{}), + make_tuple(sequence<0>{}, + sequence<1>{}, + sequence<2>{}, + sequence<3>{}, + sequence<4, 5>{}, + sequence<6>{})); constexpr auto a_lds_block_desc = transform_tensor_descriptor( - a_lds_block_desc_permuted, - make_tuple(make_pass_through_transform(number{}), + a_lds_block_desc_1, + make_tuple(make_merge_transform_v3_division_mod( + make_tuple(number{}, number{}, number{}, number{})), make_merge_transform_v3_division_mod( - make_tuple(number{}, number{}))), - make_tuple(sequence<1>{}, sequence<0, 2>{}), + make_tuple(number{}, number{}, number{}))), + make_tuple(sequence<0, 1, 3, 4>{}, sequence<2, 5, 6>{}), make_tuple(sequence<0>{}, sequence<1>{})); + // return a_lds_block_desc_permuted; return a_lds_block_desc; } template - CK_TILE_HOST_DEVICE static constexpr auto MakeF8xF4_WriteALdsBlockDescriptor() - { -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackA(); - return make_naive_tensor_descriptor(make_tuple(number{}, number{}), - make_tuple(number{}, number<1>{}), - number{}, - number<1>{}); -#else - return MakeF16xF4_ReadALdsBlockDescriptor(); -#endif - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeF8xF4_ALDS_TileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto MakeMXF4_ALDS_TileDistribution() { using TileShape = typename Problem::BlockGemmShape; static_assert(TileShape::WarpTile::at(I1) == 16, "requires XDL_N == 16"); static_assert(TileShape::BlockWarps::at(I0) == 1, "requires Wave_M == 1"); - constexpr int Repeat = TileShape::BlockWarps::at(number<1>{}); - constexpr int M0 = TileShape::WarpTile::at(I0); + constexpr int M_warps = TileShape::BlockWarps::at(number<0>{}); + constexpr int N_warps = TileShape::BlockWarps::at(number<1>{}); + constexpr int M_Lane = TileShape::WarpTile::at(I0); // 16 - constexpr int K_Lane = 64 / TileShape::WarpTile::at(I1); // 4 + constexpr int K_Lane = 64 / M_Lane; // 4 - constexpr int K2 = TileShape::WarpTile::at(I2) / K_Lane; // 128 / 4 = 32 - constexpr int XDL_PerThreadK = KBPerLoad / K2; // 32 / 32 = 1 - constexpr int K0 = K_Lane; // 4 + constexpr int K_Thread = TileShape::WarpTile::at(I2) / K_Lane; // 32 + constexpr index_t num_access_v = static_cast(wg_attr_num_access); + constexpr int K1 = K_Thread / num_access_v; // 16 return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1>>, - tuple, sequence<0, 0>>, - sequence<2>, - sequence<2>>{}); + std::conditional_t< + num_access_v == 1, + tile_distribution_encoding< + sequence, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<0, 2>>, + sequence<2>, + sequence<1>>, + tile_distribution_encoding< // + sequence, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<1, 2>>, + sequence<2, 2>, + sequence<0, 2>>>{}); } - // assum a8 scale dtype is fp32 template - CK_TILE_HOST_DEVICE static constexpr auto MakeAScaleDramTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto MakeMXFP4_BFlatDramTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; + + static_assert(TileShape::WarpTile::at(I1) == 16, "only for XDL_N == 16"); + + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t WaveSize = get_warp_size(); + constexpr index_t WaveNum = BlockSize / WaveSize; + + constexpr index_t K1 = WaveSize; // threads cnt in K dim + constexpr index_t KWavePerBlk = 1; + constexpr index_t K0 = KWavePerBlk; + + constexpr index_t NWavePerBlk = TileShape::BlockWarps::at(number<1>{}); // N_Warp + + constexpr index_t WaveRepeat = WaveNum / TileShape::flatNPerWarp; + constexpr index_t kKPerThread = 32; + constexpr index_t num_access_v = static_cast(wg_attr_num_access); + constexpr index_t K2 = kKPerThread / num_access_v; + + return make_static_tile_distribution( + std::conditional_t< // + num_access_v == 1, + tile_distribution_encoding< // + sequence, + tuple, // 4 2 + sequence>, // 1 64 32 + tuple, sequence<2>>, + tuple, sequence<1>>, + sequence<2>, + sequence<2>>, + tile_distribution_encoding< // + sequence, + tuple, // 4 2 + sequence>, // 2 1 64 16 + tuple, sequence<2>>, + tuple, sequence<2>>, + sequence<2, 2>, + sequence<0, 3>>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeMXFP4_ScaleA_DramTileDistribution() { using TileShape = typename Problem::BlockGemmShape; // ck_tile::TileFlatmmShape - constexpr int Repeat = TileShape::BlockWarps::at(number<1>{}); // 4 + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t WaveSize = get_warp_size(); + constexpr index_t WaveNum = BlockSize / WaveSize; - constexpr index_t M0 = TileShape::WarpTile::at(I0); // 16 + constexpr index_t kMPerBlock = TileShape::BlockTile::at(I0); - constexpr int K_Lane = 64 / TileShape::WarpTile::at(I1); // 4 + constexpr index_t M_Warps = TileShape::BlockWarps::at(I0); + constexpr index_t N_Warps = TileShape::BlockWarps::at(I1); - constexpr int K2 = TileShape::WarpTile::at(I2) / K_Lane / 32; // 128 / 4 = 1 - constexpr int K0 = K_Lane; // 4 + static_assert(WaveNum == M_Warps * N_Warps, "Block warps do not match block size"); + + constexpr index_t M_Lanes = TileShape::WarpTile::at(I0); + constexpr index_t K_Lanes = 64 / M_Lanes; + + // Y dimension (M) decomposition + constexpr index_t Y2 = M_Lanes; + constexpr index_t Y1 = M_Warps; + constexpr index_t Y0 = kMPerBlock / (MXdlPack * Y1 * Y2); + + // X dimension (K) decomposition + constexpr index_t X0 = K_Lanes; + constexpr index_t X1 = 1; // packed 2x2 E8M0 data into 1 int32_t for load return make_static_tile_distribution( - tile_distribution_encoding, // repeat N_warps - tuple, sequence>, - tuple, sequence<2, 1>>, - tuple, sequence<0, 0>>, + tile_distribution_encoding, // repeat N_warps + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<0, 2>>, + sequence<1, 2>, + sequence<0, 1>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeMXFP4_ScaleB_DramTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; // ck_tile::TileFlatmmShape + + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t WaveSize = get_warp_size(); + constexpr index_t WaveNum = BlockSize / WaveSize; + + constexpr index_t kNPerBlock = TileShape::BlockTile::at(I1); + + constexpr index_t M_Warps = TileShape::BlockWarps::at(I0); + constexpr index_t N_Warps = TileShape::BlockWarps::at(I1); + + static_assert(WaveNum == M_Warps * N_Warps, "Block warps do not match block size"); + + constexpr index_t N_Lanes = TileShape::WarpTile::at(I1); + constexpr index_t K_Lanes = 64 / N_Lanes; + + // Y dimension (M) decomposition + constexpr index_t Y2 = N_Lanes; + constexpr index_t Y1 = N_Warps; + constexpr index_t Y0 = kNPerBlock / (NXdlPack * Y1 * Y2); + + // X dimension (K) decomposition + constexpr index_t X0 = K_Lanes; + constexpr index_t X1 = 1; // packed 2x2 E8M0 data into 1 int32_t for load + + return make_static_tile_distribution( + tile_distribution_encoding, // ? + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<0, 2>>, + sequence<1, 2>, + sequence<0, 1>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeMXFP4_ScaleA_FlatDramTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; + + constexpr index_t M_Warp = TileShape::BlockWarps::at(number<0>{}); + constexpr index_t K_Lane = 64 / TileShape::WarpTile::at(I0); + constexpr index_t M_Lane = TileShape::WarpTile::at(I0); + constexpr index_t N_Wrap = TileShape::BlockWarps::at(number<1>{}); + constexpr index_t MWavePerBlk = M_Warp; + + return make_static_tile_distribution( + tile_distribution_encoding, // ? + tuple, // second direction + sequence>, // first direction + tuple, sequence<2, 1>>, // which direction + tuple, sequence<0, 1>>, // which index + // sequence<2>, sequence<1>>{}); } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeMXFP4_ScaleB_FlatDramTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; + + constexpr index_t N_Warp = TileShape::BlockWarps::at(number<1>{}); + constexpr index_t K_Lane = 64 / TileShape::WarpTile::at(I1); + constexpr index_t N_Lane = TileShape::WarpTile::at(I1); + constexpr index_t M_Wrap = TileShape::BlockWarps::at(number<0>{}); + constexpr index_t NWavePerBlk = N_Warp; + + return make_static_tile_distribution( + tile_distribution_encoding, // ? + tuple, // second direction + sequence>, // first direction + tuple, sequence<2, 1>>, // which direction + tuple, sequence<0, 1>>, // which index + // + sequence<2>, + sequence<1>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeA() + { + using ADataType = remove_cvref_t; + constexpr index_t APackedSize = numeric_traits::PackedSize; + return sizeof(ADataType) * + MakeMXFP4_ALdsBlockDescriptor().get_element_space_size() / APackedSize; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return GetSmemSizeA(); + } }; } // namespace ck_tile