Added check if there are regions that do not implement the RegionBranchOpInterface.

[dfki-jugr]

Summary

Add a check if regions do not implement the RegionBranchOpInterface. This is not
allowed in the current deallocation steps. Furthermore, we handle edge-cases,
where a single region is attached and the parent operation has no results.

This fixes: https://bugs.llvm.org/show_bug.cgi?id=48575

[herhut]

mlir/lib/Transforms/BufferDeallocation.cpp

83

Naming: This does not check for control flow. It checks whether interfaces are present or whether the control flow is supported.

84

Would region.walk do what you want?

87

check that operation's with at least one region implement...

I did not check this but I am a little surprised that operations with inputs do not require this interface. How is the connection between an input and the uses inside of the succesor regions made in this case?

94

nit: use empty or getNumResults() == 0.

99

Consider not aborting early but reporting more cases?

539

This could be made slightly better by moving the emitError into the checkForControlFlow helper (and maybe call it validateControlFlowIsSupported or something). Then you could emit the error on the op that fails the check, which gives the user of this a more actionable error.

[dfki-jugr]

mlir/lib/Transforms/BufferDeallocation.cpp

84

We are wondering if the walk is recursive?
If yes, we could use region.walk.

87

In general, we would need to check the inputs if there are some uses inside the successor region.
Considering our implementation, we insert deallocations after the operation such that they are independent from the inputs.

[dfki-jugr]

Addressed comments.

[herhut]

Please fix comments.

mlir/lib/Transforms/BufferDeallocation.cpp

87

Nit: the comment is wrong. This does not check if the operation has at least one region and implements the interface, it checks if all operations that have at least one region implement the interface.

91

nit: since _in that case_ the

mlir/lib/Transforms/BufferDeallocation.cpp

static void walkReturnOperations(Region *region, const FuncT &func) {
   for (Block &block : *region)
     for (Operation &operation : block) {
       // Skip non-return-like terminators.
       if (operation.hasTrait<OpTrait::ReturnLike>())
         func(&operation);
     }
 }
 
+/// Checks if all operations in a given region have at least one attached region
+/// that implements the RegionBranchOpInterface. This is not required in edge
+/// cases, where we have a single attached region and the parent operation has
+/// no results.
+static bool validateSupportedControlFlow(Region &region) {
+  bool success = true;
+  region.walk([&success](Operation *operation) {
+    auto regions = operation->getRegions();
+    // Walk over all operations in a region and check if the operation has at
+    // least one region and implements the RegionBranchOpInterface. If there
+    // is an operation that does not fulfill this condition, we cannot apply
+    // the deallocation steps. Furthermore, we accept cases, where we have a
+    // region that returns no results, since, in that case, the intra-region
+    // control flow does not affect the transformation.
+    size_t size = regions.size();
+    if (((size == 1 && !operation->getResults().empty()) || size > 1) &&
+        !dyn_cast<RegionBranchOpInterface>(operation)) {
+      operation->emitError("All operations with attached regions need to "
+                           "implement the RegionBranchOpInterface.");
+      success = false;
+    }
+  });
+  return success;
+}
+
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Backedges analysis
 //===----------------------------------------------------------------------===//
 
 /// A straight-forward program analysis which detects loop backedges induced by
 /// explicit control flow.

 struct BufferDeallocationPass : BufferDeallocationBase<BufferDeallocationPass> {
 
   void runOnFunction() override {
     // Ensure that there are supported loops only.
     Backedges backedges(getFunction());
     if (backedges.size()) {
       getFunction().emitError(
           "Structured control-flow loops are supported only.");
-      return;
+      return signalPassFailure();
+    }
+
+    // Check that the control flow structures are supported.
+    if (!validateSupportedControlFlow(getFunction().getRegion())) {
+      return signalPassFailure();
     }
 
     // Place all required temporary alloc, copy and dealloc nodes.
     BufferDeallocation deallocation(getFunction());
     deallocation.deallocate();
   }
 };
 

