mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-03 21:21:22 +00:00
initial stream-k implementation with example (#699)
* initial stream-k implementation with example * fix unexpected change in err * improve a little bit performance by reorganize pipeline. * improve perf a little bit by swizzle block idx * add profiler * update example * fix spelling * shrink karg for streamk * support dynamic buffer using memory coherence glc_slc bit from template * control memory coherence while construct dynamic buffer * update reduction for streamk(not ready yet) * Add template parameter to make_dynamic_buffer to support amd_buffer coherence setting * fix build issue * fix several bug * now result is correct, everything works (but has scratch) * remove scratch by manually reset coordinate * update device code * fix a bug in final reduce * fix something in example * update async memset * fix enum as camel case * modify coherence enum name * clean code and use atomic streamk by default * remove unused var * throw exception if have empty pointer * fix format * fix CI warning * fix type in init * modify CI error * filter out on gfx10+ * restore changed example code --------- Co-authored-by: Qianfeng Zhang <Qianfeng.Zhang@amd.com>
This commit is contained in:
carlushuang
@@ -7,6 +7,8 @@
|
||||
#include "ck/utility/number.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
#include "ck/tensor_description/multi_index_transform_helper.hpp"
|
||||
#include <limits>
|
||||
#include <stdlib.h>
|
||||
|
||||
namespace ck {
|
||||
|
||||
@@ -669,4 +671,406 @@ struct BlockToCTileMap_3DGrid_KSplit
|
||||
}
|
||||
};
|
||||
|
||||
enum StreamKReductionStrategy
|
||||
{
|
||||
Atomic = 0, // sk block use atomic to do reduction
|
||||
Reduction, // let some workgroup responsible for doing the reduction operation
|
||||
};
|
||||
|
||||
template <uint32_t MPerBlock_,
|
||||
uint32_t NPerBlock_,
|
||||
uint32_t KPerBlock_,
|
||||
StreamKReductionStrategy ReductionStrategy_ = StreamKReductionStrategy::Atomic,
|
||||
uint32_t TileSwizzleSubM_ = 8>
|
||||
struct BlockToCTileMap_GemmStreamK
|
||||
{
|
||||
static constexpr uint32_t min_k_iters_per_sk_block = 2;
|
||||
static constexpr uint32_t MPerBlock = MPerBlock_;
|
||||
static constexpr uint32_t NPerBlock = NPerBlock_;
|
||||
static constexpr uint32_t KPerBlock = KPerBlock_;
|
||||
static constexpr StreamKReductionStrategy ReductionStrategy = ReductionStrategy_;
|
||||
static constexpr uint32_t tile_swizzle_sub_m = TileSwizzleSubM_;
|
||||
|
||||
//--------------------------------------
|
||||
// pass to device
|
||||
uint32_t sk_num_blocks;
|
||||
uint32_t sk_num_big_blocks;
|
||||
uint32_t dp_start_block_idx;
|
||||
uint32_t reduction_start_block_idx;
|
||||
uint32_t k_iters_per_big_block;
|
||||
MDiv2 n_tiles;
|
||||
MDiv k_iters_per_tile;
|
||||
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 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 sk_occupancy = occupancy;
|
||||
uint32_t dp_tiles = full_dispatch_tiles;
|
||||
uint32_t sk_tiles = partial_dispatche_tiles;
|
||||
|
||||
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;
|
||||
}
|
||||
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;
|
||||
}
|
||||
|
||||
// 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<int>::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));
|
||||
reduction_start_block_idx = dp_start_block_idx + dp_num_blocks;
|
||||
|
||||
if constexpr(ReductionStrategy == StreamKReductionStrategy::Reduction)
|
||||
{
|
||||
uint32_t upper_big = math::lcm(k_iters_per_big_block, k_iters_per_tile.get());
|
||||
uint32_t upper_little = math::lcm(k_iters_per_big_block - 1, k_iters_per_tile.get());
|
||||
eqav_tiles_big = MDiv(upper_big / k_iters_per_tile.get());
|
||||
eqav_tiles_little = MDiv(upper_little / k_iters_per_tile.get());
|
||||
}
|
||||
|
||||
#if 0
|
||||
printf("cu:%d, occupancy:%d, grids:%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, "
|
||||
"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().x,
|
||||
num_tiles,
|
||||
dp_tiles,
|
||||
sk_num_big_blocks,
|
||||
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,
|
||||
reduction_start_block_idx,
|
||||
get_sk_tiles(),
|
||||
get_workspace_size(sizeof(float)));
|
||||
#endif
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_sk_total_iters() const
|
||||
{
|
||||
uint32_t sk_total_iters = sk_num_big_blocks * k_iters_per_big_block +
|
||||
(sk_num_blocks - sk_num_big_blocks) * (k_iters_per_big_block - 1);
|
||||
return sk_total_iters;
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_sk_tiles() const
|
||||
{
|
||||
// tiles for sk
|
||||
uint32_t sk_total_iters = get_sk_total_iters();
|
||||
return k_iters_per_tile.div(sk_total_iters);
|
||||
}
|
||||
|
||||
__host__ __device__ dim3 get_grid_dims() const
|
||||
{
|
||||
if constexpr(ReductionStrategy == StreamKReductionStrategy::Reduction)
|
||||
{
|
||||
return dim3(reduction_start_block_idx + get_sk_tiles(), 1, 1);
|
||||
}
|
||||
else
|
||||
return dim3(reduction_start_block_idx, 1, 1);
|
||||
}
|
||||
|
||||
__device__ uint32_t get_block_idx() const
|
||||
{
|
||||
// TODO: swizzle block index for better locality
|
||||
return __builtin_amdgcn_readfirstlane(blockIdx.x);
|
||||
}
|
||||
|
||||
__device__ void
|
||||
get_block_itr(uint32_t block_idx, uint32_t& iter_start, uint32_t& iter_end) const
|
||||
{
|
||||
if(block_idx < sk_num_big_blocks)
|
||||
{
|
||||
iter_start = block_idx * k_iters_per_big_block;
|
||||
iter_end = iter_start + k_iters_per_big_block;
|
||||
}
|
||||
else if(block_idx < sk_num_blocks)
|
||||
{
|
||||
iter_start = (sk_num_big_blocks * k_iters_per_big_block) +
|
||||
(block_idx - sk_num_big_blocks) * (k_iters_per_big_block - 1);
|
||||
iter_end = iter_start + (k_iters_per_big_block - 1);
|
||||
}
|
||||
else if(block_idx >= dp_start_block_idx)
|
||||
{
|
||||
uint32_t sk_total_iters = get_sk_total_iters();
|
||||
uint32_t dp_iters_per_block = k_iters_per_tile.get();
|
||||
iter_start = sk_total_iters + (block_idx - dp_start_block_idx) * dp_iters_per_block;
|
||||
iter_end = iter_start + dp_iters_per_block;
|
||||
}
|
||||
}
|
||||
|
||||
__device__ uint32_t get_current_iter_length(uint32_t iter_start,
|
||||
uint32_t iter_end,
|
||||
uint32_t total_iter_length) const
|
||||
{
|
||||
uint32_t iter_length_mod, iter_length_quo /*unused*/;
|
||||
k_iters_per_tile.divmod(iter_end, iter_length_quo, iter_length_mod);
|
||||
uint32_t current_iter_length = math::min(
|
||||
iter_length_mod == 0 ? (iter_end - iter_start) : iter_length_mod, total_iter_length);
|
||||
return current_iter_length;
|
||||
}
|
||||
|
||||
__device__ uint32_t get_tile_idx(uint32_t iter) const { return k_iters_per_tile.div(iter); }
|
||||
|
||||
__device__ void
|
||||
get_tile_idx_with_offset(uint32_t iter, uint32_t& tile_idx, uint32_t& iter_offset) const
|
||||
{
|
||||
k_iters_per_tile.divmod(iter, tile_idx, iter_offset);
|
||||
}
|
||||
|
||||
__device__ auto tile_to_spatial(uint32_t tile_idx, uint32_t m, uint32_t n) const
|
||||
{
|
||||
uint32_t m_tile_idx, n_tile_idx;
|
||||
uint32_t n_tiles_value = math::integer_divide_ceil(n, NPerBlock);
|
||||
n_tiles.divmod(tile_idx, n_tiles_value, m_tile_idx, n_tile_idx);
|
||||
|
||||
// swizzle tile
|
||||
uint32_t m_tiles = math::integer_divide_ceil(m, MPerBlock);
|
||||
|
||||
uint32_t tile_swizzle_sub_m_rem = m_tiles % tile_swizzle_sub_m;
|
||||
|
||||
const auto sub_m_adapt = (m_tile_idx < (m_tiles - tile_swizzle_sub_m_rem))
|
||||
? tile_swizzle_sub_m
|
||||
: tile_swizzle_sub_m_rem;
|
||||
|
||||
uint32_t m_tile_idx_sub0, m_tile_idx_sub1;
|
||||
m_tile_idx_sub0 = m_tile_idx / tile_swizzle_sub_m;
|
||||
m_tile_idx_sub1 = m_tile_idx % tile_swizzle_sub_m;
|
||||
|
||||
uint32_t tile_idx_local = n_tile_idx + m_tile_idx_sub1 * n_tiles_value;
|
||||
|
||||
uint32_t m_tile_idx_with_adapt, n_tile_idx_with_adapt;
|
||||
|
||||
n_tile_idx_with_adapt = tile_idx_local / sub_m_adapt;
|
||||
m_tile_idx_with_adapt = tile_idx_local % sub_m_adapt;
|
||||
return make_tuple(m_tile_idx_with_adapt + m_tile_idx_sub0 * tile_swizzle_sub_m,
|
||||
n_tile_idx_with_adapt);
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_workspace_size_for_acc(uint32_t acc_element_bytes) const
|
||||
{
|
||||
static constexpr uint32_t alignment = 128;
|
||||
uint32_t acc_buffer_bytes =
|
||||
MPerBlock * NPerBlock * get_total_acc_buffers() * acc_element_bytes;
|
||||
return (acc_buffer_bytes + alignment - 1) / alignment * alignment;
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_workspace_size_for_semaphore() const
|
||||
{
|
||||
return get_sk_tiles() * sizeof(uint32_t);
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_workspace_size(uint32_t acc_element_bytes) const
|
||||
{
|
||||
return get_workspace_size_for_acc(acc_element_bytes) + get_workspace_size_for_semaphore();
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_tile_intersections(uint32_t tiles_,
|
||||
const MDiv& eqav_tiles_) const
|
||||
{
|
||||
uint32_t tile_idx_ = tiles_ == 0 ? 0 : (tiles_ - 1);
|
||||
uint32_t max_eqav_tiles_ = eqav_tiles_.get() - 1;
|
||||
uint32_t quo_, rem_;
|
||||
eqav_tiles_.divmod(tile_idx_, quo_, rem_);
|
||||
return quo_ * max_eqav_tiles_ + rem_;
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_tiles_cover_sk_block(uint32_t num_sk_blocks_,
|
||||
uint32_t iters_per_sk_block_) const
|
||||
{
|
||||
return k_iters_per_tile.div(num_sk_blocks_ * iters_per_sk_block_ + k_iters_per_tile.get() -
|
||||
1);
|
||||
}
|
||||
|
||||
__host__ __device__ uint32_t get_total_acc_buffers() const
|
||||
{
|
||||
uint32_t tiles_cover_big_blocks =
|
||||
get_tiles_cover_sk_block(sk_num_big_blocks, k_iters_per_big_block);
|
||||
uint32_t tiles_cover_little_blocks =
|
||||
get_tiles_cover_sk_block(sk_num_blocks - sk_num_big_blocks, k_iters_per_big_block - 1);
|
||||
|
||||
uint32_t total_intersec_big =
|
||||
get_tile_intersections(tiles_cover_big_blocks, eqav_tiles_big);
|
||||
uint32_t total_intersec_little =
|
||||
get_tile_intersections(tiles_cover_little_blocks, eqav_tiles_little);
|
||||
|
||||
return sk_num_blocks + total_intersec_big + total_intersec_little;
|
||||
}
|
||||
|
||||
__device__ uint32_t get_acc_buffer_offset_from_tile(uint32_t tile_idx_) const
|
||||
{
|
||||
// TODO: from big to little
|
||||
uint32_t tiles_cover_big_blocks =
|
||||
get_tiles_cover_sk_block(sk_num_big_blocks, k_iters_per_big_block);
|
||||
if(tile_idx_ < tiles_cover_big_blocks)
|
||||
{
|
||||
uint32_t touched_sk_blocks =
|
||||
(tile_idx_ * k_iters_per_tile.get() + k_iters_per_big_block - 1) /
|
||||
k_iters_per_big_block;
|
||||
uint32_t current_intersec = get_tile_intersections(tile_idx_, eqav_tiles_big);
|
||||
return touched_sk_blocks + current_intersec;
|
||||
}
|
||||
else
|
||||
{
|
||||
uint32_t iters_per_little_sk_block = k_iters_per_big_block - 1;
|
||||
uint32_t tile_idx_little_reverse = get_sk_tiles() - tile_idx_;
|
||||
uint32_t touched_sk_blocks =
|
||||
(tile_idx_little_reverse * k_iters_per_tile.get() + iters_per_little_sk_block - 1) /
|
||||
iters_per_little_sk_block;
|
||||
uint32_t current_intersec =
|
||||
get_tile_intersections(tile_idx_little_reverse, eqav_tiles_little);
|
||||
return get_total_acc_buffers() - (touched_sk_blocks + current_intersec);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ uint32_t get_acc_buffer_offset_from_block(uint32_t block_idx_) const
|
||||
{
|
||||
uint32_t iters_per_big_sk_block = k_iters_per_big_block;
|
||||
uint32_t iters_per_little_sk_block = k_iters_per_big_block - 1;
|
||||
if(block_idx_ < sk_num_big_blocks)
|
||||
{
|
||||
uint32_t touched_tiles = k_iters_per_tile.div(block_idx_ * iters_per_big_sk_block +
|
||||
k_iters_per_tile.get() - 1);
|
||||
uint32_t current_intersec = get_tile_intersections(touched_tiles, eqav_tiles_big);
|
||||
return block_idx_ + current_intersec;
|
||||
}
|
||||
else
|
||||
{
|
||||
uint32_t block_idx_little_reverse = sk_num_blocks - block_idx_;
|
||||
uint32_t touched_tiles = k_iters_per_tile.div(
|
||||
block_idx_little_reverse * iters_per_little_sk_block + k_iters_per_tile.get() - 1);
|
||||
uint32_t current_intersec = get_tile_intersections(touched_tiles, eqav_tiles_little);
|
||||
return get_total_acc_buffers() - (block_idx_little_reverse + current_intersec);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
|
||||
namespace ck {
|
||||
|
||||
struct GridwiseGemmPipeline_v3
|
||||
{
|
||||
__host__ __device__ static constexpr bool IsSupported(index_t)
|
||||
{
|
||||
// TODO: improve applicability
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename AGridDesc,
|
||||
typename ABlockDesc,
|
||||
typename ABlockTransfer,
|
||||
typename AGridBuffer,
|
||||
typename ABlockBuffer,
|
||||
typename ABlockTransferStep,
|
||||
typename BGridDesc,
|
||||
typename BBlockDesc,
|
||||
typename BBlockTransfer,
|
||||
typename BGridBuffer,
|
||||
typename BBlockBuffer,
|
||||
typename BBlockTransferStep,
|
||||
typename BlockwiseGemm,
|
||||
typename CThreadBuffer>
|
||||
__device__ static void Run(const AGridDesc& a_grid_desc,
|
||||
const ABlockDesc& a_block_desc,
|
||||
ABlockTransfer& a_blockwise_copy,
|
||||
const AGridBuffer& a_grid_buf,
|
||||
ABlockBuffer& a_block_buf,
|
||||
const ABlockTransferStep& a_block_copy_step,
|
||||
const BGridDesc& b_grid_desc,
|
||||
const BBlockDesc& b_block_desc,
|
||||
BBlockTransfer& b_blockwise_copy,
|
||||
const BGridBuffer& b_grid_buf,
|
||||
BBlockBuffer& b_block_buf,
|
||||
const BBlockTransferStep& b_block_copy_step,
|
||||
const BlockwiseGemm& blockwise_gemm,
|
||||
CThreadBuffer& c_thread_buf,
|
||||
index_t num_loop)
|
||||
{
|
||||
// global read 0
|
||||
a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf);
|
||||
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf);
|
||||
|
||||
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
|
||||
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
|
||||
|
||||
// Initialize C
|
||||
c_thread_buf.Clear();
|
||||
|
||||
// LDS write 0
|
||||
a_blockwise_copy.RunWrite(a_block_desc, a_block_buf);
|
||||
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf);
|
||||
|
||||
num_loop--;
|
||||
|
||||
while(num_loop > 0)
|
||||
{
|
||||
a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf);
|
||||
block_sync_lds();
|
||||
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf);
|
||||
|
||||
blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
|
||||
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
|
||||
a_blockwise_copy.RunWrite(a_block_desc, a_block_buf);
|
||||
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf);
|
||||
|
||||
num_loop--;
|
||||
}
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user