[fir] add fir.array_modify op

[clementval]

Summary

fir.array_update is only handling intrinsic assignments.
They are two big differences with user defined assignments:

1. The LHS and RHS types may not match, this does not play well with fir.array_update that relies on both the merge and the updated element to have the same type.
2. user defined assignment has a call semantics, with potential side effects. So if a fir.array_update can hide a call, it traits would need to be updated.

Instead of hiding more semantic in the fir.array_update, introduce
a new fir.array_modify op that allows de-correlating indicating that
an array value element is modified, and how it is modified.
This allow the ArrayValueCopy pass to still perform copy elision
while not having to implement the call itself, and could in general
be used for all kind of assignments (e.g. character assignment).

Update the alias analysis to not rely on the merge arguments (since
fir.array_modify has none).
Instead, analyze what is done with the element address.
This implies adding the ability to follow the users of fir.array_modify,
as well as being able to go through fir.store that may be generated to
store the RHS value in order to pass it to a user define routine.
This is done by adding a ReachCollector class to gather all array
accesses.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Jean Perier <[e-mail]>

[schweitz]

LGTM

flang/include/flang/Optimizer/Dialect/FIROps.td

def fir_ArrayUpdateOp : fir_Op<"array_update", [AttrSizedOperandSegments,
   let assemblyFormat = [{
     $sequence `,` $merge `,` $indices (`typeparams` $typeparams^)? attr-dict
       `:` functional-type(operands, results)
   }];
 
   let verifier = "return ::verify(*this);";
 }
 
