diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td --- a/flang/include/flang/Optimizer/Dialect/FIROps.td +++ b/flang/include/flang/Optimizer/Dialect/FIROps.td @@ -2146,7 +2146,9 @@ }]; } -def fir_IfOp : region_Op<"if", [NoRegionArguments]> { +def fir_IfOp : region_Op<"if", [DeclareOpInterfaceMethods, 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 @@ 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 --- 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 @@ } } +// 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 index, llvm::ArrayRef operands, + llvm::SmallVectorImpl ®ions) { + // 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()) { + 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 operands, + llvm::SmallVectorImpl &invocationBounds) { + if (auto cond = operands[0].dyn_cast_or_null()) { + // 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 --- a/flang/test/Transforms/stack-arrays.fir +++ b/flang/test/Transforms/stack-arrays.fir @@ -48,7 +48,7 @@ // 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 @@ // 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> + } else { + fir.freemem %a : !fir.heap> + } + 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 diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -717,7 +717,7 @@ //===----------------------------------------------------------------------===// def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods, + "getRegionInvocationBounds"]>, DeclareOpInterfaceMethods, SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveMemoryEffects, NoRegionArguments]> {