[Mlir-commits] [mlir] [mlir][gpu] Sink integer arithmetic into gpu.launch (PR #221462)

[Hamza Qureshi]

https://github.com/hamzaqureshi5 created https://github.com/llvm/llvm-project/pull/221462

isLikelyAnIndexComputation() accepted constants, memref.dim, arith.select and arith.cmpi, but not integer arithmetic. An index computation such as arith.addi feeding a gpu.launch was left outside as kernel argument, even though it is cheap to recompute inside the kernel. Add arith.addi, arith.subi and arith.muli. 

Fixes #221396

>From bf63a5df6e90a50a54ba98b4a49d0b235c947ce0 Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Sat, 5 Sep 2026 21:07:03 +0500
Subject: [PATCH] [mlir][gpu] Sink integer arithmetic into gpu.launch

isLikelyAnIndexComputation() accepted constants, memref.dim, arith.select
and arith.cmpi, but no integer arithmetic. An index computation such as
arith.addi feeding a gpu.launch was therefore left outside it and became an
extra kernel argument after outlining, even though it is cheap to recompute
inside the kernel.

Add arith.addi, arith.subi and arith.muli. They are side-effect free and in
the same cost class as the comparisons and selects that are already sunk,
and memref.dim, which is also on the list, can be a load for a dynamic
dimension. Division is left out as materially more expensive to recompute
per thread.

Sinking is still only performed when it does not introduce new kernel
arguments, which extractBeneficiaryOps() already ensures.

Fixes #221396
---
 .../GPU/Transforms/KernelOutlining.cpp        |  8 ++++-
 mlir/test/Dialect/GPU/sink-ops.mlir           | 29 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 0ee3d425f66be..79c623edb7456 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -72,9 +72,15 @@ static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody,
 /// Identifies operations that are beneficial to sink into kernels. These
 /// operations may not have side-effects, as otherwise sinking (and hence
 /// duplicating them) is not legal.
+///
+/// Sinking duplicates an operation into every thread, so only operations that
+/// are cheap to recompute are considered. Simple integer arithmetic is in the
+/// same cost class as the comparisons and selects already handled here, and
+/// sinking it avoids passing the result as an additional kernel argument.
 static bool isLikelyAnIndexComputation(Operation *op) {
   return matchPattern(op, m_Constant()) ||
-         isa<memref::DimOp, arith::SelectOp, arith::CmpIOp>(op);
+         isa<memref::DimOp, arith::SelectOp, arith::CmpIOp, arith::AddIOp,
+             arith::SubIOp, arith::MulIOp>(op);
 }
 
 /// For a given operation `op`, computes whether it is beneficial to sink the
diff --git a/mlir/test/Dialect/GPU/sink-ops.mlir b/mlir/test/Dialect/GPU/sink-ops.mlir
index 8f5924f2e3265..41c251c13dc88 100644
--- a/mlir/test/Dialect/GPU/sink-ops.mlir
+++ b/mlir/test/Dialect/GPU/sink-ops.mlir
@@ -98,3 +98,32 @@ func.func @multiple_uses2(%arg0 : memref<*xf32>) {
   }
   return
 }
+
+// -----
+
+// Integer arithmetic feeding the launch is recomputed inside it, so that the
+// value does not have to be passed in as a kernel argument.
+
+// CHECK-LABEL: @index_arithmetic
+// CHECK-SAME: %[[ARG0:.*]]: index
+func.func @index_arithmetic(%arg0: index) {
+  %c1 = arith.constant 1 : index
+  %sum = arith.addi %arg0, %c1 : index
+  %diff = arith.subi %arg0, %c1 : index
+  %prod = arith.muli %arg0, %c1 : index
+  // CHECK: gpu.launch blocks
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,
+                                       %grid_z = %c1)
+             threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,
+                                        %block_z = %c1) {
+    // CHECK: %[[C1:.*]] = arith.constant 1 : index
+    // CHECK-NEXT: %[[SUM:.*]] = arith.addi %[[ARG0]], %[[C1]]
+    // CHECK-NEXT: %[[DIFF:.*]] = arith.subi %[[ARG0]], %[[C1]]
+    // CHECK-NEXT: %[[PROD:.*]] = arith.muli %[[ARG0]], %[[C1]]
+    // CHECK-NEXT: "use"(%[[ARG0]], %[[SUM]], %[[DIFF]], %[[PROD]])
+    // CHECK-NEXT: gpu.terminator
+    "use"(%arg0, %sum, %diff, %prod) : (index, index, index, index) -> ()
+    gpu.terminator
+  }
+  return
+}