[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]

https://github.com/wang-y-z updated https://github.com/llvm/llvm-project/pull/101172

>From d18a9d711b060444c1e52370b56e182d3b7938df Mon Sep 17 00:00:00 2001
From: isaacw <isaacw at nvidia.com>
Date: Tue, 30 Jul 2024 04:18:57 -0700
Subject: [PATCH 1/4] init

---
 mlir/lib/IR/Operation.cpp                     |  2 +-
 .../Pass/scf2cf-print-liveness-crash.mlir     | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Pass/scf2cf-print-liveness-crash.mlir

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b51357198b1ca..88055ef7205db 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -391,7 +391,7 @@ bool Operation::isBeforeInBlock(Operation *other) {
   // parent.
   if (!block->isOpOrderValid()) {
     block->recomputeOpOrder();
-  } else {
+  } else if (block->getOperations().size() > 1) {
     // Update the order either operation if necessary.
     updateOrderIfNecessary();
     other->updateOrderIfNecessary();
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

>From e3f2a63d4a86bc2be321c52efa1c16658c3ace04 Mon Sep 17 00:00:00 2001
From: isaacw <isaacw at nvidia.com>
Date: Fri, 2 Aug 2024 00:06:50 -0700
Subject: [PATCH 2/4] fix addNodeToList logic

---
 mlir/lib/IR/Operation.cpp | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 88055ef7205db..235d6ac9873ab 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -391,7 +391,8 @@ bool Operation::isBeforeInBlock(Operation *other) {
   // parent.
   if (!block->isOpOrderValid()) {
     block->recomputeOpOrder();
-  } else if (block->getOperations().size() > 1) {
+  // } else if (!llvm::hasSingleElement(*block)) {
+  } else {
     // Update the order either operation if necessary.
     updateOrderIfNecessary();
     other->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.

>From 2b2c1dbaa3016dc6c2dea0c88b1665676f916b01 Mon Sep 17 00:00:00 2001
From: isaacw <isaacw at nvidia.com>
Date: Fri, 2 Aug 2024 00:08:01 -0700
Subject: [PATCH 3/4] fix

---
 mlir/lib/IR/Operation.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 235d6ac9873ab..3f4bcc1f16de8 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -391,7 +391,6 @@ 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();

>From 1a20afda12d631ded34dde2960ce574187c9f788 Mon Sep 17 00:00:00 2001
From: isaacw <isaacw at nvidia.com>
Date: Thu, 8 Aug 2024 21:40:33 -0700
Subject: [PATCH 4/4] fix updateOrderIfNecessary

---
 mlir/lib/IR/Operation.cpp | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 3f4bcc1f16de8..326d1dc745230 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -406,7 +406,7 @@ void Operation::updateOrderIfNecessary() {
   assert(block && "expected valid parent");
 
   // If the order is valid for this operation there is nothing to do.
-  if (hasValidOrder())
+  if (hasValidOrder() || llvm::hasSingleElement(*block))
     return;
   Operation *blockFront = &block->front();
   Operation *blockBack = &block->back();
@@ -502,14 +502,9 @@ 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!");
-  
-  Block *curParent = getContainingBlock();
-  op->block = curParent;
-
+  op->block = getContainingBlock();
   // 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.