[mlir] Fix Regions 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] 9a8bb4bc63/mlir/lib/IR/Operation.cpp (L154)

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

mlir/include/mlir/IR/Region.h

@@ -240,6 +240,7 @@ public:
   /// 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());
   }

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: {

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

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

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();