From 5b1e24427433e27f6614de011eac0607111aeffb Mon Sep 17 00:00:00 2001 From: Harisankar Sadasivan Date: Fri, 22 Mar 2024 19:41:42 +0000 Subject: [PATCH] 2-tile sk+ DP with atomics for FP16 --- example/01_gemm/README.md | 22 + .../device/impl/device_gemm_xdl_streamk.hpp | 32 +- .../gpu/grid/block_to_ctile_map.hpp | 183 +-- .../gpu/grid/gridwise_gemm_xdlops_streamk.hpp | 1180 +++++++++-------- 4 files changed, 696 insertions(+), 721 deletions(-) diff --git a/example/01_gemm/README.md b/example/01_gemm/README.md index 226783b03b..b041f0a3c3 100644 --- a/example/01_gemm/README.md +++ b/example/01_gemm/README.md @@ -21,3 +21,25 @@ Warm up Start running 5 times... Perf: 1.19685 ms, 107.657 TFlops, 78.8501 GB/s ``` + +# Instructions for ```example_gemm_xdl_streamk``` + +## Run ```example_gemm_xdl_streamk``` +```bash +# arg1: verification (0=no, 1=yes) +# arg2: initialization (0=no init, 1=integer value, 2=decimal value) +# arg3: time kernel (0=no, 1=yes) +# arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC +# arg10: NumSKBlocks(optional, defaults to DP GEMM) +bin/example_gemm_xdl_streamk 1 2 1 3840 4096 4096 4096 4096 4096 312 +``` + +Result (MI250 @ 1700Mhz, 181TFlops peak FP16 on 1 dye) +``` +a_m_k: dim 2, lengths {3840, 4096}, strides {4096, 1} +b_k_n: dim 2, lengths {4096, 4096}, strides {4096, 1} +c_m_n: dim 2, lengths {3840, 4096}, strides {4096, 1} +Recommended grid size :312 +Perf: 1.21689 ms, 105.884 TFlops, 79.2748 GB/s, GemmXdlStreamK_RRR_B256_Vec8x2x8_128x128x4x8 + +``` diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_streamk.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_streamk.hpp index 51b8958d61..5133556d79 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_streamk.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_streamk.hpp @@ -137,6 +137,8 @@ struct DeviceGemmXdlStreamK : public DeviceGemmStreamK; int occupancy, num_cu; - hipError_t rtn; - rtn = hipOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, kernel, BlockSize, GridwiseGemm::GetSharedMemoryNumberOfByte()); - hip_check_error(rtn); - + hip_check_error( + hipOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy, kernel, BlockSize, 0)); hipDeviceProp_t dev_prop; hipDevice_t dev; - rtn = hipGetDevice(&dev); - hip_check_error(rtn); - rtn = hipGetDeviceProperties(&dev_prop, dev); - hip_check_error(rtn); + hip_check_error(hipGetDevice(&dev)); + hip_check_error(hipGetDeviceProperties(&dev_prop, dev)); num_cu = dev_prop.multiProcessorCount; + printf("Assuming full GPU availability, recommended stream-k grid size for tuning :%0d\n", + num_cu * occupancy); return Argument{p_a, p_b, @@ -318,17 +317,12 @@ struct DeviceGemmXdlStreamK : public DeviceGemmStreamK; int occupancy, num_cu; - hipError_t rtn; - rtn = hipOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, kernel, BlockSize, GridwiseGemm::GetSharedMemoryNumberOfByte()); - hip_check_error(rtn); - + hip_check_error( + hipOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy, kernel, BlockSize, 0)); hipDeviceProp_t dev_prop; hipDevice_t dev; - rtn = hipGetDevice(&dev); - hip_check_error(rtn); - rtn = hipGetDeviceProperties(&dev_prop, dev); - hip_check_error(rtn); + hip_check_error(hipGetDevice(&dev)); + hip_check_error(hipGetDeviceProperties(&dev_prop, dev)); num_cu = dev_prop.multiProcessorCount; return std::make_unique(reinterpret_cast(p_a), diff --git a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp index a89e14cbdb..883298454a 100644 --- a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp +++ b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp @@ -1010,142 +1010,69 @@ struct BlockToCTileMap_GemmStreamK MDiv eqav_tiles_big; // for reduction MDiv eqav_tiles_little; // for reduction - // MDiv tile_swizzle_sub_m_rem; - //-------------------------------------- - // prefer construct on host BlockToCTileMap_GemmStreamK(uint32_t m, uint32_t n, uint32_t k, uint32_t num_cu, uint32_t occupancy, - uint32_t sk_blocks = 0xffffffff) + uint32_t sk_blocks = 0) { + // total output tiles uint32_t num_tiles = math::integer_divide_ceil(m, MPerBlock) * math::integer_divide_ceil(n, NPerBlock); k_iters_per_tile = MDiv(math::integer_divide_ceil(k, KPerBlock)); - // one cu can hold one wg at one time, from the whole chip's point of view - // if number of wg is same as num_cu, we call it 1 dispatch - // if number of wg is 2x num_cu, we call it 2 dispatches. - // one dispatch can deliver wg same as num_cu (full dispatch), or less than num_cu (partial - // dispatch) - // - uint32_t full_dispatches = num_tiles / num_cu; - uint32_t full_dispatch_tiles = full_dispatches * num_cu; - uint32_t partial_dispatche_tiles = num_tiles - full_dispatch_tiles; + uint32_t dp_tiles, dp_num_blocks, sk_total_iters; - uint32_t sk_occupancy = occupancy; - uint32_t dp_tiles = full_dispatch_tiles; - uint32_t sk_tiles = partial_dispatche_tiles; + // default to regular DP GEMM if sk blocks == 0 + sk_num_blocks = sk_blocks; + if(sk_num_blocks == 0 || sk_num_blocks == 0xFFFFFFFF) + { + sk_num_blocks = 0; + dp_tiles = num_tiles; + sk_num_big_blocks = 0; + k_iters_per_big_block = 0; - if(full_dispatches < occupancy) - { - // in this case, we allocate all blocks as sk blocks - // sk_occupancy = occupancy - full_dispatches; - sk_occupancy = 1; // TODO: single occ seems better - dp_tiles = full_dispatch_tiles; - sk_tiles = partial_dispatche_tiles; - } - else if((occupancy > 1) && (full_dispatches % occupancy == occupancy - 1)) - { - // e.g. occupancy = 2, full_dispatches = 3, 5, 7 ... - // occupancy = 3, full_dispatches = 5, 8, 11 ... - // occupancy = 4, full_dispatches = 7, 11 ... - sk_occupancy = 1; // left 1 slot for sk occupancy - dp_tiles = full_dispatch_tiles; - sk_tiles = partial_dispatche_tiles; + dp_num_blocks = num_tiles; // all tile to be dp block + dp_start_block_idx = 0; + sk_total_iters = 0; // clear this tiles } + // 2-tile sk + DP GEMM else { - // others, we reduce 1 dispatch from dp, together with partial dispatch, - // to construct sk dispatch - sk_occupancy = occupancy - ((full_dispatches - 1) % occupancy); - dp_tiles = full_dispatch_tiles - num_cu; - sk_tiles = partial_dispatche_tiles + num_cu; + // grid size + uint32_t grid_size = occupancy * num_cu; + // check if there's enough work for DP+ stream-k + bool bigEnough = num_tiles > grid_size; + // max of 2 sk tiles per block + uint32_t sk_tiles = bigEnough ? grid_size + num_tiles % grid_size : num_tiles; + // remaining tiles are DP tiles + dp_tiles = bigEnough ? (num_tiles - sk_tiles) : 0; + + sk_total_iters = k_iters_per_tile.get() * sk_tiles; + + // k_iters_per_sk_block is the floor of avg each ck block loop over tiles. + // we need to decide how many iters for each sk block + // let m = k_iters_per_sk_block + // some of the sk block (little) will cover m iters, some (big) will cover m+1 + // we have + // 1) l + b = sk_blocks + // 2) l * m + b * (m + 1) = sk_total_iters + // => (l + b) * m + b = sk_total_iters + // => sk_blocks * m + b = sk_total_iters + // => b = sk_total_iters - m * sk_blocks + // NOTE: big could be zero + uint32_t k_iters_per_sk_block = sk_total_iters / sk_num_blocks; + sk_num_big_blocks = sk_total_iters - k_iters_per_sk_block * sk_num_blocks; + k_iters_per_big_block = k_iters_per_sk_block + 1; + + dp_num_blocks = dp_tiles; + dp_start_block_idx = (sk_num_blocks + num_cu - 1) / num_cu * num_cu; } - // uint32_t dp_iters_per_block = k_iters_per_tile.get(); - uint32_t sk_total_iters = k_iters_per_tile.get() * sk_tiles; - uint32_t dp_num_blocks = 0; - - { - uint32_t min_sk_tiles = (sk_tiles >= num_cu) ? num_cu : (sk_tiles + 1); - uint32_t max_sk_tiles = - (sk_tiles >= num_cu) ? num_cu * sk_occupancy - : math::min(num_cu, sk_total_iters / min_k_iters_per_sk_block); - - // if use dp for sk-block, how many iters do we need - uint32_t dp_for_sk_iters = k_iters_per_tile.get(); - - uint32_t best_sk_score = - std::numeric_limits::max(); // we need to find the smallest sk iters - for(uint32_t tentative_sk_blocks = min_sk_tiles; tentative_sk_blocks < max_sk_tiles; - tentative_sk_blocks++) - { - uint32_t tentative_sk_iters_per_block = - (sk_total_iters + tentative_sk_blocks - 1) / tentative_sk_blocks; - uint32_t tentative_sk_iters = tentative_sk_iters_per_block; - uint32_t sk_blocks_per_tile = (tentative_sk_blocks + sk_tiles - 1) / sk_tiles; - - // TODO: carefully adjust this parameter - // the more sk_blocks_per_tile, the worse the overhead - uint32_t cross_sk_blocks_overhead = sk_blocks_per_tile; - if(tentative_sk_blocks % sk_tiles != 0) - { - // penalty for uneven divide - cross_sk_blocks_overhead += - sk_blocks_per_tile * tentative_sk_iters_per_block / 50; - } - - uint32_t tentative_sk_score = tentative_sk_iters + cross_sk_blocks_overhead; - - if(tentative_sk_score < best_sk_score) - { - best_sk_score = tentative_sk_score; - sk_num_blocks = tentative_sk_blocks; - } - } - - if(best_sk_score >= dp_for_sk_iters) - { - sk_num_blocks = 0; - } - - // give a chance to control num of sk blocks - sk_num_blocks = sk_blocks != 0xffffffff ? sk_blocks : sk_num_blocks; - - if(sk_num_blocks == 0) - { - sk_num_big_blocks = 0; - k_iters_per_big_block = 0; - - dp_num_blocks = num_tiles; // all tile to be dp block - dp_start_block_idx = 0; - sk_total_iters = 0; // clear this tiles - } - else - { - // k_iters_per_sk_block is the floor of avg each ck block loop over tiles. - // we need to decide how many iters for each sk block - // let m = k_iters_per_sk_block - // some of the sk block (little) will cover m iters, some (big) will cover m+1 - // we have - // 1) l + b = sk_blocks - // 2) l * m + b * (m + 1) = sk_total_iters - // => (l + b) * m + b = sk_total_iters - // => sk_blocks * m + b = sk_total_iters - // => b = sk_total_iters - m * sk_blocks - // NOTE: big could be zero - uint32_t k_iters_per_sk_block = sk_total_iters / sk_num_blocks; - sk_num_big_blocks = sk_total_iters - k_iters_per_sk_block * sk_num_blocks; - k_iters_per_big_block = k_iters_per_sk_block + 1; - - dp_num_blocks = dp_tiles; - dp_start_block_idx = (sk_num_blocks + num_cu - 1) / num_cu * num_cu; - } - } - n_tiles = MDiv2(math::integer_divide_ceil(n, NPerBlock)); + n_tiles = MDiv2(math::integer_divide_ceil(n, NPerBlock)); + // using multiple blocks for parallel reduction reduction_start_block_idx = dp_start_block_idx + dp_num_blocks; if constexpr(ReductionStrategy == StreamKReductionStrategy::Reduction) @@ -1157,13 +1084,14 @@ struct BlockToCTileMap_GemmStreamK } #if 0 - printf("cu:%d, occupancy:%d, grids:%d, num_tiles:%d, dp_tiles:%d, sk_num_big_blocks:%d, " + printf("cu:%d, occupancy:%d, gridsize:%d, num_tiles:%d, dp_tiles:%d, sk_num_big_blocks:%d, " "sk_num_blocks:%d, " - "sk_total_iters:%d, dp_start_block_idx:%d, dp_iters_per_block:%d, dp_num_blocks:%d, " + "sk_total_iters:%d, dp_start_block_idx:%d, dp_num_blocks:%d, " "k_iters_per_tile:%d, k_iters_per_big_block:%d, reduction_start_block_idx:%u, " "sk_tiles:%u, workspace(acc float):%u\n", num_cu, occupancy, + // get_grid_dims(num_cu, occupancy).x, get_grid_dims().x, num_tiles, dp_tiles, @@ -1171,7 +1099,7 @@ struct BlockToCTileMap_GemmStreamK sk_num_blocks, sk_total_iters, dp_start_block_idx, - dp_iters_per_block, + dp_num_blocks, k_iters_per_tile.get(), k_iters_per_big_block, @@ -1195,7 +1123,8 @@ struct BlockToCTileMap_GemmStreamK return k_iters_per_tile.div(sk_total_iters); } - __host__ __device__ dim3 get_grid_dims() const + // __host__ __device__ constexpr dim3 get_grid_dims(int num_cu, int occupancy) const + __host__ __device__ constexpr dim3 get_grid_dims() const { if constexpr(ReductionStrategy == StreamKReductionStrategy::Reduction) { @@ -1203,6 +1132,16 @@ struct BlockToCTileMap_GemmStreamK } else return dim3(reduction_start_block_idx, 1, 1); + // return dim3(num_cu * occupancy, 1, 1); // HS + } + __host__ __device__ uint32_t total_blocks_allocated() const + { + if constexpr(ReductionStrategy == StreamKReductionStrategy::Reduction) + { + return __builtin_amdgcn_readfirstlane(reduction_start_block_idx + get_sk_tiles()); + } + else + return __builtin_amdgcn_readfirstlane(reduction_start_block_idx); } __device__ uint32_t get_block_idx() const diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_streamk.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_streamk.hpp index e9190dee29..4281a6e12f 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_streamk.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_streamk.hpp @@ -145,6 +145,7 @@ struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_streamk index_t StrideA; index_t StrideB; index_t StrideC; + index_t num_cu, occupancy; // stream-k arguments Block2CTileMap block_mapping; Argument(const FloatAB* p_a_grid_, @@ -156,8 +157,8 @@ struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_streamk index_t StrideA_, index_t StrideB_, index_t StrideC_, - uint32_t num_cu, - uint32_t occupancy, + uint32_t num_cu_, + uint32_t occupancy_, uint32_t num_sk_blocks_) : p_a_grid(p_a_grid_), p_b_grid(p_b_grid_), @@ -168,7 +169,9 @@ struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_streamk StrideA(StrideA_), StrideB(StrideB_), StrideC(StrideC_), - block_mapping(M, N, K, num_cu, occupancy, num_sk_blocks_) + num_cu(num_cu_), + occupancy(occupancy_), + block_mapping(M, N, K, num_cu_, occupancy_, num_sk_blocks_) { } @@ -452,16 +455,16 @@ struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_streamk Block2CTileMap block_mapping, void* __restrict__ p_shared_block) { - uint32_t m = M; - uint32_t n = N; - uint32_t k = K; - uint32_t pad_m = (m + MPerBlock - 1) / MPerBlock * MPerBlock; - uint32_t pad_n = (n + NPerBlock - 1) / NPerBlock * NPerBlock; - uint32_t pad_k = (k + KPerBlock - 1) / KPerBlock * KPerBlock; - uint32_t stride_a = StrideA; - uint32_t stride_b = StrideB; - uint32_t stride_c = StrideC; - + uint32_t m = M; + uint32_t n = N; + uint32_t k = K; + uint32_t pad_m = (m + MPerBlock - 1) / MPerBlock * MPerBlock; + uint32_t pad_n = (n + NPerBlock - 1) / NPerBlock * NPerBlock; + uint32_t pad_k = (k + KPerBlock - 1) / KPerBlock * KPerBlock; + uint32_t stride_a = StrideA; + uint32_t stride_b = StrideB; + uint32_t stride_c = StrideC; + uint32_t block_idx = block_mapping.get_block_idx(); const auto a_k0_m_k1_grid_desc = MakeAGridDescriptor_K0_M_K1(m, pad_m, k, pad_k, stride_a); const auto b_k0_n_k1_grid_desc = MakeBGridDescriptor_K0_N_K1(k, pad_k, n, pad_n, stride_b); const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N(m, pad_m, n, pad_n, stride_c); @@ -520,623 +523,640 @@ struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_streamk // gridwise GEMM pipeline const auto gridwise_gemm_pipeline = GridwiseGemmPipeline_v3(); - - uint32_t block_idx = block_mapping.get_block_idx(); - bool is_sk_block = block_idx < block_mapping.sk_num_blocks; - bool is_dp_block = block_idx >= block_mapping.dp_start_block_idx && - block_idx < block_mapping.reduction_start_block_idx; - bool is_reduction_block = block_idx >= block_mapping.reduction_start_block_idx; - bool is_padding_block = block_idx >= block_mapping.sk_num_blocks && - block_idx < block_mapping.dp_start_block_idx; - uint32_t iter_start, iter_end; - block_mapping.get_block_itr(block_idx, iter_start, iter_end); - uint32_t total_iter_length = iter_end - iter_start; - - if(is_padding_block) - return; - uint32_t* p_semaphore = reinterpret_cast(reinterpret_cast(p_workspace) + block_mapping.get_workspace_size_for_acc(sizeof(FloatAcc))); - if constexpr(Block2CTileMap::ReductionStrategy == StreamKReductionStrategy::Reduction) + // offset for last acc buffer of this block + uint32_t block_acc_offset = + (block_mapping.get_acc_buffer_offset_from_block(block_idx + 1) - 1) * MPerBlock * + NPerBlock; + uint32_t iter_start, iter_end; + bool is_sk_block, is_dp_block, is_padding_block, is_reduction_block; + + uint32_t total_iter_length; + +#pragma unroll + // stream-k: for new work for all the persistent blocks. + for(; block_idx < block_mapping.total_blocks_allocated(); block_idx += gridDim.x) { - if(is_reduction_block) + + is_sk_block = block_idx < block_mapping.sk_num_blocks; + is_dp_block = block_idx >= block_mapping.dp_start_block_idx && + block_idx < block_mapping.reduction_start_block_idx; + + is_padding_block = block_idx >= block_mapping.sk_num_blocks && + block_idx < block_mapping.dp_start_block_idx; + if(is_padding_block) { - // descriptors - constexpr auto cluster_length_reduce = GetClusterLengthReduction(); - constexpr auto reduce_desc = make_cluster_descriptor(cluster_length_reduce); - const auto reduce_thread_cluster_idx = - reduce_desc.CalculateBottomIndex(make_multi_index(get_thread_local_1d_id())); - const auto thread_m_cluster_id = reduce_thread_cluster_idx[I0]; - const auto thread_n_cluster_id = reduce_thread_cluster_idx[I1]; + continue; + } - constexpr auto MReduceIters = - math::integer_divide_ceil(Number{}, cluster_length_reduce.At(I0)); - constexpr auto NReduceIters = math::integer_divide_ceil( - Number{}, - cluster_length_reduce.At(I1) * - Number{}); + block_mapping.get_block_itr(block_idx, iter_start, iter_end); + total_iter_length = iter_end - iter_start; - constexpr auto acc_thread_buf_load_desc = make_naive_tensor_descriptor_packed( - make_tuple(I1, Number{})); - constexpr auto acc_thread_buf_store_desc = make_naive_tensor_descriptor_packed( - make_tuple(I1, I1, I1, Number{})); + if constexpr(Block2CTileMap::ReductionStrategy == StreamKReductionStrategy::Reduction) + { + is_reduction_block = block_idx >= block_mapping.reduction_start_block_idx; + if(is_reduction_block) + { + // descriptors + constexpr auto cluster_length_reduce = GetClusterLengthReduction(); + constexpr auto reduce_desc = make_cluster_descriptor(cluster_length_reduce); + const auto reduce_thread_cluster_idx = reduce_desc.CalculateBottomIndex( + make_multi_index(get_thread_local_1d_id())); + const auto thread_m_cluster_id = reduce_thread_cluster_idx[I0]; + const auto thread_n_cluster_id = reduce_thread_cluster_idx[I1]; - constexpr auto c_partial_acc_block_m_n = GetPartialAccBlockDescriptor(); + constexpr auto MReduceIters = math::integer_divide_ceil( + Number{}, cluster_length_reduce.At(I0)); + constexpr auto NReduceIters = math::integer_divide_ceil( + Number{}, + cluster_length_reduce.At(I1) * + Number{}); - constexpr auto partial_acc_load_step_n = make_multi_index( - 0, cluster_length_reduce.At(I1) * CBlockTransferScalarPerVector_NWaveNPerXDL); - constexpr auto partial_acc_load_step_n_reverse = - make_multi_index(0, - -1 * cluster_length_reduce.At(I1).value * (NReduceIters - 1) * - CBlockTransferScalarPerVector_NWaveNPerXDL); - constexpr auto partial_acc_load_step_m = - make_multi_index(cluster_length_reduce.At(I0), 0); + constexpr auto acc_thread_buf_load_desc = make_naive_tensor_descriptor_packed( + make_tuple(I1, Number{})); + constexpr auto acc_thread_buf_store_desc = + make_naive_tensor_descriptor_packed(make_tuple( + I1, I1, I1, Number{})); - constexpr auto partial_acc_store_step_n = make_multi_index( - 0, - 0, - 0, - cluster_length_reduce.At(I1) * CBlockTransferScalarPerVector_NWaveNPerXDL); - constexpr auto partial_acc_store_step_n_reverse = - make_multi_index(0, - 0, - 0, - -1 * cluster_length_reduce.At(I1).value * (NReduceIters - 1) * - CBlockTransferScalarPerVector_NWaveNPerXDL); - constexpr auto partial_acc_store_step_m = - make_multi_index(0, cluster_length_reduce.At(I0), 0, 0); + constexpr auto c_partial_acc_block_m_n = GetPartialAccBlockDescriptor(); - StaticBuffer - parcial_acc_buf; - StaticBuffer - acc_buf; + constexpr auto partial_acc_load_step_n = make_multi_index( + 0, + cluster_length_reduce.At(I1) * CBlockTransferScalarPerVector_NWaveNPerXDL); + constexpr auto partial_acc_load_step_n_reverse = make_multi_index( + 0, + -1 * cluster_length_reduce.At(I1).value * (NReduceIters - 1) * + CBlockTransferScalarPerVector_NWaveNPerXDL); + constexpr auto partial_acc_load_step_m = + make_multi_index(cluster_length_reduce.At(I0), 0); - // start to compute - auto reduction_idx = blockIdx.x - block_mapping.reduction_start_block_idx; - auto spatial_idx = block_mapping.tile_to_spatial(reduction_idx, m, n); + constexpr auto partial_acc_store_step_n = make_multi_index( + 0, + 0, + 0, + cluster_length_reduce.At(I1) * CBlockTransferScalarPerVector_NWaveNPerXDL); + constexpr auto partial_acc_store_step_n_reverse = make_multi_index( + 0, + 0, + 0, + -1 * cluster_length_reduce.At(I1).value * (NReduceIters - 1) * + CBlockTransferScalarPerVector_NWaveNPerXDL); + constexpr auto partial_acc_store_step_m = + make_multi_index(0, cluster_length_reduce.At(I0), 0, 0); - workgroup_barrier wg_barrier(p_semaphore); + StaticBuffer + parcial_acc_buf; + StaticBuffer + acc_buf; - uint32_t tile_acc_offset_start = - block_mapping.get_acc_buffer_offset_from_tile(reduction_idx); - uint32_t tile_acc_offset_end = - block_mapping.get_acc_buffer_offset_from_tile(reduction_idx + 1); + // start to compute + auto reduction_idx = block_idx - block_mapping.reduction_start_block_idx; + auto spatial_idx = block_mapping.tile_to_spatial(reduction_idx, m, n); - auto acc_load = ThreadwiseTensorSliceTransfer_v2< - FloatAcc, // SrcData, - FloatAcc, // DstData, - decltype(c_partial_acc_block_m_n), // SrcDesc, - decltype(acc_thread_buf_load_desc), // DstDesc, - Sequence<1, CBlockTransferScalarPerVector_NWaveNPerXDL>, // SliceLengths, - Sequence<0, 1>, // DimAccessOrder, - 1, // SrcVectorDim, - CBlockTransferScalarPerVector_NWaveNPerXDL, // SrcScalarPerVector, - 1, // SrcScalarStrideInVector, - false // SrcResetCoordinateAfterRun, - >{c_partial_acc_block_m_n, - make_multi_index(thread_m_cluster_id, - thread_n_cluster_id * - CBlockTransferScalarPerVector_NWaveNPerXDL)}; + workgroup_barrier wg_barrier(p_semaphore); - auto acc_store = ThreadwiseTensorSliceTransfer_v1r3< - FloatAcc, // SrcData, - FloatC, // DstData, - decltype(acc_thread_buf_store_desc), // SrcDesc, - decltype(c_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc, - CElementwiseOperation, // ElementwiseOperation, - Sequence<1, 1, 1, CBlockTransferScalarPerVector_NWaveNPerXDL>, // SliceLengths, - Sequence<0, 1, 2, 3>, // DimAccessOrder, - 3, // DstVectorDim, - CBlockTransferScalarPerVector_NWaveNPerXDL, // DstScalarPerVector, - InMemoryDataOperationEnum::Set, // InMemoryDataOperationEnum DstInMemOp, - 1, // DstScalarStrideInVector, - false // DstResetCoordinateAfterRun, - >{c_grid_desc_mblock_mperblock_nblock_nperblock, - make_multi_index(__builtin_amdgcn_readfirstlane(spatial_idx[I0]), - thread_m_cluster_id, - __builtin_amdgcn_readfirstlane(spatial_idx[I1]), - thread_n_cluster_id * - CBlockTransferScalarPerVector_NWaveNPerXDL), - CElementwiseOperation{}}; + uint32_t tile_acc_offset_start = + block_mapping.get_acc_buffer_offset_from_tile(reduction_idx); + uint32_t tile_acc_offset_end = + block_mapping.get_acc_buffer_offset_from_tile(reduction_idx + 1); - // block synchronization - wg_barrier.wait_eq(reduction_idx, tile_acc_offset_end - tile_acc_offset_start); + auto acc_load = ThreadwiseTensorSliceTransfer_v2< + FloatAcc, // SrcData, + FloatAcc, // DstData, + decltype(c_partial_acc_block_m_n), // SrcDesc, + decltype(acc_thread_buf_load_desc), // DstDesc, + Sequence<1, CBlockTransferScalarPerVector_NWaveNPerXDL>, // SliceLengths, + Sequence<0, 1>, // DimAccessOrder, + 1, // SrcVectorDim, + CBlockTransferScalarPerVector_NWaveNPerXDL, // SrcScalarPerVector, + 1, // SrcScalarStrideInVector, + false // SrcResetCoordinateAfterRun, + >{c_partial_acc_block_m_n, + make_multi_index(thread_m_cluster_id, + thread_n_cluster_id * + CBlockTransferScalarPerVector_NWaveNPerXDL)}; + + auto acc_store = ThreadwiseTensorSliceTransfer_v1r3< + FloatAcc, // SrcData, + FloatC, // DstData, + decltype(acc_thread_buf_store_desc), // SrcDesc, + decltype(c_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc, + CElementwiseOperation, // ElementwiseOperation, + Sequence<1, + 1, + 1, + CBlockTransferScalarPerVector_NWaveNPerXDL>, // SliceLengths, + Sequence<0, 1, 2, 3>, // DimAccessOrder, + 3, // DstVectorDim, + CBlockTransferScalarPerVector_NWaveNPerXDL, // DstScalarPerVector, + InMemoryDataOperationEnum::Set, // InMemoryDataOperationEnum DstInMemOp, + 1, // DstScalarStrideInVector, + false // DstResetCoordinateAfterRun, + >{c_grid_desc_mblock_mperblock_nblock_nperblock, + make_multi_index(__builtin_amdgcn_readfirstlane(spatial_idx[I0]), + thread_m_cluster_id, + __builtin_amdgcn_readfirstlane(spatial_idx[I1]), + thread_n_cluster_id * + CBlockTransferScalarPerVector_NWaveNPerXDL), + CElementwiseOperation{}}; + + // block synchronization + wg_barrier.wait_eq(reduction_idx, tile_acc_offset_end - tile_acc_offset_start); #if 0 if(threadIdx.x == 0) { - printf("bid:%d, rid:%d, os:%d,%d, spatial:%d,%d\n", static_cast(blockIdx.x), + printf("bid:%d, rid:%d, os:%d,%d, spatial:%d,%d\n", static_cast(block_idx), reduction_idx, __builtin_amdgcn_readfirstlane(tile_acc_offset_start), __builtin_amdgcn_readfirstlane(tile_acc_offset_end), __builtin_amdgcn_readfirstlane(spatial_idx[I0]), __builtin_amdgcn_readfirstlane(spatial_idx[I1])); } #endif - using Accumulation = ck::detail:: - AccumulateWithNanCheck; + using Accumulation = ck::detail:: + AccumulateWithNanCheck; - for(int i_m = 0; i_m < MReduceIters; i_m++) - { - static_for<0, NReduceIters, 1>{}([&](auto i_n_reduce) { - acc_buf.Clear(); - for(auto i = tile_acc_offset_start; i < tile_acc_offset_end; i++) - { - auto c_partial_acc_buf = - make_dynamic_buffer( - reinterpret_cast(p_workspace) + - i * c_partial_acc_block_m_n.GetElementSpaceSize(), - c_partial_acc_block_m_n.GetElementSpaceSize()); - - acc_load.Run(c_partial_acc_block_m_n, - c_partial_acc_buf, - acc_thread_buf_load_desc, - make_tuple(I0, I0), - parcial_acc_buf); - - static_for<0, CBlockTransferScalarPerVector_NWaveNPerXDL, 1>{}( - [&](auto i_vec) { - constexpr auto offset = - acc_thread_buf_load_desc.CalculateOffset( - make_tuple(0, i_vec)); - Accumulation::Calculate(acc_buf(Number{}), - parcial_acc_buf[Number{}]); - }); - } - - if(thread_n_cluster_id * CBlockTransferScalarPerVector_NWaveNPerXDL < - NPerBlock) - { - acc_store.Run(acc_thread_buf_store_desc, - make_tuple(I0, I0, I0, I0), - acc_buf, - c_grid_desc_mblock_mperblock_nblock_nperblock, - c_grid_buf); - } - if constexpr(NReduceIters != 1) - { - if constexpr(i_n_reduce != (NReduceIters - 1)) - { - acc_load.MoveSrcSliceWindow(c_partial_acc_block_m_n, - partial_acc_load_step_n); - acc_store.MoveDstSliceWindow( - c_grid_desc_mblock_mperblock_nblock_nperblock, - partial_acc_store_step_n); - } - else - { - acc_load.MoveSrcSliceWindow(c_partial_acc_block_m_n, - partial_acc_load_step_n_reverse); - acc_store.MoveDstSliceWindow( - c_grid_desc_mblock_mperblock_nblock_nperblock, - partial_acc_store_step_n_reverse); - } - } - }); + for(int i_m = 0; i_m < MReduceIters; i_m++) { - acc_load.MoveSrcSliceWindow(c_partial_acc_block_m_n, - partial_acc_load_step_m); - acc_store.MoveDstSliceWindow(c_grid_desc_mblock_mperblock_nblock_nperblock, - partial_acc_store_step_m); - } - } - return; - } - } - - // offset for last acc buffer of this block - uint32_t block_acc_offset = - (block_mapping.get_acc_buffer_offset_from_block(block_idx + 1) - 1) * MPerBlock * - NPerBlock; - - while(true) - { - uint32_t current_iter_length = __builtin_amdgcn_readfirstlane( - block_mapping.get_current_iter_length(iter_start, iter_end, total_iter_length)); - uint32_t tile_idx, iter_offset; - block_mapping.get_tile_idx_with_offset(iter_end - 1, tile_idx, iter_offset); - iter_offset = __builtin_amdgcn_readfirstlane(iter_offset - current_iter_length + 1); - auto spatial_idx = block_mapping.tile_to_spatial(tile_idx, m, n); - - const index_t m_block_data_idx_on_grid = - __builtin_amdgcn_readfirstlane(spatial_idx[I0] * MPerBlock); - - const index_t n_block_data_idx_on_grid = - __builtin_amdgcn_readfirstlane(spatial_idx[I1] * NPerBlock); - - const index_t k0_block_data_idx_on_grid = - __builtin_amdgcn_readfirstlane(iter_offset * K0PerBlock); - - // A matrix blockwise copy - auto a_blockwise_copy = - ThreadGroupTensorSliceTransfer_v4r1, - ABlockTransferThreadClusterLengths_K0_M_K1, - ABlockTransferThreadClusterArrangeOrder, - FloatAB, - FloatAB, - decltype(a_k0_m_k1_grid_desc), - decltype(a_block_desc_k0_m_k1), - ABlockTransferSrcAccessOrder, - Sequence<1, 0, 2>, - ABlockTransferSrcVectorDim, - 2, - ABlockTransferSrcScalarPerVector, - ABlockTransferDstScalarPerVector_K1, - 1, - 1, - AThreadTransferSrcResetCoordinateAfterRun, - true>( - a_k0_m_k1_grid_desc, - make_multi_index(k0_block_data_idx_on_grid, m_block_data_idx_on_grid, 0), - a_element_op, - a_block_desc_k0_m_k1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}); - - // B matrix blockwise copy - auto b_blockwise_copy = - ThreadGroupTensorSliceTransfer_v4r1, - BBlockTransferThreadClusterLengths_K0_N_K1, - BBlockTransferThreadClusterArrangeOrder, - FloatAB, - FloatAB, - decltype(b_k0_n_k1_grid_desc), - decltype(b_block_desc_k0_n_k1), - BBlockTransferSrcAccessOrder, - Sequence<1, 0, 2>, - BBlockTransferSrcVectorDim, - 2, - BBlockTransferSrcScalarPerVector, - BBlockTransferDstScalarPerVector_K1, - 1, - 1, - BThreadTransferSrcResetCoordinateAfterRun, - true>( - b_k0_n_k1_grid_desc, - make_multi_index(k0_block_data_idx_on_grid, n_block_data_idx_on_grid, 0), - b_element_op, - b_block_desc_k0_n_k1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}); - - const index_t num_k_block_main_loop = current_iter_length; - - gridwise_gemm_pipeline.Run(a_k0_m_k1_grid_desc, - a_block_desc_k0_m_k1, - a_blockwise_copy, - a_grid_buf, - a_block_buf, - a_block_slice_copy_step, - b_k0_n_k1_grid_desc, - b_block_desc_k0_n_k1, - b_blockwise_copy, - b_grid_buf, - b_block_buf, - b_block_slice_copy_step, - blockwise_gemm, - c_thread_buf, - num_k_block_main_loop); - - // output: register to global memory - { - constexpr index_t MWave = MPerBlock / (MRepeat * MPerXDL); - constexpr index_t NWave = NPerBlock / (NRepeat * NPerXDL); - - constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc = - blockwise_gemm.GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); - - constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc = - blockwise_gemm.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); - - constexpr auto M0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I0); - constexpr auto N0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I1); - constexpr auto M1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I2); - constexpr auto N1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I3); - constexpr auto M2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I4); - constexpr auto M3 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I5); - constexpr auto M4 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I6); - constexpr auto N2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I7); - - constexpr auto c_block_desc_mblock_mpershuffle_nblock_npershuffle = - GetCBlockDescriptor_MBlock_MPerShuffle_NBlock_NPerShuffle(); - - constexpr auto c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle = - GetCBlockDescriptor_MShuffleRepeat_MPerShuffle_NShuffleRepeat_NPerShuffle(); - - auto c_block_buf = make_dynamic_buffer( - reinterpret_cast(p_shared_block), - c_block_desc_mblock_mpershuffle_nblock_npershuffle.GetElementSpaceSize()); - - auto c_partial_acc_buf = - make_dynamic_buffer( - reinterpret_cast(p_workspace) + block_acc_offset, - c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle - .GetElementSpaceSize()); - - constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2 = transform_tensor_descriptor( - c_block_desc_mblock_mpershuffle_nblock_npershuffle, - make_tuple(make_freeze_transform(I0), // freeze mblock - make_unmerge_transform( - make_tuple(CShuffleMRepeatPerShuffle, - M1, - M2, - M3, - M4)), // M1 = MWave, M2 * M3 * M4 = MPerXDL - make_freeze_transform(I0), // freeze nblock - make_unmerge_transform( - make_tuple(CShuffleNRepeatPerShuffle, - N1, - N2))), // M1 = MWave, M2 * M3 * M4 = MPerXDL - make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), - make_tuple(Sequence<>{}, - Sequence<0, 2, 4, 5, 6>{}, - Sequence<>{}, - Sequence<1, 3, 7>{})); - - // calculate origin of thread output tensor on global memory - // blockwise GEMM c matrix starting index - const auto c_thread_mtx_on_block = - blockwise_gemm.CalculateCThreadOriginDataIndex(I0, I0, I0, I0); - - const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0]; - const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1]; - - const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor = - make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))), - make_tuple(Sequence<0, 1, 2, 3, 4>{}), - make_tuple(Sequence<0>{})); - - const auto m_thread_data_on_block_idx = - m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex( - make_multi_index(m_thread_data_on_block)); - - const auto n_thread_data_on_block_to_n0_n1_n2_adaptor = - make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(N0, N1, N2))), - make_tuple(Sequence<0, 1, 2>{}), - make_tuple(Sequence<0>{})); - - const auto n_thread_data_on_block_idx = - n_thread_data_on_block_to_n0_n1_n2_adaptor.CalculateBottomIndex( - make_multi_index(n_thread_data_on_block)); - - // VGPR to LDS - auto c_thread_copy_vgpr_to_lds = ThreadwiseTensorSliceTransfer_v1r3< - FloatAcc, - FloatCShuffle, - decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc), - decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2), - ck::tensor_operation::element_wise::PassThrough, - Sequence, - Sequence<0, 1, 2, 3, 4, 5, 6, 7>, - 7, - 1, - InMemoryDataOperationEnum::Set, - 1, - true>{c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, - make_multi_index(0, - 0, - m_thread_data_on_block_idx[I1], - n_thread_data_on_block_idx[I1], - m_thread_data_on_block_idx[I2], - m_thread_data_on_block_idx[I3], - m_thread_data_on_block_idx[I4], - n_thread_data_on_block_idx[I2]), - ck::tensor_operation::element_wise::PassThrough{}}; - - // LDS to global - auto c_block_copy_lds_to_global = ThreadGroupTensorSliceTransfer_v6r1r2< - ThisThreadBlock, // index_t BlockSize, - CElementwiseOperation, // ElementwiseOperation, - // InMemoryDataOperationEnum::Set, // DstInMemOp, - Sequence<1, - CShuffleMRepeatPerShuffle * MWave * MPerXDL, - 1, - CShuffleNRepeatPerShuffle * NWave * NPerXDL>, // BlockSliceLengths, - CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, - Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, - FloatCShuffle, // typename SrcData, - FloatC, // typename DstData, - decltype(c_block_desc_mblock_mpershuffle_nblock_npershuffle), - decltype(c_grid_desc_mblock_mperblock_nblock_nperblock), - Sequence<0, 1, 2, 3>, // typename DimAccessOrder, - 3, // index_t VectorDim, - CBlockTransferScalarPerVector_NWaveNPerXDL, // index_t ScalarPerVector, - false, // bool ThreadTransferSrcResetCoordinateAfterRun, - false> // bool ThreadTransferDstResetCoordinateAfterRun - {c_block_desc_mblock_mpershuffle_nblock_npershuffle, - make_multi_index(0, 0, 0, 0), - c_grid_desc_mblock_mperblock_nblock_nperblock, - make_multi_index(__builtin_amdgcn_readfirstlane(spatial_idx[I0]), - 0, - __builtin_amdgcn_readfirstlane(spatial_idx[I1]), - 0), - c_element_op}; - - // LDS to global partial acc - auto c_block_copy_lds_to_partial_acc = ThreadGroupTensorSliceTransfer_v6r1r2< - ThisThreadBlock, // index_t BlockSize, - CElementwiseOperation, // ElementwiseOperation, - // InMemoryDataOperationEnum::Set, // DstInMemOp, - Sequence<1, - CShuffleMRepeatPerShuffle * MWave * MPerXDL, - 1, - CShuffleNRepeatPerShuffle * NWave * NPerXDL>, // BlockSliceLengths, - CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, - Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, - FloatCShuffle, // typename SrcData, - FloatCShuffle, // typename DstData, - decltype(c_block_desc_mblock_mpershuffle_nblock_npershuffle), - decltype(c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle), - Sequence<0, 1, 2, 3>, // typename DimAccessOrder, - 3, // index_t VectorDim, - CBlockTransferScalarPerVector_NWaveNPerXDL, // index_t ScalarPerVector, - false, // bool ThreadTransferSrcResetCoordinateAfterRun, => need to be false, - // othre wise has scratch - false> // bool ThreadTransferDstResetCoordinateAfterRun, => need to be false, - // othre wise has scratch - {c_block_desc_mblock_mpershuffle_nblock_npershuffle, - make_multi_index(0, 0, 0, 0), - c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle, - make_multi_index(0, 0, 0, 0), - c_element_op}; - - constexpr auto mxdlperwave_forward_step = - make_multi_index(0, CShuffleMRepeatPerShuffle * MWave * MPerXDL, 0, 0); - constexpr auto nxdlperwave_forward_step = - make_multi_index(0, 0, 0, CShuffleNRepeatPerShuffle * NWave * NPerXDL); - constexpr auto nxdlperwave_backward_step = - make_multi_index(0, 0, 0, -CShuffleNRepeatPerShuffle * NWave * NPerXDL); - - static_for<0, MRepeat, CShuffleMRepeatPerShuffle>{}([&](auto mxdlperwave_iter) { - constexpr auto mxdlperwave = mxdlperwave_iter; - - static_for<0, NRepeat, CShuffleNRepeatPerShuffle>{}([&](auto nxdlperwave_iter) { - constexpr bool nxdlperwave_forward_sweep = - (mxdlperwave % (2 * CShuffleMRepeatPerShuffle) == 0); - - constexpr index_t nxdlperwave_value = - nxdlperwave_forward_sweep - ? nxdlperwave_iter - : (NRepeat - nxdlperwave_iter - CShuffleNRepeatPerShuffle); - - constexpr auto nxdlperwave = Number{}; - - // make sure it's safe to do ds_write - block_sync_lds(); - - // VGPR to LDS - c_thread_copy_vgpr_to_lds.Run( - c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc, - make_tuple(mxdlperwave, nxdlperwave, I0, I0, I0, I0, I0, I0), - c_thread_buf, - c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, - c_block_buf); - - // make sure it's safe to do ds_read - block_sync_lds(); - - c_block_copy_lds_to_global.SetSrcSliceOrigin( - c_block_desc_mblock_mpershuffle_nblock_npershuffle, - make_tuple(0, 0, 0, 0)); - - // LDS to global - if(is_dp_block) - c_block_copy_lds_to_global.template Run( - c_block_desc_mblock_mpershuffle_nblock_npershuffle, - c_block_buf, - c_grid_desc_mblock_mperblock_nblock_nperblock, - c_grid_buf); - else if(is_sk_block) - { - if constexpr(Block2CTileMap::ReductionStrategy == - StreamKReductionStrategy::Reduction) + static_for<0, NReduceIters, 1>{}([&](auto i_n_reduce) { + acc_buf.Clear(); + for(auto i = tile_acc_offset_start; i < tile_acc_offset_end; i++) { - // constexpr offset - c_block_copy_lds_to_partial_acc.SetSrcSliceOrigin( + auto c_partial_acc_buf = + make_dynamic_buffer( + reinterpret_cast(p_workspace) + + i * c_partial_acc_block_m_n.GetElementSpaceSize(), + c_partial_acc_block_m_n.GetElementSpaceSize()); + + acc_load.Run(c_partial_acc_block_m_n, + c_partial_acc_buf, + acc_thread_buf_load_desc, + make_tuple(I0, I0), + parcial_acc_buf); + + static_for<0, CBlockTransferScalarPerVector_NWaveNPerXDL, 1>{}( + [&](auto i_vec) { + constexpr auto offset = + acc_thread_buf_load_desc.CalculateOffset( + make_tuple(0, i_vec)); + Accumulation::Calculate(acc_buf(Number{}), + parcial_acc_buf[Number{}]); + }); + } + + if(thread_n_cluster_id * CBlockTransferScalarPerVector_NWaveNPerXDL < + NPerBlock) + { + acc_store.Run(acc_thread_buf_store_desc, + make_tuple(I0, I0, I0, I0), + acc_buf, + c_grid_desc_mblock_mperblock_nblock_nperblock, + c_grid_buf); + } + if constexpr(NReduceIters != 1) + { + if constexpr(i_n_reduce != (NReduceIters - 1)) + { + acc_load.MoveSrcSliceWindow(c_partial_acc_block_m_n, + partial_acc_load_step_n); + acc_store.MoveDstSliceWindow( + c_grid_desc_mblock_mperblock_nblock_nperblock, + partial_acc_store_step_n); + } + else + { + acc_load.MoveSrcSliceWindow(c_partial_acc_block_m_n, + partial_acc_load_step_n_reverse); + acc_store.MoveDstSliceWindow( + c_grid_desc_mblock_mperblock_nblock_nperblock, + partial_acc_store_step_n_reverse); + } + } + }); + { + acc_load.MoveSrcSliceWindow(c_partial_acc_block_m_n, + partial_acc_load_step_m); + acc_store.MoveDstSliceWindow( + c_grid_desc_mblock_mperblock_nblock_nperblock, + partial_acc_store_step_m); + } + } + return; + } + } + while(true) + { + uint32_t current_iter_length = __builtin_amdgcn_readfirstlane( + block_mapping.get_current_iter_length(iter_start, iter_end, total_iter_length)); + uint32_t tile_idx, iter_offset; + block_mapping.get_tile_idx_with_offset(iter_end - 1, tile_idx, iter_offset); + iter_offset = __builtin_amdgcn_readfirstlane(iter_offset - current_iter_length + 1); + auto spatial_idx = block_mapping.tile_to_spatial(tile_idx, m, n); + + const index_t m_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(spatial_idx[I0] * MPerBlock); + + const index_t n_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(spatial_idx[I1] * NPerBlock); + + const index_t k0_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(iter_offset * K0PerBlock); + + // A matrix blockwise copy + auto a_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1< + ThisThreadBlock, + AElementwiseOperation, + ck::tensor_operation::element_wise::PassThrough, + InMemoryDataOperationEnum::Set, + Sequence, + ABlockTransferThreadClusterLengths_K0_M_K1, + ABlockTransferThreadClusterArrangeOrder, + FloatAB, + FloatAB, + decltype(a_k0_m_k1_grid_desc), + decltype(a_block_desc_k0_m_k1), + ABlockTransferSrcAccessOrder, + Sequence<1, 0, 2>, + ABlockTransferSrcVectorDim, + 2, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_K1, + 1, + 1, + AThreadTransferSrcResetCoordinateAfterRun, + true>(a_k0_m_k1_grid_desc, + make_multi_index(k0_block_data_idx_on_grid, m_block_data_idx_on_grid, 0), + a_element_op, + a_block_desc_k0_m_k1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + + // B matrix blockwise copy + auto b_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1< + ThisThreadBlock, + BElementwiseOperation, + ck::tensor_operation::element_wise::PassThrough, + InMemoryDataOperationEnum::Set, + Sequence, + BBlockTransferThreadClusterLengths_K0_N_K1, + BBlockTransferThreadClusterArrangeOrder, + FloatAB, + FloatAB, + decltype(b_k0_n_k1_grid_desc), + decltype(b_block_desc_k0_n_k1), + BBlockTransferSrcAccessOrder, + Sequence<1, 0, 2>, + BBlockTransferSrcVectorDim, + 2, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_K1, + 1, + 1, + BThreadTransferSrcResetCoordinateAfterRun, + true>(b_k0_n_k1_grid_desc, + make_multi_index(k0_block_data_idx_on_grid, n_block_data_idx_on_grid, 0), + b_element_op, + b_block_desc_k0_n_k1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + + const index_t num_k_block_main_loop = current_iter_length; + + gridwise_gemm_pipeline.Run(a_k0_m_k1_grid_desc, + a_block_desc_k0_m_k1, + a_blockwise_copy, + a_grid_buf, + a_block_buf, + a_block_slice_copy_step, + b_k0_n_k1_grid_desc, + b_block_desc_k0_n_k1, + b_blockwise_copy, + b_grid_buf, + b_block_buf, + b_block_slice_copy_step, + blockwise_gemm, + c_thread_buf, + num_k_block_main_loop); + + // output: register to global memory + { + constexpr index_t MWave = MPerBlock / (MRepeat * MPerXDL); + constexpr index_t NWave = NPerBlock / (NRepeat * NPerXDL); + + constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc = + blockwise_gemm.GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); + + constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc = + blockwise_gemm.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); + + constexpr auto M0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I0); + constexpr auto N0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I1); + constexpr auto M1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I2); + constexpr auto N1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I3); + constexpr auto M2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I4); + constexpr auto M3 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I5); + constexpr auto M4 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I6); + constexpr auto N2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I7); + + constexpr auto c_block_desc_mblock_mpershuffle_nblock_npershuffle = + GetCBlockDescriptor_MBlock_MPerShuffle_NBlock_NPerShuffle(); + + constexpr auto c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle = + GetCBlockDescriptor_MShuffleRepeat_MPerShuffle_NShuffleRepeat_NPerShuffle(); + + auto c_block_buf = make_dynamic_buffer( + reinterpret_cast(p_shared_block), + c_block_desc_mblock_mpershuffle_nblock_npershuffle.GetElementSpaceSize()); + + auto c_partial_acc_buf = + make_dynamic_buffer( + reinterpret_cast(p_workspace) + block_acc_offset, + c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle + .GetElementSpaceSize()); + + constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2 = + transform_tensor_descriptor( + c_block_desc_mblock_mpershuffle_nblock_npershuffle, + make_tuple(make_freeze_transform(I0), // freeze mblock + make_unmerge_transform( + make_tuple(CShuffleMRepeatPerShuffle, + M1, + M2, + M3, + M4)), // M1 = MWave, M2 * M3 * M4 = MPerXDL + make_freeze_transform(I0), // freeze nblock + make_unmerge_transform( + make_tuple(CShuffleNRepeatPerShuffle, + N1, + N2))), // M1 = MWave, M2 * M3 * M4 = MPerXDL + make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), + make_tuple(Sequence<>{}, + Sequence<0, 2, 4, 5, 6>{}, + Sequence<>{}, + Sequence<1, 3, 7>{})); + + // calculate origin of thread output tensor on global memory + // blockwise GEMM c matrix starting index + const auto c_thread_mtx_on_block = + blockwise_gemm.CalculateCThreadOriginDataIndex(I0, I0, I0, I0); + + const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0]; + const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1]; + + const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor = + make_single_stage_tensor_adaptor( + make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))), + make_tuple(Sequence<0, 1, 2, 3, 4>{}), + make_tuple(Sequence<0>{})); + + const auto m_thread_data_on_block_idx = + m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex( + make_multi_index(m_thread_data_on_block)); + + const auto n_thread_data_on_block_to_n0_n1_n2_adaptor = + make_single_stage_tensor_adaptor( + make_tuple(make_merge_transform(make_tuple(N0, N1, N2))), + make_tuple(Sequence<0, 1, 2>{}), + make_tuple(Sequence<0>{})); + + const auto n_thread_data_on_block_idx = + n_thread_data_on_block_to_n0_n1_n2_adaptor.CalculateBottomIndex( + make_multi_index(n_thread_data_on_block)); + + // VGPR to LDS + auto c_thread_copy_vgpr_to_lds = ThreadwiseTensorSliceTransfer_v1r3< + FloatAcc, + FloatCShuffle, + decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc), + decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2), + ck::tensor_operation::element_wise::PassThrough, + Sequence, + Sequence<0, 1, 2, 3, 4, 5, 6, 7>, + 7, + 1, + InMemoryDataOperationEnum::Set, + 1, + true>{c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, + make_multi_index(0, + 0, + m_thread_data_on_block_idx[I1], + n_thread_data_on_block_idx[I1], + m_thread_data_on_block_idx[I2], + m_thread_data_on_block_idx[I3], + m_thread_data_on_block_idx[I4], + n_thread_data_on_block_idx[I2]), + ck::tensor_operation::element_wise::PassThrough{}}; + + // LDS to global + auto c_block_copy_lds_to_global = ThreadGroupTensorSliceTransfer_v6r1r2< + ThisThreadBlock, // index_t BlockSize, + CElementwiseOperation, // ElementwiseOperation, + // InMemoryDataOperationEnum::Set, // DstInMemOp, + Sequence<1, + CShuffleMRepeatPerShuffle * MWave * MPerXDL, + 1, + CShuffleNRepeatPerShuffle * NWave * NPerXDL>, // BlockSliceLengths, + CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, + FloatCShuffle, // typename SrcData, + FloatC, // typename DstData, + decltype(c_block_desc_mblock_mpershuffle_nblock_npershuffle), + decltype(c_grid_desc_mblock_mperblock_nblock_nperblock), + Sequence<0, 1, 2, 3>, // typename DimAccessOrder, + 3, // index_t VectorDim, + CBlockTransferScalarPerVector_NWaveNPerXDL, // index_t ScalarPerVector, + false, // bool ThreadTransferSrcResetCoordinateAfterRun, + false> // bool ThreadTransferDstResetCoordinateAfterRun + {c_block_desc_mblock_mpershuffle_nblock_npershuffle, + make_multi_index(0, 0, 0, 0), + c_grid_desc_mblock_mperblock_nblock_nperblock, + make_multi_index(__builtin_amdgcn_readfirstlane(spatial_idx[I0]), + 0, + __builtin_amdgcn_readfirstlane(spatial_idx[I1]), + 0), + c_element_op}; + + // LDS to global partial acc + auto c_block_copy_lds_to_partial_acc = ThreadGroupTensorSliceTransfer_v6r1r2< + ThisThreadBlock, // index_t BlockSize, + CElementwiseOperation, // ElementwiseOperation, + // InMemoryDataOperationEnum::Set, // DstInMemOp, + Sequence<1, + CShuffleMRepeatPerShuffle * MWave * MPerXDL, + 1, + CShuffleNRepeatPerShuffle * NWave * NPerXDL>, // BlockSliceLengths, + CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, + FloatCShuffle, // typename SrcData, + FloatCShuffle, // typename DstData, + decltype(c_block_desc_mblock_mpershuffle_nblock_npershuffle), + decltype(c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle), + Sequence<0, 1, 2, 3>, // typename DimAccessOrder, + 3, // index_t VectorDim, + CBlockTransferScalarPerVector_NWaveNPerXDL, // index_t ScalarPerVector, + false, // bool ThreadTransferSrcResetCoordinateAfterRun, => need to be + // false, othre wise has scratch + false> // bool ThreadTransferDstResetCoordinateAfterRun, => need to be + // false, othre wise has scratch + {c_block_desc_mblock_mpershuffle_nblock_npershuffle, + make_multi_index(0, 0, 0, 0), + c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle, + make_multi_index(0, 0, 0, 0), + c_element_op}; + + constexpr auto mxdlperwave_forward_step = + make_multi_index(0, CShuffleMRepeatPerShuffle * MWave * MPerXDL, 0, 0); + constexpr auto nxdlperwave_forward_step = + make_multi_index(0, 0, 0, CShuffleNRepeatPerShuffle * NWave * NPerXDL); + constexpr auto nxdlperwave_backward_step = + make_multi_index(0, 0, 0, -CShuffleNRepeatPerShuffle * NWave * NPerXDL); + + static_for<0, MRepeat, CShuffleMRepeatPerShuffle>{}([&](auto mxdlperwave_iter) { + constexpr auto mxdlperwave = mxdlperwave_iter; + + static_for<0, NRepeat, CShuffleNRepeatPerShuffle>{}( + [&](auto nxdlperwave_iter) { + constexpr bool nxdlperwave_forward_sweep = + (mxdlperwave % (2 * CShuffleMRepeatPerShuffle) == 0); + + constexpr index_t nxdlperwave_value = + nxdlperwave_forward_sweep + ? nxdlperwave_iter + : (NRepeat - nxdlperwave_iter - CShuffleNRepeatPerShuffle); + + constexpr auto nxdlperwave = Number{}; + + // make sure it's safe to do ds_write + block_sync_lds(); + + // VGPR to LDS + c_thread_copy_vgpr_to_lds.Run( + c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc, + make_tuple(mxdlperwave, nxdlperwave, I0, I0, I0, I0, I0, I0), + c_thread_buf, + c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, + c_block_buf); + + // make sure it's safe to do ds_read + block_sync_lds(); + + c_block_copy_lds_to_global.SetSrcSliceOrigin( c_block_desc_mblock_mpershuffle_nblock_npershuffle, make_tuple(0, 0, 0, 0)); - c_block_copy_lds_to_partial_acc.SetDstSliceOrigin( - c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle, - make_tuple(mxdlperwave.value, 0, nxdlperwave.value, 0)); + // LDS to global + if(is_dp_block) + c_block_copy_lds_to_global + .template Run( + c_block_desc_mblock_mpershuffle_nblock_npershuffle, + c_block_buf, + c_grid_desc_mblock_mperblock_nblock_nperblock, + c_grid_buf); + else if(is_sk_block) + { + if constexpr(Block2CTileMap::ReductionStrategy == + StreamKReductionStrategy::Reduction) + { + // constexpr offset + c_block_copy_lds_to_partial_acc.SetSrcSliceOrigin( + c_block_desc_mblock_mpershuffle_nblock_npershuffle, + make_tuple(0, 0, 0, 0)); - c_block_copy_lds_to_partial_acc - .template Run( - c_block_desc_mblock_mpershuffle_nblock_npershuffle, - c_block_buf, - c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle, - c_partial_acc_buf); - } - else if constexpr(Block2CTileMap::ReductionStrategy == - StreamKReductionStrategy::Atomic) - { - c_block_copy_lds_to_global - .template Run( - c_block_desc_mblock_mpershuffle_nblock_npershuffle, - c_block_buf, + c_block_copy_lds_to_partial_acc.SetDstSliceOrigin( + c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle, + make_tuple(mxdlperwave.value, 0, nxdlperwave.value, 0)); + + c_block_copy_lds_to_partial_acc.template Run< + decltype(c_block_buf), + decltype(c_partial_acc_buf), + InMemoryDataOperationEnum::Set>( + c_block_desc_mblock_mpershuffle_nblock_npershuffle, + c_block_buf, + c_block_desc_mshuffle_mpershuffle_nshuffle_npershuffle, + c_partial_acc_buf); + } + else if constexpr(Block2CTileMap::ReductionStrategy == + StreamKReductionStrategy::Atomic) + { + c_block_copy_lds_to_global + .template Run( + c_block_desc_mblock_mpershuffle_nblock_npershuffle, + c_block_buf, + c_grid_desc_mblock_mperblock_nblock_nperblock, + c_grid_buf); + } + } + + // move on nxdlperwave dimension + if constexpr(nxdlperwave_forward_sweep && + (nxdlperwave < NRepeat - CShuffleNRepeatPerShuffle)) + { + c_block_copy_lds_to_global.MoveDstSliceWindow( c_grid_desc_mblock_mperblock_nblock_nperblock, - c_grid_buf); - } - } + nxdlperwave_forward_step); + } + else if constexpr((!nxdlperwave_forward_sweep) && (nxdlperwave > 0)) + { + c_block_copy_lds_to_global.MoveDstSliceWindow( + c_grid_desc_mblock_mperblock_nblock_nperblock, + nxdlperwave_backward_step); + } + }); - // move on nxdlperwave dimension - if constexpr(nxdlperwave_forward_sweep && - (nxdlperwave < NRepeat - CShuffleNRepeatPerShuffle)) + // move on mxdlperwave dimension + if constexpr(mxdlperwave < MRepeat - CShuffleMRepeatPerShuffle) { c_block_copy_lds_to_global.MoveDstSliceWindow( c_grid_desc_mblock_mperblock_nblock_nperblock, - nxdlperwave_forward_step); - } - else if constexpr((!nxdlperwave_forward_sweep) && (nxdlperwave > 0)) - { - c_block_copy_lds_to_global.MoveDstSliceWindow( - c_grid_desc_mblock_mperblock_nblock_nperblock, - nxdlperwave_backward_step); + mxdlperwave_forward_step); } }); - // move on mxdlperwave dimension - if constexpr(mxdlperwave < MRepeat - CShuffleMRepeatPerShuffle) + if constexpr(Block2CTileMap::ReductionStrategy == + StreamKReductionStrategy::Reduction) { - c_block_copy_lds_to_global.MoveDstSliceWindow( - c_grid_desc_mblock_mperblock_nblock_nperblock, - mxdlperwave_forward_step); + if(is_sk_block) + { + // increase the counter for this tile + workgroup_barrier wg_barrier(p_semaphore); + wg_barrier.inc(tile_idx); + } } - }); + } + + // exit condition + iter_end -= current_iter_length; + if(iter_end <= iter_start) + break; if constexpr(Block2CTileMap::ReductionStrategy == StreamKReductionStrategy::Reduction) { - if(is_sk_block) - { - // increase the counter for this tile - workgroup_barrier wg_barrier(p_semaphore); - wg_barrier.inc(tile_idx); - } + block_acc_offset -= MPerBlock * NPerBlock; } + // make sure next loop LDS is ready for use + block_sync_lds(); } - - // exit condition - iter_end -= current_iter_length; - if(iter_end <= iter_start) - break; - - if constexpr(Block2CTileMap::ReductionStrategy == StreamKReductionStrategy::Reduction) - { - block_acc_offset -= MPerBlock * NPerBlock; - } - // make sure next loop LDS is ready for use - block_sync_lds(); } }