[Mlir-commits] [mlir] 850b2c6 - [mlir] Fix `Region`s `takeBody` method if the region is not empty

[Markus Böck]

Author: Markus Böck
Date: 2022-04-21T15:32:59+02:00
New Revision: 850b2c6b3c73a48cce05c163c49fcea89491a5e1

URL: https://github.com/llvm/llvm-project/commit/850b2c6b3c73a48cce05c163c49fcea89491a5e1
DIFF: https://github.com/llvm/llvm-project/commit/850b2c6b3c73a48cce05c163c49fcea89491a5e1.diff

LOG: [mlir] Fix `Region`s `takeBody` method if the region is not empty

The current implementation of takeBody first clears the Region, before then taking ownership of the blocks of the other regions. The issue here however, is that when clearing the region, it does not take into account references of operations to each other. In particular, blocks are deleted from front to back, and operations within a block are very likely to be deleted despite still having uses, causing an assertion to trigger [0].

This patch fixes that issue by simply calling dropAllReferences()before clearing the blocks.

[0] https://github.com/llvm/llvm-project/blob/9a8bb4bc635de9d56706262083c15eb1e0cf3e87/mlir/lib/IR/Operation.cpp#L154

Differential Revision: https://reviews.llvm.org/D123913

Added: 
    mlir/test/IR/test-take-body.mlir
    mlir/test/lib/IR/TestRegions.cpp

Modified: 
    mlir/include/mlir/IR/Region.h
    mlir/test/lib/IR/CMakeLists.txt
    mlir/tools/mlir-opt/mlir-opt.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index 671a981a54502..6e4b141cdc510 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -240,6 +240,7 @@ class Region {
   /// Takes body of another region (that region will have no body after this
   /// operation completes).  The current body of this region is cleared.
   void takeBody(Region &other) {
+    dropAllReferences();
     blocks.clear();
     blocks.splice(blocks.end(), other.getBlocks());
   }

diff  --git a/mlir/test/IR/test-take-body.mlir b/mlir/test/IR/test-take-body.mlir
new file mode 100644
index 0000000000000..d965bb93a926b
--- /dev/null
+++ b/mlir/test/IR/test-take-body.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s --test-take-body -split-input-file
+
+func @foo() {
+  %0 = "test.foo"() : () -> i32
+  cf.br ^header
+	
+^header:
+  cf.br ^body
+
+^body:
+  "test.use"(%0) : (i32) -> ()
+  cf.br ^header
+}
+
+func private @bar() {
+  return
+}
+
+// CHECK-LABEL: func @foo
+// CHECK-NEXT: return
+
+// CHECK-LABEL: func private @bar()
+// CHECK-NOT: {

diff  --git a/mlir/test/lib/IR/CMakeLists.txt b/mlir/test/lib/IR/CMakeLists.txt
index a8b54d4d2d4f5..8b519538719e5 100644
--- a/mlir/test/lib/IR/CMakeLists.txt
+++ b/mlir/test/lib/IR/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_library(MLIRTestIR
   TestSideEffects.cpp
   TestSlicing.cpp
   TestSymbolUses.cpp
+  TestRegions.cpp
   TestTypes.cpp
   TestVisitors.cpp
   TestVisitorsGeneric.cpp

diff  --git a/mlir/test/lib/IR/TestRegions.cpp b/mlir/test/lib/IR/TestRegions.cpp
new file mode 100644
index 0000000000000..e850691631be2
--- /dev/null
+++ b/mlir/test/lib/IR/TestRegions.cpp
@@ -0,0 +1,45 @@
+//===- TestRegions.cpp - Pass to test Region's methods --------------------===//
+//
+// 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 "TestDialect.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// This is a test pass that tests Region's takeBody method by making the first
+/// function take the body of the second.
+struct TakeBodyPass
+    : public PassWrapper<TakeBodyPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TakeBodyPass)
+
+  StringRef getArgument() const final { return "test-take-body"; }
+  StringRef getDescription() const final { return "Test Region's takeBody"; }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+
+    SmallVector<func::FuncOp> functions =
+        llvm::to_vector(module.getOps<func::FuncOp>());
+    if (functions.size() != 2) {
+      module.emitError("Expected only two functions in test");
+      signalPassFailure();
+      return;
+    }
+
+    functions[0].getBody().takeBody(functions[1].getBody());
+  }
+};
+
+} // namespace
+
+namespace mlir {
+void registerRegionTestPasses() { PassRegistration<TakeBodyPass>(); }
+} // namespace mlir

diff  --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 6a82811ee85ab..83a1c30af56de 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -37,6 +37,7 @@ void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
+void registerRegionTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
@@ -128,6 +129,7 @@ void registerTestPasses() {
   registerSideEffectTestPasses();
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
+  registerRegionTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();