[flang][hlfir] Handle box and non constant must_free in end_associate.

Differential Revision: https://reviews.llvm.org/D142700
llvm · Jan 27, 2023 · 36e1890 · 36e1890
1 parent 86e6a8c
commit 36e1890
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 12 deletions.
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
@@ -331,24 +331,25 @@ struct AssociateOpConversion
   }
 };
 
-static void genFreeIfMustFree(mlir::Location loc,
-                              mlir::ConversionPatternRewriter &rewriter,
+static void genFreeIfMustFree(mlir::Location loc, fir::FirOpBuilder &builder,
                               mlir::Value var, mlir::Value mustFree) {
   auto genFree = [&]() {
-    if (var.getType().isa<fir::BaseBoxType>())
-      TODO(loc, "unbox");
-    if (!var.getType().isa<fir::HeapType>())
-      var = rewriter.create<fir::ConvertOp>(
-          loc, fir::HeapType::get(fir::unwrapRefType(var.getType())), var);
-    rewriter.create<fir::FreeMemOp>(loc, var);
+    // fir::FreeMemOp operand type must be a fir::HeapType.
+    mlir::Type heapType = fir::HeapType::get(
+        hlfir::getFortranElementOrSequenceType(var.getType()));
+    if (var.getType().isa<fir::BaseBoxType, fir::BoxCharType>())
+      var = builder.create<fir::BoxAddrOp>(loc, heapType, var);
+    else if (!var.getType().isa<fir::HeapType>())
+      var = builder.create<fir::ConvertOp>(loc, heapType, var);
+    builder.create<fir::FreeMemOp>(loc, var);
   };
   if (auto cstMustFree = fir::getIntIfConstant(mustFree)) {
     if (*cstMustFree != 0)
       genFree();
-    // else, nothing to do.
+    // else, mustFree is false, nothing to do.
     return;
   }
-  TODO(loc, "conditional free");
+  builder.genIfThen(loc, mustFree).genThen(genFree).end();
 }
 
 struct EndAssociateOpConversion
@@ -360,7 +361,9 @@ struct EndAssociateOpConversion
   matchAndRewrite(hlfir::EndAssociateOp endAssociate, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = endAssociate->getLoc();
-    genFreeIfMustFree(loc, rewriter, adaptor.getVar(), adaptor.getMustFree());
+    auto module = endAssociate->getParentOfType<mlir::ModuleOp>();
+    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    genFreeIfMustFree(loc, builder, adaptor.getVar(), adaptor.getMustFree());
     rewriter.eraseOp(endAssociate);
     return mlir::success();
   }
@@ -378,9 +381,11 @@ struct DestroyOpConversion
     mlir::Location loc = destroy->getLoc();
     mlir::Value bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     if (!fir::isa_trivial(bufferizedExpr.getType())) {
+      auto module = destroy->getParentOfType<mlir::ModuleOp>();
+      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
       mlir::Value mustFree = getBufferizedExprMustFreeFlag(adaptor.getExpr());
       mlir::Value firBase = hlfir::Entity(bufferizedExpr).getFirBase();
-      genFreeIfMustFree(loc, rewriter, firBase, mustFree);
+      genFreeIfMustFree(loc, builder, firBase, mustFree);
     }
     rewriter.eraseOp(destroy);
     return mlir::success();

diff --git a/flang/test/HLFIR/associate-codegen.fir b/flang/test/HLFIR/associate-codegen.fir
@@ -79,6 +79,43 @@ func.func @associate_char(%arg0: !fir.boxchar<1> ) {
 // CHECK-NOT: fir.freemem
 
 
+func.func @test_end_associate_box(%var: !fir.box<!fir.array<?xf64>>) {
+  %true = arith.constant 1 : i1
+  hlfir.end_associate %var, %true : !fir.box<!fir.array<?xf64>>, i1
+  return
+}
+// CHECK-LABEL:   func.func @test_end_associate_box(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?xf64>>) {
+// CHECK:  %[[VAL_1:.*]] = arith.constant true
+// CHECK:  %[[VAL_2:.*]] = fir.box_addr %[[VAL_0]] : (!fir.box<!fir.array<?xf64>>) -> !fir.heap<!fir.array<?xf64>>
+// CHECK:  fir.freemem %[[VAL_2]] : !fir.heap<!fir.array<?xf64>>
+
+
+func.func @test_end_associate_boxchar(%var: !fir.boxchar<2>) {
+  %true = arith.constant 1 : i1
+  hlfir.end_associate %var, %true : !fir.boxchar<2>, i1
+  return
+}
+// CHECK-LABEL:   func.func @test_end_associate_boxchar(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.boxchar<2>) {
+// CHECK:  %[[VAL_1:.*]] = arith.constant true
+// CHECK:  %[[VAL_2:.*]] = fir.box_addr %[[VAL_0]] : (!fir.boxchar<2>) -> !fir.heap<!fir.char<2,?>>
+// CHECK:  fir.freemem %[[VAL_2]] : !fir.heap<!fir.char<2,?>>
+
+
+func.func @test_end_associate_box_dynamic(%var: !fir.box<!fir.array<?xf64>>, %must_free: i1) {
+  hlfir.end_associate %var, %must_free : !fir.box<!fir.array<?xf64>>, i1
+  return
+}
+// CHECK-LABEL:   func.func @test_end_associate_box_dynamic(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?xf64>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: i1) {
+// CHECK:  fir.if %[[VAL_1]] {
+// CHECK:    %[[VAL_2:.*]] = fir.box_addr %[[VAL_0]] : (!fir.box<!fir.array<?xf64>>) -> !fir.heap<!fir.array<?xf64>>
+// CHECK:    fir.freemem %[[VAL_2]] : !fir.heap<!fir.array<?xf64>>
+// CHECK:  }
+
+
 func.func private @take_i4(!fir.ref<i32>)
 func.func private @take_r4(!fir.ref<f32>)
 func.func private @take_l4(!fir.ref<!fir.logical<4>>)