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

Fix build errors

Browse Source
This commit is contained in:
Manish Kumar
2025-11-25 14:51:55 +00:00
parent 7aead9064b
commit 8b6c11b490
7 changed files with 232 additions and 347 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
example/ck_tile/17_grouped_gemm/CMakeLists.txt
View File

@@ -2,6 +2,7 @@ add_executable(tile_example_grouped_gemm EXCLUDE_FROM_ALL grouped_gemm.cpp)
add_executable(tile_example_quant_grouped_gemm EXCLUDE_FROM_ALL quant_grouped_gemm.cpp)
add_executable(tile_example_grouped_gemm_preshuffle EXCLUDE_FROM_ALL grouped_gemm_preshuffle.cpp)
add_executable(tile_example_grouped_gemm_multi_d EXCLUDE_FROM_ALL grouped_gemm_multi_d.cpp)
add_executable(tile_example_grouped_gemm_persistent_async EXCLUDE_FROM_ALL grouped_gemm_persistent_async.cpp)
set(EXAMPLE_GEMM_COMPILE_OPTIONS)
if(CK_USE_OCP_FP8)
list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
@@ -9,4 +10,5 @@ endif()
target_compile_options(tile_example_grouped_gemm PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
target_compile_options(tile_example_grouped_gemm_preshuffle PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
target_compile_options(tile_example_grouped_gemm_multi_d PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
target_compile_options(tile_example_quant_grouped_gemm PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
target_compile_options(tile_example_quant_grouped_gemm PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
target_compile_options(tile_example_grouped_gemm_persistent_async PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})

34
example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.cpp
View File

@@ -1,10 +1,12 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#include "run_grouped_gemm_example.inc"
#include "persistent_async_scheduler.hpp"
#include "persistent_async_utils.hpp"
#include "ck_tile/core/utility/gemm_validation.hpp"
#include <hip/hip_runtime.h>
#include "grouped_gemm.hpp"
#include "grouped_gemm_persistent_async.hpp"
/**
* @brief Helper to allocate and initialize chunk signals
@@ -13,7 +15,7 @@
* @param stream HIP stream for async operations
* @return Device pointer to chunk signals array
*/
static uint32_t* allocate_chunk_signals(int num_chunks, hipStream_t stream)
[[maybe_unused]] static uint32_t* allocate_chunk_signals(int num_chunks, hipStream_t stream)
{
uint32_t* signals_device = nullptr;
@@ -34,7 +36,7 @@ static uint32_t* allocate_chunk_signals(int num_chunks, hipStream_t stream)
* @param chunk_idx Index of chunk to signal
* @param stream HIP stream for async operations
*/
static void signal_chunk_ready(uint32_t* signals, int chunk_idx, hipStream_t stream)
[[maybe_unused]] static void signal_chunk_ready(uint32_t* signals, int chunk_idx, hipStream_t stream)
{
uint32_t ready = 1;
ck_tile::hip_check_error(hipMemcpyAsync(
@@ -43,7 +45,7 @@ static void signal_chunk_ready(uint32_t* signals, int chunk_idx, hipStream_t str
int main(int argc, char* argv[])
{
auto arg_parser = create_args();
auto [result, arg_parser] = create_args(argc, argv);
// Add async-specific arguments
arg_parser.insert(
@@ -52,12 +54,28 @@ int main(int argc, char* argv[])
"tile_idx_pivot_m", "0", "Pivot offset for M dimension (for hotspot spreading)");
arg_parser.insert("enable_async", "1", "Enable async input signaling (0=disabled, 1=enabled)");
auto result = arg_parser.parse(argc, argv);
// TO-DO Add example
if(!result)
return -1;
/*TO-DO
// Parse async-specific arguments
const bool enable_async = arg_parser.get_int("enable_async") != 0;
const ck_tile::index_t tiles_per_chunk_m = arg_parser.get_int("tiles_per_chunk_m");
const ck_tile::index_t tile_idx_pivot_m = arg_parser.get_int("tile_idx_pivot_m");
const std::string a_layout = arg_parser.get_str("a_layout");
const std::string b_layout = arg_parser.get_str("b_layout");
const std::string data_type = arg_parser.get_str("prec");
auto res = invoke_grouped_gemm_persistent_async<ck_tile::half_t>(
a_layout, b_layout, data_type, arg_parser,
, tiles_per_chunk_m, tile_idx_pivot_m);
*/
return 0;
}

315
example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.hpp
View File

@@ -1,279 +1,108 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "gemm_utils.hpp"
#include "persistent_async_scheduler.hpp"
#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp"
#include "ck_tile/ops/epilogue.hpp"
/**
* @brief Invoker for Persistent Async GEMM
*
* This invoker implements persistent GEMM with asynchronous input readiness.
* It extends the standard GEMM with support for:
* - Chunk-based async input signaling
* - Producer-consumer synchronization
* - Pivot-based tile traversal
*/
struct GemmPersistentAsyncInvoker
{
template <typename GemmConfig,
typename ADataType,
typename BDataType,
typename DsDataType,
typename AccDataType,
typename CDataType,
typename ALayout,
typename BLayout,
typename DsLayout,
typename CLayout,
bool Persistent,
typename CDEElementWise>
static float gemm(const ck_tile::GemmHostArgs& args,
const ck_tile::stream_config& s,
const ck_tile::PersistentAsyncArgs& async_args)
template <typename GroupedGemKernelParam, typename ADataType, typename BDataType, typename AccDataType, typename DsDataType, typename CDataType, typename DsLayout, typename ALayout, typename BLayout, typename CLayout>
void invoke_grouped_gemm_persistent(const ck_tile::stream_config& s,
const ck_tile::index_t num_groups,
void* kargs_ptr,
bool splitk)
{
constexpr bool TransposeC = false;
constexpr bool DoubleSmemBuffer = false;
static_assert(Persistent, "This invoker only supports persistent GEMM.");
constexpr int kBlockPerCu = 1;
constexpr ck_tile::index_t TileParitionerGroupNum = 8;
constexpr ck_tile::index_t TileParitionerM01 = 4;
using GemmShape = ck_tile::TileGemmShape<
ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
ck_tile::sequence<GemmConfig::M_Warp, GemmConfig::N_Warp, GemmConfig::K_Warp>,
ck_tile::sequence<GemmConfig::M_Warp_Tile,
GemmConfig::N_Warp_Tile,
GemmConfig::K_Warp_Tile>>;
using TilePartitioner =
ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
GemmConfig::TilePartitionerGroupNum,
GemmConfig::TilePartitionerM01>;
using Traits = ck_tile::TileGemmTraits<GemmConfig::kPadM,
GemmConfig::kPadN,
GemmConfig::kPadK,
ALayout,
BLayout,
ELayout,
GemmConfig::NumWaveGroups>;
using GemmShape =
ck_tile::TileGemmShape<ck_tile::sequence<GroupedGemKernelParam::M_Tile,
GroupedGemKernelParam::N_Tile,
GroupedGemKernelParam::K_Tile>,
ck_tile::sequence<GroupedGemKernelParam::M_Warp,
GroupedGemKernelParam::N_Warp,
GroupedGemKernelParam::K_Warp>,
ck_tile::sequence<GroupedGemKernelParam::M_Warp_Tile,
GroupedGemKernelParam::N_Warp_Tile,
GroupedGemKernelParam::K_Warp_Tile>>;
using TilePartitioner = ck_tile::
GemmSpatiallyLocalTilePartitioner<GemmShape, TileParitionerGroupNum, TileParitionerM01>;
using GemmUniversalTraits =
ck_tile::TileGemmUniversalTraits<GemmConfig::kPadM,
GemmConfig::kPadN,
GemmConfig::kPadK,
GemmConfig::DoubleSmemBuffer,
ALayout,
BLayout,
CLayout,
GemmConfig::TransposeC,
GemmConfig::UseStructuredSparsity,
true, // Persistent = true
GemmConfig::NumWaveGroups,
GemmConfig::Preshuffle>;
ck_tile::PersistentTileGemmUniversalTraits<GroupedGemKernelParam::kPadM,
GroupedGemKernelParam::kPadN,
GroupedGemKernelParam::kPadK,
DoubleSmemBuffer,
ALayout,
BLayout,
CLayout,
TransposeC>;
using GemmPipelineProblem =
ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
using BaseGemmPipeline = typename PipelineTypeTraits<
GemmConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
const ck_tile::index_t k_grain = args.k_batch * GemmConfig::K_Tile;
const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * GemmConfig::K_Tile;
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
float ave_time{0};
const auto Run = [&](const auto has_hot_loop_,
const auto tail_number_,
const auto memory_operation_) {
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
constexpr auto scheduler = GemmConfig::Scheduler;
const auto Run = [&](const auto memory_operation_) {
constexpr auto scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
constexpr auto memory_operation = memory_operation_.value;
// We create the GEMM pipeline without specifying hotloop or tailnumber.
// These are automatically run inside the kernel based on the given input data.
using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
BDataType,
AccDataType,
GemmShape,
GemmUniversalTraits,
scheduler,
has_hot_loop_v,
tail_number_v>;
using GemmPipeline = typename PipelineTypeTraits<
GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
using WorkspaceType = ck_tile::remove_cvref_t<typename GemmConfig::WorkspaceType>;
scheduler>;
using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<UniversalGemmProblem>;
using GemmEpilogue = ck_tile::CShuffleEpilogue<
ck_tile::CShuffleEpilogueProblem<ADataType,
BDataType,
DsDataType,
AccDataType,
WorkspaceType,
CDataType,
DsLayout,
ELayout,
CDEElementWise,
CLayout,
ck_tile::element_wise::PassThrough,
TilePartitioner::MPerBlock,
TilePartitioner::NPerBlock,
GemmConfig::M_Warp,
GemmConfig::N_Warp,
GemmConfig::M_Warp_Tile,
GemmConfig::N_Warp_Tile,
GemmConfig::K_Warp_Tile,
GroupedGemKernelParam::M_Warp,
GroupedGemKernelParam::N_Warp,
GroupedGemKernelParam::M_Warp_Tile,
GroupedGemKernelParam::N_Warp_Tile,
GroupedGemKernelParam::K_Warp_Tile,
UniversalGemmProblem::TransposeC,
memory_operation,
GemmConfig::NumWaveGroups>>;
using GemmKernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
ck_tile::DeviceMem ws_m_n_dev_buf(args.M * args.N * sizeof(WorkspaceType));
ck_tile::GemmHostArgs ws_args = ck_tile::GemmHostArgs(args);
auto c_ptr = ws_args.c_ptr;
ws_args.c_ptr = ws_m_n_dev_buf.GetDeviceBuffer();
// Add persistent async arguments to ws_args
ws_args.chunk_signals = async_args.chunk_signals;
ws_args.tiles_per_chunk_m = async_args.tiles_per_chunk_m;
auto gemm_kargs = GemmKernel::MakeKernelArgs(ws_args);
const dim3 grids = Persistent ? GemmKernel::MaxOccupancyGridSize(s)
: GemmKernel::GridSize(args.M, args.N, args.k_batch);
const dim3 blocks = GemmKernel::BlockSize();
if(!GemmKernel::IsSupportedArgument(gemm_kargs))
{
throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n");
}
using XElementwiseOperation = ck_tile::element_wise::UnaryConvert;
using BlockTile = ck_tile::sequence<2048>;
using BlockWarps = ck_tile::sequence<8>;
using WarpTile = ck_tile::sequence<64>;
using ElementwiseShape =
ck_tile::ElementWiseShape<BlockWarps, BlockTile, WarpTile, WorkspaceType>;
using Problem = ck_tile::ElementWisePipelineProblem<WorkspaceType,
WorkspaceType,
CDataType,
ElementwiseShape,
XElementwiseOperation>;
using ElementwiseKernel =
ck_tile::ElementWiseKernel<Problem, ck_tile::ElementWiseDefaultPolicy>;
ck_tile::index_t total_elements = 1;
std::vector<ck_tile::index_t> shape = {args.M, args.N};
for(auto d : shape)
total_elements *= d;
const ck_tile::index_t kBlockSize = ElementwiseKernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = 1;
constexpr ck_tile::index_t elements_per_block = BlockTile::at(ck_tile::number<0>{});
ck_tile::index_t kGridSize =
(total_elements + elements_per_block - 1) / elements_per_block;
auto input_tensors = ck_tile::make_tuple(static_cast<WorkspaceType*>(ws_args.c_ptr));
auto input_size = ck_tile::make_tuple(args.M, args.N);
// Check if the kernel configuration is supported
if(!ElementwiseKernel::IsSupportedArgument(input_size))
{
throw std::runtime_error(
"Wrong! Elementwise arguments not supported! Skipping gemm!\n");
}
memory_operation>>;
using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
const dim3 blocks = Kernel::BlockSize();
const dim3 grids = Kernel::MaxOccupancyGridSize(s);
if(s.log_level_ > 0)
{
std::cout << "Launching Persistent Async GEMM kernel:\n"
<< " Kernel: " << Kernel::GetName() << '\n'
<< " Shape: " << GemmShape::GetName() << '\n'
<< " Problem: " << UniversalGemmProblem::GetName() << '\n'
<< " Pipeline: " << GemmPipeline::GetName() << '\n'
<< " Grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}\n"
<< " Blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
<< "}\n"
<< " Async Args:\n"
<< " tiles_per_chunk_m: " << async_args.tiles_per_chunk_m << '\n'
<< " tile_idx_pivot_m: " << async_args.tile_idx_pivot_m << '\n'
<< " chunk_signals: "
<< (async_args.chunk_signals ? "enabled" : "disabled") << std::endl;
std::cout << "Launching kernel: " << Kernel::GetName()
<< " with args:" << " grid: {" << grids.x << ", " << grids.y << ", "
<< grids.z << "}" << ", blocks: {" << blocks.x << ", " << blocks.y << ", "
<< blocks.z << "}" << std::endl;
}
// Declare rotating_mem_ptr here so it stays in scope until it is needed
std::unique_ptr<ck_tile::RotatingMemWrapper<ADataType, BDataType>> rotating_mem_ptr;
std::function<void()> preprocess;
auto clear_gemm_output = [&]() {
if(args.k_batch > 1)
hipGetErrorString(hipMemsetAsync(
ws_args.c_ptr, 0, args.M * args.N * sizeof(WorkspaceType), s.stream_id_));
};
if(s.flush_cache_)
{
std::cout << "Flushing cache..." << std::endl;
ck_tile::HostTensor<ADataType> a_m(ck_tile::host_tensor_descriptor(
args.M, args.K, args.stride_A, is_row_major(ALayout{})));
ck_tile::HostTensor<BDataType> b_n(ck_tile::host_tensor_descriptor(
args.K, args.N, args.stride_B, is_row_major(BLayout{})));
auto size_a_buffer = a_m.get_element_space_size_in_bytes();
auto size_b_buffer = b_n.get_element_space_size_in_bytes();
rotating_mem_ptr =
std::make_unique<ck_tile::RotatingMemWrapper<ADataType, BDataType>>(
gemm_kargs.as_ptr[0],
gemm_kargs.bs_ptr[0],
s.rotating_count_,
size_a_buffer,
size_b_buffer);
rotating_mem_ptr->Print();
preprocess = [&]() {
ck_tile::flush_icache();
rotating_mem_ptr->Next();
clear_gemm_output();
};
}
else
{
preprocess = clear_gemm_output;
}
ave_time = ck_tile::launch_kernel_time_mask(
s,
preprocess,
ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
GemmKernel{}, grids, blocks, 0, gemm_kargs),
ck_tile::make_kernel<kBlockPerCu>(ElementwiseKernel{},
kGridSize,
kBlockSize,
0,
input_size,
ck_tile::make_tuple(args.N, 1), // Input Stride
ck_tile::make_tuple(args.N, 1), // Output Stride
input_tensors,
static_cast<CDataType*>(c_ptr)));
return ave_time;
ck_tile::launch_kernel(s,
ck_tile::make_kernel<kBlockPerCu>(
Kernel{},
grids,
blocks,
0,
ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
num_groups));
};
const auto RunSplitK = [&](const auto has_hot_loop_, const auto tail_number_) {
if(args.k_batch == 1)
{
return Run(has_hot_loop_, tail_number_, MemoryOpSet{});
}
else
{
return Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
}
};
if(splitk)
{
Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::atomic_add>{});
}
else
{
return ave_time = BaseGemmPipeline::TailHandler(RunSplitK, has_hot_loop, tail_num);
Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::set>{});
}
}
};

19
example/ck_tile/17_grouped_gemm/persistent_async_scheduler.hpp
View File

@@ -16,19 +16,28 @@ namespace ck_tile {
*/
struct PersistentAsyncArgs
{
/// Number of M tiles per chunk (granularity of async readiness signaling)
index_t tiles_per_chunk_m = 0;
/// Device pointer to global chunk readiness flags (1 = ready, 0 = not ready)
uint32_t* chunk_signals = nullptr;
/// Pivot offset for M dimension (for hotspot spreading in tile scheduling)
index_t tile_idx_pivot_m = 0;
PersistentAsyncArgs(index_t tiles_per_chunk_m_,
uint32_t* chunk_signals_,
index_t tile_idx_pivot_m_,
bool enable_async_)
/// Enable/disable async input signaling (false = disabled, true = enabled)
bool enable_async = false;
CK_TILE_HOST_DEVICE PersistentAsyncArgs() = default;
CK_TILE_HOST_DEVICE PersistentAsyncArgs(index_t tiles_per_chunk_m_,
uint32_t* chunk_signals_,
index_t tile_idx_pivot_m_,
bool enable_async_ = false)
: tiles_per_chunk_m(tiles_per_chunk_m_),
chunk_signals(chunk_signals_),
tile_idx_pivot_m(tile_idx_pivot_m_)
tile_idx_pivot_m(tile_idx_pivot_m_),
enable_async(enable_async_)
{
}
};

63
example/ck_tile/17_grouped_gemm/persistent_async_utils.hpp Executable file
View File

@@ -0,0 +1,63 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
namespace ck_tile {
/**
* @brief Safe iteration boundary fence for persistent kernels
*
* This function ensures memory consistency between iterations in persistent loops by:
* - Waiting for all vector memory operations to complete (vmcnt=0)
* - Waiting for all LDS/GDS operations to complete (lgkmcnt=0)
* - Synchronizing all workgroup threads via barrier
*
* This prevents race conditions when reusing LDS or moving to the next tile.
*/
CK_TILE_DEVICE static void iteration_boundary_fence()
{
__builtin_amdgcn_s_waitcnt(0);
__builtin_amdgcn_s_waitcnt(0);
__builtin_amdgcn_s_barrier();
}
/**
* @brief Wait for chunk readiness signal (producer-consumer synchronization)
*
* This function implements producer-consumer synchronization for async input readiness:
* - One lane polls the chunk_signals[chunk_idx] flag with acquire semantics
* - When signal becomes ready (value == 1), a workgroup barrier releases all threads
*
* @param chunk_signals Device pointer to global chunk readiness flags array
* @param chunk_idx Index of the chunk to wait for
*
* @note Only lane 0 performs the polling to minimize global memory traffic
* @note Uses acquire semantics to ensure proper memory ordering
*/
CK_TILE_DEVICE static void wait_chunk_signal(const uint32_t* chunk_signals, index_t chunk_idx)
{
// Only lane 0 polls the signal to minimize global memory traffic
if(threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0)
{
volatile const uint32_t* signal_ptr = chunk_signals + chunk_idx;
// Poll until chunk is ready (signal == 1)
// Use acquire semantics for proper memory ordering
uint32_t signal_value;
do {
signal_value = __builtin_nontemporal_load(signal_ptr);
__builtin_amdgcn_s_sleep(1); // Brief sleep to reduce contention
} while(signal_value == 0);
// Memory fence with acquire semantics
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
}
// Barrier to release all threads in the workgroup
__builtin_amdgcn_s_barrier();
}
} // namespace ck_tile

16
include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp Normal file → Executable file
View File

@@ -37,9 +37,7 @@ struct GemmHostArgs
index_t K_,
index_t stride_A_,
index_t stride_B_,
index_t stride_E_,
uint32_t* chunk_signals_ = nullptr,
index_t tiles_per_chunk_m_ = 0)
index_t stride_E_)
: a_ptr(a_ptr_),
b_ptr(b_ptr_),
e_ptr(e_ptr_),
@@ -49,9 +47,7 @@ struct GemmHostArgs
stride_A(stride_A_),
stride_B(stride_B_),
stride_E(stride_E_),
k_batch(k_batch_),
chunk_signals(chunk_signals_),
tiles_per_chunk_m(tiles_per_chunk_m_)
k_batch(k_batch_)
{
}
@@ -76,10 +72,6 @@ struct GemmHostArgs
};
index_t k_batch;
// Persistent async arguments
uint32_t* chunk_signals;
index_t tiles_per_chunk_m;
};
template <typename TilePartitioner_, typename GemmPipeline_, typename EpiloguePipeline_>
@@ -161,9 +153,7 @@ struct GemmKernel
{hostArgs.stride_A},
{hostArgs.stride_B},
{/*hostArgs.stride_Ds*/},
hostArgs.stride_E,
hostArgs.chunk_signals,
hostArgs.tiles_per_chunk_m));
hostArgs.stride_E));
}
CK_TILE_HOST static auto

