[fir] add fir.array_modify op

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.

Reviewed By: schweitz

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

Co-authored-by: Valentin Clement <clementval@gmail.com>

diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index c2e42ad..b7e1e2c 100644 (file)
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -1484,6 +1484,67 @@ def fir_ArrayUpdateOp : fir_Op<"array_update", [AttrSizedOperandSegments,
   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]> {
 

diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index ed13f93..94e8d62 100644 (file)
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -497,6 +497,18 @@ static mlir::LogicalResult verify(fir::ArrayUpdateOp op) {
 }
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 

diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir
index e64cff7..5f1588f 100644 (file)
--- a/flang/test/Fir/fir-ops.fir
+++ b/flang/test/Fir/fir-ops.fir
@@ -636,6 +636,16 @@ func @test_misc_ops(%arr1 : !fir.ref<!fir.array<?x?xf32>>, %m : index, %n : inde
   %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
 }
 

diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir
index a2b367c..f9f0a63 100644 (file)
--- a/flang/test/Fir/invalid.fir
+++ b/flang/test/Fir/invalid.fir
@@ -592,3 +592,17 @@ func @test_misc_ops(%arr1 : !fir.ref<!fir.array<?x?xf32>>, %m : index, %n : inde
   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
+}