[Mlir-commits] [mlir] [MLIR][Operation] Fix `isBeforeInBlock` crash bug mentioned in https://github.com/llvm/llvm-project/issues/60909 (PR #101172)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: wang-y-z (wang-y-z)

<details>
<summary>Changes</summary>

# summary
This MR fix `isBeforeInBlock` crash bug mentioned in https://github.com/llvm/llvm-project/issues/60909

# Trigger condition
1. A block only have one operation.
2. `block->isOpOrderValid()` is true.
3. call: `op->isBeforeInBlock(op)`, compared with op itself.

Will crash on `assert(blockFront != blockBack && "expected more than one operation");`

# Case study
Simplified repro case in `mlir/test/Pass/scf2cf-print-liveness-crash.mlir`
When put `-convert-scf-to-cf -test-print-liveness` together in one cmd line, the first pass will work normally and crash on the second pass.
Details please refer  https://github.com/llvm/llvm-project/issues/60909

# Solutions
option1. in `isBeforeInBlock`, check if block only have one operation before step into `updateOrderIfNecessary`, if have only one, it must return false
option2. in `isBeforeInBlock`, check if `this == other`, if true return false

both OK for me.

# Other
> Why the `block->isOpOrderValid()` will be true? Why not find the root cause of this phenomenon? 

I think `isBeforeInBlock` as a fundamental API in operation.cpp should should not make too many assumptions about the state of `block->isOpOrderValid()`

Let's think about this in IR:
```
op1 {
^bb0:
   op2
}
```
When `block->isOpOrderValid() is true`, op2->isBeforeInBlock(op2) will crash on `updateOrderIfNecessary`
When `block->isOpOrderValid() is false`, `op2->isBeforeInBlock(op2)` will step into `recomputeOpOrder` and make `parentValidOpOrderPair.setInt(true);` , if user call `op2->isBeforeInBlock(op2)` again, it will crash on `updateOrderIfNecessary` too
So, I think fix on `isBeforeInBlock` is closer to the essence of the issue. 

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


2 Files Affected:

- (modified) mlir/lib/IR/Operation.cpp (+6-1) 
- (added) mlir/test/Pass/scf2cf-print-liveness-crash.mlir (+19) 


``````````diff
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b51357198b1ca..235d6ac9873ab 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -391,6 +391,7 @@ bool Operation::isBeforeInBlock(Operation *other) {
   // parent.
   if (!block->isOpOrderValid()) {
     block->recomputeOpOrder();
+  // } else if (!llvm::hasSingleElement(*block)) {
   } else {
     // Update the order either operation if necessary.
     updateOrderIfNecessary();
@@ -502,10 +503,14 @@ Block *llvm::ilist_traits<::mlir::Operation>::getContainingBlock() {
 /// keep the block pointer up to date.
 void llvm::ilist_traits<::mlir::Operation>::addNodeToList(Operation *op) {
   assert(!op->getBlock() && "already in an operation block!");
-  op->block = getContainingBlock();
+  
+  Block *curParent = getContainingBlock();
+  op->block = curParent;
 
   // Invalidate the order on the operation.
   op->orderIndex = Operation::kInvalidOrderIdx;
+  // Invalidate the ordering of the parent block.
+  curParent->invalidateOpOrder();
 }
 
 /// This is a trait method invoked when an operation is removed from a block.
diff --git a/mlir/test/Pass/scf2cf-print-liveness-crash.mlir b/mlir/test/Pass/scf2cf-print-liveness-crash.mlir
new file mode 100644
index 0000000000000..665e085cfedee
--- /dev/null
+++ b/mlir/test/Pass/scf2cf-print-liveness-crash.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(convert-scf-to-cf), func.func(test-print-liveness))"
+
+module {
+  func.func @for_if_for(%arg0: index, %arg1: index, %arg2: index, %arg3: i1) {
+    %cst = arith.constant dense<0.000000e+00> : tensor<128x32xf16>
+    %0 = scf.for %arg4 = %arg0 to %arg1 step %arg2 iter_args(%arg5 = %cst) -> (tensor<128x32xf16>) {
+      %1 = scf.if %arg3 -> (tensor<128x32xf16>) {
+        scf.yield %arg5 : tensor<128x32xf16>
+      } else {
+        %2 = scf.for %arg6 = %arg0 to %arg1 step %arg2 iter_args(%arg7 = %arg5) -> (tensor<128x32xf16>) {
+          scf.yield %arg7 : tensor<128x32xf16>
+        }
+        scf.yield %2 : tensor<128x32xf16>
+      }
+      scf.yield %1 : tensor<128x32xf16>
+    }
+    return
+  }
+}
\ No newline at end of file

``````````

</details>


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