[Mlir-commits] [mlir] [mlir] Fix crash in dataflow test passes when func op carries a tag (PR #203276)

[llvmlistbot at llvm.org]

https://github.com/lijinpei-amd created https://github.com/llvm/llvm-project/pull/203276

The `test-foo-analysis` and `test-staged-analyses` passes walk the function and look up the analysis state at every tagged op. The walk also visits the anchor func op itself, but the analyses only model program points within the function body, so the state at the func op's program point is null. When a `tag` attribute is placed on the func op, the lookup returns null and the state assertion fails.

Skip the anchor op in both walks.

Fixes #203221
Fixes #203222

>From 9bd78524dd8fae4d94c139cc2ca10bf373063f36 Mon Sep 17 00:00:00 2001
From: Li Jinpei <jinpli at amd.com>
Date: Thu, 11 Jun 2026 22:02:43 +0800
Subject: [PATCH] [mlir] Fix crash in dataflow test passes when func op carries
 a tag

The `test-foo-analysis` and `test-staged-analyses` passes walk the function
and look up the analysis state at every tagged op. The walk also visits the
anchor func op itself, but the analyses only model program points within the
function body, so the state at the func op's program point is null. When a
`tag` attribute is placed on the func op, the lookup returns null and the
state assertion fails.

Skip the anchor op in both walks.

Fixes #203221
Fixes #203222

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
 mlir/test/Analysis/DataFlow/test-staged-analyses.mlir |  8 ++++++++
 mlir/test/Analysis/test-foo-analysis.mlir             | 10 ++++++++++
 mlir/test/lib/Analysis/TestDataFlowFramework.cpp      |  4 ++++
 3 files changed, 22 insertions(+)

diff --git a/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir b/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
index 5da7ad8bf1274..861cac20811f2 100644
--- a/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
+++ b/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
@@ -48,3 +48,11 @@ func.func @requires_staged_bar() {
   // CHECK: "test.branch"()[^bb{{[0-9]+}}] {bar_state = true, foo = 2 : ui64, foo_state = 5 : i64, tag = "annotate"} : () -> ()
   "test.branch"() [^bb1] {tag = "annotate", foo = 2 : ui64} : () -> ()
 }
+
+// CHECK-LABEL: func.func @tag_on_func()
+// CHECK-NOT: foo_state
+func.func @tag_on_func() attributes {tag = "annotate"} {
+  // CHECK: "test.foo"() {bar_state = true, foo = 1 : ui64, foo_state = 1 : i64, tag = "annotate"} : () -> ()
+  "test.foo"() {tag = "annotate", foo = 1 : ui64} : () -> ()
+  return
+}
diff --git a/mlir/test/Analysis/test-foo-analysis.mlir b/mlir/test/Analysis/test-foo-analysis.mlir
index edd153d433878..782cf478a20ba 100644
--- a/mlir/test/Analysis/test-foo-analysis.mlir
+++ b/mlir/test/Analysis/test-foo-analysis.mlir
@@ -93,3 +93,13 @@ func.func @test_double_loop() -> () {
   "test.foo"() {tag = "end"} : () -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: function: @test_tag_on_func
+// CHECK-NOT: on_func ->
+func.func @test_tag_on_func() attributes {tag = "on_func"} {
+  // CHECK: a -> 0
+  "test.foo"() {tag = "a"} : () -> ()
+  return
+}
diff --git a/mlir/test/lib/Analysis/TestDataFlowFramework.cpp b/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
index 9af7e205aaee9..54b3ae02bdd1f 100644
--- a/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
+++ b/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
@@ -319,6 +319,8 @@ void TestFooAnalysisPass::runOnOperation() {
   os << "function: @" << func.getSymName() << "\n";
 
   func.walk([&](Operation *op) {
+    if (op == func.getOperation())
+      return;
     auto tag = op->getAttrOfType<StringAttr>(kTagAttrName);
     if (!tag)
       return;
@@ -342,6 +344,8 @@ void TestStagedAnalysesPass::runOnOperation() {
     return signalPassFailure();
 
   func.walk([&](Operation *op) {
+    if (op == func.getOperation())
+      return;
     if (!op->hasAttr(kTagAttrName))
       return;