From 8b185e872e1172eba2444ccb49648469598d72e8 Mon Sep 17 00:00:00 2001 From: Thrupti Raj Lakshmana Gowda Date: Mon, 27 Oct 2025 09:15:34 -0500 Subject: [PATCH 01/17] Ck tile engine preshuffle (#2919) * Partial Progress : Preshuffle working code for datatype * Partial Progress : Preshuffle Cleanup * Working code for default config with min max step * Partial Progress : PermuteN implemented in validation * Partial Progress : PermuteN changes in Preshuffle * CK Tile Engine Preshuffle Complete * CK TILE ENGINE : Preshuffle Layout validation * CK Tile Engine Preshuffle Validation * Preshuffle Validation check * CK Tile Engine Preshuffle : Fixing Validation Cases * Addressing PR review Comments * Changes in config * Addressing Review Comments * Adding additional architecture in Jenkins * Partial Progress : Selective Datatype and layouts * Limited datatypes and layouts * Addressing CI errors * Datatype updates * Datatype updates * Datatype changes to Preshuffle * Addressing Review Comments * Addressing Review Comments * Datatype changes * Changes to Cmake * Update on Jenkins * Formatting with precommit * Ruff Formatting --- Jenkinsfile | 10 +- .../configs/simple_test_config.json | 6 +- tile_engine/ops/gemm/CMakeLists.txt | 2 +- .../ops/gemm_preshuffle/CMakeLists.txt | 13 +- .../commons/validation_utils.py | 148 ++++++++++++- .../configs/default_config.json | 81 ++++---- .../configs/user_provided_config.json | 22 +- ...ffle.hpp => gemm_preshuffle_benchmark.hpp} | 8 + ...p => gemm_preshuffle_benchmark_single.cpp} | 15 +- .../gemm_preshuffle_common.hpp | 76 +++---- .../gemm_preshuffle_instance_builder.py | 194 +++++++++++------- .../gemm_preshuffle_profiler.hpp | 27 ++- 12 files changed, 393 insertions(+), 209 deletions(-) rename tile_engine/ops/gemm_preshuffle/{benchmark_gemm_preshuffle.hpp => gemm_preshuffle_benchmark.hpp} (97%) rename tile_engine/ops/gemm_preshuffle/{benchmark_gemm_preshuffle_single.cpp => gemm_preshuffle_benchmark_single.cpp} (89%) diff --git a/Jenkinsfile b/Jenkinsfile index 7a8574df05..b89d6fb657 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -1488,7 +1488,7 @@ pipeline { -D GEMM_LAYOUT="rcr;rrr;crr;ccr" \ -D GEMM_MULTI_D_DATATYPE="fp16" \ -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8" \ + -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ -DCMAKE_CXX_FLAGS=" -O3 " .. && \ ninja -j64 benchmark_gemm_all && \ @@ -1528,7 +1528,7 @@ pipeline { -D GEMM_LAYOUT="rcr;rrr;crr;ccr" \ -D GEMM_MULTI_D_DATATYPE="fp16" \ -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8" \ + -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ -DCMAKE_CXX_FLAGS=" -O3 " .. && \ ninja -j64 benchmark_gemm_all && \ @@ -1570,11 +1570,7 @@ pipeline { -DCMAKE_CXX_FLAGS=" -O3 " .. && \ ninja -j64 benchmark_gemm_all && \ python3 ../tile_engine/ops/gemm/gemm_benchmark.py . --problem-sizes "1024,1024,1024" \ - --warmup 5 --repeat 5 --verbose --json results.json && \ - ninja -j64 benchmark_gemm_fp16_rcr && \ - ninja -j64 benchmark_gemm_fp16_rrr && \ - ninja -j64 benchmark_gemm_fp16_crr && \ - ninja -j64 benchmark_gemm_fp16_ccr """ + --warmup 5 --repeat 5 --verbose --json results.json """ } steps{ buildHipClangJobAndReboot(setup_args:setup_args, no_reboot:true, build_type: 'Release', execute_cmd: execute_args) diff --git a/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json b/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json index c80210b963..a4f32a1907 100644 --- a/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json +++ b/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json @@ -1,6 +1,4 @@ { - "problem": { - }, "tile_config": { "tile_m": { "values": [ @@ -85,5 +83,7 @@ false ] } - } + }, + "k_block_per_cu": 1, + "permute_n": false } diff --git a/tile_engine/ops/gemm/CMakeLists.txt b/tile_engine/ops/gemm/CMakeLists.txt index 0e0ca41c9a..1eb49c0c7f 100644 --- a/tile_engine/ops/gemm/CMakeLists.txt +++ b/tile_engine/ops/gemm/CMakeLists.txt @@ -1,5 +1,5 @@ set(GEMM_DATATYPE "fp8;fp16" CACHE STRING "List of datatypes for GEMM (semicolon-separated)") -set(GEMM_LAYOUT "rcr" CACHE STRING "List of layout for GEMM (semicolon-separated)") +set(GEMM_LAYOUT "rcr;rrr;crr;ccr" CACHE STRING "List of layout for GEMM (semicolon-separated)") set(GEMM_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)") option(ENABLE_CCACHE_GEMM "Enable ccache for GEMM ops compilation" OFF) diff --git a/tile_engine/ops/gemm_preshuffle/CMakeLists.txt b/tile_engine/ops/gemm_preshuffle/CMakeLists.txt index 972ad9d0db..d80d2661d1 100644 --- a/tile_engine/ops/gemm_preshuffle/CMakeLists.txt +++ b/tile_engine/ops/gemm_preshuffle/CMakeLists.txt @@ -1,4 +1,4 @@ -set(GEMM_PRESHUFFLE_DATATYPE "fp8;fp16" CACHE STRING "List of datatypes for GEMM Preshuffle (semicolon-separated)") +set(GEMM_PRESHUFFLE_DATATYPE "fp16;fp8" CACHE STRING "List of datatypes for GEMM Preshuffle (semicolon-separated)") set(GEMM_PRESHUFFLE_LAYOUT "rcr" CACHE STRING "List of layout for GEMM Preshuffle (semicolon-separated)") set(GEMM_PRESHUFFLE_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)") option(ENABLE_CCACHE_GEMM_PRESHUFFLE "Enable ccache for GEMM Preshuffle ops compilation" OFF) @@ -65,7 +65,7 @@ function(create_individual_gemm_preshuffle_target datatype layout trait tile_con # Create the executable add_executable(${target_name} EXCLUDE_FROM_ALL - ${GEMM_PRESHUFFLE_SOURCE_DIR}/benchmark_gemm_preshuffle_single.cpp + ${GEMM_PRESHUFFLE_SOURCE_DIR}/gemm_preshuffle_benchmark_single.cpp ${instance_header} ) @@ -176,7 +176,7 @@ function(build_individual_gemm_preshuffle_targets datatype layout) OUTPUT_VARIABLE list_output ERROR_VARIABLE list_error ) - + if(NOT ret EQUAL 0) message(FATAL_ERROR "Failed to list kernels for ${datatype} ${layout}: ${list_error}") endif() @@ -273,10 +273,10 @@ else() endforeach() # Create trait-based collection targets - # These are common trait components used across all GEMM kernels - set(GEMM_PRESHUFFLE_PIPELINES "preshufflev1;preshufflev2") + # These are common trait components used across all GEMM Preshuffle kernels + set(GEMM_PRESHUFFLE_PIPELINES "preshufflev2") set(GEMM_PRESHUFFLE_EPILOGUES "default;cshuffle") - set(GEMM_PRESHUFFLE_SCHEDULERS "intrawave;interwave;default") + set(GEMM_PRESHUFFLE_SCHEDULERS "default") foreach(pipeline IN LISTS GEMM_PRESHUFFLE_PIPELINES) add_custom_target(benchmark_gemm_preshuffle_${pipeline}_pipeline) @@ -291,7 +291,6 @@ else() endforeach() # Build individual targets for each datatype/layout combination - foreach(dt IN LISTS GEMM_PRESHUFFLE_DATATYPE) foreach(l IN LISTS GEMM_PRESHUFFLE_LAYOUT) build_individual_gemm_preshuffle_targets(${dt} ${l}) diff --git a/tile_engine/ops/gemm_preshuffle/commons/validation_utils.py b/tile_engine/ops/gemm_preshuffle/commons/validation_utils.py index 454e26a7b5..b38ff5dffb 100644 --- a/tile_engine/ops/gemm_preshuffle/commons/validation_utils.py +++ b/tile_engine/ops/gemm_preshuffle/commons/validation_utils.py @@ -32,7 +32,6 @@ WARP_TILE_SUPPORTED_COMBINATIONS = { [16, 16, 16], [32, 32, 16], [16, 16, 32], - [4, 64, 16], [64, 4, 16], ], "bf16_bf16_bf16": [ @@ -40,7 +39,6 @@ WARP_TILE_SUPPORTED_COMBINATIONS = { [16, 16, 16], [32, 32, 16], [16, 16, 32], - [4, 64, 16], [64, 4, 16], ], "fp8_fp8_fp16": [[32, 32, 16], [32, 32, 32]], @@ -52,7 +50,6 @@ WARP_TILE_SUPPORTED_COMBINATIONS = { [16, 16, 16], [32, 32, 16], [16, 16, 32], - [4, 64, 16], [64, 4, 16], ], "bf16_bf16_bf16": [ @@ -60,7 +57,6 @@ WARP_TILE_SUPPORTED_COMBINATIONS = { [16, 16, 16], [32, 32, 16], [16, 16, 32], - [4, 64, 16], [64, 4, 16], ], "fp8_fp8_fp16": [[32, 32, 16], [32, 32, 32], [16, 16, 32], [16, 16, 64]], @@ -73,7 +69,6 @@ WARP_TILE_SUPPORTED_COMBINATIONS = { [16, 16, 16], [32, 32, 16], [16, 16, 32], - [4, 64, 16], [64, 4, 16], ], "bf16_bf16_bf16": [ @@ -81,7 +76,6 @@ WARP_TILE_SUPPORTED_COMBINATIONS = { [16, 16, 16], [32, 32, 16], [16, 16, 32], - [4, 64, 16], [64, 4, 16], ], "fp8_fp8_fp16": [ @@ -122,6 +116,12 @@ def element_size(data_type: str) -> float: def is_trait_combination_valid(pipeline: str, epilogue: str, scheduler: str) -> bool: """Check if a trait combination is valid.""" + if pipeline not in ["preshufflev2"]: + raise ValueError("Accepted pipeline values are: ['preshufflev2']") + if epilogue not in ["default", "cshuffle"]: + return ValueError("Accepted epilogue values are: ['default', 'cshuffle']") + if scheduler not in ["default"]: + return ValueError("Accepted scheduler values are: ['default']") return (pipeline, epilogue, scheduler) not in TRAIT_UNSUPPORTED_COMBINATIONS @@ -173,7 +173,7 @@ def validate_lds_capacity( matrix_b_size = (tile_n * tile_k) * element_size(b_datatype) total_tile_in_lds = matrix_a_size + matrix_b_size - max_tile_size = 2**15 if pipeline == "compv4" else 2**16 + max_tile_size = 2**15 if pipeline in ["preshufflev2", "compv4"] else 2**16 if total_tile_in_lds > max_tile_size: error_msg = ( @@ -266,6 +266,35 @@ def is_tile_config_valid( if warp_k * warp_tile_k > tile_k: return False + # Validate vector load alignment + m_iter_per_warp = tile_m / (warp_m * warp_tile_m) + vector_valid, vector_error = validate_vector_load_alignment( + warp_tile_m, + warp_tile_k, + a_datatype, + m_iter_per_warp, + wave_size=64, + vector_load_size=16, + ) + if not vector_valid: + logging.debug(f"Vector load alignment failed: {vector_error}") + return False + + # Validate M0, M1, M2 configuration for matrix A row-major layout + m0_m1_m2_valid, m0_m1_m2_error = validate_m0_m1_m2_configuration( + tile_m, + tile_k, + warp_m, + warp_n, + warp_k, + a_datatype, + vector_load_size=16, + warp_size=64, + ) + if not m0_m1_m2_valid: + logging.debug(f"M0/M1/M2 configuration validation failed: {m0_m1_m2_error}") + return False + # Validate warp configuration if not validate_warp_configuration(warp_m, warp_n, warp_k): logging.debug( @@ -318,12 +347,117 @@ def is_tile_config_valid( return True +def validate_vector_load_alignment( + wg_m: int, + wg_k: int, + a_datatype: str, + m_iter_per_warp: int, + wave_size: int, + vector_load_size: int, +) -> Tuple[bool, str]: + try: + # Calculate the memory access pattern size + a_element_size = element_size(a_datatype) + access_size = (wg_m * wg_k * a_element_size * m_iter_per_warp) / wave_size + + # Check if it's aligned to vector load size + if access_size % vector_load_size != 0: + error_msg = ( + f"Vector load alignment violation: " + f"({wg_m} * {wg_k} * {a_element_size} * {m_iter_per_warp} / {wave_size}) " + f"% {vector_load_size} = {access_size % vector_load_size} != 0. " + f"Access size: {access_size} bytes" + ) + return False, error_msg + + return True, "" + + except Exception as e: + return False, f"Error in vector load validation: {str(e)}" + + +def validate_m0_m1_m2_configuration( + tile_m: int, + tile_k: int, + warp_m: int, + warp_n: int, + warp_k: int, + a_datatype: str, + vector_load_size: int = 16, + warp_size: int = 64, +) -> Tuple[bool, str]: + """ + Validate M0, M1, M2 configuration for matrix A row-major layout. + This ensures proper memory access pattern alignment. + """ + try: + # Validation for A as row-major + MPerBlock = tile_m + + # Calculate K1 using element size + K1 = vector_load_size / element_size(a_datatype) + + # Check if K1 is valid (must be integer) + if K1 != int(K1): + return ( + False, + f"K1 = {K1} is not an integer. vector_load_size({vector_load_size}) must be divisible by element_size({a_datatype})", + ) + K1 = int(K1) + + # Calculate K0 + if tile_k % K1 != 0: + return False, f"tile_k({tile_k}) must be divisible by K1({K1})" + K0 = tile_k // K1 + + # Calculate M2 + if warp_size % K0 != 0: + return False, f"warp_size({warp_size}) must be divisible by K0({K0})" + M2 = warp_size // K0 + + # Calculate number of warps and block size + NumWarps = warp_m * warp_n * warp_k + BlockSize = NumWarps * warp_size + + # Calculate M0 (assuming get_warp_size() returns warp_size) + M0 = BlockSize // warp_size # This should equal NumWarps + + # Calculate M1 + if (M2 * M0) == 0: + return False, f"M2({M2}) * M0({M0}) cannot be zero" + + if MPerBlock % (M2 * M0) != 0: + return ( + False, + f"MPerBlock({MPerBlock}) must be divisible by M2({M2}) * M0({M0}) = {M2 * M0}", + ) + M1 = MPerBlock // (M2 * M0) + + # Validate the assertion: M0 * M1 * M2 == MPerBlock + calculated_m_per_block = M0 * M1 * M2 + if calculated_m_per_block != MPerBlock: + error_msg = ( + f"Incorrect M0, M1, M2 configuration! " + f"M0({M0}) * M1({M1}) * M2({M2}) = {calculated_m_per_block} != MPerBlock({MPerBlock}). " + f"Configuration: K0={K0}, K1={K1}, NumWarps={NumWarps}, BlockSize={BlockSize}" + ) + return False, error_msg + + return True, "" + + except ZeroDivisionError as e: + return False, f"Division by zero in M0/M1/M2 calculation: {str(e)}" + except Exception as e: + return False, f"Error in M0/M1/M2 validation: {str(e)}" + + # [TODO] Handle this while moving code to commons Add more datatype to this function if needed def get_dtype_string(datatype: str) -> str: """Get C++ type string for datatype""" dtype_map = { "fp16": "ck_tile::fp16_t", "fp8": "ck_tile::fp8_t", + "bf8": "ck_tile::bf8_t", "bf16": "ck_tile::bf16_t", "fp32": "float", "fp64": "double", diff --git a/tile_engine/ops/gemm_preshuffle/configs/default_config.json b/tile_engine/ops/gemm_preshuffle/configs/default_config.json index d4c3537c65..4606cf0c27 100644 --- a/tile_engine/ops/gemm_preshuffle/configs/default_config.json +++ b/tile_engine/ops/gemm_preshuffle/configs/default_config.json @@ -1,62 +1,72 @@ { "tile_config": { "tile_m": { - "values": [ - 128 - ] + "max": 256, + "min": 64, + "step": 64 }, "tile_n": { - "values": [ - 128 - ] + "max": 256, + "min": 64, + "step": 64 }, "tile_k": { - "values": [ - 128 - ] + "max": 256, + "min": 64, + "step": 64 }, "warp_m": { - "values": [ - 1 - ] + "values": [ + 4, + 2, + 1 + ] }, "warp_n": { - "values": [ - 4 - ] + "values": [ + 4, + 2, + 1 + ] }, "warp_k": { - "values": [ - 1 - ] + "values": [ + 1 + ] }, "warp_tile_m": { - "values": [ - 16 - ] + "values": [ + 4, + 16, + 32 + ] }, "warp_tile_n": { - "values": [ - 16 - ] + "values": [ + 16, + 32, + 64 + ] }, "warp_tile_k": { - "values": [ - 16,32 - ] + "values": [ + 8, + 16, + 32, + 64, + 128 + ] } }, "trait_config": { "pipeline": { "values": [ - "preshufflev1", "preshufflev2" ] }, "scheduler": { "values": [ - "interwave", - "intrawave" + "default" ] }, "epilogue": { @@ -81,11 +91,12 @@ ] }, "persistent": { - "values": [ - true, - false - ] + "values": [ + true, + false + ] } }, - "k_block_per_cu": 2 + "k_block_per_cu": 1, + "permute_n": true } \ No newline at end of file diff --git a/tile_engine/ops/gemm_preshuffle/configs/user_provided_config.json b/tile_engine/ops/gemm_preshuffle/configs/user_provided_config.json index c0fc1f6cf8..cf7c79462e 100644 --- a/tile_engine/ops/gemm_preshuffle/configs/user_provided_config.json +++ b/tile_engine/ops/gemm_preshuffle/configs/user_provided_config.json @@ -2,27 +2,27 @@ "tile_config": { "tile_m": { "values": [ - 128 + 64 ] }, "tile_n": { "values": [ - 128 + 64 ] }, "tile_k": { "values": [ - 64 + 192 ] }, "warp_m": { "values": [ - 1 + 2 ] }, "warp_n": { "values": [ - 4 + 2 ] }, "warp_k": { @@ -42,7 +42,7 @@ }, "warp_tile_k": { "values": [ - 16,32 + 32 ] } }, @@ -54,12 +54,13 @@ }, "scheduler": { "values": [ - "intrawave" + "default" ] }, "epilogue": { "values": [ - "default" + "default", + "cshuffle" ] }, "pad_m": { @@ -79,9 +80,10 @@ }, "persistent": { "values": [ - false + true ] } }, - "k_block_per_cu": 8 + "k_block_per_cu": 1, + "permute_n": false } \ No newline at end of file diff --git a/tile_engine/ops/gemm_preshuffle/benchmark_gemm_preshuffle.hpp b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_benchmark.hpp similarity index 97% rename from tile_engine/ops/gemm_preshuffle/benchmark_gemm_preshuffle.hpp rename to tile_engine/ops/gemm_preshuffle/gemm_preshuffle_benchmark.hpp index 74fccf6bf2..77a9f26527 100644 --- a/tile_engine/ops/gemm_preshuffle/benchmark_gemm_preshuffle.hpp +++ b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_benchmark.hpp @@ -23,6 +23,14 @@ inline constexpr auto get_metric_name(Metric m) } } +struct KernelConfig +{ + std::tuple tile_dims; + std::tuple warp_dims; + std::tuple warp_tile_dims; + bool permuteN; +}; + struct GemmProblem { int split_k_; diff --git a/tile_engine/ops/gemm_preshuffle/benchmark_gemm_preshuffle_single.cpp b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_benchmark_single.cpp similarity index 89% rename from tile_engine/ops/gemm_preshuffle/benchmark_gemm_preshuffle_single.cpp rename to tile_engine/ops/gemm_preshuffle/gemm_preshuffle_benchmark_single.cpp index 152e27e77e..1f03d1cf9b 100644 --- a/tile_engine/ops/gemm_preshuffle/benchmark_gemm_preshuffle_single.cpp +++ b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_benchmark_single.cpp @@ -75,7 +75,7 @@ inline auto create_args(int argc, char* argv[]) return std::make_tuple(result, arg_parser); } -void benchmark_gemm_preshuffle_single(const ck_tile::ArgParser& arg_parser) +void benchmark_single(const ck_tile::ArgParser& arg_parser) { // Use DataTypeTraits to get the actual type names from the generated header // The generated header defines ADataType, BDataType, AccDataType, CDataType @@ -124,9 +124,16 @@ void benchmark_gemm_preshuffle_single(const ck_tile::ArgParser& arg_parser) try { // Create a lambda that wraps the kernel launch - std::tuple warp_tile_dims = std::make_tuple( SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK); + std::tuple tile_dims = + std::make_tuple(SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK); + std::tuple warp_dims = std::make_tuple(SelectedKernel::WarpPerBlock_M, + SelectedKernel::WarpPerBlock_N, + SelectedKernel::WarpPerBlock_K); + bool permuteN = SelectedKernel::PermuteN; + + KernelConfig config{tile_dims, warp_dims, warp_tile_dims, permuteN}; auto kernel_func = [](const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& stream) { @@ -134,7 +141,7 @@ void benchmark_gemm_preshuffle_single(const ck_tile::ArgParser& arg_parser) }; // Benchmark the kernel - profiler.benchmark(gemm_problem, kernel_func, warp_tile_dims); + profiler.benchmark(gemm_problem, kernel_func, config); // Select best instance based on metric profiler.select_best_instance(static_cast(arg_parser.get_int("metric"))); @@ -153,7 +160,7 @@ int main(int argc, char* argv[]) if(!result) return EXIT_FAILURE; - benchmark_gemm_preshuffle_single(parser); + benchmark_single(parser); return 0; } catch(const std::exception& e) diff --git a/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_common.hpp b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_common.hpp index 4fb98dc3c2..09ec895ab5 100644 --- a/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_common.hpp +++ b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_common.hpp @@ -75,58 +75,6 @@ constexpr auto is_row_major(Layout) return ck_tile::bool_constant>{}; } -// // Permutation function for pk_int4_t -// template -// void permute_vectors_i4x4_b(Tensor& tensor) -// { -// const ck_tile::index_t K = tensor.get_length(0); -// const ck_tile::index_t N = tensor.get_length(1); -// // vector pk_i4x4 permute -// for(int i = 0; i < N; i++) -// { -// for(int j = 0; j < K; j += 8) -// { -// int8_t input[8]; - -// for(int k = 0; k < 4; k++) -// { -// int8_t i4x2 = tensor(j + k * 2, i).data; -// input[k * 2 + 0] = (i4x2 >> 4) & 0xf; -// input[k * 2 + 1] = (i4x2 >> 0) & 0xf; -// } - -// // permute 01234567->20643175 -// { -// int8_t hi = input[2]; -// int8_t lo = input[0]; -// int8_t i4x2 = (hi << 4) | lo; -// tensor(j + 0, i) = i4x2; -// } - -// { -// int8_t hi = input[6]; -// int8_t lo = input[4]; -// int8_t i4x2 = (hi << 4) | lo; -// tensor(j + 2, i) = i4x2; -// } - -// { -// int8_t hi = input[3]; -// int8_t lo = input[1]; -// int8_t i4x2 = (hi << 4) | lo; -// tensor(j + 4, i) = i4x2; -// } - -// { -// int8_t hi = input[7]; -// int8_t lo = input[5]; -// int8_t i4x2 = (hi << 4) | lo; -// tensor(j + 6, i) = i4x2; -// } -// } -// } -// } - // Structure to hold kernel traits for dispatcher struct KernelTraits { @@ -211,3 +159,27 @@ auto shuffle_b(const ck_tile::HostTensor& t, std::copy(t.begin(), t.end(), t_view.begin()); return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); } + +template +auto shuffle_b_permuteN(const ck_tile::HostTensor& t, + ck_tile::index_t N_Warp_Tile, + ck_tile::index_t K_Warp_Tile, + ck_tile::index_t N_Tile, + ck_tile::index_t N_Warp) +{ + assert(t.get_lengths().size() == 2); + + int n_ = t.get_lengths()[1]; + int k_ = t.get_lengths()[0]; + int divisor = N_Warp_Tile == 32 ? 2 : 4; + int NRepeat = N_Tile / N_Warp_Tile / N_Warp; + ck_tile::HostTensor t_view({n_ / N_Tile, + N_Warp, + N_Warp_Tile, + NRepeat, + k_ / K_Warp_Tile, + divisor, + K_Warp_Tile / divisor}); + std::copy(t.begin(), t.end(), t_view.begin()); + return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 5, 2, 6}); +} diff --git a/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_instance_builder.py b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_instance_builder.py index e6e075cb36..1d4b027716 100644 --- a/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_instance_builder.py +++ b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_instance_builder.py @@ -95,67 +95,87 @@ class GemmPreshuffleKernelBuilder: def _get_tile_configs(self, fast_mode=False): """Get tile configurations for the current datatype and layout""" - if "tile_configs" in self.config: - # Old format - return ( - self.config["tile_configs"].get(self.datatype, {}).get(self.layout, []) + + tile_config = self.config["tile_config"] + + # Generate values in the config if default range is given + if tile_config.get("tile_m").get("values") is None: + tile_config.get("tile_m")["values"] = self._generate_values( + tile_config.get("tile_m").get("min"), + tile_config.get("tile_m").get("max"), + tile_config.get("tile_m").get("step"), + ) + if tile_config.get("tile_n").get("values") is None: + tile_config.get("tile_n")["values"] = self._generate_values( + tile_config.get("tile_n").get("min"), + tile_config.get("tile_n").get("max"), + tile_config.get("tile_n").get("step"), + ) + if tile_config.get("tile_k").get("values") is None: + tile_config.get("tile_k")["values"] = self._generate_values( + tile_config.get("tile_k").get("min"), + tile_config.get("tile_k").get("max"), + tile_config.get("tile_k").get("step"), ) - elif "tile_config" in self.config: - # New format - generate combinations from individual parameter values - tile_config = self.config["tile_config"] - # Get all possible values for each parameter - tile_m_values = tile_config.get("tile_m", {}).get("values", [256]) - tile_n_values = tile_config.get("tile_n", {}).get("values", [256]) - tile_k_values = tile_config.get("tile_k", {}).get("values", [32]) - warp_m_values = tile_config.get("warp_m", {}).get("values", [2]) - warp_n_values = tile_config.get("warp_n", {}).get("values", [2]) - warp_k_values = tile_config.get("warp_k", {}).get("values", [1]) - warp_tile_m_values = tile_config.get("warp_tile_m", {}).get("values", [32]) - warp_tile_n_values = tile_config.get("warp_tile_n", {}).get("values", [32]) - warp_tile_k_values = tile_config.get("warp_tile_k", {}).get("values", [32]) + # Get all possible values for each parameter + tile_m_values = tile_config.get("tile_m").get("values") + tile_n_values = tile_config.get("tile_n").get("values") + tile_k_values = tile_config.get("tile_k").get("values") + warp_m_values = tile_config.get("warp_m").get("values") + warp_n_values = tile_config.get("warp_n").get("values") + warp_k_values = tile_config.get("warp_k").get("values") + warp_tile_m_values = tile_config.get("warp_tile_m").get("values") + warp_tile_n_values = tile_config.get("warp_tile_n").get("values") + warp_tile_k_values = tile_config.get("warp_tile_k").get("values") - # Generate all combinations - configs = [] - for tile_m in tile_m_values: - for tile_n in tile_n_values: - for tile_k in tile_k_values: - for warp_m in warp_m_values: - for warp_n in warp_n_values: - for warp_k in warp_k_values: - for warp_tile_m in warp_tile_m_values: - for warp_tile_n in warp_tile_n_values: - for warp_tile_k in warp_tile_k_values: - # Validate configuration - if self._validate_tile_config( - tile_m, - tile_n, - tile_k, - warp_m, - warp_n, - warp_k, - warp_tile_m, - warp_tile_n, - warp_tile_k, - fast_mode=fast_mode, - ): - configs.append( - { - "tile_m": tile_m, - "tile_n": tile_n, - "tile_k": tile_k, - "warp_m": warp_m, - "warp_n": warp_n, - "warp_k": warp_k, - "warp_tile_m": warp_tile_m, - "warp_tile_n": warp_tile_n, - "warp_tile_k": warp_tile_k, - } - ) - return configs - else: - # Fallback to default - return [] + # Generate all combinations + configs = [] + for tile_m in tile_m_values: + for tile_n in tile_n_values: + for tile_k in tile_k_values: + for warp_m in warp_m_values: + for warp_n in warp_n_values: + for warp_k in warp_k_values: + for warp_tile_m in warp_tile_m_values: + for warp_tile_n in warp_tile_n_values: + for warp_tile_k in warp_tile_k_values: + # Validate configuration + if self._validate_tile_config( + tile_m, + tile_n, + tile_k, + warp_m, + warp_n, + warp_k, + warp_tile_m, + warp_tile_n, + warp_tile_k, + fast_mode=fast_mode, + ): + configs.append( + { + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "warp_m": warp_m, + "warp_n": warp_n, + "warp_k": warp_k, + "warp_tile_m": warp_tile_m, + "warp_tile_n": warp_tile_n, + "warp_tile_k": warp_tile_k, + } + ) + return configs + + def _generate_values(self, min_val, max_val, step): + """Generate a list of values from min to max with the given step""" + values = [] + val = min_val + while val <= max_val: + values.append(val) + val += step + return values def _generate_trait_combinations(self): """Generate all combinations of traits""" @@ -270,6 +290,12 @@ class GemmPreshuffleKernelBuilder: return True else: + # Validate preshuffle specific constraints + if self.config.get("permute_n"): + valid = (tile_n / warp_tile_n / warp_n) % 2 == 0 + if not valid: + return False + # Full validation for generation # Determine data types for validation a_datatype = self.datatype @@ -299,7 +325,7 @@ class GemmPreshuffleKernelBuilder: ) def _generate_kernel_instance( - self, tile_config, trait_combo, k_block_per_cu, is_header=True + self, tile_config, trait_combo, k_block_per_cu, permute_n, is_header=True ): """Generate a single kernel instance""" ( @@ -349,9 +375,9 @@ class GemmPreshuffleKernelBuilder: acc_type = "float" # Determine output type - c_type = get_dtype_string(self.datatype) + c_type = self.datatype if self.datatype in ["fp8", "bf8"]: - c_type = "ck_tile::fp16_t" + c_type = "fp16" # Determine layouts based on self.layout a_layout, b_layout, c_layout = get_abc_layouts(self.layout) @@ -374,7 +400,7 @@ class GemmPreshuffleKernelBuilder: using ADataType = {get_dtype_string(self.datatype)}; using BDataType = {get_dtype_string(self.datatype)}; using AccDataType = {acc_type}; -using CDataType = {c_type}; +using CDataType = {get_dtype_string(c_type)}; using ALayout = {a_layout}; using BLayout = {b_layout}; @@ -408,6 +434,8 @@ struct SelectedKernel {{ static constexpr bool Preshuffle = true; static constexpr ck_tile::index_t NumWaveGroups = 1; + static constexpr bool PermuteN = {"true" if permute_n else "false"}; + // Tile shape using TileShape = ck_tile::TileGemmShape< ck_tile::sequence, @@ -485,7 +513,10 @@ struct SelectedKernel {{ WarpTileK, // KPerXdl_ TransposeC, // isCTransposed_ memory_operation, // MemoryOperation_ - NumWaveGroups>; // kNumWaveGroups_ + NumWaveGroups, // kNumWaveGroups_ + false, // FixedVectorSize_ + 1, // VectorSizeC_ + PermuteN>; // isPermuteN_ using GemmEpilogue = ck_tile::CShuffleEpilogue; """ @@ -580,6 +611,7 @@ struct SelectedKernel {{ tile_configs = self._get_tile_configs() trait_combos = self._generate_trait_combinations() k_block_per_cu = self.config.get("k_block_per_cu") + permute_n = self.config.get("permute_n") # Prepare work items for parallel processing work_items = [] @@ -590,6 +622,7 @@ struct SelectedKernel {{ tile_config, trait_combo, k_block_per_cu, + permute_n, self.working_path, self.datatype, self.layout, @@ -681,21 +714,29 @@ struct SelectedKernel {{ def _generate_single_kernel_individual(work_item): """Worker function to generate a single individual kernel file""" - tile_config, trait_combo, k_block_per_cu, working_path, datatype, layout = work_item + ( + tile_config, + trait_combo, + k_block_per_cu, + permute_n, + working_path, + datatype, + layout, + ) = work_item # Create a temporary builder instance for this worker builder = GemmPreshuffleKernelBuilder(working_path, datatype, layout) try: kernel_name, instance_code = builder._generate_kernel_instance( - tile_config, trait_combo, k_block_per_cu + tile_config, trait_combo, k_block_per_cu, permute_n ) - # Create simplified filename without the "gemm_" prefix - # Remove "gemm_" from the beginning of kernel_name for the filename + # Create simplified filename without the "gemm_preshuffle_" prefix + # Remove "gemm_preshuffle_" from the beginning of kernel_name for the filename simplified_name = kernel_name - if simplified_name.startswith("gemm_"): - simplified_name = simplified_name[5:] # Remove "gemm_" prefix + if simplified_name.startswith("gemm_preshuffle_"): + simplified_name = simplified_name[16:] # Remove "gemm_preshuffle_" prefix # Write individual header file header_file = working_path / f"gemm_single_{simplified_name}.hpp" @@ -727,7 +768,7 @@ def main(): parser.add_argument( "--layout", required=True, - choices=["rcr", "rrr", "ccr", "crr"], + choices=["rcr"], help="Matrix layout", ) parser.add_argument("--config_json", required=True, help="Configuration JSON file") @@ -735,7 +776,9 @@ def main(): "--num_workers", type=int, help="Number of parallel workers (default: auto)" ) parser.add_argument( - "--gen_individual", action="store_true", help="Generate individual kernel files" + "--gen_all_individual", + action="store_true", + help="Generate individual kernel files", ) parser.add_argument( "--gen_single", action="store_true", help="Generate a single kernel file" @@ -763,7 +806,7 @@ def main(): assert len(layout_parts) == 3, ( f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr' where r stands for row major and c stands for column major)" ) - assert layout_parts[0] == "r" and layout_parts[1] == "c", ( + assert layout_parts[0] in ["r"] and layout_parts[1] in ["c"], ( f"Invalid matrix_a layout : {layout_parts[0]} or matrix_b layout: {layout_parts[1]} (matrix_a must be 'r' for row major and matrix_b must be 'c' for column major as it is the only supported layout for preshuffle)" ) assert layout_parts[2] == "r", ( @@ -816,10 +859,11 @@ def main(): ) k_block_per_cu = builder.config.get("k_block_per_cu") + permute_n = builder.config.get("permute_n") # Generate the kernel kernel_name, instance_code = builder._generate_kernel_instance( - tile_config, trait_combo, k_block_per_cu + tile_config, trait_combo, k_block_per_cu, permute_n ) # Write the file @@ -835,13 +879,13 @@ def main(): print(f"Generated {header_file}") - elif args.gen_individual: + elif args.gen_all_individual: # Generate all individual kernel files builder.run(args.num_workers) pass else: parser.error( - "Must specify one of: --list_kernels, --gen_individual, or --gen_single" + "Must specify one of: --list_kernels, --gen_all_individual, or --gen_single" ) diff --git a/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_profiler.hpp b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_profiler.hpp index 4f2a929ba0..7d212c934c 100644 --- a/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_profiler.hpp +++ b/tile_engine/ops/gemm_preshuffle/gemm_preshuffle_profiler.hpp @@ -2,7 +2,7 @@ #include "ck_tile/host/device_prop.hpp" #include "ck_tile/ops/gemm.hpp" -#include "benchmark_gemm_preshuffle.hpp" +#include "gemm_preshuffle_benchmark.hpp" class GemmProfiler { @@ -17,7 +17,7 @@ class GemmProfiler void benchmark(GemmProblem& gemm_problem, std::function kernel_func, - const std::tuple& warp_tile_dims) + KernelConfig& config) { // Create a vector with a single callable that returns both name and time std::vector(ck_tile::GemmHostArgs&, @@ -30,13 +30,13 @@ class GemmProfiler return std::make_tuple(std::string(KERNEL_NAME), time); }); - benchmark(gemm_problem, callables, warp_tile_dims); + benchmark(gemm_problem, callables, config); } void benchmark(GemmProblem& gemm_problem, std::vector( ck_tile::GemmHostArgs&, const ck_tile::stream_config&)>>& callables, - const std::tuple& warp_tile_dims) + KernelConfig& config) { const ALayout layout_a = ALayout{}; const BLayout layout_b = BLayout{}; @@ -110,11 +110,22 @@ class GemmProfiler for(const auto& callable : callables) { - ck_tile::index_t N_Warp_Tile = std::get<1>(warp_tile_dims); - ck_tile::index_t K_Warp_Tile = std::get<2>(warp_tile_dims); + ck_tile::index_t N_Warp_Tile = std::get<1>(config.warp_tile_dims); + ck_tile::index_t K_Warp_Tile = std::get<2>(config.warp_tile_dims); + ck_tile::index_t N_Tile = std::get<1>(config.tile_dims); + ck_tile::index_t N_Warp = std::get<1>(config.warp_dims); + + ck_tile::HostTensor b_shuffle_host = [&]() { + if(config.permuteN) + { + return shuffle_b_permuteN(b_k_n, N_Warp_Tile, K_Warp_Tile, N_Tile, N_Warp); + } + else + { + return shuffle_b(b_k_n, N_Warp_Tile, K_Warp_Tile); + } + }(); - ck_tile::HostTensor b_shuffle_host = - shuffle_b(b_k_n, N_Warp_Tile, K_Warp_Tile); b_k_n_dev_buf.ToDevice(b_shuffle_host.data()); ck_tile::GemmHostArgs gemm_args = { From a1ce64374f91a0a8982b61a3bf2594fec3c75d60 Mon Sep 17 00:00:00 2001 From: andrew clark Date: Mon, 27 Oct 2025 08:24:36 -0600 Subject: [PATCH 02/17] Jenkins Alerts Notifications (#3086) * Testing minimal pipeline * Update Jenkinsfile * Testing webhook * Testing webhook * Testing webhook * Testing build log output * Testing log retrieval * Testing * Testing pattern matching * Fixing regex * Testing error detection * Testing log formatting Including additional context around log failure. * Testing notification message format * Update Jenkinsfile * Notification formatting * Testing secure interpolation * Testing string interpolation * Notification format * Fixing markdown * Testing markdown * Testing markdown * Revert "Testing markdown" This reverts commit adeb6d2d55d2e2a4b1817c3e82b1f0095bd25aeb. * Testing different markdown format * Revert "Testing different markdown format" This reverts commit bf5406a1cd25f55208a08211472a758a018a9d2a. * Testing markdown * Testing markdown * Testing markdown * Testing markdown * Testing markdown * Testing notification * Testing notification * Testing notification * Testing failure mode * Testing failure mode * Adding new patterns and tests * Commenting * Stage name fix * Moving to notification on failure only * Fixing notification format * Testing env vars * Testing build url redirect * Testing no log errors * Testing no errors case * Integrating into primary jenkinsfile * Updating notification message Removed emoji from message --- Jenkinsfile | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/Jenkinsfile b/Jenkinsfile index b89d6fb657..9acbbeeca2 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -12,6 +12,14 @@ def show_node_info() { """ } +// Error patterns to scan build logs for specific failure types and send detailed notifications. +def failurePatterns = [ + [pattern: /login attempt to .* failed with status: 401 Unauthorized/, description: "Docker registry authentication failed"], + [pattern: /docker login failed/, description: "Docker login failed"], + [pattern: /HTTP request sent .* 404 Not Found/, description: "HTTP request failed with 404"], + [pattern: /cat: .* No such file or directory/, description: "GPU not found"], +] + class Version { int major, minor, patch @Override @@ -1849,4 +1857,36 @@ pipeline { } } } + post { + failure { + node(rocmnode("nogpu")) { + script { + // Get the build log. + def buildLog = sh(script: 'wget -q --no-check-certificate -O - ' + BUILD_URL + 'consoleText', returnStdout: true) + // Check for patterns in the log. + def foundPatterns = [] + for (patternMap in failurePatterns) { + def result = checkForPattern(patternMap.pattern, buildLog) + if (result.found) { + foundPatterns.add([ + description: patternMap.description, + matchedLine: result.matchedLine, + context: result.context + ]) + } + } + // Send a notification for each matched failure pattern. + for (patternMap in foundPatterns) { + withCredentials([string(credentialsId: 'ck_ci_errors_webhook_url', variable: 'WEBHOOK_URL')]) { + sh ''' + curl -X POST "${WEBHOOK_URL}" \ + -H 'Content-Type: application/json' \ + -d '{"text": "\\n\\n**Build Failed**\\n\\n**Issues detected:** ''' + patternMap.description + '''\\n\\n**Log context:**\\n```\\n''' + patternMap.context.replace("'", "\\'") + '''\\n```\\n\\n**Job:** ''' + env.JOB_NAME + '''\\n\\n**Build:** #''' + env.BUILD_NUMBER + '''\\n\\n**URL:** ''' + env.RUN_DISPLAY_URL + '''"}' + ''' + } + } + } + } + } + } } From 06973b1cf4987b5f2e7fc1fe504b56df58edaf1f Mon Sep 17 00:00:00 2001 From: Enrico Degregori <73224202+EnricoDeg@users.noreply.github.com> Date: Mon, 27 Oct 2025 16:09:02 +0100 Subject: [PATCH 03/17] Fix multi-abd tests bug (#3099) --- .../profiler/profile_gemm_multi_abd_impl.hpp | 80 +++++++------------ 1 file changed, 28 insertions(+), 52 deletions(-) diff --git a/profiler/include/profiler/profile_gemm_multi_abd_impl.hpp b/profiler/include/profiler/profile_gemm_multi_abd_impl.hpp index 46745fd02b..51922fde33 100644 --- a/profiler/include/profiler/profile_gemm_multi_abd_impl.hpp +++ b/profiler/include/profiler/profile_gemm_multi_abd_impl.hpp @@ -188,66 +188,42 @@ bool profile_gemm_multi_abd_impl(int do_verification, EDataType, remove_cvref_t>>::type; - auto get_a_matrix = [&]() -> auto { - // in case of pass through we avoid allocating a new - // tensor and copying values - if constexpr(is_same_v) + Tensor a_m_k({M, K}); + for(int m = 0; m < M; ++m) + { + for(int k = 0; k < K; ++k) { - return as_m_k(Number<0>{}); + // result + auto data_refs1 = ck::tie(a_m_k(m, k)); + // inputs + auto data_refs2 = + generate_tie([&](auto i) -> auto& { return as_m_k(Number{})(m, k); }, + Number{}); + auto data_refs = concat_tuple_of_refs(data_refs1, data_refs2); + unpack(a_element_op, data_refs); } - else - { - Tensor a_m_k({M, K}); - for(int m = 0; m < M; ++m) - { - for(int k = 0; k < K; ++k) - { - // result - auto data_refs1 = ck::tie(a_m_k(m, k)); - // inputs - auto data_refs2 = - generate_tie([&](auto i) -> auto& { return as_m_k(Number{})(m, k); }, - Number{}); - auto data_refs = concat_tuple_of_refs(data_refs1, data_refs2); - unpack(a_element_op, data_refs); - } - } - return a_m_k; - } - }; + } using BComputeType = typename std::conditional<(NumBTensor > 1), EDataType, remove_cvref_t>>::type; - auto get_b_matrix = [&]() -> auto { - // in case of pass through we avoid allocating a new - // tensor and copying values - if constexpr(is_same_v) + Tensor b_k_n({K, N}); + for(int k = 0; k < K; ++k) + { + for(int n = 0; n < N; ++n) { - return bs_k_n(Number<0>{}); + // result + auto data_refs1 = ck::tie(b_k_n(k, n)); + // inputs + auto data_refs2 = + generate_tie([&](auto i) -> auto& { return bs_k_n(Number{})(k, n); }, + Number{}); + auto data_refs = concat_tuple_of_refs(data_refs1, data_refs2); + unpack(b_element_op, data_refs); } - else - { - Tensor b_k_n({K, N}); - for(int k = 0; k < K; ++k) - { - for(int n = 0; n < N; ++n) - { - // result - auto data_refs1 = ck::tie(b_k_n(k, n)); - // inputs - auto data_refs2 = - generate_tie([&](auto i) -> auto& { return bs_k_n(Number{})(k, n); }, - Number{}); - auto data_refs = concat_tuple_of_refs(data_refs1, data_refs2); - unpack(b_element_op, data_refs); - } - } - return b_k_n; - } - }; + } using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm Date: Mon, 27 Oct 2025 08:29:15 -0700 Subject: [PATCH 04/17] Add .cline* files to .gitignore (#3101) Developers who use cline on the code base need to ignore .cline* directories like .cline_storage and .clinerules. Using a wildcard to ignore any other cline-related directories. --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index bcc5888b7f..6641e5bc58 100644 --- a/.gitignore +++ b/.gitignore @@ -36,6 +36,9 @@ tags # Editors .vscode +# Cline +.cline* + # build-in-source directory (see exceptions below) build* From 054fdb765cd74c0f7bbb6561ea58713df82ed85f Mon Sep 17 00:00:00 2001 From: arai713 <67439843+arai713@users.noreply.github.com> Date: Mon, 27 Oct 2025 09:14:17 -0700 Subject: [PATCH 05/17] [CK_TILE] Stream-K operator() Reboot (#3064) * Persistent Stream-K Kernel Implementation This change implements an operator() function in the reboot::StreamKKernel class that is enabled when the Persistent flag is set to true. In this case, the data-parallel portion and the Stream-K portion of the kernel are fully persistent. The changes were made in the reboot namespace. A future PR will remove the old Stream-K kernel class and remove the reboot namespace. * Unit Tests for Persistent Stream-K Kernel This change contains the inital test suite for the Persitent Stream-K Kernel. The files contain "reboot" in the name; a future PR will remove tests for the old Stream-K Kernel and remove the "reboot" naming. A future commit will add tests for the non-persistent kernel. Also added estimate_num_wgs_per_tile to the StreamKTilePartitionerBase class. This allows us to estimate the number of accumulations done per macro tile in C to use during validation when computing relative and absolute tolerance. * Adding implementation for the Non-Persistent Stream-K kernel This code is adding the operator() function for the Non-Persistent Stream-K kernel. Persistency of the kernel is determined through a template argument. The Non-Persistent kernel will allocate additional workgroups for the data parallel section, leading to a different structure for processing the data parallel and Stream-K sections. There has been an addition to the TilePartitioner to get access to the whether Persistent has been set to true or false in the StreamKKernel. * Adding in the tests for the Non-Persistent Stream-K kernel * Refactor Stream-K Reboot Unit Tests This commit makes the following changes: - Update test cases to determine M, N, and K based on the number of CUs. This ensures that each test case is one of Edge Case, SK Only, DP Only, or DP + 2 Tile SK regardless of the architecture. - Since the DP + 2 Tile SK test case takes long to run, this change moves this case into a separate .inc file and labels it as an extended test. - Since the extended test takes > 30 seconds to run, this test is added to the list of regression tests. * Fix spelling errors in comments for test cases Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Changes based on review Removed const volatile for typenames Set up alias for is_tuple_t Naming changes for clarity: GemmCommon -> BaseGemm Moved std::enable_if_t out of template parameters and changed to a return type for operator() Added constructor for StreamKKernelArgs to clarify UniversalGemm inheritance --------- Co-authored-by: Emily Martins Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../ops/gemm/kernel/streamk_gemm_kernel.hpp | 472 ++++++++++++++++++ .../kernel/streamk_gemm_tile_partitioner.hpp | 7 + .../streamk_gemm_tile_partitioner_impl.hpp | 21 + test/CMakeLists.txt | 1 + test/ck_tile/gemm_streamk/CMakeLists.txt | 12 + ...gemm_streamk_reboot_bf16_nonpersistent.cpp | 19 + ...st_gemm_streamk_reboot_bf16_persistent.cpp | 19 + ...gemm_streamk_reboot_fp16_nonpersistent.cpp | 19 + ...st_gemm_streamk_reboot_fp16_persistent.cpp | 19 + ...gemm_streamk_reboot_bf16_nonpersistent.cpp | 19 + ...st_gemm_streamk_reboot_bf16_persistent.cpp | 19 + ...gemm_streamk_reboot_fp16_nonpersistent.cpp | 19 + ...st_gemm_streamk_reboot_fp16_persistent.cpp | 19 + ...est_gemm_streamk_reboot_extended_cases.inc | 24 + .../test_gemm_streamk_reboot_smoke_cases.inc | 47 ++ .../test_gemm_streamk_reboot_types.hpp | 56 +++ .../test_gemm_streamk_reboot_util.cpp | 10 + .../test_gemm_streamk_reboot_util.hpp | 283 +++++++++++ .../test_streamk_tile_partitioner.cpp | 20 + .../test_streamk_tile_partitioner_common.hpp | 17 + 20 files changed, 1122 insertions(+) create mode 100644 test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp create mode 100644 test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_persistent.cpp create mode 100644 test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp create mode 100644 test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_persistent.cpp create mode 100644 test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp create mode 100644 test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_persistent.cpp create mode 100644 test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp create mode 100644 test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_persistent.cpp create mode 100644 test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_extended_cases.inc create mode 100644 test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_smoke_cases.inc create mode 100644 test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_types.hpp create mode 100644 test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.cpp create mode 100644 test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.hpp diff --git a/include/ck_tile/ops/gemm/kernel/streamk_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/streamk_gemm_kernel.hpp index 58bce4795f..915aebd1e6 100644 --- a/include/ck_tile/ops/gemm/kernel/streamk_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/streamk_gemm_kernel.hpp @@ -8,6 +8,478 @@ #include "ck_tile/host/concat.hpp" namespace ck_tile { +namespace reboot { + +/// @brief The Stream K GEMM kernel host arguments. +/// +/// @par Overview +/// This structure is passed to @ref StreamKKernel "StreamKKernel" when creating the kernel +/// arguments object. It contains all necessary information required to build proper kernel +/// arguments and launch the kernel on GPU. This structure defines the GEMM problem +/// configuration by stating all required information like M,N,K sizes and respective strides. +struct StreamKHostArgs : public ck_tile::UniversalGemmHostArgs<> +{ + CK_TILE_HOST explicit StreamKHostArgs(const void* a_ptr_, + const void* b_ptr_, + void* c_ptr_, + index_t M_, + index_t N_, + index_t K_, + index_t stride_A_, + index_t stride_B_, + index_t stride_C_, + StreamKReductionStrategy reduction_strategy_) + : UniversalGemmHostArgs<>({a_ptr_}, + {b_ptr_}, + {/*ds_ptr*/}, + c_ptr_, + /*k_batch_ =*/1, + M_, + N_, + K_, + {stride_A_}, + {stride_B_}, + {/*stride_Ds_*/}, + stride_C_), + reduction_strategy{reduction_strategy_} + { + } + + ck_tile::StreamKReductionStrategy reduction_strategy; +}; + +/// @brief The Stream K GEMM kernel class. +/// +/// @par Overview +/// This class is responsible for the Stream-K kernel, making use of UniversalGemm. +// The main kernel functions are the operator() functions. There is one for Persistent +// and one for Non-Persistent data parallel sections of the Stream-K algorithm. +// +// Both the Non-Persistent and Persistent kernels make use of `BaseGemm()` and +// `StreamKGemm()`. `BaseGemm()` computes offsets into the A,B,C tensors, then calls +// `RunGemm()` which runs the GEMM pipeline and epilogue. `StreamKGemm()` performs the +// main Stream-K algorithm. Each iteration of the Stream-K loop calls `BaseGemm()`. +template +struct StreamKKernel +{ + /// @brief Inject the UniversalGemmKernel base class to support execution of all necessary + /// functions. + using UniversalGemmKernel = + UniversalGemmKernel; + + static constexpr index_t kBlockSize = UniversalGemmKernel::kBlockSize; + static constexpr bool PersistentDP = UniversalGemmKernel::PersistentKernel; + + using TilePartitioner = TilePartitioner_; + using GemmPipeline = GemmPipeline_; + using EpiloguePipeline = EpiloguePipeline_; + + static_assert( + TilePartitioner::PERSISTENT == PersistentDP, + "Persistent flag from TilePartitioner must match Persistent flag from UniversalGemm."); + + /// @brief Specify the layout configurations for A, B, and C + using ALayout = typename GemmPipeline::ALayout; + using BLayout = typename GemmPipeline::BLayout; + using CLayout = typename GemmPipeline::CLayout; + + /// @brief Specify the data type configurations for A, B, and C + using ADataType = typename GemmPipeline::ADataType; + using BDataType = typename GemmPipeline::BDataType; + using CDataType = typename EpiloguePipeline::ODataType; + + template + static constexpr bool is_tuple_v = is_detected::value; + + /// @brief ALayout and ADataType are expected to be scalars, not a tuple. + static_assert(!is_tuple_v && !is_tuple_v, + "ALayout and ADataType must be scalars."); + + /// @brief BLayout and BDataType are expected to be scalars, not a tuple. + static_assert(!is_tuple_v && !is_tuple_v, + "BLayout and BDataType must be scalars."); + + /// @brief CLayout and CDataType are expected to be scalars, not a tuple. + static_assert(!is_tuple_v && !is_tuple_v, + "CLayout and CDataType must be scalars."); + + struct StreamKKernelArgs : ck_tile::UniversalGemmKernelArgs<> + { + StreamKKernelArgs(const StreamKHostArgs& host_args, index_t grid) + : UniversalGemmKernelArgs{host_args.as_ptr, + host_args.bs_ptr, + host_args.ds_ptr, + host_args.e_ptr, + host_args.M, + host_args.N, + host_args.K, + host_args.stride_As, + host_args.stride_Bs, + host_args.stride_Ds, + host_args.stride_E, + host_args.k_batch}, + reduction_strategy{host_args.reduction_strategy}, + // The workspace pointer is set to nullptr because we must first + // instantiate the TilePartitioner to get the necessary size + workspace_ptr{nullptr}, + tile_partitioner{TilePartitioner{host_args.M, host_args.N, host_args.K, grid}} + + { + } + + /// @brief The strategy used by work groups to compute final results in C tensor. + StreamKReductionStrategy reduction_strategy; + /// @brief A pointer to a buffer in device memory for accumulating partial via reduction + /// strategy. + void* workspace_ptr; + /// @brief An instance of the TilePartioner class for assisting with mapping workgroups to + /// the C tensor. + TilePartitioner tile_partitioner; + }; + + using KernelArgs = StreamKKernelArgs; + using Kernel = StreamKKernel; + + [[nodiscard]] CK_TILE_HOST static const std::string GetName() + { + // clang-format off + using P_ = GemmPipeline; + using WarpTile = typename P_::BlockGemmShape::WarpTile; + + return concat('_', "streamk", gemm_prec_str(), + concat('x', P_::MPerBlock, P_::NPerBlock, P_::KPerBlock), + concat('x', WarpTile::at(number<0>{}), WarpTile::at(number<1>{}), WarpTile::at(number<2>{})), + concat('x', P_::GetVectorSizeA(), P_::GetVectorSizeB(), P_::GetVectorSizeC()), + concat('x', P_::kPadM, P_::kPadN, P_::kPadK)); + // clang-format on + } + + /// @brief Compute the grid size for the Stream K kernel using the tile_partitioner. + /// @return The grid size. + CK_TILE_HOST static auto GridSize(const TilePartitioner& tile_partitioner) -> dim3 + { + return tile_partitioner.grid_size(); + } + + /// @brief Get the maximum occupancy grid size for the persistent kernel on the current device. + /// @return The maximum occupancy grid size. + /// @note This function queries the maximum occupancy of the kernel using + /// `hipOccupancyMaxActiveBlocksPerMultiprocessor`. + CK_TILE_HOST static auto MaxOccupancyGridSize(const stream_config& s) -> dim3 + { + return UniversalGemmKernel::MaxOccupancyGridSize(s); + } + + CK_TILE_HOST static constexpr auto BlockSize() -> dim3 + { + return UniversalGemmKernel::BlockSize(); + } + + /// @brief Constructs kernel arguments for the Stream-K kernel. + /// @param host_args Stream-K host arguments. + /// @param num_cu Number of compute units (CUs). The default is the number of CUs on the device. + /// The caller may select their own to assist with test reproducibility, etc. + /// @param occupancy The maximum number of active blocks per CU for this kernel. The caller may + /// select their own to assist with test reproducibility, etc. + /// @return The kernel arguments for Stream-K. + CK_TILE_HOST static StreamKKernelArgs MakeKernelArgs(const StreamKHostArgs& host_args, + int num_cu = NumCU(), + int occupancy = Occupancy()) + { + const index_t grid = num_cu * occupancy; + + return StreamKKernelArgs{host_args, grid}; + } + + template + CK_TILE_DEVICE static void + RunGemm(const std::array& as_ptr, + const std::array& bs_ptr, + const std::array& ds_ptr, + CDataType* c_ptr, + void* smem_ptr_0, + const typename UniversalGemmKernel::KernelArgs& kargs, + const index_t num_loop, + const index_t block_idx_m, + const index_t block_idx_n, + const index_t k_size) + { + // Create Gemm tensor views, pad views and tile windows + const auto& gemm_tensor_views_tuple = + UniversalGemmKernel::template MakeGemmTensorViews( + as_ptr, bs_ptr, ds_ptr, c_ptr, kargs, k_size); + + const auto& gemm_pad_views = UniversalGemmKernel::MakeGemmPadViews(gemm_tensor_views_tuple); + auto gemm_tile_windows = + UniversalGemmKernel::MakeGemmTileWindows(gemm_pad_views, block_idx_m, block_idx_n); + + // Run GEMM cooperatively by whole workgroup. + const auto& as_block_window = gemm_tile_windows.at(UniversalGemmKernel::I0); + const auto& bs_block_window = gemm_tile_windows.at(UniversalGemmKernel::I1); + const auto& ds_block_window = gemm_tile_windows.at(UniversalGemmKernel::I2); + + // Since num_loop can vary per WG and per iteration of the Stream-K while loop, we compute + // has_hot_loop and tail_num here. This is a similar pattern used by grouped GEMM. In this + // case, we call the GemmPipeline's operator() function that takes both has_hot_loop and + // tail_num. + const bool has_hot_loop = GemmPipeline::BlockHasHotloop(num_loop); + const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto& c_block_tile = GemmPipeline{}(as_block_window[UniversalGemmKernel::I0], + bs_block_window[UniversalGemmKernel::I0], + num_loop, + has_hot_loop, + tail_num, + smem_ptr_0); + + if(UseDefaultScheduler || (get_warp_id() == 0)) + { + // Run Epilogue Pipeline + auto& c_block_window = gemm_tile_windows.at(UniversalGemmKernel::I3); + + EpiloguePipeline{}(c_block_window, c_block_tile, ds_block_window, smem_ptr_0); + } + } + + CK_TILE_HOST static bool IsSupportedArgument(const StreamKKernelArgs& kargs) + { + if(kargs.reduction_strategy == StreamKReductionStrategy::Reduction) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("CK Tile Stream-K only supports the atomic reduction strategy."); + } + return false; + } + return UniversalGemmKernel::IsSupportedArgument(kargs); + } + + /// @brief Computes the buffer size needed to store accumulation results for Stream K. + /// @return The buffer size needed. + CK_TILE_HOST static uint32_t GetWorkSpaceSize(const StreamKKernelArgs& kargs) + { + return kargs.tile_partitioner.GetWorkSpaceSize(sizeof(CDataType)); + } + + /// @brief Sets the kargs' current workspace_ptr to the given workspace_ptr. + /// @note Assumes that the given workspace_ptr points to allocated device memory. + CK_TILE_HOST static void SetWorkSpacePointer(StreamKKernelArgs& kargs, void* workspace_ptr) + { + kargs.workspace_ptr = workspace_ptr; + } + + /// @brief Computes offsets into A, B, and C tensors then runs the GEMM pipeline and epilogue. + /// @param kargs Stream-K kernel arguments. + /// @param tile_idx The 1D tile index in the C tensor for this workgroup. + /// @param num_loop The number of iterations (at the macro tile level) in the K dimension this + /// workgroup will perform in the C tile. + /// @param i_k_a The K offset in the A tensor. + /// @param i_k_b The K offset in the B tensor. + /// @param k_size The portion of the K dimension this workgroup processes in the assigned + /// `tile_idx`. + /// @param smem_ptr_0 Pointer to LDS. + CK_TILE_DEVICE void BaseGemm(StreamKKernelArgs& kargs, + index_t tile_idx, + index_t num_loop, + index_t i_k_a, + index_t i_k_b, + index_t k_size, + void* smem_ptr_0) const + { + const auto c_macro_tile_idx = kargs.tile_partitioner.get_output_tile_index(tile_idx); + index_t i_m = c_macro_tile_idx[UniversalGemmKernel::I0] * TilePartitioner::MPerBlock; + index_t i_n = c_macro_tile_idx[UniversalGemmKernel::I1] * TilePartitioner::NPerBlock; + + const ADataType* a_ptr = static_cast(kargs.as_ptr[0]) + i_k_a; + const BDataType* b_ptr = static_cast(kargs.bs_ptr[0]) + i_k_b; + CDataType* c_ptr = static_cast(kargs.e_ptr); + + // Run the GEMM pipeline and Epilogue. + RunGemm( + {a_ptr}, {b_ptr}, {/*ds_ptr*/}, c_ptr, smem_ptr_0, kargs, num_loop, i_m, i_n, k_size); + } + + /// @brief Runs the main Stream-K algorithm. + /// @param kargs Stream-K kernel arguments. + /// @param cta_idx The current Stream-K workgroup's index. + /// @param smem_ptr_0 Pointer to LDS. + /// @note It is assumed that the first Stream-K workgroup has a `cta_idx` of zero. If a + /// non-persistent data-parallel (DP) section is used, then a Stream-K workgroup's `cta_idx` + /// should be something like `blockIdx.x` minus number of DP workgroups. + CK_TILE_DEVICE void + StreamKGemm(StreamKKernelArgs& kargs, index_t cta_idx, void* smem_ptr_0) const + { + index_t iter_start, iter_end; + kargs.tile_partitioner.get_iter_boundaries(iter_start, iter_end, cta_idx); + + while(iter_start < iter_end) + { + // Get the 1D tile index in the C tensor that this workgroup will work in for this + // iteration of the loop. + index_t tile_idx = + amd_wave_read_first_lane(kargs.tile_partitioner.get_tile_index(iter_start)); + + // Get the start and end boundaries for the current tile. + index_t tile_iter_start, tile_iter_end; + kargs.tile_partitioner.get_tile_boundaries(tile_iter_start, tile_iter_end, tile_idx); + + // Get the start and end iteration within the current tile for the workgroup. + index_t local_iter_start = amd_wave_read_first_lane( + kargs.tile_partitioner.get_local_iter(iter_start, tile_iter_start)); + index_t local_iter_end = + amd_wave_read_first_lane(kargs.tile_partitioner.get_local_iter_end( + tile_iter_start, iter_end, tile_iter_end)); + + // Get the iteration length. + index_t num_loop_sk = local_iter_end - local_iter_start; + + // Determine the total size along the K dimension the workgroup is using in this + // iteration (used to construct tensor views). + index_t k_size = num_loop_sk * TilePartitioner::KPerBlock; + + // Get the K offsets for the A and B tensors + auto [i_k_a, i_k_b] = GetKOffsets( + local_iter_start, kargs.stride_As[0], kargs.stride_Bs[0]); + + if constexpr(TilePartitioner::ReductionStrategy == StreamKReductionStrategy::Atomic) + { + BaseGemm(kargs, tile_idx, num_loop_sk, i_k_a, i_k_b, k_size, smem_ptr_0); + } + else + { + // TODO: Apply reduction logic. + } + + // Prepare for next Stream-K loop iteration. + iter_start = tile_iter_end; + block_sync_lds(); + } + } + + /// @brief Entry point for the Stream-K Kernel with non-persistent DP. + /// + /// @par Overview + /// For the Non-Persistent kernel, each data parallel workgroup will + /// compute the results for their assigned macro-tile by calling `BaseGemm()`. + /// The Stream-K workgroups will do their assigned work by calling + /// `StreamKGemm()`, which calls `BaseGemm()` in the Stream-K loop. + template + CK_TILE_DEVICE typename std::enable_if_t operator()(StreamKKernelArgs kargs) const + { + // Allocate LDS + __shared__ char smem_ptr_0[UniversalGemmKernel::GetSmemSize()]; + + index_t block_idx = ck_tile::get_block_1d_id(); + index_t dp_num_loop = kargs.tile_partitioner.get_iters_per_tile(); + index_t dp_ctas = kargs.tile_partitioner.get_dp_ctas(); + bool is_dp_ctas = block_idx < kargs.tile_partitioner.get_dp_ctas(); + + // Check if at the data parallel section + if(is_dp_ctas) + { + BaseGemm(kargs, block_idx, dp_num_loop, 0, 0, kargs.K, smem_ptr_0); + } + else + { + // Stream-K + StreamKGemm(kargs, block_idx - dp_ctas, smem_ptr_0); + } + } + + /// @brief Entry point for the Stream-K Kernel with persistent DP. + /// + /// @par Overview + /// For the Persistent kernel, each workgroup will first compute their + /// assigned data-parallel tiles. Each data parallel tile will be computed + /// by calling `BaseGemm()`. Then the workgroups will proceed with the + /// Stream-K portion by calling `StreamKGemm()`, which calls `BaseGemm()` + /// in the Stream-K loop. + template + CK_TILE_DEVICE typename std::enable_if_t operator()(StreamKKernelArgs kargs) const + { + // Allocate LDS + __shared__ char smem_ptr_0[UniversalGemmKernel::GetSmemSize()]; + + index_t block_idx = ck_tile::get_block_1d_id(); + index_t dp_num_loop = kargs.tile_partitioner.get_iters_per_tile(); + + // Data-parallel section + for(index_t tile_idx = block_idx; tile_idx < kargs.tile_partitioner.get_dp_tiles(); + tile_idx += kargs.tile_partitioner.get_grid()) + { + BaseGemm(kargs, tile_idx, dp_num_loop, 0, 0, kargs.K, smem_ptr_0); + } + + // Stream-K section + StreamKGemm(kargs, block_idx, smem_ptr_0); + } + + private: + /// @brief Computes the K offsets in the A and B tensors given iter_offset, where iter_offset is + /// the starting macro tile index in the K dimension for the workgroup. + /// @return A tuple containing the offsets into the A and B tensors accounting for the layouts + /// of A and B. + /// @note The default case is that A is assumed to be row major and B is assumed to be column + /// major. + template + CK_TILE_DEVICE static tuple + GetKOffsets(index_t iter_offset, index_t stride_a, index_t stride_b) + { + index_t stride_offset_a; + index_t stride_offset_b; + if constexpr(std::is_same_v) + { + stride_offset_a = stride_a; + } + else + { + stride_offset_a = 1; + } + + if constexpr(std::is_same_v) + { + stride_offset_b = stride_b; + } + else + { + stride_offset_b = 1; + } + + index_t base_offset = iter_offset * TilePartitioner::KPerBlock; + + return make_tuple(base_offset * stride_offset_a, base_offset * stride_offset_b); + } + + CK_TILE_HOST static int NumCU() + { + hipDeviceProp_t dev_prop; + hipDevice_t dev; + hip_check_error(hipGetDevice(&dev)); + hip_check_error(hipGetDeviceProperties(&dev_prop, dev)); + int num_cu = dev_prop.multiProcessorCount; + + return num_cu; + } + + /// @brief Computes the occupancy (i.e. maximum number of active blocks per CU) for the kernel + /// @return The occupancy + /// @note This function queries the maximum occupancy of the kernel using + /// `hipOccupancyMaxActiveBlocksPerMultiprocessor`. + CK_TILE_HOST static int Occupancy() + { + int occupancy; + + // Since occupancy of 1 is valid for stream k, we set min_num_block_per_cu to 1 + constexpr int min_block_per_cu = 1; + const auto kernel = kentry; + + hip_check_error( + hipOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy, kernel, kBlockSize, 0)); + + return occupancy; + } +}; +} // namespace reboot /// @brief The Stream K GEMM kernel host arguments. /// diff --git a/include/ck_tile/ops/gemm/kernel/streamk_gemm_tile_partitioner.hpp b/include/ck_tile/ops/gemm/kernel/streamk_gemm_tile_partitioner.hpp index 1962f3518a..e98c60e5f0 100644 --- a/include/ck_tile/ops/gemm/kernel/streamk_gemm_tile_partitioner.hpp +++ b/include/ck_tile/ops/gemm/kernel/streamk_gemm_tile_partitioner.hpp @@ -186,6 +186,11 @@ struct StreamKTilePartitionerBase */ CK_TILE_HOST_DEVICE index_t get_n() const noexcept; + /** + * @brief Returns an estimate of the number of workgroups writing to the same macro tile in C. + */ + CK_TILE_HOST index_t estimate_num_wgs_per_tile() const noexcept; + protected: index_t num_tiles_; index_t grid_; @@ -246,6 +251,7 @@ struct StreamKTilePartitioner_v2::get_n() c return n_; } +template +CK_TILE_HOST index_t +StreamKTilePartitionerBase::estimate_num_wgs_per_tile() + const noexcept +{ + // In the case of non-atomic reduction or data-parallel only, there will always be 1 workgroup + // writing final results to a given macro tile in C. + int num_wgs_per_tile = 1; + + // Otherwise, for atomics, multiple workgroups may be writing to the same macro tile in C. + if(sk_ctas_ > 0 && ReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic) + { + ck_tile::index_t iters_per_sk_cta_non_zero = ck_tile::max(iters_per_sk_cta_, 1); + // Estimate the number of workgroups per macro tile. + num_wgs_per_tile = (iters_per_tile_ / iters_per_sk_cta_non_zero) + + ((iters_per_tile_ % iters_per_sk_cta_non_zero) != 0); + } + + return std::max(num_wgs_per_tile, 1); +} + template diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 47b776f401..810ae8d231 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -46,6 +46,7 @@ set(REGRESSION_TESTS test_ck_tile_fmha_fwd_bf16 test_ck_tile_fmha_fwd_fp16 test_ck_tile_fmha_fwd_fp8 + test_ck_tile_streamk_reboot_extended ) function(add_test_executable TEST_NAME) diff --git a/test/ck_tile/gemm_streamk/CMakeLists.txt b/test/ck_tile/gemm_streamk/CMakeLists.txt index 331118da59..eba411e271 100644 --- a/test/ck_tile/gemm_streamk/CMakeLists.txt +++ b/test/ck_tile/gemm_streamk/CMakeLists.txt @@ -117,6 +117,18 @@ if(GPU_TARGETS MATCHES "gfx9") # #${CMAKE_CURRENT_SOURCE_DIR}/extended_tests/mem/bf16_ccc_mem_128x128x32_2x2x1_32x32x16_NonPersistent.cpp # ) add_gtest_executable(test_ck_tile_streamk_tile_partitioner test_streamk_tile_partitioner.cpp) + add_gtest_executable(test_ck_tile_streamk_reboot_smoke + ${CMAKE_CURRENT_SOURCE_DIR}/smoke_tests/test_gemm_streamk_reboot_fp16_persistent.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/smoke_tests/test_gemm_streamk_reboot_bf16_persistent.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/smoke_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/smoke_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp + test_gemm_streamk_reboot_util.cpp) + add_gtest_executable(test_ck_tile_streamk_reboot_extended + ${CMAKE_CURRENT_SOURCE_DIR}/extended_tests/test_gemm_streamk_reboot_fp16_persistent.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/extended_tests/test_gemm_streamk_reboot_bf16_persistent.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/extended_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/extended_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp + test_gemm_streamk_reboot_util.cpp) else() message(DEBUG "Skipping test_ck_tile_streamk tests for current target") endif() diff --git a/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp new file mode 100644 index 0000000000..eb4478f3d6 --- /dev/null +++ b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootBf16NonPersistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootBf16NonPersistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootBf16NonPersistent, KernelTypesStreamKBf16NonPersistent); + +#include "test_gemm_streamk_reboot_extended_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_persistent.cpp b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_persistent.cpp new file mode 100644 index 0000000000..c42ada1a98 --- /dev/null +++ b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_bf16_persistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootBf16Persistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootBf16Persistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootBf16Persistent, KernelTypesStreamKBf16Persistent); + +#include "test_gemm_streamk_reboot_extended_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp new file mode 100644 index 0000000000..664c16a5e6 --- /dev/null +++ b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootFp16NonPersistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootFp16NonPersistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootFp16NonPersistent, KernelTypesStreamKFp16NonPersistent); + +#include "test_gemm_streamk_reboot_extended_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_persistent.cpp b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_persistent.cpp new file mode 100644 index 0000000000..39c79b4180 --- /dev/null +++ b/test/ck_tile/gemm_streamk/extended_tests/test_gemm_streamk_reboot_fp16_persistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootFp16Persistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootFp16Persistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootFp16Persistent, KernelTypesStreamKFp16Persistent); + +#include "test_gemm_streamk_reboot_extended_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp new file mode 100644 index 0000000000..0c1813fb65 --- /dev/null +++ b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_nonpersistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootBf16NonPersistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootBf16NonPersistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootBf16NonPersistent, KernelTypesStreamKBf16NonPersistent); + +#include "test_gemm_streamk_reboot_smoke_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_persistent.cpp b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_persistent.cpp new file mode 100644 index 0000000000..e78092c4ba --- /dev/null +++ b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_bf16_persistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootBf16Persistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootBf16Persistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootBf16Persistent, KernelTypesStreamKBf16Persistent); + +#include "test_gemm_streamk_reboot_smoke_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp new file mode 100644 index 0000000000..5e6118bd0c --- /dev/null +++ b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_nonpersistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootFp16NonPersistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootFp16NonPersistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootFp16NonPersistent, KernelTypesStreamKFp16NonPersistent); + +#include "test_gemm_streamk_reboot_smoke_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_persistent.cpp b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_persistent.cpp new file mode 100644 index 0000000000..9f9c8f8234 --- /dev/null +++ b/test/ck_tile/gemm_streamk/smoke_tests/test_gemm_streamk_reboot_fp16_persistent.cpp @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_streamk_reboot_types.hpp" +#include "test_gemm_streamk_reboot_util.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileStreamKRebootFp16Persistent : public TestCkTileStreamKReboot +{ +}; + +#define TEST_SUITE_NAME TestCkTileStreamKRebootFp16Persistent + +TYPED_TEST_SUITE(TestCkTileStreamKRebootFp16Persistent, KernelTypesStreamKFp16Persistent); + +#include "test_gemm_streamk_reboot_smoke_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_extended_cases.inc b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_extended_cases.inc new file mode 100644 index 0000000000..8b6522bd75 --- /dev/null +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_extended_cases.inc @@ -0,0 +1,24 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +TYPED_TEST(TEST_SUITE_NAME, StreamK_DP2TSK) +{ + const ck_tile::index_t num_cu = get_cu_count(); + constexpr ck_tile::index_t M_Tile = std::tuple_element_t<7, TypeParam>::value; + constexpr ck_tile::index_t N_Tile = std::tuple_element_t<8, TypeParam>::value; + + // For DP 2-Tile SK, there are 2 important terms: + // Term 1: (M_Tile * num_cu * 2) - This ensures we have at least 2 cycles that will fully + // saturate all CUs. This assumes tile sizes are large enough such that occupancy is 1. + // Term 2: (M_Tile * 2) - This ensures we have 1 cycle that does not fully saturate all CUs + // (i.e., we will have remainder tiles). This guarantees we have 1 full tile cycle plus + // remainder tiles for the 2 Tile SK portion; the rest of the tiles will fully saturate all CUs + // for the DP portion. + ck_tile::index_t M = (M_Tile * num_cu * 2) + (M_Tile * 2); + ck_tile::index_t N = N_Tile; + ck_tile::index_t K = 2048; + + this->Run(M, N, K); +} diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_smoke_cases.inc b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_smoke_cases.inc new file mode 100644 index 0000000000..d714b3446c --- /dev/null +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_smoke_cases.inc @@ -0,0 +1,47 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +TYPED_TEST(TEST_SUITE_NAME, StreamK_EdgeCase) +{ + ck_tile::index_t M = 256; + ck_tile::index_t N = 256; + ck_tile::index_t K = 256; + + this->Run(M, N, K); +} + +TYPED_TEST(TEST_SUITE_NAME, StreamK_DPOnly) +{ + const ck_tile::index_t num_cu = get_cu_count(); + constexpr ck_tile::index_t M_Tile = std::tuple_element_t<7, TypeParam>::value; + constexpr ck_tile::index_t N_Tile = std::tuple_element_t<8, TypeParam>::value; + constexpr ck_tile::index_t K_Tile = std::tuple_element_t<9, TypeParam>::value; + + // For DP only, we ensure that the number of tiles is a multiple of the number of CUs. This + // assumes tile sizes are large enough such that occupancy is 1. + ck_tile::index_t M = M_Tile * num_cu; + ck_tile::index_t N = N_Tile; + ck_tile::index_t K = K_Tile; + + this->Run(M, N, K); +} + +TYPED_TEST(TEST_SUITE_NAME, StreamK_SKOnly) +{ + const ck_tile::index_t num_cu = get_cu_count(); + constexpr ck_tile::index_t M_Tile = std::tuple_element_t<7, TypeParam>::value; + constexpr ck_tile::index_t N_Tile = std::tuple_element_t<8, TypeParam>::value; + constexpr ck_tile::index_t K_Tile = std::tuple_element_t<9, TypeParam>::value; + + // For SK only, we have 4 macro tiles in C. But, we need to make sure there is enough work along + // the K dimension to avoid falling into the edge case. Thus, we always have at least num_cu + // macro tiles in the K dimension. This assumes tile sizes are large enough such that occupancy + // is 1. + ck_tile::index_t M = M_Tile * 2; + ck_tile::index_t N = N_Tile * 2; + ck_tile::index_t K = K_Tile * num_cu; + + this->Run(M, N, K); +} diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_types.hpp b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_types.hpp new file mode 100644 index 0000000000..1db53ddd64 --- /dev/null +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_types.hpp @@ -0,0 +1,56 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include +#include + +#include "gtest/gtest.h" + +#include "ck_tile/host.hpp" + +using F16 = ck_tile::half_t; +using F32 = float; +using BF16 = ck_tile::bf16_t; + +using Row = ck_tile::tensor_layout::gemm::RowMajor; +using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + +using Persistent = std::true_type; +using NonPersistent = std::false_type; + +using I32 = ck_tile::number<32>; +using I256 = ck_tile::number<256>; + +// clang-format off +using KernelTypesStreamKFp16Persistent = ::testing::Types< +// ALayout BLayout CLayout ADataType BDataType AccDataType CDataType M_MacroTile N_MacroTile K_MacroTile Persistent + + std::tuple< Row, Row, Row, F16, F16, F32, F16, I256, I256, I32, Persistent>, + std::tuple< Row, Col, Row, F16, F16, F32, F16, I256, I256, I32, Persistent>, + std::tuple< Col, Col, Row, F16, F16, F32, F16, I256, I256, I32, Persistent>, + std::tuple< Col, Row, Row, F16, F16, F32, F16, I256, I256, I32, Persistent> +>; + +using KernelTypesStreamKBf16Persistent = ::testing::Types< + std::tuple< Row, Row, Row, BF16, BF16, F32, BF16, I256, I256, I32, Persistent>, + std::tuple< Row, Col, Row, BF16, BF16, F32, BF16, I256, I256, I32, Persistent>, + std::tuple< Col, Col, Row, BF16, BF16, F32, BF16, I256, I256, I32, Persistent>, + std::tuple< Col, Row, Row, BF16, BF16, F32, BF16, I256, I256, I32, Persistent> +>; + +using KernelTypesStreamKFp16NonPersistent = ::testing::Types< +// ALayout BLayout CLayout ADataType BDataType AccDataType CDataType M_MacroTile N_MacroTile K_MacroTile Persistent + + std::tuple< Row, Row, Row, F16, F16, F32, F16, I256, I256, I32, NonPersistent>, + std::tuple< Row, Col, Row, F16, F16, F32, F16, I256, I256, I32, NonPersistent>, + std::tuple< Col, Col, Row, F16, F16, F32, F16, I256, I256, I32, NonPersistent>, + std::tuple< Col, Row, Row, F16, F16, F32, F16, I256, I256, I32, NonPersistent> +>; + +using KernelTypesStreamKBf16NonPersistent = ::testing::Types< + std::tuple< Row, Row, Row, BF16, BF16, F32, BF16, I256, I256, I32, NonPersistent>, + std::tuple< Row, Col, Row, BF16, BF16, F32, BF16, I256, I256, I32, NonPersistent>, + std::tuple< Col, Col, Row, BF16, BF16, F32, BF16, I256, I256, I32, NonPersistent>, + std::tuple< Col, Row, Row, BF16, BF16, F32, BF16, I256, I256, I32, NonPersistent> +>; +// clang-format on diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.cpp b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.cpp new file mode 100644 index 0000000000..39a92d622d --- /dev/null +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.cpp @@ -0,0 +1,10 @@ +#include "test_gemm_streamk_reboot_util.hpp" + +ck_tile::index_t get_cu_count() +{ + hipDeviceProp_t dev_prop; + hipDevice_t dev; + ck_tile::hip_check_error(hipGetDevice(&dev)); + ck_tile::hip_check_error(hipGetDeviceProperties(&dev_prop, dev)); + return dev_prop.multiProcessorCount; +} diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.hpp b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.hpp new file mode 100644 index 0000000000..85863989b0 --- /dev/null +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_reboot_util.hpp @@ -0,0 +1,283 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include + +#include "ck_tile/host.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" + +template +auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeType = + std::conditional_t; + // Calculate thresholds + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + + // The logic below may need to become more advanced once bugs in Stream-K Tile Partitioner are + // resolved. Because the number of WGs contributing to a macro tile in C may not be the same for + // all macro tiles in C. + + // Calculate error due to more than 1 WG contributing to the same macro tile in C + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} + +ck_tile::index_t get_cu_count(); + +template +class TestCkTileStreamKReboot : public ::testing::Test +{ + protected: + using ALayout = std::tuple_element_t<0, Tuple>; + using BLayout = std::tuple_element_t<1, Tuple>; + using CLayout = std::tuple_element_t<2, Tuple>; + using ADataType = std::tuple_element_t<3, Tuple>; + using BDataType = std::tuple_element_t<4, Tuple>; + using AccDataType = std::tuple_element_t<5, Tuple>; + using CDataType = std::tuple_element_t<6, Tuple>; + using DsLayout = ck_tile::tuple<>; + using DsDataType = ck_tile::tuple<>; + static constexpr ck_tile::index_t M_Tile = std::tuple_element_t<7, Tuple>::value; + static constexpr ck_tile::index_t N_Tile = std::tuple_element_t<8, Tuple>::value; + static constexpr ck_tile::index_t K_Tile = std::tuple_element_t<9, Tuple>::value; + static constexpr bool Persistent = std::tuple_element_t<10, Tuple>::value; + + template + ck_tile::index_t invoke_streamk(const ck_tile::reboot::StreamKHostArgs& args, + const ck_tile::stream_config& s) + { + constexpr ck_tile::index_t M_Warp = 2; + constexpr ck_tile::index_t N_Warp = 2; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = 32; + constexpr ck_tile::index_t N_Warp_Tile = 32; + constexpr ck_tile::index_t K_Warp_Tile = 16; + + constexpr bool kPadM = PadM; + constexpr bool kPadN = PadN; + constexpr bool kPadK = PadK; + constexpr bool preshuffle = Preshuffle; + + constexpr bool DoubleSmemBuffer = false; + constexpr int kBlockPerCu = 1; + constexpr bool StructuredSparsity = false; + constexpr bool NumWaveGroup = 1; + + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = + ck_tile::StreamKTilePartitioner_v2; + + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + + const auto Run = [&](const auto memory_operation_) { + constexpr auto memory_operation = memory_operation_.value; + constexpr auto scheduler = ck_tile::GemmPipelineScheduler::Intrawave; + + // We create the GEMM pipeline without specifying has_hot_loop or tail_num. + // This is because num_loop can vary (a) per WG and (b) per iteration of the Stream-K + // while loop. Instead, has_hot_loop and tail_num are determined in the Stream-K + // Kernel's RunGemm function. This is a similar pattern used by grouped GEMM. + using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; + // For initial testing, we will just test with one pipeline. + // More extensive testing is coming later and will test other pipelines. + using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem; + + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + M_Warp, + N_Warp, + M_Warp_Tile, + N_Warp_Tile, + K_Warp_Tile, + UniversalGemmProblem::TransposeC, + memory_operation>>; + + using Kernel = + ck_tile::reboot::StreamKKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + + if(!Kernel::IsSupportedArgument(kargs)) + { + EXPECT_TRUE(false); + } + + dim3 grid_dims = Kernel::GridSize(kargs.tile_partitioner); + dim3 block_dims = Kernel::BlockSize(); + + ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grid_dims, block_dims, 0, kargs)); + + return kargs.tile_partitioner.estimate_num_wgs_per_tile(); + }; + + return Run(ck_tile::integral_constant{}); + } + + public: + void Run(ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + ck_tile::StreamKReductionStrategy reduction_strategy = + ck_tile::StreamKReductionStrategy::Atomic, + ck_tile::index_t stride_A = 0, + ck_tile::index_t stride_B = 0, + ck_tile::index_t stride_C = 0) + { + // Since M, N, and K will vary depending on the number of CUs, we print it here to + // facilitate test output readability. + std::cout << "M: " << M << ", N: " << N << ", K: " << K << std::endl; + + using namespace ck_tile::literals; + + if(reduction_strategy == ck_tile::StreamKReductionStrategy::Reduction) + { + throw std::runtime_error("Reduction Strategy is current unsupported!\n"); + } + + auto f_host_tensor_descriptor = [](std::size_t row, + std::size_t col, + std::size_t stride, + auto layout) { + if constexpr(std::is_same_v) + { + return ck_tile::HostTensorDescriptor({row, col}, {stride, 1_uz}); + } + else + { + return ck_tile::HostTensorDescriptor({row, col}, {1_uz, stride}); + } + }; + + auto f_get_default_stride = + [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { + if(stride == 0) + { + if constexpr(std::is_same_v) + { + return col; + } + else + { + return row; + } + } + else + return stride; + }; + + stride_A = f_get_default_stride(M, K, stride_A, ALayout{}); + stride_B = f_get_default_stride(K, N, stride_B, BLayout{}); + stride_C = f_get_default_stride(M, N, stride_C, CLayout{}); + + ck_tile::HostTensor a_m_k(f_host_tensor_descriptor(M, K, stride_A, ALayout{})); + ck_tile::HostTensor b_k_n(f_host_tensor_descriptor(K, N, stride_B, BLayout{})); + ck_tile::HostTensor c_m_n_dev_result( + f_host_tensor_descriptor(M, N, stride_C, CLayout{})); + + ck_tile::FillUniformDistributionIntegerValue{-5, 5, /*seed*/ 11939}(a_m_k); + ck_tile::FillUniformDistributionIntegerValue{-5, 5, /*seed*/ 11940}(b_k_n); + + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); + + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + + ck_tile::reboot::StreamKHostArgs args{a_m_k_dev_buf.GetDeviceBuffer(), + b_k_n_dev_buf.GetDeviceBuffer(), + c_m_n_dev_buf.GetDeviceBuffer(), + M, + N, + K, + stride_A, + stride_B, + stride_C, + reduction_strategy}; + + ck_tile::index_t num_accumulations_per_tile = + invoke_streamk( + args, ck_tile::stream_config{nullptr, false, 0, 0, 1}); + + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); + + ck_tile::HostTensor c_m_n_host_ref( + f_host_tensor_descriptor(M, N, stride_C, CLayout{})); + c_m_n_host_ref.SetZero(); + + ck_tile::reference_gemm( + a_m_k, b_k_n, c_m_n_host_ref); + + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = calculate_rtol_atol( + K, num_accumulations_per_tile, max_accumulated_value); + + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass); + }; +}; diff --git a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp index 89d72d844b..9028f7bf10 100644 --- a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp +++ b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp @@ -77,6 +77,26 @@ TEST(StreamKTilePartitionerBaseGetWorkSpaceSize, ReductionStrategy) expected_partials_size + expected_flags_size); } +TEST(StreamKTilePartitionerBaseEstimateNumWgsPerTile, EstimateNumWgsPerTileLowerValue) +{ + using Config = StreamKTilePartitionerBaseConfigDP2TileSK; + + ck_tile::StreamKTilePartitionerBase tile_partitioner{ + Config::M, Config::N, Config::K, Config::GRID}; + + EXPECT_EQ(tile_partitioner.estimate_num_wgs_per_tile(), 1); +} + +TEST(StreamKTilePartitionerBaseEstimateNumWgsPerTile, EstimateNumWgsPerTileEqualValue) +{ + using Config = StreamKTilePartitionerBaseConfigSKOnlyWith2WgsPerSKTile; + + ck_tile::StreamKTilePartitionerBase tile_partitioner{ + Config::M, Config::N, Config::K, Config::GRID}; + + EXPECT_EQ(tile_partitioner.estimate_num_wgs_per_tile(), 2); +} + TEST(StreamKTilePartitionerBaseGetLocalIter, GetLocalIter) { // Types diff --git a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp index 4fc654a7ea..eb62f4253b 100644 --- a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp +++ b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp @@ -194,6 +194,23 @@ struct StreamKTilePartitionerBaseConfigDP2TileSK : public StreamKTilePartitioner ck_tile::sequence>; }; +struct StreamKTilePartitionerBaseConfigSKOnlyWith2WgsPerSKTile + : public StreamKTilePartitionerBaseConfig +{ + static constexpr ck_tile::index_t M = 16; + static constexpr ck_tile::index_t N = 4; + static constexpr ck_tile::index_t K = 16; + static constexpr ck_tile::index_t GRID = 8; + + static constexpr ck_tile::index_t M_TILE = 4; + static constexpr ck_tile::index_t N_TILE = 4; + static constexpr ck_tile::index_t K_TILE = 8; + + using GemmShape = ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; +}; + struct StreamKTilePartitionerBaseConfigDPOnly : public StreamKTilePartitionerBaseConfig { static constexpr ck_tile::index_t M = 12; From 5c1974065e6c0a95c5db6e09903f078ec7a386be Mon Sep 17 00:00:00 2001 From: Johannes Graner Date: Mon, 27 Oct 2025 18:43:09 +0100 Subject: [PATCH 06/17] [CK_TILE] Add conv fwd + bias + clamp example (#3012) * Implement argument passing to element-wise functions for fwd convolution * Add files for fwd + bias + clamp example * Implement Bias * Implement Clamp * Elementwise function composition * Composition unit test * Implement fwd + bias + clamp example * Simplify argument passing and composition * elfunc -> bias_and_clamp * Rename function to specify example * Move element-wise function instantiation to kernel * Make bias a runtime tensor * No ugly namespace aliasing * Initialize element-wise function on host * Remove function initialization helper, simplify Compose initialization * Remove unintended LSP compatibility patch * Clean up includes and unused code * Switch names in cshuffle epilogue * Move CDElementwise to conv traits * Re-add required include * Initialize bias in same way as other tensors * Better type specification for ds pointer * Disable 1D convolution * Add warning for non-group-constant bias --- .../20_grouped_convolution/CMakeLists.txt | 3 + ...grouped_convolution_forward_bias_clamp.cpp | 58 ++++ .../grouped_convolution_forward_invoker.hpp | 13 +- ...ped_convolution_fwd_bias_clamp_example.inc | 301 ++++++++++++++++++ .../run_grouped_convolution_fwd_example.inc | 14 +- .../reference/reference_grouped_conv_fwd.hpp | 33 +- .../unary_element_wise_operation.hpp | 66 ++++ .../ops/epilogue/cshuffle_epilogue.hpp | 10 +- .../grouped_convolution_forward_kernel.hpp | 24 +- .../utils/grouped_convolution_utils.hpp | 29 +- .../elementwise/test_elementwise_1d.cpp | 14 +- 11 files changed, 524 insertions(+), 41 deletions(-) create mode 100644 example/ck_tile/20_grouped_convolution/grouped_convolution_forward_bias_clamp.cpp create mode 100644 example/ck_tile/20_grouped_convolution/run_grouped_convolution_fwd_bias_clamp_example.inc diff --git a/example/ck_tile/20_grouped_convolution/CMakeLists.txt b/example/ck_tile/20_grouped_convolution/CMakeLists.txt index 10332137e2..e9614061e1 100644 --- a/example/ck_tile/20_grouped_convolution/CMakeLists.txt +++ b/example/ck_tile/20_grouped_convolution/CMakeLists.txt @@ -4,6 +4,9 @@ list(APPEND EXAMPLE_CONV_COMPILE_OPTIONS -mllvm -enable-noalias-to-md-conversion add_executable(tile_example_grouped_conv_fwd EXCLUDE_FROM_ALL grouped_convolution_forward.cpp) target_compile_options(tile_example_grouped_conv_fwd PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) +add_executable(tile_example_grouped_conv_fwd_bias_clamp EXCLUDE_FROM_ALL grouped_convolution_forward_bias_clamp.cpp) +target_compile_options(tile_example_grouped_conv_fwd_bias_clamp PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + add_executable(tile_example_grouped_conv_bwd_weight EXCLUDE_FROM_ALL grouped_convolution_backward_weight.cpp) target_compile_options(tile_example_grouped_conv_bwd_weight PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_bias_clamp.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_bias_clamp.cpp new file mode 100644 index 0000000000..ed215cb178 --- /dev/null +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_bias_clamp.cpp @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include +#include +#include +#include +#include + +#include "ck_tile/host.hpp" +#include "grouped_convolution_utils.hpp" +#include "grouped_convolution_forward_invoker.hpp" +#include "run_grouped_convolution_fwd_bias_clamp_example.inc" + +template