mlir/test/Transforms/buffer-deallocation.mlir

-// RUN: mlir-opt -buffer-deallocation -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -verify-diagnostics -buffer-deallocation -split-input-file %s | FileCheck %s
 
 // This file checks the behaviour of BufferDeallocation pass for moving and
 // inserting missing DeallocOps in their correct positions. Furthermore,
 // copies and their corresponding AllocOps are inserted.
 
 // Test Case:
 //    bb0
 //   /   \

 // CHECK-NEXT: dealloc %[[VAL_7]]
 
 // -----
 
 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
 // The BufferDeallocation transformation should fail on this explicit
 // control-flow loop since they are not supported.
 
-// CHECK-LABEL: func @loop_dynalloc
+// expected-error@+1 {{Structured control-flow loops are supported only}}
 func @loop_dynalloc(
   %arg0 : i32,
   %arg1 : i32,
   %arg2: memref<?xf32>,
   %arg3: memref<?xf32>) {
   %const0 = constant 0 : i32
   br ^loopHeader(%const0, %arg2 : i32, memref<?xf32>)
 
^loopBody(%val : i32, %buff2: memref<?xf32>):
   %alloc1 = alloc(%size) : memref<?xf32>
   br ^loopHeader(%inc, %alloc1 : i32, memref<?xf32>)
 
 ^exit(%buff3 : memref<?xf32>):
   test.copy(%buff3, %arg3) : (memref<?xf32>, memref<?xf32>)
   return
 }
 
-// expected-error@+1 {{Structured control-flow loops are supported only}}
-
 // -----
 
 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
 // The BufferDeallocation transformation should fail on this explicit
 // control-flow loop since they are not supported.
 
-// CHECK-LABEL: func @do_loop_alloc
+// expected-error@+1 {{Structured control-flow loops are supported only}}
 func @do_loop_alloc(
   %arg0 : i32,
   %arg1 : i32,
   %arg2: memref<2xf32>,
   %arg3: memref<2xf32>) {
   %const0 = constant 0 : i32
   br ^loopBody(%const0, %arg2 : i32, memref<2xf32>)
 
cond_br %lessThan,
     ^loopBody(%i, %buff : i32, memref<2xf32>),
     ^exit(%buff : memref<2xf32>)
 
 ^exit(%buff3 : memref<2xf32>):
   test.copy(%buff3, %arg3) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
-// expected-error@+1 {{Structured control-flow loops are supported only}}
-
 // -----
 
 // CHECK-LABEL: func @assumingOp(
 func @assumingOp(
   %arg0: !shape.witness,
   %arg2: memref<2xf32>,
   %arg3: memref<2xf32>) {
   // Confirm the alloc will be dealloc'ed in the block.

 //      CHECK: %[[ASSUMING_RESULT:.*]] = shape.assuming %[[ARG0]]
 // CHECK-NEXT:    %[[TMP_ALLOC:.*]] = alloc()
 // CHECK-NEXT:    %[[RETURNING_ALLOC:.*]] = alloc()
 // CHECK-NEXT:    linalg.copy(%[[TMP_ALLOC]], %[[RETURNING_ALLOC]])
 // CHECK-NEXT:    dealloc %[[TMP_ALLOC]]
 // CHECK-NEXT:    shape.assuming_yield %[[RETURNING_ALLOC]]
 //      CHECK: test.copy(%[[ASSUMING_RESULT:.*]], %[[ARG2]])
 // CHECK-NEXT: dealloc %[[ASSUMING_RESULT]]
+
+// -----
+
+// Test Case: The op "test.bar" does not implement the RegionBranchOpInterface.
+// This is not allowed in buffer deallocation.
+
+func @noRegionBranchOpInterface() {
+// expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
+  %0 = "test.bar"() ( {
+// expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
+    %1 = "test.bar"() ( {
+      "test.yield"() : () -> ()
+    }) : () -> (i32)
+    "test.yield"() : () -> ()
+  }) : () -> (i32)
+  "test.terminator"() : () -> ()
+}