[flang-commits] [flang] [flang][FIR] add canonicalization pattern for fir.if returning OPTIONAL (PR #205353)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: jeanPerier

<details>
<summary>Changes</summary>

Lowering is generating patterns when forwarding OPTIONAL in calls that looks like:

```
   %present = fir.is_present %var : (T) -> i1
    %if_result = fir.if %present -> (T) {
      fir.result %var : T
    } else {
      %absent = fir.absent T
      fir.result %absent : T
    }
```

This specific pattern is a no-op and `%var` can be used directly. The lowering logic that generates such patterns is inside non trivial compiler code that has to deal with more complex scenarios where the code inside the fir.if is more complex. Add a FIR pattern to canonicalize such code to help with later analysis (like aliasing).

---
Full diff: https://github.com/llvm/llvm-project/pull/205353.diff


3 Files Affected:

- (modified) flang/include/flang/Optimizer/Dialect/FIROps.td (+2) 
- (modified) flang/lib/Optimizer/Dialect/FIROps.cpp (+59) 
- (added) flang/test/Fir/present-absent-if-fold.fir (+66) 


``````````diff
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 35a87c29d0cb6..c6e4d1b3b4d11 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -2560,6 +2560,8 @@ def fir_IfOp
         "bool":$withElseRegion)>
   ];
 
