[flang][hlfir] add allocatable assignment semantic to hlfir.assign

[clementval]

LGTM

[PeteSteinfeld]

All builds and tests correctly and looks good.

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

 
 def fir_AssignOp : hlfir_Op<"assign", [MemoryEffects<[MemWrite]>]> {
   let summary = "Assign an expression or variable value to a Fortran variable";
 
   let description = [{
     Assign rhs to lhs following Fortran intrinsic assignments rules.
     The operation deals with inserting a temporary if the lhs and rhs
     may overlap.
+    The optional "realloc" flag allows indicating that this assignment
+    has the Fortran 95 semantics for assignments to a whole allocatable.
+    In such case, the left hand side must be an allocatable that may be
+    unallocated or allocated with a different type and shape than the right
+    hand side. It will be allocated or re-allocated as needed during the
+    assignment.
+    When "realloc" is set and this is a character assignment, the optional
+    flag "keep_lhs_length_if_realloc" indicates that the character
+    left hand side should retain its length after the assignment. If the
+    right hand side has a different length, truncation and padding will
+    occur. This covers the case of explicit and assumed length character
+    allocatables.
+    Otherwise, the left hand side will be allocated or reallocated to match the
+    right hand side length if they differ. This covers the case of deferred
+    length character allocatables.
   }];
 
   let arguments = (ins AnyFortranEntity:$rhs,
-                   Arg<AnyFortranVariable, "", [MemWrite]>:$lhs);
+                   Arg<AnyFortranVariable, "", [MemWrite]>:$lhs,
+                   UnitAttr:$realloc,
+                   UnitAttr:$keep_lhs_length_if_realloc);
 
   let assemblyFormat = [{
-    $rhs `to` $lhs attr-dict `:` type(operands)
+    $rhs `to` $lhs (`realloc` $realloc^)?
+    (`keep_lhs_len` $keep_lhs_length_if_realloc^)?
+    attr-dict `:` type(operands)
   }];
