From e231cfb3dcb752ee4eedc714315f9dc8f0e34403 Mon Sep 17 00:00:00 2001 From: Thrupti Raj Lakshmana Gowda Date: Tue, 13 Jan 2026 18:20:30 -0600 Subject: [PATCH] [CK TILE ENGINE] CI fix for Basic Tile Engine (#3554) * memory op changes * memory op changes * Fixing TILE_ENGINE_BASIC in Tile Engine * Removing gfx90a from Tile Engine Run * [CK TILE ENGINE] increasing ci configs for BASIC case * Setting RUN_TILE_ENGINE_BASIC_TESTS to ON by default --------- Co-authored-by: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> [ROCm/composable_kernel commit: 51027474afe07ba069123f37798867270d59ac12] --- Jenkinsfile | 40 +++----------- tile_engine/ops/gemm/gemm_instance_builder.py | 52 +++++-------------- .../configs/default_ci_config.json | 7 ++- .../configs/default_ci_config.json | 1 + .../configs/default_ci_config.json | 7 ++- 5 files changed, 31 insertions(+), 76 deletions(-) diff --git a/Jenkinsfile b/Jenkinsfile index 87e0d03752..9c670183fd 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -1201,8 +1201,8 @@ pipeline { description: "Run the ck_tile FMHA tests (default: OFF)") booleanParam( name: "RUN_TILE_ENGINE_BASIC_TESTS", - defaultValue: false, - description: "Run the tile_engine_basic tests (default: OFF)") + defaultValue: true, + description: "Run the tile_engine_basic tests (default: ON)") booleanParam( name: "RUN_TILE_ENGINE_GEMM_TESTS", defaultValue: false, @@ -1650,7 +1650,10 @@ pipeline { -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ -D GEMM_PRESHUFFLE_CONFIG_FILE="default_ci_config.json" .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all """ + ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ + python3 ../tile_engine/ops/gemm/gemm_universal/gemm_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ } steps{ buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) @@ -1667,37 +1670,6 @@ pipeline { } parallel { - stage("Run TILE_ENGINE_GEMM Tests on gfx90a") - { - when { - beforeAgent true - expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } - } - agent{ label rocmnode("gfx90a") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx90a" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_STREAMK_DATATYPE="fp8;fp16" \ - -D GEMM_STREAMK_LAYOUT="rcr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all benchmark_gemm_streamk_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } - steps{ - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) - cleanWs() - } - } stage("Run TILE_ENGINE_GEMM Tests on gfx942") { when { diff --git a/tile_engine/ops/gemm/gemm_instance_builder.py b/tile_engine/ops/gemm/gemm_instance_builder.py index 089f968649..9c60c565de 100644 --- a/tile_engine/ops/gemm/gemm_instance_builder.py +++ b/tile_engine/ops/gemm/gemm_instance_builder.py @@ -643,40 +643,31 @@ struct SelectedKernel {{ using GemmPipeline = {pipeline_impl_map.get(pipeline)};""" - # Runfunction body - instance_code += """ - - const auto Run = [&](const auto memory_operation_) {""" - # Scheduler initialization if self.kernel_name_prefix in ["gemm_universal"]: instance_code += f""" - constexpr auto scheduler = {scheduler_type_map.get(scheduler)};""" - - # Memory operation - instance_code += """ - [[maybe_unused]] constexpr auto memory_operation = memory_operation_.value;""" + constexpr auto scheduler = {scheduler_type_map.get(scheduler)};""" # UniversalGemmProblem if self.kernel_name_prefix in ["gemm_universal"]: instance_code += """ - using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem< - ADataType, - BDataType, - AccDataType, - TileShape, - ck_tile::TileGemmUniversalTraits, - scheduler>;""" + using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem< + ADataType, + BDataType, + AccDataType, + TileShape, + ck_tile::TileGemmUniversalTraits, + scheduler>;""" # GemmPipeline if self.kernel_name_prefix in ["gemm_universal"]: instance_code += f""" - using GemmPipeline = {pipeline_impl_map.get(pipeline)};""" + using GemmPipeline = {pipeline_impl_map.get(pipeline)};""" # Epilogue instance_code += self.populate_epilogue(epilogue) @@ -748,23 +739,8 @@ struct SelectedKernel {{ ck_tile::make_kernel(GemmKernel{{}}, grids, blocks, 0, kargs)); return ave_time; - }};""" - - # Run SplitK handler - - instance_code += """ - - float ave_time = 0.f; - if(args.k_batch == 1) { - ave_time = Run(ck_tile::integral_constant{}); - } else { - ave_time = Run(ck_tile::integral_constant{}); - } - return ave_time; - } -}; + }} +}}; """ return instance_code diff --git a/tile_engine/ops/gemm/gemm_multi_d/configs/default_ci_config.json b/tile_engine/ops/gemm/gemm_multi_d/configs/default_ci_config.json index 2df04d0ac1..0698786c33 100644 --- a/tile_engine/ops/gemm/gemm_multi_d/configs/default_ci_config.json +++ b/tile_engine/ops/gemm/gemm_multi_d/configs/default_ci_config.json @@ -49,7 +49,9 @@ "trait_config": { "pipeline": { "values": [ - "compv4" + "compv3", + "compv4", + "mem" ] }, "scheduler": { @@ -60,7 +62,8 @@ }, "epilogue": { "values": [ - "cshuffle" + "cshuffle", + "default" ] }, "pad_m": { diff --git a/tile_engine/ops/gemm/gemm_preshuffle/configs/default_ci_config.json b/tile_engine/ops/gemm/gemm_preshuffle/configs/default_ci_config.json index b32d42dabf..868debad3f 100644 --- a/tile_engine/ops/gemm/gemm_preshuffle/configs/default_ci_config.json +++ b/tile_engine/ops/gemm/gemm_preshuffle/configs/default_ci_config.json @@ -59,6 +59,7 @@ }, "epilogue": { "values": [ + "default", "cshuffle" ] }, diff --git a/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json b/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json index 2dd8230edc..38376a410b 100644 --- a/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json +++ b/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json @@ -49,7 +49,9 @@ "trait_config": { "pipeline": { "values": [ - "compv4" + "compv3", + "compv4", + "mem" ] }, "scheduler": { @@ -60,7 +62,8 @@ }, "epilogue": { "values": [ - "cshuffle" + "cshuffle", + "default" ] }, "pad_m": {