From 43f34f58349ae178fd1c95d6a73c6858f35f2ea1 Mon Sep 17 00:00:00 2001
From: Julian Gross <julian.gross@dfki.de>
Date: Wed, 13 Jan 2021 11:13:41 +0100
Subject: [PATCH] Added check if there are regions that do not implement the
 RegionBranchOpInterface.

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

Differential Revision: https://reviews.llvm.org/D94586
---
 mlir/lib/Transforms/BufferDeallocation.cpp    | 32 ++++++++++++++++++++++++++-
 mlir/test/Transforms/buffer-deallocation.mlir | 27 ++++++++++++++++------
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Transforms/BufferDeallocation.cpp b/mlir/lib/Transforms/BufferDeallocation.cpp
index bdd8515..a415569 100644
--- a/mlir/lib/Transforms/BufferDeallocation.cpp
+++ b/mlir/lib/Transforms/BufferDeallocation.cpp
@@ -76,6 +76,31 @@ static void walkReturnOperations(Region *region, const FuncT &func) {
     }
 }
 
+/// 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 {
 
 //===----------------------------------------------------------------------===//
@@ -506,7 +531,12 @@ struct BufferDeallocationPass : BufferDeallocationBase<BufferDeallocationPass> {
     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.
diff --git a/mlir/test/Transforms/buffer-deallocation.mlir b/mlir/test/Transforms/buffer-deallocation.mlir
index f61d501..90b39b2 100644
--- a/mlir/test/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Transforms/buffer-deallocation.mlir
@@ -1,4 +1,4 @@
-// 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,
@@ -1094,7 +1094,7 @@ func @loop_nested_alloc(
 // 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,
@@ -1121,15 +1121,13 @@ func @loop_dynalloc(
   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,
@@ -1155,8 +1153,6 @@ func @do_loop_alloc(
   return
 }
 
-// expected-error@+1 {{Structured control-flow loops are supported only}}
-
 // -----
 
 // CHECK-LABEL: func @assumingOp(
@@ -1193,3 +1189,20 @@ func @assumingOp(
 // 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"() : () -> ()
+}
-- 
2.7.4