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

[jeanPerier]

Summary

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.

[clementval]

LGTM

[jeanPerier]

Rebase to incorporate windows build fix.

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

def hlfir_EndAssociateOp : hlfir_Op<"end_associate", []> {
   let assemblyFormat = [{
     $var `,` $must_free attr-dict `:` type(operands)
   }];
 
   let builders = [OpBuilder<(ins "hlfir::AssociateOp":$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]> {
   let summary = "synthetic op to prevent reassociation";
 
   let description = [{
     Same as fir.reassoc, except it accepts hlfir.expr arguments.
   }];

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
-#include <tuple>
 #include <optional>
+#include <tuple>
 
 //===----------------------------------------------------------------------===//
 // DeclareOp
 //===----------------------------------------------------------------------===//
 
 /// Given a FIR memory type, and information about non default lower bounds, get
 /// the related HLFIR variable type.
 mlir::Type hlfir::DeclareOp::getHLFIRVariableType(mlir::Type inputType,
return build(builder, result, associate.getFirBase(),
                associate.getMustFreeStrorageFlag());
 }
 
 //===----------------------------------------------------------------------===//
 // AsExprOp
 //===----------------------------------------------------------------------===//
 
 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>()) {
     typeShape.append(seqType.getShape().begin(), seqType.getShape().end());
     type = seqType.getEleTy();
   }
 
   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);
 }
 
 //===----------------------------------------------------------------------===//
 // ElementalOp
 //===----------------------------------------------------------------------===//
 
 void hlfir::ElementalOp::build(mlir::OpBuilder &builder,
                                mlir::OperationState &odsState,

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

struct AsExprOpConversion : public mlir::OpConversionPattern<hlfir::AsExprOp> {
   explicit AsExprOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::AsExprOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::AsExprOp asExpr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     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);
     mlir::Value bufferizedExpr =
         packageBufferizedExpr(loc, builder, temp, cleanup);
     rewriter.replaceOp(asExpr, bufferizedExpr);
     return mlir::success();
   }

flang/test/HLFIR/as_expr-codegen.fir

 // CHECK:    %[[VAL_4:.*]] = fir.shape %[[VAL_1]], %[[VAL_3]]#1 : (index, index) -> !fir.shape<2>
 // CHECK:    %[[VAL_5:.*]] = fir.allocmem !fir.array<10x?xi32>, %[[VAL_3]]#1 {bindc_name = ".tmp", uniq_name = ""}
 // CHECK:    %[[VAL_6:.*]] = arith.constant true
 // CHECK:    %[[VAL_7:.*]]:2 = hlfir.declare %[[VAL_5]](%[[VAL_4]]) {uniq_name = ".tmp"} : (!fir.heap<!fir.array<10x?xi32>>, !fir.shape<2>) -> (!fir.box<!fir.array<10x?xi32>>, !fir.heap<!fir.array<10x?xi32>>)
 // CHECK:    hlfir.assign %[[VAL_0]] to %[[VAL_7]]#0 : !fir.class<!fir.array<10x?xi32>>, !fir.box<!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>

flang/test/HLFIR/as_expr.fir

 
 func.func @array_expr_2(%arg0: !fir.ref<!fir.array<10xi32>>) {
   %0 = hlfir.as_expr %arg0 : (!fir.ref<!fir.array<10xi32>>) -> !hlfir.expr<10xi32>
   return
 }
 // 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>