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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lijinpei-amd

<details>
<summary>Changes</summary>

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

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


3 Files Affected:

- (modified) mlir/test/Analysis/DataFlow/test-staged-analyses.mlir (+8) 
- (modified) mlir/test/Analysis/test-foo-analysis.mlir (+10) 
- (modified) mlir/test/lib/Analysis/TestDataFlowFramework.cpp (+4) 


``````````diff
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;
 

``````````

</details>


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