Skip to content

[AIROCMLIR-1067] Extend problemKey with supportsSplitK #2453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
justinrosner wants to merge 6 commits into
base: develop
Choose a base branch
from
Open

[AIROCMLIR-1067] Extend problemKey with supportsSplitK #2453

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
40a3e20
Extend problemKey with supportsSplitK
justinrosner Aug 17, 2026
61a195f
Add perfRunner changes
justinrosner Aug 17, 2026
20c15ea
Minor tweaks
justinrosner Aug 17, 2026
6e01dbc
Attend to AI review comments
justinrosner Aug 17, 2026
3e9b4fd
Attend to more review comments
justinrosner Aug 17, 2026
7cd05e5
Yet another round of AI review comments
justinrosner Aug 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1738,6 +1738,24 @@ static LogicalResult getTuningProblemStr(rock::RockGemmWrapperInterface gemmIF,
// since it can store each field separately.
// Currently serialize the problem in MIOpenDriver command friendly format
LogicalResult getTuningProblemStr(ModuleOp mod, SmallVectorImpl<char> &out) {
auto serializeWithSplitKSupport = [&](auto tuningOp) -> LogicalResult {
if (failed(getTuningProblemStr(tuningOp, out)))
return failure();

// Legality must be judged on the function holding the op we just
// serialized. The ModuleOp overload always inspects the module's first
// function, which is not necessarily that one.
auto func = tuningOp->template getParentOfType<func::FuncOp>();
if (!func)
return failure();

llvm::raw_svector_ostream problemOS(out);
problemOS << " -supportsSplitK "
Comment thread
justinrosner marked this conversation as resolved.
<< (succeeded(rock::testFusionLegalitySplitK(func)) ? "true"
Comment thread
justinrosner marked this conversation as resolved.
: "false");
return success();
};

{
rock::RockGemmWrapperInterface gemmIF;
WalkResult findPrimary =
Expand All @@ -1746,7 +1764,7 @@ LogicalResult getTuningProblemStr(ModuleOp mod, SmallVectorImpl<char> &out) {
return WalkResult::interrupt();
});
if (findPrimary.wasInterrupted())
return getTuningProblemStr(gemmIF, out);
return serializeWithSplitKSupport(gemmIF);
}
{
rock::RockGemmGemmWrapperInterface gemmGemmOp;
Expand All @@ -1756,7 +1774,7 @@ LogicalResult getTuningProblemStr(ModuleOp mod, SmallVectorImpl<char> &out) {
return WalkResult::interrupt();
});
if (findGemmGemm.wasInterrupted())
return getTuningProblemStr(gemmGemmOp, out);
return serializeWithSplitKSupport(gemmGemmOp);
Comment thread
justinrosner marked this conversation as resolved.
}
return failure();
}
Expand Down
2 changes: 2 additions & 0 deletions mlir/test/fusion/fusability-conv-add-max.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
// CHECK-SPLITK: fusible:0
// RUN: rocmlir-gen -emit-module-fusibility-for=v3:16,16,4,16,16,1,1,1,2,1,1 - < %s | FileCheck %s --check-prefixes=CHECK-NONSPLITK
// CHECK-NONSPLITK: fusible:1
// RUN: rocmlir-gen --emit-tuning-key - < %s | FileCheck %s --check-prefix=CHECK-TUNING-KEY
// CHECK-TUNING-KEY: -supportsSplitK false
module {
func.func @mlir_convolution_add_relu(%arg0: memref<64x1x1x1xf32>, %arg1: memref<1x256x56x56xf32>, %arg2: memref<64x256x1x1xf32>, %arg3: memref<1x64x56x56xf32>) attributes {rock.enable_splitk_for_tuning, rock.kernel, mhal.arch = "amdgcn-amd-amdhsa:gfx90a:sramecc+:xnack-"} {
%cst = arith.constant 0.000000e+00 : f32
Expand Down
2 changes: 2 additions & 0 deletions mlir/test/fusion/fusability-dot-add.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
// CHECK-SPLITK: fusible:1
// RUN: rocmlir-gen -emit-module-fusibility-for=v3:16,16,4,16,16,1,1,1,2,1,1 - < %s | FileCheck %s --check-prefixes=CHECK-NONSPLITK
// CHECK-NONSPLITK: fusible:1
// RUN: rocmlir-gen --emit-tuning-key - < %s | FileCheck %s --check-prefix=CHECK-TUNING-KEY
// CHECK-TUNING-KEY: -supportsSplitK true
module {
func.func @mlir_dot_add(%arg0: memref<1x2x320xf32>, %arg1: memref<1x2x1280xf32>, %arg2: memref<1x1280x320xf32>, %arg3: memref<1x2x320xf32>) attributes {rock.enable_splitk_for_tuning, rock.kernel, mhal.arch = "amdgcn-amd-amdhsa:gfx90a:sramecc+:xnack-"} {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<1x2x320xf32>
Expand Down
2 changes: 2 additions & 0 deletions mlir/test/fusion/fusability-dot-mul.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
// CHECK-SPLITK: fusible:1
// RUN: rocmlir-gen -emit-module-fusibility-for=v3:16,16,4,16,16,1,1,1,2,1,1 - < %s | FileCheck %s --check-prefixes=CHECK-NONSPLITK
// CHECK-NONSPLITK: fusible:1
// RUN: rocmlir-gen --emit-tuning-key - < %s | FileCheck %s --check-prefix=CHECK-TUNING-KEY
// CHECK-TUNING-KEY: -supportsSplitK true
module {
func.func @mlir_dot_mul(%arg0: memref<1x2x320xf32>, %arg1: memref<1x2x1280xf32>, %arg2: memref<1x1280x320xf32>, %arg3: memref<1x2x320xf32>) attributes {rock.enable_splitk_for_tuning, rock.kernel, mhal.arch = "amdgcn-amd-amdhsa:gfx90a:sramecc+:xnack-"} {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<1x2x320xf32>
Expand Down
4 changes: 3 additions & 1 deletion mlir/test/rocmlir-gen/gemm-misc-options.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,9 @@
// RUN: rocmlir-gen --emit-tuning-key -p -t fp8_fp8 --arch gfx1201 | FileCheck %s --check-prefix=CONVOCPFP8
// CONVOCPFP8: amdgcn-amd-amdhsa:gfx1201 {{.*}} convfp8_fp8 -F 1 -f GNC01 -I NGC01 -O NGC01 -n 128 -c 8 -H 32 -W 32 -k 128 -y 3 -x 3 -p 0 -q 0 -u 1 -v 1 -l 1 -j 1 -g 1
// RUN: rocmlir-gen --arch gfx908 --operation gemm -p --emit-tuning-key | FileCheck %s --check-prefix=GEMM
// GEMM: amdgcn-amd-amdhsa:gfx908 {{.*}} -t f32 -out_datatype f32 -transA false -transB false -g 1 -m 1024 -n 512 -k 769
// GEMM: amdgcn-amd-amdhsa:gfx908 {{.*}} -t f32 -out_datatype f32 -transA false -transB false -g 1 -m 1024 -n 512 -k 769 -supportsSplitK true
Comment thread
justinrosner marked this conversation as resolved.
// RUN: rocmlir-gen --arch gfx908 --operation gemm -t i8 -p --emit-tuning-key | FileCheck %s --check-prefix=GEMM_I8
// GEMM_I8: amdgcn-amd-amdhsa:gfx908 {{.*}} -t i8 -out_datatype i32 -transA false -transB false -g 1 -m 1024 -n 512 -k 769 -supportsSplitK false{{$}}
// RUN: rocmlir-gen --emit-tuning-key -p -t fp8_fp8 --arch gfx950 | FileCheck %s --check-prefix=CONVOCPFP8_GFX950
// CONVOCPFP8_GFX950: amdgcn-amd-amdhsa:gfx950 {{.*}} convfp8_fp8 -F 1 -f GNC01 -I NGC01 -O NGC01 -n 128 -c 8 -H 32 -W 32 -k 128 -y 3 -x 3 -p 0 -q 0 -u 1 -v 1 -l 1 -j 1 -g 1
// RUN: rocmlir-gen --arch gfx942 --operation gemm -p --num_cu 40 --num_chiplets 20 | FileCheck %s --check-prefix=NUM_CHIPLETS
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/rocmlir-gen/problem-key.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
// RUN: rocmlir-gen --arch gfx942 --operation attention -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f32 -g 1 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_1
// CHECK_1: -t f32 -transQ false -transK false -transV false -transO false -causal false -return_lse false -split_kv 1 -num_heads_q 1 -num_heads_kv 1 -g 1 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false
// CHECK_1: -t f32 -transQ false -transK false -transV false -transO false -causal false -return_lse false -split_kv 1 -num_heads_q 1 -num_heads_kv 1 -g 1 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false -supportsSplitK true
// RUN: rocmlir-gen --arch gfx942 --operation attention -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f16 -g 4 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_2
// CHECK_2: -t f16 -transQ false -transK false -transV false -transO false -causal false -return_lse false -split_kv 1 -num_heads_q 1 -num_heads_kv 1 -g 4 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false
// RUN: rocmlir-gen --arch gfx942 --operation attention -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t i8 -g 8 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_3
Expand Down
Loading
Loading