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

[zero9178]

Summary

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

[rriddle]

mlir/test/lib/IR/TestTakeBody.cpp

1 ↗

(On Diff #423297)

Can this file be more generally for Region related tests? Otherwise, we'll have a huge number of test files potentially.

42 ↗

(On Diff #423297)

This should probably have Region in the name. If this file is for Region tests in general, this could just be registerRegionTestPasses().

[zero9178]

Address review comments

mlir/include/mlir/IR/Region.h

public:
   void cloneInto(Region *dest, BlockAndValueMapping &mapper);
   /// Clone this region into 'dest' before the given position in 'dest'.
   void cloneInto(Region *dest, Region::iterator destPos,
                  BlockAndValueMapping &mapper);
 
   /// 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());
   }
 
   /// Returns 'block' if 'block' lies in this region, or otherwise finds the
   /// ancestor of 'block' that lies in this region. Returns nullptr if the
   /// latter fails.
   Block *findAncestorBlockInRegion(Block &block);

mlir/test/IR/test-take-body.mlir

+// 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

add_mlir_library(MLIRTestIR
   TestOpaqueLoc.cpp
   TestOperationEquals.cpp
   TestPrintDefUse.cpp
   TestPrintInvalid.cpp
   TestPrintNesting.cpp
   TestSideEffects.cpp
   TestSlicing.cpp
   TestSymbolUses.cpp
+  TestRegions.cpp
   TestTypes.cpp
   TestVisitors.cpp
   TestVisitorsGeneric.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTestDialect
   )
 
 target_include_directories(MLIRTestIR
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test
   )

mlir/test/lib/IR/TestRegions.cpp

+//===- 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

 void registerConvertToTargetEnvPass();
 void registerCloneTestPasses();
 void registerPassManagerTestPass();
 void registerPrintSpirvAvailabilityPass();
 void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
+void registerRegionTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLoopPermutationPass();
 void registerTestMatchers();
 void registerTestOperationEqualPass();
void registerTestPasses() {
   registerCloneTestPasses();
   registerConvertToTargetEnvPass();
   registerPassManagerTestPass();
   registerPrintSpirvAvailabilityPass();
   registerShapeFunctionTestPasses();
   registerSideEffectTestPasses();
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
+  registerRegionTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLoopPermutationPass();
   registerTestMatchers();
   registerTestOperationEqualPass();