From 25db7cb70c312b956240b3d65c3bc72020ac9409 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier@nvidia.com>
Date: Wed, 1 Mar 2023 10:43:16 +0100
Subject: [PATCH] [flang][hlfir] Implement hlfir.declare optional codegen

The hlfir fir.box with the local lower bounds and type parameters
must be generated conditionally when the entity is optional.

Differential Revision: https://reviews.llvm.org/D144962
---
 .../Optimizer/HLFIR/Transforms/ConvertToFIR.cpp    | 64 +++++++++++++++-------
 flang/test/HLFIR/declare-codegen.fir               | 22 ++++++++
 2 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
index ce09ef0..160a77c 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
@@ -286,26 +286,52 @@ public:
     mlir::Value hlfirBase;
     mlir::Type hlfirBaseType = declareOp.getBase().getType();
     if (hlfirBaseType.isa<fir::BaseBoxType>()) {
-      // Need to conditionally rebox/embox for optional.
-      if (mlir::cast<fir::FortranVariableOpInterface>(declareOp.getOperation())
-              .isOptional())
-        TODO(loc, "converting hlfir declare of optional box to fir");
-      if (!firBase.getType().isa<fir::BaseBoxType>()) {
-        llvm::SmallVector<mlir::Value> typeParams;
-        auto maybeCharType =
-            fir::unwrapSequenceType(fir::unwrapPassByRefType(hlfirBaseType))
-                .dyn_cast<fir::CharacterType>();
-        if (!maybeCharType || maybeCharType.hasDynamicLen())
-          typeParams.append(declareOp.getTypeparams().begin(),
-                            declareOp.getTypeparams().end());
-        hlfirBase = rewriter.create<fir::EmboxOp>(
-            loc, hlfirBaseType, firBase, declareOp.getShape(),
-            /*slice=*/mlir::Value{}, typeParams);
+      auto module = declareOp->getParentOfType<mlir::ModuleOp>();
+      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+      // Helper to generate the hlfir fir.box with the local lower bounds and
+      // type parameters.
+      auto genHlfirBox = [&]() -> mlir::Value {
+        if (!firBase.getType().isa<fir::BaseBoxType>()) {
+          llvm::SmallVector<mlir::Value> typeParams;
+          auto maybeCharType =
+              fir::unwrapSequenceType(fir::unwrapPassByRefType(hlfirBaseType))
+                  .dyn_cast<fir::CharacterType>();
+          if (!maybeCharType || maybeCharType.hasDynamicLen())
+            typeParams.append(declareOp.getTypeparams().begin(),
+                              declareOp.getTypeparams().end());
+          return builder.create<fir::EmboxOp>(
+              loc, hlfirBaseType, firBase, declareOp.getShape(),
+              /*slice=*/mlir::Value{}, typeParams);
+        } else {
+          // Rebox so that lower bounds are correct.
+          return builder.create<fir::ReboxOp>(loc, hlfirBaseType, firBase,
+                                              declareOp.getShape(),
+                                              /*slice=*/mlir::Value{});
+        }
+      };
+      if (!mlir::cast<fir::FortranVariableOpInterface>(declareOp.getOperation())
+               .isOptional()) {
+        hlfirBase = genHlfirBox();
       } else {
-        // Rebox so that lower bounds are correct.
-        hlfirBase = rewriter.create<fir::ReboxOp>(loc, hlfirBaseType, firBase,
-                                                  declareOp.getShape(),
-                                                  /*slice=*/mlir::Value{});
+        // Need to conditionally rebox/embox the optional: the input fir.box
+        // may be null and the rebox would be illegal. It is also important to
+        // preserve the optional aspect: the hlfir fir.box should be null if
+        // the entity is absent so that later fir.is_present on the hlfir base
+        // are valid.
+        mlir::Value isPresent =
+            builder.create<fir::IsPresentOp>(loc, builder.getI1Type(), firBase);
+        hlfirBase = builder
+                        .genIfOp(loc, {hlfirBaseType}, isPresent,
+                                 /*withElseRegion=*/true)
+                        .genThen([&] {
+                          builder.create<fir::ResultOp>(loc, genHlfirBox());
+                        })
+                        .genElse([&]() {
+                          mlir::Value absent =
+                              builder.create<fir::AbsentOp>(loc, hlfirBaseType);
+                          builder.create<fir::ResultOp>(loc, absent);
+                        })
+                        .getResults()[0];
       }
     } else if (hlfirBaseType.isa<fir::BoxCharType>()) {
       assert(declareOp.getTypeparams().size() == 1 &&
diff --git a/flang/test/HLFIR/declare-codegen.fir b/flang/test/HLFIR/declare-codegen.fir
index 601ad4c..3e80a52 100644
--- a/flang/test/HLFIR/declare-codegen.fir
+++ b/flang/test/HLFIR/declare-codegen.fir
@@ -178,3 +178,25 @@ func.func @array_declare_unlimited_polymorphic_boxaddr(%arg0: !fir.ref<!fir.clas
 // CHECK-LABEL: func.func @array_declare_unlimited_polymorphic_boxaddr(
 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.class<!fir.ptr<!fir.array<?x?xnone>>>>) {
 // CHECK:         %[[VAL_1:.*]] = fir.declare %[[VAL_0]] {uniq_name = "x"} : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?x?xnone>>>>) -> !fir.ref<!fir.class<!fir.ptr<!fir.array<?x?xnone>>>>
+
+
+func.func @test_optional_declare(%arg0: !fir.box<!fir.array<?xi32>>) {
+  %c42 = arith.constant 42 : index
+  %0 = fir.shift %c42 : (index) -> !fir.shift<1>
+  %1:2 = hlfir.declare %arg0(%0) {fortran_attrs = #fir.var_attrs<optional>, uniq_name = "x"} : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>)
+  return
+}
+
+// CHECK-LABEL:   func.func @test_optional_declare(
+// CHECK-SAME:      %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) {
+// CHECK:  %[[VAL_1:.*]] = arith.constant 42 : index
+// CHECK:  %[[VAL_2:.*]] = fir.shift %[[VAL_1]] : (index) -> !fir.shift<1>
+// CHECK:  %[[VAL_3:.*]] = fir.declare %[[VAL_0]](%[[VAL_2]]) {fortran_attrs = #fir.var_attrs<optional>, uniq_name = "x"} : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>) -> !fir.box<!fir.array<?xi32>>
+// CHECK:  %[[VAL_4:.*]] = fir.is_present %[[VAL_3]] : (!fir.box<!fir.array<?xi32>>) -> i1
+// CHECK:  %[[VAL_5:.*]] = fir.if %[[VAL_4]] -> (!fir.box<!fir.array<?xi32>>) {
+// CHECK:    %[[VAL_6:.*]] = fir.rebox %[[VAL_3]](%[[VAL_2]]) : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>) -> !fir.box<!fir.array<?xi32>>
+// CHECK:    fir.result %[[VAL_6]] : !fir.box<!fir.array<?xi32>>
+// CHECK:  } else {
+// CHECK:    %[[VAL_7:.*]] = fir.absent !fir.box<!fir.array<?xi32>>
+// CHECK:    fir.result %[[VAL_7]] : !fir.box<!fir.array<?xi32>>
+// CHECK:  }
-- 
2.7.4