[flang-commits] [flang] 1b0ec39 - [Flang] Give fir.if RegionBranchOpInterface

[Sacha Ballantyne via flang-commits]

Author: Sacha Ballantyne
Date: 2023-03-09T12:19:37Z
New Revision: 1b0ec3981c7eeecaf478997e7d94d4d0a8d7ad60

URL: https://github.com/llvm/llvm-project/commit/1b0ec3981c7eeecaf478997e7d94d4d0a8d7ad60
DIFF: https://github.com/llvm/llvm-project/commit/1b0ec3981c7eeecaf478997e7d94d4d0a8d7ad60.diff

LOG: [Flang] Give fir.if RegionBranchOpInterface

fir.if currently isn't treated as a 'proper' conditional, so passes are unable to determine which regions are executed at times.

This patch gives fir.if this interface, which shouldn't do too much on its own but should allow future changes to take advantage
for various purposes

Reviewed By: vzakhari

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

Added: 
    

Modified: 
    flang/include/flang/Optimizer/Dialect/FIROps.td
    flang/lib/Optimizer/Dialect/FIROps.cpp
    flang/test/Transforms/stack-arrays.fir

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 1fd653b6076b9..2bc4ec0401d52 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -2146,7 +2146,9 @@ def fir_DoLoopOp : region_Op<"do_loop",
   }];
 }
 
-def fir_IfOp : region_Op<"if", [NoRegionArguments]> {
+def fir_IfOp : region_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
+    "getRegionInvocationBounds"]>, RecursiveMemoryEffects,
+    NoRegionArguments]> {
   let summary = "if-then-else conditional operation";
   let description = [{
     Used to conditionally execute operations. This operation is the FIR
@@ -2168,7 +2170,7 @@ def fir_IfOp : region_Op<"if", [NoRegionArguments]> {
 
   let regions = (region
     SizedRegion<1>:$thenRegion,
-    AnyRegion:$elseRegion
+    MaxSizedRegion<1>:$elseRegion
   );
 
   let skipDefaultBuilders = 1;

diff  --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 671936f8186fd..c62a7bc1174c9 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -19,6 +19,7 @@
 #include "flang/Optimizer/Support/Utils.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
@@ -3403,6 +3404,59 @@ void fir::IfOp::build(mlir::OpBuilder &builder, mlir::OperationState &result,
   }
 }
 
+// These 2 functions copied from scf.if implementation.
+
+/// Given the region at `index`, or the parent operation if `index` is None,
+/// return the successor regions. These are the regions that may be selected
+/// during the flow of control. `operands` is a set of optional attributes that
+/// correspond to a constant value for each operand, or null if that operand is
+/// not a constant.
+void fir::IfOp::getSuccessorRegions(
+    std::optional<unsigned> index, llvm::ArrayRef<mlir::Attribute> operands,
+    llvm::SmallVectorImpl<mlir::RegionSuccessor> &regions) {
+  // The `then` and the `else` region branch back to the parent operation.
+  if (index) {
+    regions.push_back(mlir::RegionSuccessor(getResults()));
+    return;
+  }
+
+  // Don't consider the else region if it is empty.
+  mlir::Region *elseRegion = &this->getElseRegion();
+  if (elseRegion->empty())
+    elseRegion = nullptr;
+
+  // Otherwise, the successor is dependent on the condition.
+  bool condition;
+  if (auto condAttr = operands.front().dyn_cast_or_null<mlir::IntegerAttr>()) {
+    condition = condAttr.getValue().isOne();
+  } else {
+    // If the condition isn't constant, both regions may be executed.
+    regions.push_back(mlir::RegionSuccessor(&getThenRegion()));
+    // If the else region does not exist, it is not a viable successor.
+    if (elseRegion)
+      regions.push_back(mlir::RegionSuccessor(elseRegion));
+    return;
+  }
+
+  // Add the successor regions using the condition.
+  regions.push_back(
+      mlir::RegionSuccessor(condition ? &getThenRegion() : elseRegion));
+}
+
+void fir::IfOp::getRegionInvocationBounds(
+    llvm::ArrayRef<mlir::Attribute> operands,
+    llvm::SmallVectorImpl<mlir::InvocationBounds> &invocationBounds) {
+  if (auto cond = operands[0].dyn_cast_or_null<mlir::BoolAttr>()) {
+    // If the condition is known, then one region is known to be executed once
+    // and the other zero times.
+    invocationBounds.emplace_back(0, cond.getValue() ? 1 : 0);
+    invocationBounds.emplace_back(0, cond.getValue() ? 0 : 1);
+  } else {
+    // Non-constant condition. Each region may be executed 0 or 1 times.
+    invocationBounds.assign(2, {0, 1});
+  }
+}
+
 mlir::ParseResult fir::IfOp::parse(mlir::OpAsmParser &parser,
                                    mlir::OperationState &result) {
   result.regions.reserve(2);

diff  --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index 7b73bf0930743..d470ea704be48 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -48,7 +48,7 @@ func.func @dfa1(%arg0: !fir.ref<!fir.logical<4>> {fir.bindc_name = "cond"}) {
 // CHECK-NEXT:   return
 // CHECK-NEXT: }
 
-// Check scf.if (fir.if is not considered a branch operation)
+// Check scf.if
 func.func @dfa2(%arg0: i1) {
   %a = fir.allocmem !fir.array<1xi8>
   scf.if %arg0 {
@@ -66,6 +66,24 @@ func.func @dfa2(%arg0: i1) {
 // CHECK-NEXT:  return
 // CHECK-NEXT:  }
 
+// Check freemem in both regions
+func.func @dfa3(%arg0: i1) {
+  %a = fir.allocmem !fir.array<1xi8>
+  fir.if %arg0 {
+    fir.freemem %a : !fir.heap<!fir.array<1xi8>>
+  } else {
+    fir.freemem %a : !fir.heap<!fir.array<1xi8>>
+  }
+  return
+}
+// CHECK:     func.func @dfa3(%arg0: i1) {
+// CHECK-NEXT:  %[[MEM:.*]] = fir.alloca !fir.array<1xi8>
+// CHECK-NEXT:  fir.if %arg0 {
+// CHECK-NEXT:  } else {
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+// CHECK-NEXT:  }
+
 // check the alloca is placed after all operands become available
 func.func @placement1() {
   // do some stuff with other ssa values