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

[Ayush Kumar Gaur]

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

### 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

>From 9982fb621ae7d873f51106f7293642a5314a1909 Mon Sep 17 00:00:00 2001
From: Ayush3941 <ayushkgaur1 at gmail.com>
Date: Sat, 14 Feb 2026 11:53:47 -0500
Subject: [PATCH] [mlir][DFA] Treat gpu.kernel functions as externally callable

---
 .../Analysis/DataFlow/DeadCodeAnalysis.cpp    | 10 +++++++
 mlir/test/Transforms/remove-dead-values.mlir  | 29 +++++++++++++++++++
 2 files changed, 39 insertions(+)

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