[Mlir-commits] [mlir] 064cbec - [MLIR][GPU] Make sure to propagate known cluster sizes in kernel outlining (#176894)

[llvmlistbot at llvm.org]

Author: Adam Paszke
Date: 2026-01-21T11:10:40+01:00
New Revision: 064cbec2c1997e76faee2a1468409ce427bacc7a

URL: https://github.com/llvm/llvm-project/commit/064cbec2c1997e76faee2a1468409ce427bacc7a
DIFF: https://github.com/llvm/llvm-project/commit/064cbec2c1997e76faee2a1468409ce427bacc7a.diff

LOG: [MLIR][GPU] Make sure to propagate known cluster sizes in kernel outlining (#176894)

Otherwise, the changes from #174404 don't kick in.

Added: 
    

Modified: 
    mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
    mlir/test/Dialect/GPU/outlining.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 97adad64d78c4..adfe365ad9039 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -211,6 +211,10 @@ static gpu::GPUFuncOp outlineKernelFuncImpl(gpu::LaunchOp launchOp,
   if (auto gridBounds =
           maybeConstantDimsAttr(launchOp.getGridSizeOperandValues()))
     outlinedFunc.setKnownGridSizeAttr(gridBounds);
+  if (auto clusterSize = launchOp.getClusterSizeOperandValues()) {
+    if (auto clusterBounds = maybeConstantDimsAttr(*clusterSize))
+      outlinedFunc.setKnownClusterSizeAttr(clusterBounds);
+  }
 
   IRMapping map;
 

diff  --git a/mlir/test/Dialect/GPU/outlining.mlir b/mlir/test/Dialect/GPU/outlining.mlir
index e14521ab9bb5c..2e5d12936ae65 100644
--- a/mlir/test/Dialect/GPU/outlining.mlir
+++ b/mlir/test/Dialect/GPU/outlining.mlir
@@ -485,6 +485,7 @@ func.func @launch_cluster() {
 // CHECK-NEXT: gpu.func @launch_cluster_kernel
 // CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)
 // CHECK-SAME: known_block_size = array<i32: 20, 24, 28>
+// CHECK-SAME: known_cluster_size = array<i32: 1, 2, 1>
 // CHECK-SAME: known_grid_size = array<i32: 8, 12, 16>
 // CHECK-NEXT: %[[BID:.*]] = gpu.block_id x
 // CHECK-NEXT: = gpu.block_id y