+def fir_ArrayModifyOp : fir_Op<"array_modify", [AttrSizedOperandSegments,
+    NoSideEffect]> {
+  let summary = "Get an address for an array value to modify it.";
+
+  let description = [{
+    Modify the value of an element in an array value through actions done
+    on the returned address. A new array value is also
+    returned where all element values of the input array are identical except
+    for the selected element which is the value after the modification done
+    on the element address.
+
+    ```fortran
+      real :: a(n)
+      ...
+      ! Elemental user defined assignment from type(SomeType) to real.
+      a = value_of_some_type
+    ```
+
+    One can use `fir.array_modify` to update the (implied) value of `a(i)`
+    in an array expression as shown above.
+
+    ```mlir
+      %s = fir.shape %n : (index) -> !fir.shape<1>
+      // Load the entire array 'a'.
+      %v = fir.array_load %a(%s) : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>)
+          -> !fir.array<?xf32>
+      // Update the value of one of the array value's elements with a user
+      // defined assignment from %rhs.
+      %new = fir.do_loop %i = ... (%inner = %v) {
+        %rhs = ...
+        %addr, %r = fir.array_modify %inner, %i, %j : (!fir.array<?xf32>,
+      index) -> fir.ref<f32>, !fir.array<?xf32>
+        fir.call @user_def_assign(%addr, %rhs) (fir.ref<f32>,
+      fir.ref<!fir.type<SomeType>>) -> ()
+        fir.result %r : !fir.ref<!fir.array<?xf32>>
+      }
+      fir.array_merge_store %v, %new to %a : !fir.ref<!fir.array<?xf32>>
+    ```
+
+    An array value modification behaves as if a mapping function from the indices
+    to the new value has been added, replacing the previous mapping. These
+    mappings can be added to the ssa-value, but will not be materialized in
+    memory until the `fir.array_merge_store` is performed.
+  }];
+
+  let arguments = (ins
+    fir_SequenceType:$sequence,
+    Variadic<AnyCoordinateType>:$indices,
+    Variadic<AnyIntegerType>:$typeparams
+  );
+
+  let results = (outs fir_ReferenceType, fir_SequenceType);
+
+  let assemblyFormat = [{
+    $sequence `,` $indices (`typeparams` $typeparams^)? attr-dict
+      `:` functional-type(operands, results)
+  }];
+
+  let verifier = [{ return ::verify(*this); }];
+}
+
 def fir_ArrayMergeStoreOp : fir_Op<"array_merge_store",
     [AttrSizedOperandSegments]> {
 
   let summary = "Store merged array value to memory.";
 
   let description = [{
     Store a merged array value to memory.
 

flang/lib/Optimizer/Dialect/FIROps.cpp

if (indSize == arrTy.getDimension() &&
     return op.emitOpError("merged value does not have element type");
   auto ty = validArraySubobject(op);
   if (!ty || ty != ::adjustedElementType(op.merge().getType()))
     return op.emitOpError("merged value and/or indices do not type check");
   return mlir::success();
 }
 
 //===----------------------------------------------------------------------===//
+// ArrayModifyOp
+//===----------------------------------------------------------------------===//
+
+static mlir::LogicalResult verify(fir::ArrayModifyOp op) {
+  auto arrTy = op.sequence().getType().cast<fir::SequenceType>();
+  auto indSize = op.indices().size();
+  if (indSize < arrTy.getDimension())
+    return op.emitOpError("number of indices must match array dimension");
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
 // BoxAddrOp
 //===----------------------------------------------------------------------===//
 
 mlir::OpFoldResult fir::BoxAddrOp::fold(llvm::ArrayRef<mlir::Attribute> opnds) {
   if (auto v = val().getDefiningOp()) {
     if (auto box = dyn_cast<fir::EmboxOp>(v))
       return box.memref();
     if (auto box = dyn_cast<fir::EmboxCharOp>(v))

flang/test/Fir/fir-ops.fir

func @test_misc_ops(%arr1 : !fir.ref<!fir.array<?x?xf32>>, %m : index, %n : index, %o : index, %p : index) {
   // CHECK: [[AV2:%.*]] = fir.array_update [[AV1]], [[FVAL]], [[I10]], [[J20]] : (!fir.array<?x?xf32>, f32, index, index) -> !fir.array<?x?xf32>
   // CHECK: fir.array_merge_store [[AV1]], [[AV2]] to [[ARR1]] : !fir.array<?x?xf32>, !fir.array<?x?xf32>, !fir.ref<!fir.array<?x?xf32>>
   %s = fir.shape_shift %m, %n, %o, %p : (index, index, index, index) -> !fir.shapeshift<2>
   %av1 = fir.array_load %arr1(%s) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.array<?x?xf32>
   %f = fir.array_fetch %av1, %i10, %j20 : (!fir.array<?x?xf32>, index, index) -> f32
   %av2 = fir.array_update %av1, %f, %i10, %j20 : (!fir.array<?x?xf32>, f32, index, index) -> !fir.array<?x?xf32>
   fir.array_merge_store %av1, %av2 to %arr1 : !fir.array<?x?xf32>, !fir.array<?x?xf32>, !fir.ref<!fir.array<?x?xf32>>
 
+  // CHECK: [[AV3:%.*]] = fir.array_load [[ARR1]]([[SHAPE]]) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.array<?x?xf32>
+  // CHECK: [[FVAL2:%.*]] = fir.array_fetch [[AV3]], [[I10]], [[J20]] : (!fir.array<?x?xf32>, index, index) -> f32
+  // CHECK: [[AV4:%.*]]:2 = fir.array_modify [[AV3]], [[I10]], [[J20]] : (!fir.array<?x?xf32>, index, index) -> (!fir.ref<f32>, !fir.array<?x?xf32>)
+  // CHECK: fir.store [[FVAL2]] to [[AV4]]#0 : !fir.ref<f32>
+  // CHECK: fir.array_merge_store [[AV3]], [[AV4]]#1 to [[ARR1]] : !fir.array<?x?xf32>, !fir.array<?x?xf32>, !fir.ref<!fir.array<?x?xf32>>
+  %av3 = fir.array_load %arr1(%s) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.array<?x?xf32>
+  %f2 = fir.array_fetch %av3, %i10, %j20 : (!fir.array<?x?xf32>, index, index) -> f32
+  %addr, %av4 = fir.array_modify %av3, %i10, %j20 : (!fir.array<?x?xf32>, index, index) -> (!fir.ref<f32>, !fir.array<?x?xf32>)
+  fir.store %f2 to %addr : !fir.ref<f32>
+  fir.array_merge_store %av3, %av4 to %arr1 : !fir.array<?x?xf32>, !fir.array<?x?xf32>, !fir.ref<!fir.array<?x?xf32>>
   return
 }
 
 // CHECK-LABEL: @test_shift
 func @test_shift(%arg0: !fir.box<!fir.array<?xf32>>) -> !fir.ref<f32> {
   %c4 = constant 4 : index
   %c100 = constant 100 : index
   // CHECK: fir.shift %{{.*}} : (index) -> !fir.shift<1>

flang/test/Fir/invalid.fir

func @test_misc_ops(%arr1 : !fir.ref<!fir.array<?x?xf32>>, %m : index, %n : index, %o : index, %p : index) {
   %s = fir.shape_shift %m, %n, %o, %p : (index, index, index, index) -> !fir.shapeshift<2>
   %av1 = fir.array_load %arr1(%s) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.array<?x?xf32>
   %f = fir.array_fetch %av1, %m, %n : (!fir.array<?x?xf32>, index, index) -> f32
   %av2 = fir.array_update %av1, %f, %m, %n : (!fir.array<?x?xf32>, f32, index, index) -> !fir.array<?x?xf32>
   // expected-error@+1 {{'fir.array_merge_store' op operand #0 must be result of a fir.array_load op}}
   fir.array_merge_store %av2, %av2 to %arr1 : !fir.array<?x?xf32>, !fir.array<?x?xf32>, !fir.ref<!fir.array<?x?xf32>>
   return
 }
+
+// -----
+
+func @bad_array_modify(%arr1 : !fir.ref<!fir.array<?x?xf32>>, %m : index, %n : index, %o : index, %p : index, %f : f32) {
+  %i10 = constant 10 : index
+  %j20 = constant 20 : index
+  %s = fir.shape_shift %m, %n, %o, %p : (index, index, index, index) -> !fir.shapeshift<2>
+  %av1 = fir.array_load %arr1(%s) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.array<?x?xf32>
+  // expected-error@+1 {{'fir.array_modify' op number of indices must match array dimension}}
+  %addr, %av2 = fir.array_modify %av1, %i10 : (!fir.array<?x?xf32>, index) -> (!fir.ref<f32>, !fir.array<?x?xf32>)
+  fir.store %f to %addr : !fir.ref<f32>
+  fir.array_merge_store %av1, %av2 to %arr1 : !fir.array<?x?xf32>, !fir.array<?x?xf32>, !fir.ref<!fir.array<?x?xf32>>
+  return
+}