+
+  let extraClassDeclaration = [{
+    /// Does this assignment have the Fortran 95 semantics of assignments
+    /// to a whole allocatable?
+    bool isAllocatableAssignment() {
+      return getRealloc();
+    }
+    /// Is the assignment left hand side a whole allocatable character
+    /// that should retain its length after the assignment?
+    bool mustKeepLhsLengthInAllocatableAssignment() {
+      return getKeepLhsLengthIfRealloc();
+    }
+  }];
+
+  let hasVerifier = 1;
 }
 
 def hlfir_DesignateOp : hlfir_Op<"designate", [AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<fir_FortranVariableOpInterface>]> {
   let summary = "Designate a Fortran variable";
 
   let description = [{
     This operation represents a Fortran "part-ref", except that it can embed a

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

 #include "llvm/ADT/TypeSwitch.h"
 #include <optional>
 #include <tuple>
 
 //===----------------------------------------------------------------------===//
 // DeclareOp
 //===----------------------------------------------------------------------===//
 
+/// Is this a fir.[ref/ptr/heap]<fir.[box/class]<fir.heap<T>>> type?
+static bool isAllocatableBoxRef(mlir::Type type) {
+  fir::BaseBoxType boxType =
+      fir::dyn_cast_ptrEleTy(type).dyn_cast_or_null<fir::BaseBoxType>();
+  return boxType && boxType.getEleTy().isa<fir::HeapType>();
+}
+
+mlir::LogicalResult hlfir::AssignOp::verify() {
+  mlir::Type lhsType = getLhs().getType();
+  if (isAllocatableAssignment() && !isAllocatableBoxRef(lhsType))
+    return emitOpError("lhs must be an allocatable when `realloc` is set");
+  if (mustKeepLhsLengthInAllocatableAssignment() &&
+      !(isAllocatableAssignment() &&
+        hlfir::getFortranElementType(lhsType).isa<fir::CharacterType>()))
+    return emitOpError("`realloc` must be set and lhs must be a character "
+                       "allocatable when `keep_lhs_length_if_realloc` is set");
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
+// 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,
                                                   bool hasExplicitLowerBounds) {
   mlir::Type type = fir::unwrapRefType(inputType);
   if (type.isa<fir::BaseBoxType>())
     return inputType;
   if (auto charType = type.dyn_cast<fir::CharacterType>())

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

matchAndRewrite(hlfir::AssignOp assignOp,
     }
     auto [rhsExv, rhsCleanUp] =
         hlfir::translateToExtendedValue(loc, builder, rhs);
     auto [lhsExv, lhsCleanUp] =
         hlfir::translateToExtendedValue(loc, builder, lhs);
     assert(!lhsCleanUp && !rhsCleanUp &&
            "variable to fir::ExtendedValue must not require cleanup");
 
-    if (lhs.isArray()) {
+    auto emboxRHS = [&](fir::ExtendedValue &rhsExv) -> mlir::Value {
       // There may be overlap between lhs and rhs. The runtime is able to detect
       // and to make a copy of the rhs before modifying the lhs if needed.
       // The code below relies on this and does not do any compile time alias
       // analysis.
       const bool rhsIsValue = fir::isa_trivial(fir::getBase(rhsExv).getType());
       if (rhsIsValue) {
         // createBox can only be called for fir::ExtendedValue that are
         // already in memory. Place the integer/real/complex/logical scalar
         // in memory (convert to the LHS type so that i1 are allocated in
         // a proper Fortran logical storage).
         mlir::Type lhsValueType = lhs.getFortranElementType();
         mlir::Value rhsVal =
             builder.createConvert(loc, lhsValueType, fir::getBase(rhsExv));
         mlir::Value temp = builder.create<fir::AllocaOp>(loc, lhsValueType);
         builder.create<fir::StoreOp>(loc, rhsVal, temp);
         rhsExv = temp;
       }
+      return fir::getBase(builder.createBox(loc, rhsExv));
+    };
 
+    if (assignOp.isAllocatableAssignment()) {
+      // Whole allocatable assignment: use the runtime to deal with the
+      // reallocation.
+      mlir::Value from = emboxRHS(rhsExv);
+      mlir::Value to = fir::getBase(lhsExv);
+      if (assignOp.mustKeepLhsLengthInAllocatableAssignment()) {
+        // Indicate the runtime that it should not reallocate in case of length
+        // mismatch, and that it should use the LHS explicit/assumed length if
+        // allocating/reallocation the LHS.
+        TODO(loc, "assignment to explicit length whole allocatable");
+      } else if (lhs.isPolymorphic()) {
+        // Indicate the runtime that the LHS must have the RHS dynamic type
+        // after the assignment.
+        TODO(loc, "assignment to whole polymorphic entity");
+      } else {
+        fir::runtime::genAssign(builder, loc, to, from);
+      }
+    } else if (lhs.isArray()) {
       // Use the runtime for simplicity. An optimization pass will be added to
       // inline array assignment when profitable.
-      auto to = fir::getBase(builder.createBox(loc, lhsExv));
-      auto from = fir::getBase(builder.createBox(loc, rhsExv));
-
+      mlir::Value from = emboxRHS(rhsExv);
+      mlir::Value to = fir::getBase(builder.createBox(loc, lhsExv));
+      // This is not a whole allocatable assignment: the runtime will not
+      // reallocate and modify "toMutableBox" even if it is taking it by
+      // reference.
       auto toMutableBox = builder.createTemporary(loc, to.getType());
-      // As per 10.2.1.2 point 1 (1) polymorphic variables must be allocatable.
-      // It is assumed here that they have been reallocated with the dynamic
-      // type and that the mutableBox will not be modified.
       builder.create<fir::StoreOp>(loc, to, toMutableBox);
       fir::runtime::genAssign(builder, loc, toMutableBox, from);
     } else {
       // Assume overlap does not matter for scalar (dealt with memmove for
       // characters).
       // This is not true if this is a derived type with "recursive" allocatable
       // components, in which case an overlap would matter because the LHS
       // reallocation, if any, may modify the RHS component value before it is

flang/test/HLFIR/assign-codegen.fir

func.func @test_i1_scalar_to_array(%lhs: !fir.box<!fir.array<?x!fir.logical<4>>>, %rhs: i1) {
   return
 }
 // CHECK-LABEL:   func.func @test_i1_scalar_to_array(
 // CHECK:  %[[VAL_5:.*]] = fir.convert %{{.*}} : (i1) -> !fir.logical<4>
 // CHECK:  %[[VAL_6:.*]] = fir.alloca !fir.logical<4>
 // CHECK:  %[[VAL_7:.*]] = fir.embox %[[VAL_6]] : (!fir.ref<!fir.logical<4>>) -> !fir.box<!fir.logical<4>>
 // CHECK:  %[[VAL_12:.*]] = fir.convert %[[VAL_7]] : (!fir.box<!fir.logical<4>>) -> !fir.box<none>
 // CHECK:  %[[VAL_14:.*]] = fir.call @_FortranAAssign(%{{.*}}, %[[VAL_12]]
+
+func.func @alloc_assign(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, %arg1: !fir.box<!fir.array<?xi32>>) {
+  hlfir.assign %arg1 to %arg0 realloc : !fir.box<!fir.array<?xi32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  return
+}
+// CHECK-LABEL:  func.func @alloc_assign(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>,
+// CHECK-SAME: %[[VAL_1:.*]]: !fir.box<!fir.array<?xi32>>) {
+// CHECK:  %[[VAL_2:.*]] = fir.convert %[[VAL_0]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<none>>
+// CHECK:  %[[VAL_3:.*]] = fir.convert %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
+// CHECK:  fir.call @_FortranAAssign(%[[VAL_2]], %[[VAL_3]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) -> none

flang/test/HLFIR/assign.fir

 func.func @array_4(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.box<!fir.array<?xi32>>) {
   hlfir.assign %arg1 to %arg0 : !fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>
   return
 }
 // CHECK-LABEL:   func.func @array_4(
 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>,
 // CHECK-SAME:    %[[VAL_1:.*]]: !fir.box<!fir.array<?xi32>>) {
 // CHECK:  hlfir.assign %[[VAL_1]] to %[[VAL_0]] : !fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>
+
+func.func @alloc_assign(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, %arg1: !fir.box<!fir.array<?xi32>>) {
+  hlfir.assign %arg1 to %arg0 realloc : !fir.box<!fir.array<?xi32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  return
+}
+// CHECK-LABEL:   func.func @alloc_assign(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: !fir.box<!fir.array<?xi32>>) {
+// CHECK:  hlfir.assign %[[VAL_1]] to %[[VAL_0]] realloc : !fir.box<!fir.array<?xi32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+
+func.func @alloc_assign_keep_lhs_len(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,?>>>>>, %arg1: !fir.box<!fir.array<?x!fir.char<1,?>>>) {
+  hlfir.assign %arg1 to %arg0 realloc keep_lhs_len : !fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,?>>>>>
+  return
+}
+// CHECK-LABEL:   func.func @alloc_assign_keep_lhs_len(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,?>>>>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>>) {
+// CHECK:  hlfir.assign %[[VAL_1]] to %[[VAL_0]] realloc keep_lhs_len : !fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,?>>>>>

flang/test/HLFIR/invalid.fir

 }
 
 // -----
 func.func @bad_matmul8(%arg0: !hlfir.expr<2xi32>, %arg1: !hlfir.expr<2x3xi32>) {
   // expected-error@+1 {{'hlfir.matmul' op incorrect result shape}}
   %0 = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<2xi32>, !hlfir.expr<2x3xi32>) -> !hlfir.expr<1x3xi32>
   return
 }
+
+// -----
+func.func @bad_assign_1(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.box<!fir.array<?xi32>>) {
+  // expected-error@+1 {{'hlfir.assign' op lhs must be an allocatable when `realloc` is set}}
+  hlfir.assign %arg1 to %arg0 realloc : !fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>
+  return
+}
+
+// -----
+func.func @bad_assign_2(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, %arg1: !fir.box<!fir.array<?xi32>>) {
+  // expected-error@+1 {{'hlfir.assign' op `realloc` must be set and lhs must be a character allocatable when `keep_lhs_length_if_realloc` is set}}
+  hlfir.assign %arg1 to %arg0 realloc keep_lhs_len : !fir.box<!fir.array<?xi32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  return
+}