[Mlir-commits] [mlir] [mlir] Report block argument removal from runRegionDCE (PR #232211)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hideto Ueno (uenoku)

<details>
<summary>Changes</summary>

runRegionDCE reports whether it changed the IR, but deleteDeadness did not include
dead non-entry block arguments in its change flag.

This caused GreedyRewriterDriver to not report a change even when DCE happened.

The line `EXPECT_TRUE(succeeded(runRegionDCE(...)))` in the unittest would fail without the RegionUtils change.

Assisted-by: gpt 6 sol



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


3 Files Affected:

- (modified) mlir/lib/Transforms/Utils/RegionUtils.cpp (+5-2) 
- (modified) mlir/unittests/Transforms/CMakeLists.txt (+1) 
- (added) mlir/unittests/Transforms/RegionUtils.cpp (+78) 


``````````diff
diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index cc90df06918000..9cc20b54f8d849 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -473,8 +473,11 @@ static LogicalResult deleteDeadness(RewriterBase &rewriter,
     // The entry block has an unknown contract with their enclosing block, so
     // skip it.
     for (Block &block : llvm::drop_begin(region.getBlocks(), 1)) {
-      block.eraseArguments(
-          [&](BlockArgument arg) { return !liveMap.wasProvenLive(arg); });
+      block.eraseArguments([&](BlockArgument arg) {
+        bool isDead = !liveMap.wasProvenLive(arg);
+        erasedAnything |= isDead;
+        return isDead;
+      });
     }
   }
   return success(erasedAnything);
diff --git a/mlir/unittests/Transforms/CMakeLists.txt b/mlir/unittests/Transforms/CMakeLists.txt
index dc5920087b505e..c0129be9b13fd3 100644
--- a/mlir/unittests/Transforms/CMakeLists.txt
+++ b/mlir/unittests/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
 add_mlir_unittest(MLIRTransformsTests
   Canonicalizer.cpp
   DialectConversion.cpp
+  RegionUtils.cpp
 )
 mlir_target_link_libraries(MLIRTransformsTests
   PRIVATE
diff --git a/mlir/unittests/Transforms/RegionUtils.cpp b/mlir/unittests/Transforms/RegionUtils.cpp
new file mode 100644
index 00000000000000..c7f6286c9dc6a6
--- /dev/null
+++ b/mlir/unittests/Transforms/RegionUtils.cpp
@@ -0,0 +1,78 @@
+//===- RegionUtils.cpp - Region utility unit tests ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/RegionUtils.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Parser/Parser.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+class TestBranchOp
+    : public Op<TestBranchOp, OpTrait::VariadicOperands, OpTrait::ZeroResults,
+                OpTrait::OneSuccessor, OpTrait::IsTerminator,
+                BranchOpInterface::Trait> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestBranchOp)
+
+  using Op::Op;
+
+  static StringRef getOperationName() { return "test.br"; }
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
+
+  SuccessorOperands getSuccessorOperands(unsigned index) {
+    assert(index == 0 && "invalid successor index");
+    return SuccessorOperands(MutableOperandRange(getOperation()));
+  }
+};
+
+class TestDialect : public Dialect {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDialect)
+
+  static StringRef getDialectNamespace() { return "test"; }
+
+  explicit TestDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context, TypeID::get<TestDialect>()) {
+    addOperations<TestBranchOp>();
+    allowUnknownOperations();
+  }
+};
+
+TEST(RegionUtilsTest, ReportsDeadBlockArgumentRemoval) {
+  MLIRContext context;
+  context.getOrLoadDialect<TestDialect>();
+
+  auto module = parseSourceString<ModuleOp>(R"mlir(
+    "test.container"() ({
+    ^bb0(%arg0: i32):
+      "test.br"(%arg0)[^bb1] : (i32) -> ()
+    ^bb1(%unused: i32):
+      "test.br"(%unused)[^bb1] : (i32) -> ()
+    }) : () -> ()
+  )mlir",
+                                            &context);
+  ASSERT_TRUE(module);
+
+  Operation *container = &module->getBody()->front();
+  Region &region = container->getRegion(0);
+  ASSERT_EQ(region.getBlocks().size(), 2u);
+  ASSERT_EQ(region.back().getNumArguments(), 1u);
+
+  IRRewriter rewriter(&context);
+  EXPECT_TRUE(succeeded(runRegionDCE(rewriter, container->getRegions())));
+  EXPECT_EQ(region.back().getNumArguments(), 0u);
+  EXPECT_EQ(region.front().getTerminator()->getNumOperands(), 0u);
+  EXPECT_EQ(region.back().getTerminator()->getNumOperands(), 0u);
+}
+} // namespace

``````````

</details>


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