From ee5dcdc275b136172a5bfa85a098ae506a836c85 Mon Sep 17 00:00:00 2001
From: Menooker <myjisgreat@live.cn>
Date: Sun, 22 Jun 2025 22:03:36 +0800
Subject: [PATCH] [mlir] fix assertion failure in remove-dead-values (#144849)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Simple IR patterns will trigger assertion error:

```
  func.func @test_zero_operands(%I: memref<10xindex>, %I2: memref<10xf32>) {
    %v0 = arith.constant 0 : index
    %result = memref.alloca_scope -> index {
      %c = arith.addi %v0, %v0 : index
      memref.store %c, %I[%v0] : memref<10xindex>
      memref.alloca_scope.return %c: index
    }
    func.return
  }
```

with error: `mlir/include/mlir/IR/Operation.h:988:
mlir::detail::OperandStorage& mlir::Operation::getOperandStorage():
Assertion `hasOperandStorage && "expected operation to have operand
storage"' failed.`

This PR will fix this issue.

---------

Co-authored-by: Andrzej Warzyński <andrzej.warzynski@gmail.com>
Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     |  5 +++--
 mlir/test/Transforms/remove-dead-values.mlir | 23 ++++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index ad21ce8f1870..0a02a01977ca 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -750,8 +750,9 @@ static void cleanUpDeadVals(RDVFinalCleanupList &list) {
   }
 
   // 4. Operands
-  for (auto &o : list.operands) {
-    o.op->eraseOperands(o.nonLive);
+  for (OperationToCleanup &o : list.operands) {
+    if (o.op->getNumOperands() > 0)
+      o.op->eraseOperands(o.nonLive);
   }
 
   // 5. Results
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index e55a9160f5b3..3af95db3c0e2 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -511,6 +511,28 @@ module {
 // CHECK: linalg.yield %[[yield]] : f32
 // CHECK-NOT: arith.subf
 
+
+// -----
+
+// check that ops with zero operands are correctly handled
+
+module {
+  func.func @test_zero_operands(%I: memref<10xindex>, %I2: memref<10xf32>) {
+    %v0 = arith.constant 0 : index
+    %result = memref.alloca_scope -> index {
+      %c = arith.addi %v0, %v0 : index
+      memref.store %c, %I[%v0] : memref<10xindex>
+      memref.alloca_scope.return %c: index
+    }
+    func.return
+  }
+}
+
+// CHECK-LABEL: func @test_zero_operands
+// CHECK: memref.alloca_scope
+// CHECK: memref.store
+// CHECK-NOT: memref.alloca_scope.return
+
 // -----
 
 // CHECK-LABEL: func.func @test_atomic_yield
@@ -525,3 +547,4 @@ func.func @test_atomic_yield(%I: memref<10xf32>, %idx : index) {
   }
   func.return
 }
+