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

[Adam Paszke]

https://github.com/apaszke created https://github.com/llvm/llvm-project/pull/176894

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

>From bef18ff954e893586a92be04b977600047631327 Mon Sep 17 00:00:00 2001
From: Adam Paszke <adam.paszke at gmail.com>
Date: Tue, 20 Jan 2026 11:01:44 +0000
Subject: [PATCH] [MLIR][GPU] Make sure to propagate known cluster sizes in
 kernel outlining

Otherwise, the changes from #174404 don't kick in.
---
 mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp | 4 ++++
 mlir/test/Dialect/GPU/outlining.mlir                | 1 +
 2 files changed, 5 insertions(+)

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