[Mlir-commits] [mlir] [mlir][DFA] Treat gpu.kernel functions as externally callable (PR #181491)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Ayush Kumar Gaur (Ayush3941)

<details>
<summary>Changes</summary>

### What

Fix --remove-dead-values crash on GPU kernels: private gpu.func with kernel could be treated as non-entry, leaving users with null operands and tripping matchPattern assert.

### Why

DeadCodeAnalysis didn’t recognize gpu.func kernels as externally callable entry points, so RDV erased “dead” producers, dropAllUses() nulled operands, and later folding/patterns crashed.

### Fix

Mark gpu.func with isKernel() as having unknown predecessors (externally callable) in DeadCodeAnalysis so the kernel body stays executable/live. Add a regression test covering RDV + NVVM lowering.

Fixes #<!-- -->181303

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


2 Files Affected:

- (modified) mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp (+10) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+29) 


``````````diff
diff --git a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
index 936b0c678f20c..60537c244db59 100644
--- a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
@@ -187,6 +187,16 @@ void DeadCodeAnalysis::initializeSymbolCallables(Operation *top) {
       if (!symbol)
         continue;
 
+      // Treat GPU kernels as externally callable entry points.
+      Operation *op = callable.getOperation();
+      if (op->getName().getStringRef() == "gpu.func" &&
+          op->hasAttr("gpu.kernel")) {
+        auto *state =
+            getOrCreate<PredecessorState>(getProgramPointAfter(callable));
+        propagateIfChanged(state, state->setHasUnknownPredecessors());
+        continue;
+      }
+
       // Public symbol callables or those for which we can't see all uses have
       // potentially unknown callsites.
       if (symbol.isPublic() || (!allUsesVisible && symbol.isNested())) {
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index ae83eac0c376f..8ac7be559ca84 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -796,3 +796,32 @@ func.func @scf_while_dead_iter_args() -> i32 {
   }
   return %result#0 : i32
 }
+
+// -----
+
+  // CHECK-LABEL: gpu.func @k
+  // CHECK: scf.if
+  // CHECK: gpu.barrier
+  // CHECK-NOT: arith.addi
+  // CHECK: gpu.return
+
+  // NVVM-LABEL: llvm.func @k
+  // NVVM: scf.if
+  // NVVM: nvvm.barrier0
+  // NVVM-NOT: arith.addi
+  // NVVM: llvm.return
+  module {
+    gpu.module @m {
+      gpu.func @k(%data: memref<?xi32>) kernel attributes {sym_visibility = "private"} {
+        %c0 = arith.constant 0 : index
+        %len = memref.dim %data, %c0 : memref<?xi32>
+        %p = arith.cmpi ult, %c0, %len : index
+        scf.if %p {
+          gpu.barrier
+        } else {
+          arith.addi %c0, %c0 : index
+        }
+        gpu.return
+      }
+    }
+  }

``````````

</details>


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