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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hamza Qureshi (hamzaqureshi5)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/221462.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp (+7-1) 
- (modified) mlir/test/Dialect/GPU/sink-ops.mlir (+29) 


``````````diff
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
+}

``````````

</details>


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