[flang][hlfir] Add move semantics to hlfir.as_expr.

hlfir.as_expr allows turning an array, character, or derived type
variable into a value when it the usage require an hlfir.expr (e.g,
when returning the element value inside and hlfir.elemental).

The default implementation of this operation in bufferization is to
make a copy of the variable into a temporary buffer.
This adds a time and memory overhead in cases where such copy is not
needed because the variable is already a temporary that was created
in lowering to compute the expression value, and the "as_expr" is
the sole usage of the variable.

This is for instance the case for many transformational intrinsics
that do not have hlfir.expr operation (at least for now, but some may
never benefit from having one) and must be implemented "on memory"
in lowering.

This patch adds a way to "move" the variable storage along its value.
It allows the bufferization to re-use the variable storage for the
hlfir.expr created by hlfir.as_expr, and in exchange, the
responsibility of deallocating the buffer (if the variable was heap
allocated) if passed along to the hlfir.expr, and will need to be
done after the last hlfir.expr usage.

Differential Revision: https://reviews.llvm.org/D141832

diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index 88acf78..5433cda 100644 (file)
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
@@ -328,21 +328,37 @@ def hlfir_EndAssociateOp : hlfir_Op<"end_associate", []> {
 }
 
 def hlfir_AsExprOp : hlfir_Op<"as_expr", []> {
-  let summary = "Take the value of an array, character or derived expression";
+  let summary = "Take the value of an array, character or derived variable";
 
   let description = [{
-    Take the value of an array, character or derived expression.
+    Take the value of an array, character or derived variable.
+    In general, this operation will lead to a copy of the variable
+    in the bufferization pass if it was not transformed.
+
+    However, if it is known that the variable storage will not be used anymore
+    afterwards, the variable storage ownership can be passed to the hlfir.expr
+    by providing the $must_free argument that is a boolean that indicates if
+    the storage must be freed (when it was allocated on the heap).
+    This allows Fortran lowering to build some expression value in memory when
+    there is no adequate hlfir operation, and to promote the result to an
+    hlfir.expr value without paying the price of introducing a copy.
   }];
 
-  let arguments = (ins AnyFortranVariable:$var);
+  let arguments = (ins AnyFortranVariable:$var,
+                       Optional<I1>:$must_free);
   let results = (outs hlfir_ExprType);
 
+  let extraClassDeclaration = [{
+      // Is this a "move" ?
+      bool isMove() { return getMustFree() != mlir::Value{}; }
+  }];
+
   let assemblyFormat = [{
-    $var attr-dict `:` functional-type(operands, results)
+    $var (`move` $must_free^)? attr-dict `:` functional-type(operands, results)
   }];
 
 
-  let builders = [OpBuilder<(ins "mlir::Value":$var)>];
+  let builders = [OpBuilder<(ins "mlir::Value":$var, CArg<"mlir::Value", "{}">:$must_free)>];
 }
 
 def hlfir_NoReassocOp : hlfir_Op<"no_reassoc", [NoMemoryEffect, SameOperandsAndResultType]> {

diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 6c00777..ecdc490 100644 (file)
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -18,8 +18,8 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
-#include <tuple>
 #include <optional>
+#include <tuple>
 
 //===----------------------------------------------------------------------===//
 // DeclareOp
@@ -456,7 +456,8 @@ void hlfir::EndAssociateOp::build(mlir::OpBuilder &builder,
 //===----------------------------------------------------------------------===//
 
 void hlfir::AsExprOp::build(mlir::OpBuilder &builder,
-                            mlir::OperationState &result, mlir::Value var) {
+                            mlir::OperationState &result, mlir::Value var,
+                            mlir::Value mustFree) {
   hlfir::ExprType::Shape typeShape;
   mlir::Type type = getFortranElementOrSequenceType(var.getType());
   if (auto seqType = type.dyn_cast<fir::SequenceType>()) {
@@ -466,7 +467,7 @@ void hlfir::AsExprOp::build(mlir::OpBuilder &builder,
 
   auto resultType = hlfir::ExprType::get(builder.getContext(), typeShape, type,
                                          /*isPolymorphic: TODO*/ false);
-  return build(builder, result, resultType, var);
+  return build(builder, result, resultType, var, mustFree);
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
index 97b08ae..7bcfe86 100644 (file)
--- a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
@@ -147,6 +147,14 @@ struct AsExprOpConversion : public mlir::OpConversionPattern<hlfir::AsExprOp> {
     mlir::Location loc = asExpr->getLoc();
     auto module = asExpr->getParentOfType<mlir::ModuleOp>();
     fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    if (asExpr.isMove()) {
+      // Move variable storage for the hlfir.expr buffer.
+      mlir::Value bufferizedExpr = packageBufferizedExpr(
+          loc, builder, adaptor.getVar(), adaptor.getMustFree());
+      rewriter.replaceOp(asExpr, bufferizedExpr);
+      return mlir::success();
+    }
+    // Otherwise, create a copy in a new buffer.
     hlfir::Entity source = hlfir::Entity{adaptor.getVar()};
     auto [temp, cleanup] = createTempFromMold(loc, builder, source);
     builder.create<hlfir::AssignOp>(loc, source, temp);

diff --git a/flang/test/HLFIR/as_expr-codegen.fir b/flang/test/HLFIR/as_expr-codegen.fir
index 7f7c310..c37c624 100644 (file)
--- a/flang/test/HLFIR/as_expr-codegen.fir
+++ b/flang/test/HLFIR/as_expr-codegen.fir
@@ -71,3 +71,14 @@ func.func @shape_from_box(%arg0 : !fir.class<!fir.array<10x?xi32>>) {
 // CHECK:    %[[VAL_8:.*]] = fir.undefined tuple<!fir.box<!fir.array<10x?xi32>>, i1>
 // CHECK:    %[[VAL_9:.*]] = fir.insert_value %[[VAL_8]], %[[VAL_6]], [1 : index] : (tuple<!fir.box<!fir.array<10x?xi32>>, i1>, i1) -> tuple<!fir.box<!fir.array<10x?xi32>>, i1>
 // CHECK:    %[[VAL_10:.*]] = fir.insert_value %[[VAL_9]], %[[VAL_7]]#0, [0 : index] : (tuple<!fir.box<!fir.array<10x?xi32>>, i1>, !fir.box<!fir.array<10x?xi32>>) -> tuple<!fir.box<!fir.array<10x?xi32>>, i1>
+
+func.func @test_move(%arg0 : !fir.ref<!fir.array<10x20xi32>>, %must_free: i1) {
+  %expr = hlfir.as_expr %arg0 move %must_free: (!fir.ref<!fir.array<10x20xi32>>, i1) -> !hlfir.expr<10x20xi32>
+  return
+}
+// CHECK-LABEL:   func.func @test_move(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.array<10x20xi32>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: i1) {
+// CHECK:    %[[VAL_2:.*]] = fir.undefined tuple<!fir.ref<!fir.array<10x20xi32>>, i1>
+// CHECK:    %[[VAL_3:.*]] = fir.insert_value %[[VAL_2]], %[[VAL_1]], [1 : index] : (tuple<!fir.ref<!fir.array<10x20xi32>>, i1>, i1) -> tuple<!fir.ref<!fir.array<10x20xi32>>, i1>
+// CHECK:    %[[VAL_4:.*]] = fir.insert_value %[[VAL_3]], %[[VAL_0]], [0 : index] : (tuple<!fir.ref<!fir.array<10x20xi32>>, i1>, !fir.ref<!fir.array<10x20xi32>>) -> tuple<!fir.ref<!fir.array<10x20xi32>>, i1>

diff --git a/flang/test/HLFIR/as_expr.fir b/flang/test/HLFIR/as_expr.fir
index c04e248..d47059e 100644 (file)
--- a/flang/test/HLFIR/as_expr.fir
+++ b/flang/test/HLFIR/as_expr.fir
@@ -33,3 +33,12 @@ func.func @array_expr_2(%arg0: !fir.ref<!fir.array<10xi32>>) {
 // CHECK-LABEL: func.func @array_expr_2(
 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) {
 // CHECK:   hlfir.as_expr %[[VAL_0]] : (!fir.ref<!fir.array<10xi32>>) -> !hlfir.expr<10xi32>
+
+func.func @array_expr_move(%arg0: !fir.ref<!fir.array<10xi32>>, %must_free: i1) {
+  %0 = hlfir.as_expr %arg0 move %must_free : (!fir.ref<!fir.array<10xi32>>, i1) -> !hlfir.expr<10xi32>
+  return
+}
+// CHECK-LABEL: func.func @array_expr_move(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: i1) {
+// CHECK:   hlfir.as_expr %[[VAL_0]] move %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, i1) -> !hlfir.expr<10xi32>