[Flang] Give fir.if RegionBranchOpInterface

[SBallantyne]

Summary

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

[SBallantyne]

Cleanup changes

[tblah]

flang/lib/Optimizer/Dialect/FIROps.cpp

22

nit: this can go

3431

In Fortran/Flang are if conditions strictly true == 1 or is true any non-zero value?

3448

I know you just copied this code, but it strange that this casts to a BoolAttr whereas the other to an IntegerAttr. If flang allows true as any non-zero value, you will need to make sure that is respected in any cast to BoolAttr.

flang/lib/Optimizer/Transforms/StackArrays.cpp

402 ↗

(On Diff #501864)

nit: remove this

flang/test/Transforms/stack-arrays.fir

51

Please could you make this a new test instead of changing the existing one. Reviewers on the stack arrays pass wanted to make sure it worked on scf.if so we should keep the test for that. Furthermore, this is testing different behavior than the old test (a valid conversion vs ensuring we don't modify an invalid conversion).

[tblah]

flang/lib/Optimizer/Dialect/FIROps.cpp

22

I was commenting on the old version of the patch where this was commented out. As you have uncommitted it, I presume it is actually a required include.

[SBallantyne]

Add scf.if test back

[SBallantyne]

flang/lib/Optimizer/Dialect/FIROps.cpp

3431

As fir.if takes an I1, it should be fine to check for 1 instead of 0 at this point, any integer != 0 should have been converted to 1 by now.

3448

Similar to above, while it is inconsistent that we use cast to bool vs casting to int it seems like BoolAttr follows the same rules as flang that any non-zero value will be true.

[tblah]

Thanks for your changes. LGTM but wait for somebody else to look over it

flang/test/Transforms/stack-arrays.fir

70

nit: I suspect you won't be allowed two functions named dfa2

[jeanPerier]

It makes sense to me to add this interface to fir.if and LGTM. Just adding @vzakhari here to ensure I am not missing something.

[kiranchandramohan]

flang/include/flang/Optimizer/Dialect/FIROps.td

2150

2150

We probably have to add RecursiveMemoryEffects as well.

2173

Nit: In a subsequent patch you can change this to MaxSizedRegion<1>.

[SBallantyne]

Amend test name, add RecursiveMemoryEffects trait and limit fir.if to 1 else region

[SBallantyne]

Clang format has changed in between patches, appease it

[vzakhari]

One minor thing, otherwise, looks good.

flang/include/flang/Optimizer/Dialect/FIROps.td

2150

nit: getNumRegionInvocations does not seem to be an interface method of RegionBranchOpInterface. Moreover, you do not seem to override it, so this mention is probably redundant (here and in scf.if).

[SBallantyne]

Remove redundant getNumRegionInvocations from scf.if and fir.if

[SBallantyne]

Remove change to scf

[vzakhari]

Thanks!

flang/include/flang/Optimizer/Dialect/FIROps.td

let extraClassDeclaration = [{
 
     mlir::BlockArgument iterArgToBlockArg(mlir::Value iterArg);
     void resultToSourceOps(llvm::SmallVectorImpl<mlir::Value> &results,
                            unsigned resultNum);
     mlir::Value blockArgToSourceOp(unsigned blockArgNum);
   }];
 }
 
-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
     dialect's version of `loop.if`.
 
     ```mlir
       %56 = ... : i1
       %78 = ... : !fir.ref<!T>
       fir.if %56 {
         fir.store %76 to %78 : !fir.ref<!T>
       } else {
         fir.store %77 to %78 : !fir.ref<!T>
       }
     ```
   }];
 
   let arguments = (ins I1:$condition);
   let results = (outs Variadic<AnyType>:$results);
 
   let regions = (region
     SizedRegion<1>:$thenRegion,
-    AnyRegion:$elseRegion
+    MaxSizedRegion<1>:$elseRegion
   );
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "mlir::Value":$cond, "bool":$withElseRegion)>,
     OpBuilder<(ins "mlir::TypeRange":$resultTypes, "mlir::Value":$cond,
         "bool":$withElseRegion)>
   ];

flang/lib/Optimizer/Dialect/FIROps.cpp

 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Support/FIRContext.h"
 #include "flang/Optimizer/Support/KindMapping.h"
 #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"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
void fir::IfOp::build(mlir::OpBuilder &builder, mlir::OperationState &result,
   mlir::Region *elseRegion = result.addRegion();
   if (withElseRegion) {
     elseRegion->push_back(new mlir::Block());
     if (resultTypes.empty())
       IfOp::ensureTerminator(*elseRegion, builder, result.location);
   }
 }
 
+// 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);
   mlir::Region *thenRegion = result.addRegion();
   mlir::Region *elseRegion = result.addRegion();
 
   auto &builder = parser.getBuilder();
   mlir::OpAsmParser::UnresolvedOperand cond;

flang/test/Transforms/stack-arrays.fir

 // CHECK-NEXT:   %[[BOOL:.*]] = fir.convert %[[LOGICAL]] : (!fir.logical<4>) -> i1
 // CHECK-NEXT:   fir.if %[[BOOL]] {
 // CHECK-NEXT:     fir.freemem %[[MEM]] : !fir.heap<!fir.array<?xi32>>
 // CHECK-NEXT:   } else {
 // CHECK-NEXT:   }
 // 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 {
     fir.freemem %a : !fir.heap<!fir.array<1xi8>>
   } else {
   }
   return
 }
 // CHECK:     func.func @dfa2(%arg0: i1) {
 // CHECK-NEXT:  %[[MEM:.*]] = fir.allocmem !fir.array<1xi8>
 // CHECK-NEXT:  scf.if %arg0 {
 // CHECK-NEXT:    fir.freemem %[[MEM]] : !fir.heap<!fir.array<1xi8>>
 // CHECK-NEXT:  } else {
 // CHECK-NEXT:  }
 // 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
   %1 = arith.constant 1 : index
   %2 = arith.constant 2 : index
   %3 = arith.addi %1, %2 : index
   // operand is now available
   %4 = fir.allocmem !fir.array<?xi32>, %3