ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 01:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

merge M grouped flatmm

Browse Source
This commit is contained in:
lalala-sh
2025-07-30 07:55:09 +00:00
parent 1b493fac62
commit c4aa2fef46
9 changed files with 1797 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/ck_tile/ops/flatmm/kernel/flatmm_kernel.hpp Executable file → Normal file
View File

@@ -244,7 +244,9 @@ struct FlatmmKernel
static_assert(DsLayout::size() == DsDataType::size(),
"The size of DsLayout and DsDataType should be the same");
// using KernelArgs = FlatmmKernelArgs<DsLayout::size()>;
template<class ScaleM, class ScaleN>
using KernelArgs = FlatmmKernelArgs<ScaleM, ScaleN, DsLayout::size()>;
[[nodiscard]] CK_TILE_HOST static const std::string GetName()
{
@@ -751,7 +753,7 @@ struct FlatmmKernel
CK_TILE_DEVICE void operator()(FlatmmKernelArgs<ScaleM, ScaleN, DsDataType::size()> kargs,
int partition_idx = blockIdx.x) const
{
const auto [iM, iN] = TilePartitioner{kargs.M, kargs.N}.GetOutputTileIndex(blockIdx.x);
const auto [iM, iN] = TilePartitioner{kargs.M, kargs.N}.GetOutputTileIndex(partition_idx);
const index_t i_m = __builtin_amdgcn_readfirstlane(iM * TilePartitioner::MPerBlock);
const index_t i_n = __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock);