From 8b6c11b490170a58d0742c1f9a135fe7552dcd35 Mon Sep 17 00:00:00 2001 From: Manish Kumar Date: Tue, 25 Nov 2025 14:51:55 +0000 Subject: [PATCH] Fix build errors --- .../ck_tile/17_grouped_gemm/CMakeLists.txt | 4 +- .../grouped_gemm_persistent_async.cpp | 34 +- .../grouped_gemm_persistent_async.hpp | 315 ++++-------------- .../persistent_async_scheduler.hpp | 19 +- .../persistent_async_utils.hpp | 63 ++++ .../ck_tile/ops/gemm/kernel/gemm_kernel.hpp | 16 +- .../ops/gemm/kernel/universal_gemm_kernel.hpp | 128 +++---- 7 files changed, 232 insertions(+), 347 deletions(-) create mode 100755 example/ck_tile/17_grouped_gemm/persistent_async_utils.hpp mode change 100644 => 100755 include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp mode change 100644 => 100755 include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp diff --git a/example/ck_tile/17_grouped_gemm/CMakeLists.txt b/example/ck_tile/17_grouped_gemm/CMakeLists.txt index bbfb2df006..ef135a5f9f 100755 --- a/example/ck_tile/17_grouped_gemm/CMakeLists.txt +++ b/example/ck_tile/17_grouped_gemm/CMakeLists.txt @@ -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}) \ No newline at end of file +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}) diff --git a/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.cpp b/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.cpp index abc9ac70ae..450ceaa5eb 100755 --- a/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.cpp +++ b/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.cpp @@ -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 +#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( + a_layout, b_layout, data_type, arg_parser, + , tiles_per_chunk_m, tile_idx_pivot_m); + + + + */ + + return 0; } diff --git a/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.hpp b/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.hpp index d9202a10d3..814f45900e 100755 --- a/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.hpp +++ b/example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.hpp @@ -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 - static float gemm(const ck_tile::GemmHostArgs& args, - const ck_tile::stream_config& s, - const ck_tile::PersistentAsyncArgs& async_args) +template + 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, - ck_tile::sequence, - ck_tile::sequence>; - - using TilePartitioner = - ck_tile::GemmSpatiallyLocalTilePartitioner; - - using Traits = ck_tile::TileGemmTraits; + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + using TilePartitioner = ck_tile:: + GemmSpatiallyLocalTilePartitioner; using GemmUniversalTraits = - ck_tile::TileGemmUniversalTraits; + ck_tile::PersistentTileGemmUniversalTraits; - using GemmPipelineProblem = - ck_tile::GemmPipelineProblem; - - using BaseGemmPipeline = typename PipelineTypeTraits< - GemmConfig::Pipeline>::template UniversalGemmPipeline; - - 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; - - using GemmPipeline = typename PipelineTypeTraits< - GemmConfig::Pipeline>::template GemmPipeline; - - using WorkspaceType = ck_tile::remove_cvref_t; + scheduler>; + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem>; - - using GemmKernel = ck_tile::GemmKernel; - - 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; - using Problem = ck_tile::ElementWisePipelineProblem; - using ElementwiseKernel = - ck_tile::ElementWiseKernel; - - ck_tile::index_t total_elements = 1; - std::vector 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(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; + 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> rotating_mem_ptr; - std::function 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 a_m(ck_tile::host_tensor_descriptor( - args.M, args.K, args.stride_A, is_row_major(ALayout{}))); - ck_tile::HostTensor 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>( - 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( - GemmKernel{}, grids, blocks, 0, gemm_kargs), - ck_tile::make_kernel(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(c_ptr))); - - return ave_time; + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + 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{}); + } + else + { - return ave_time = BaseGemmPipeline::TailHandler(RunSplitK, has_hot_loop, tail_num); + Run(ck_tile::integral_constant{}); + } } -}; diff --git a/example/ck_tile/17_grouped_gemm/persistent_async_scheduler.hpp b/example/ck_tile/17_grouped_gemm/persistent_async_scheduler.hpp index 65837de947..b98f655dcf 100755 --- a/example/ck_tile/17_grouped_gemm/persistent_async_scheduler.hpp +++ b/example/ck_tile/17_grouped_gemm/persistent_async_scheduler.hpp @@ -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_) { } }; diff --git a/example/ck_tile/17_grouped_gemm/persistent_async_utils.hpp b/example/ck_tile/17_grouped_gemm/persistent_async_utils.hpp new file mode 100755 index 0000000000..9c6c381196 --- /dev/null +++ b/example/ck_tile/17_grouped_gemm/persistent_async_utils.hpp @@ -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 diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp old mode 100644 new mode 100755 index 8dba7f9792..d632b1596c --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -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 @@ -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 diff --git a/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp old mode 100644 new mode 100755 index b50197f7d4..dde42cdab2 --- a/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp @@ -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& stride_As_, const std::array& stride_Bs_, const std::array& 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)