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,8 @@ }]; } -def fir_IfOp : region_Op<"if", [NoRegionArguments]> { +def fir_IfOp : region_Op<"if", [DeclareOpInterfaceMethods,NoRegionArguments]> { let summary = "if-then-else conditional operation"; let description = [{ Used to conditionally execute operations. This operation is the FIR 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,58 @@ } } +// 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,19 +48,19 @@ // CHECK-NEXT: return // CHECK-NEXT: } -// Check scf.if (fir.if is not considered a branch operation) +// Check fir.if can be properly read func.func @dfa2(%arg0: i1) { %a = fir.allocmem !fir.array<1xi8> - scf.if %arg0 { + fir.if %arg0 { fir.freemem %a : !fir.heap> } else { + fir.freemem %a : !fir.heap> } 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> +// CHECK-NEXT: %[[MEM:.*]] = fir.alloca !fir.array<1xi8> +// CHECK-NEXT: fir.if %arg0 { // CHECK-NEXT: } else { // CHECK-NEXT: } // CHECK-NEXT: return