+  let hasCanonicalizer = 1;
+
   let extraClassDeclaration = [{
     mlir::OpBuilder getThenBodyBuilder() {
       assert(!getThenRegion().empty() && "Unexpected empty 'where' region.");
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 1d61989eba6cf..ba047e71d6aa3 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -5690,6 +5690,65 @@ void fir::IfOp::resultToSourceOps(llvm::SmallVectorImpl<mlir::Value> &results,
     results.push_back(term->getOperand(resultNum));
 }
 
+// Fold away a fir.if that only forwards an optional argument or returns
+// fir.absent when it is not present:
+//
+//   %present = fir.is_present %var : (T) -> i1
+//   %r = fir.if %present -> (T) {
+//     fir.result %var : T
+//   } else {
+//     %absent = fir.absent T
+//     fir.result %absent : T
+//   }
+//
+// The result is always %var: optional arguments already encode presence.
+struct FoldPresentAbsentIfOp : public mlir::OpRewritePattern<fir::IfOp> {
+  using mlir::OpRewritePattern<fir::IfOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::IfOp ifOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (ifOp.getNumResults() != 1)
+      return mlir::failure();
+
+    auto isPresentOp = ifOp.getCondition().getDefiningOp<fir::IsPresentOp>();
+    if (!isPresentOp)
+      return mlir::failure();
+
+    mlir::Value optionalVal = isPresentOp.getVal();
+    mlir::Type resultType = ifOp.getResult(0).getType();
+    if (optionalVal.getType() != resultType)
+      return mlir::failure();
+
+    mlir::Block &thenBlock = ifOp.getThenRegion().front();
+    if (thenBlock.getOperations().size() != 1)
+      return mlir::failure();
+    auto thenResult = mlir::dyn_cast<fir::ResultOp>(thenBlock.getTerminator());
+    if (!thenResult || thenResult.getNumOperands() != 1 ||
+        thenResult.getOperand(0) != optionalVal)
+      return mlir::failure();
+
+    if (ifOp.getElseRegion().empty())
+      return mlir::failure();
+    mlir::Block &elseBlock = ifOp.getElseRegion().front();
+    if (elseBlock.getOperations().size() > 2)
+      return mlir::failure();
+    auto elseResult = mlir::dyn_cast<fir::ResultOp>(elseBlock.getTerminator());
+    if (!elseResult || elseResult.getNumOperands() != 1)
+      return mlir::failure();
+    if (!elseResult.getOperand(0).getDefiningOp<fir::AbsentOp>())
+      return mlir::failure();
+
+    rewriter.replaceOp(ifOp, optionalVal);
+    return mlir::success();
+  }
+};
+
+void fir::IfOp::getCanonicalizationPatterns(mlir::RewritePatternSet &patterns,
+                                            mlir::MLIRContext *context) {
+  patterns.add<FoldPresentAbsentIfOp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // BoxOffsetOp
 //===----------------------------------------------------------------------===//
diff --git a/flang/test/Fir/present-absent-if-fold.fir b/flang/test/Fir/present-absent-if-fold.fir
new file mode 100644
index 0000000000000..19fd7859f4a23
--- /dev/null
+++ b/flang/test/Fir/present-absent-if-fold.fir
@@ -0,0 +1,66 @@
+// RUN: fir-opt --canonicalize %s | FileCheck %s
+
+// CHECK-LABEL: func.func @fold_present_absent_if_box(
+// CHECK-SAME:     %[[VAR:.*]]: !fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>> {
+// CHECK-NOT: fir.if
+// CHECK-NOT: fir.is_present
+// CHECK-NOT: fir.absent
+// CHECK: return %[[VAR]] : !fir.box<!fir.array<?xf32>>
+func.func @fold_present_absent_if_box(
+    %var: !fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>> {
+  %present = fir.is_present %var : (!fir.box<!fir.array<?xf32>>) -> i1
+  %if_result = fir.if %present -> (!fir.box<!fir.array<?xf32>>) {
+    fir.result %var : !fir.box<!fir.array<?xf32>>
+  } else {
+    %absent = fir.absent !fir.box<!fir.array<?xf32>>
+    fir.result %absent : !fir.box<!fir.array<?xf32>>
+  }
+  return %if_result : !fir.box<!fir.array<?xf32>>
+}
+
+// CHECK-LABEL: func.func @fold_present_absent_if_ref(
+// CHECK-SAME:     %[[VAR:.*]]: !fir.ref<i32>) -> !fir.ref<i32> {
+// CHECK-NOT: fir.if
+// CHECK: return %[[VAR]] : !fir.ref<i32>
+func.func @fold_present_absent_if_ref(%var: !fir.ref<i32>) -> !fir.ref<i32> {
+  %present = fir.is_present %var : (!fir.ref<i32>) -> i1
+  %if_result = fir.if %present -> (!fir.ref<i32>) {
+    fir.result %var : !fir.ref<i32>
+  } else {
+    %absent = fir.absent !fir.ref<i32>
+    fir.result %absent : !fir.ref<i32>
+  }
+  return %if_result : !fir.ref<i32>
+}
+
+func.func private @side_effect() -> ()
+
+// CHECK-LABEL: func.func @no_fold_call_in_then(
+// CHECK: fir.is_present
+// CHECK: fir.if
+func.func @no_fold_call_in_then(%var: !fir.ref<i32>) -> !fir.ref<i32> {
+  %present = fir.is_present %var : (!fir.ref<i32>) -> i1
+  %if_result = fir.if %present -> (!fir.ref<i32>) {
+    fir.call @side_effect() : () -> ()
+    fir.result %var : !fir.ref<i32>
+  } else {
+    %absent = fir.absent !fir.ref<i32>
+    fir.result %absent : !fir.ref<i32>
+  }
+  return %if_result : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @no_fold_call_in_else(
+// CHECK: fir.is_present
+// CHECK: fir.if
+func.func @no_fold_call_in_else(%var: !fir.ref<i32>) -> !fir.ref<i32> {
+  %present = fir.is_present %var : (!fir.ref<i32>) -> i1
+  %if_result = fir.if %present -> (!fir.ref<i32>) {
+    fir.result %var : !fir.ref<i32>
+  } else {
+    %absent = fir.absent !fir.ref<i32>
+    fir.call @side_effect() : () -> ()
+    fir.result %absent : !fir.ref<i32>
+  }
+  return %if_result : !fir.ref<i32>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/205353