128
include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp Normal file → Executable file
View File

@@ -16,59 +16,59 @@
namespace ck_tile {
/**
* @brief Wait for a signal to become ready with acquire semantics
*
* Producer-only wait: One lane polls chunk_signals[chunk_idx] with acquire semantics,
* then a workgroup barrier releases everyone.
*
* @param signal_addr Pointer to the signal location in device memory
*/
CK_TILE_DEVICE static inline void wait_signal(uint32_t* signal_addr)
{
// Only one thread in the workgroup polls the signal
if(threadIdx.x == 0)
{
uint32_t ready = 0;
while(!ready)
{
// Load with acquire semantics using AMD intrinsics
// glc (globally coherent) ensures visibility across the system
asm volatile("flat_load_dword %0, %1 glc\n\t"
"s_waitcnt vmcnt(0)"
: "=v"(ready)
: "v"(signal_addr)
: "memory");
// /**
// * @brief Wait for a signal to become ready with acquire semantics
// *
// * Producer-only wait: One lane polls chunk_signals[chunk_idx] with acquire semantics,
// * then a workgroup barrier releases everyone.
// *
// * @param signal_addr Pointer to the signal location in device memory
// */
// CK_TILE_DEVICE static void wait_signal(uint32_t* signal_addr)
// {
// // Only one thread in the workgroup polls the signal
// if(threadIdx.x == 0)
// {
// uint32_t ready = 0;
// while(!ready)
// {
// // Load with acquire semantics using AMD intrinsics
// // glc (globally coherent) ensures visibility across the system
// asm volatile("flat_load_dword %0, %1 glc\n\t"
// "s_waitcnt vmcnt(0)"
// : "=v"(ready)
// : "v"(signal_addr)
// : "memory");
// Add a small delay to reduce memory traffic
if(!ready)
{
__builtin_amdgcn_s_sleep(1);
}
}
}
// // Add a small delay to reduce memory traffic
// if(!ready)
// {
// __builtin_amdgcn_s_sleep(1);
// }
// }
// }
// Workgroup barrier to release all threads after signal is ready
__builtin_amdgcn_s_barrier();
}
// // Workgroup barrier to release all threads after signal is ready
// __builtin_amdgcn_s_barrier();
// }
/**
* @brief Fence for safe iteration boundaries in persistent loops
*
* Ensures all memory operations are complete before reusing LDS or moving to next tile.
* Uses s_waitcnt vmcnt=0, lgkmcnt=0 + s_barrier.
*/
CK_TILE_DEVICE static inline void iteration_boundary_fence()
{
// Wait for all vector memory operations (global memory loads/stores)
__builtin_amdgcn_s_waitcnt_vmcnt(0);
// /**
// * @brief Fence for safe iteration boundaries in persistent loops
// *
// * Ensures all memory operations are complete before reusing LDS or moving to next tile.
// * Uses s_waitcnt vmcnt=0, lgkmcnt=0 + s_barrier.
// */
// CK_TILE_DEVICE static void iteration_boundary_fence()
// {
// // Wait for all vector memory operations (global memory loads/stores)
// __builtin_amdgcn_s_waitcnt(0);
// Wait for all LDS operations
__builtin_amdgcn_s_waitcnt_lgkmcnt(0);
// // Wait for all LDS operations
// __builtin_amdgcn_s_waitcnt(0);
// Synchronize all threads in the workgroup
__builtin_amdgcn_s_barrier();
}
// // Synchronize all threads in the workgroup
// __builtin_amdgcn_s_barrier();
// }
/// @brief The Universal GEMM kernel host arguments.
///
@@ -95,9 +95,7 @@ struct UniversalGemmHostArgs
const std::array<index_t, NumATensor>& stride_As_,
const std::array<index_t, NumBTensor>& stride_Bs_,
const std::array<index_t, NumDTensor>& stride_Ds_,
index_t stride_E_,
uint32_t* chunk_signals_ = nullptr,
index_t tiles_per_chunk_m_ = 0)
index_t stride_E_)
: as_ptr(as_ptr_),
bs_ptr(bs_ptr_),
ds_ptr(ds_ptr_),
@@ -109,9 +107,7 @@ struct UniversalGemmHostArgs
stride_Bs(stride_Bs_),
stride_Ds(stride_Ds_),
stride_E(stride_E_),
k_batch(k_batch_),
chunk_signals(chunk_signals_),
tiles_per_chunk_m(tiles_per_chunk_m_)
k_batch(k_batch_)
{
}
@@ -136,10 +132,6 @@ struct UniversalGemmHostArgs
};
index_t k_batch;
// Persistent async arguments
uint32_t* chunk_signals;
index_t tiles_per_chunk_m;
};
/// @brief The GEMM kernel device arguments.
@@ -174,11 +166,6 @@ struct UniversalGemmKernelArgs
index_t stride_E;
index_t k_batch;
/// @brief Pointer to chunk signals for async producer-consumer synchronization.
/// chunk_signals[i] == 1 indicates that chunk i is ready.
uint32_t* chunk_signals;
/// @brief Number of M tiles per chunk for async input signaling.
index_t tiles_per_chunk_m;
};
/// @brief The Universal GEMM kernel template.
@@ -381,9 +368,7 @@ struct UniversalGemmKernel
hostArgs.stride_Bs,
hostArgs.stride_Ds,
hostArgs.stride_E,
hostArgs.k_batch,
hostArgs.chunk_signals,
hostArgs.tiles_per_chunk_m};
hostArgs.k_batch};
}
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
@@ -1211,13 +1196,6 @@ struct UniversalGemmKernel
const index_t i_m = amd_wave_read_first_lane(iM * TilePartitioner::MPerBlock);
const index_t i_n = amd_wave_read_first_lane(iN * TilePartitioner::NPerBlock);
// Producer-consumer synchronization: wait for chunk to be ready
if(kargs.chunk_signals != nullptr && kargs.tiles_per_chunk_m > 0)
{
const index_t chunk_idx = iM / kargs.tiles_per_chunk_m;
wait_signal(kargs.chunk_signals + chunk_idx);
}
// Get the SplitK offset for this block
const auto k_batch = amd_wave_read_first_lane(block_id / num_tiles);
const SplitKBatchOffset splitk_batch_offset(kargs, k_batch);
@@ -1284,10 +1262,6 @@ struct UniversalGemmKernel
}
}
// Safe iteration boundary: ensure all memory operations complete
// before reusing LDS or moving to next tile
iteration_boundary_fence();
// Advance to the next work item
block_id += grid_size;
if(block_id >= num_work)