[mlir] Fix remove-dead-values pass throws error when module has a name

[Wheest]

This fixes the error described in #107870, which I also independently experienced.

This implements the logic change suggested by @joker-eph.

This requires replacing one of the tests:

// The IR remains untouched because of the presence of a non-function-like
// symbol op (module @dont_touch_unacceptable_ir).
//
// expected-error @+1 {{cannot optimize an IR with non-function symbol ops, non-call symbol user ops or branch ops}}
module @dont_touch_unacceptable_ir {
  func.func @has_cleanable_simple_op(%arg0 : i32) {
    %non_live = arith.addi %arg0, %arg0 : i32
    return
  }
}

Since this asserts that named modules are not allowed, which is the behaviour we are trying to fix. It is not clear to me why this assertion would be needed.

Anecdotally, I experienced the error going from StableHLO (which has named modules), and lowering.
Fixes #107870.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Perry Gibson (Wheest)

Changes

This fixes the error described in #107870, which I also independently experienced.

This implements the logic change suggested by @joker-eph.

This requires replacing one of the tests:

// The IR remains untouched because of the presence of a non-function-like
// symbol op (module @<!-- -->dont_touch_unacceptable_ir).
//
// expected-error @+1 {{cannot optimize an IR with non-function symbol ops, non-call symbol user ops or branch ops}}
module @<!-- -->dont_touch_unacceptable_ir {
  func.func @<!-- -->has_cleanable_simple_op(%arg0 : i32) {
    %non_live = arith.addi %arg0, %arg0 : i32
    return
  }
}

Since this asserts that named modules are not allowed, which is the behaviour we are trying to fix. It is not clear to me why this assertion would be needed.

Anecdotally, I experienced the error going from StableHLO (which has named modules), and lowering.

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+1)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+8-7)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 055256903a1522..539797151dbbb7 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -576,6 +576,7 @@ void RemoveDeadValues::runOnOperation() {
   // all symbol ops present in the IR are function-like, and all symbol user ops
   // present in the IR are call-like.
   WalkResult acceptableIR = module->walk([&](Operation *op) {
+    if (op == module) return  WalkResult::advance();
     if (isa<BranchOpInterface>(op) ||
         (isa<SymbolOpInterface>(op) && !isa<FunctionOpInterface>(op)) ||
         (isa<SymbolUserOpInterface>(op) && !isa<CallOpInterface>(op))) {
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 69426fdb620832..4622b6afa942de 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -1,13 +1,14 @@
 // RUN: mlir-opt %s -remove-dead-values -split-input-file -verify-diagnostics | FileCheck %s
 
-// The IR remains untouched because of the presence of a non-function-like
-// symbol op (module @dont_touch_unacceptable_ir).
+// -----
+
+// Dead values are removed from the IR even if the module has a name
 //
-// expected-error @+1 {{cannot optimize an IR with non-function symbol ops, non-call symbol user ops or branch ops}}
-module @dont_touch_unacceptable_ir {
-  func.func @has_cleanable_simple_op(%arg0 : i32) {
-    %non_live = arith.addi %arg0, %arg0 : i32
-    return
+module @named_module_acceptable {
+  func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
+    %0 = tensor.empty() : tensor<10xbf16>
+    // CHECK-NOT: %[[C:.*]] = tensor.empty[[C:.*]]
+    return %arg0 : tensor<10xf32>
   }
 }
 

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[joker-eph]

Do we have still a test with a non-function symbol op inside the module?

[Wheest]

The 2nd test checks for isa<BranchOpInterface>(op) and raises the error, but no we don't have a test for the "non-function symbol op inside the module" condition (I don't think we had one for this case before).

In terms of non-function symbol ops, would a global constant make sense? I've added this test which meets the requirements:

module {
// expected-error @+1 {{cannot optimize an IR with non-function symbol ops, non-call symbol user ops or branch ops}}
  memref.global "private" constant @__dont_touch_unacceptable_ir : memref<i32> = dense<0>
  func.func @main(%arg0: i32) -> i32 {
    return %arg0 : i32
  }
}

[joker-eph]

Looks good, with one thing on the test coverage.

[joker-eph]

LG, thanks

[CoTinker]

LGTM.