[flang] Update ArrayValueCopy to support array_amend and array_access

This patch update the array value copy pass to support fir-array_amend
and fir.array_access.

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

Reviewed By: PeteSteinfeld, schweitz

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

Co-authored-by: Jean Perier <jperier@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

flang/include/flang/Optimizer/Builder/Array.h

@@ -0,0 +1,27 @@
+//===-- Array.h -------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_ARRAY_H
+#define FORTRAN_OPTIMIZER_BUILDER_ARRAY_H
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+
+namespace fir::factory {
+
+/// Return true if and only if the extents are those of an assumed-size array.
+/// An assumed-size array in Fortran is declared with `*` as the upper bound of
+/// the last dimension of the array. Lowering converts the asterisk to an
+/// undefined value.
+inline bool isAssumedSize(const llvm::SmallVectorImpl<mlir::Value> &extents) {
+  return !extents.empty() &&
+         mlir::isa_and_nonnull<UndefOp>(extents.back().getDefiningOp());
+}
+
+} // namespace fir::factory
+
+#endif // FORTRAN_OPTIMIZER_BUILDER_ARRAY_H

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -371,6 +371,12 @@ public:
   /// Generate code testing \p addr is a null address.
   mlir::Value genIsNull(mlir::Location loc, mlir::Value addr);
 
+  /// Compute the extent of (lb:ub:step) as max((ub-lb+step)/step, 0). See
+  /// Fortran 2018 9.5.3.3.2 section for more details.
+  mlir::Value genExtentFromTriplet(mlir::Location loc, mlir::Value lb,
+                                   mlir::Value ub, mlir::Value step,
+                                   mlir::Type type);
+
 private:
   const KindMapping &kindMap;
 };

flang/include/flang/Optimizer/Dialect/FIROpsSupport.h

@@ -83,6 +83,10 @@ static constexpr llvm::StringRef getHostAssocAttrName() {
   return "fir.host_assoc";
 }
 
+/// Does the function, \p func, have a host-associations tuple argument?
+/// Some internal procedures may have access to host procedure variables.
+bool hasHostAssociationArgument(mlir::FuncOp func);
+
 /// Tell if \p value is:
 ///   - a function argument that has attribute \p attributeName
 ///   - or, the result of fir.alloca/fir.allocamem op that has attribute \p

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -507,6 +507,23 @@ mlir::Value fir::FirOpBuilder::genIsNull(mlir::Location loc, mlir::Value addr) {
                                   mlir::arith::CmpIPredicate::eq);
 }
 
+mlir::Value fir::FirOpBuilder::genExtentFromTriplet(mlir::Location loc,
+                                                    mlir::Value lb,
+                                                    mlir::Value ub,
+                                                    mlir::Value step,
+                                                    mlir::Type type) {
+  auto zero = createIntegerConstant(loc, type, 0);
+  lb = createConvert(loc, type, lb);
+  ub = createConvert(loc, type, ub);
+  step = createConvert(loc, type, step);
+  auto diff = create<mlir::arith::SubIOp>(loc, ub, lb);
+  auto add = create<mlir::arith::AddIOp>(loc, diff, step);
+  auto div = create<mlir::arith::DivSIOp>(loc, add, step);
+  auto cmp = create<mlir::arith::CmpIOp>(loc, mlir::arith::CmpIPredicate::sgt,
+                                         div, zero);
+  return create<mlir::arith::SelectOp>(loc, cmp, div, zero);
+}
+
 //===--------------------------------------------------------------------===//
 // ExtendedValue inquiry helper implementation
 //===--------------------------------------------------------------------===//

flang/lib/Optimizer/Dialect/FIROps.cpp

@@ -3258,6 +3258,15 @@ fir::GlobalOp fir::createGlobalOp(mlir::Location loc, mlir::ModuleOp module,
   return result;
 }
 
+bool fir::hasHostAssociationArgument(mlir::FuncOp func) {
+  if (auto allArgAttrs = func.getAllArgAttrs())
+    for (auto attr : allArgAttrs)
+      if (auto dict = attr.template dyn_cast_or_null<mlir::DictionaryAttr>())
+        if (dict.get(fir::getHostAssocAttrName()))
+          return true;
+  return false;
+}
+
 bool fir::valueHasFirAttribute(mlir::Value value,
                                llvm::StringRef attributeName) {
   // If this is a fir.box that was loaded, the fir attributes will be on the

flang/test/Fir/array-value-copy.fir

@@ -104,12 +104,12 @@ func @conversion_with_temporary(%arr0 : !fir.ref<!fir.array<10xi32>>) {
 // CHECK-LABEL: func @conversion_with_temporary(
 // CHECK-SAME:                                  %[[ARR0:.*]]: !fir.ref<!fir.array<10xi32>>)
 // Allocation of temporary array.
-// CHECK:         %[[TEMP:.*]] = fir.allocmem !fir.array<10xi32>, %{{.*}}
+// CHECK:         %[[TEMP:.*]] = fir.allocmem !fir.array<10xi32>
 // Copy of original array to temp.
 // CHECK:         fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 // CHECK:           %[[COOR0:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
-// CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
 // CHECK:           %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
 // CHECK:           fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:         }
 // Perform the assignment i = i(10:1:-1) using the temporary array.
@@ -125,8 +125,8 @@ func @conversion_with_temporary(%arr0 : !fir.ref<!fir.array<10xi32>>) {
 // Copy the result back to the original array.
 // CHECK:         fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 // CHECK:           %[[COOR0:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
-// CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0:.*]] : !fir.ref<i32>
 // CHECK:           %[[COOR1:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0:.*]] : !fir.ref<i32>
 // CHECK:           fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:         }
 // Free temporary array.
@@ -176,7 +176,7 @@ func @conversion_with_temporary_multidim(%0: !fir.ref<!fir.array<10x5xi32>>) {
 // CHECK-SAME:                                           %[[ARR0:.*]]: !fir.ref<!fir.array<10x5xi32>>) {
 // CHECK:         %[[CST10:.*]] = arith.constant 10 : index
 // CHECK:         %[[CST5:.*]] = arith.constant 5 : index
-// CHECK:         %[[TEMP:.*]] = fir.allocmem !fir.array<10x5xi32>, %c10, %c5
+// CHECK:         %[[TEMP:.*]] = fir.allocmem !fir.array<10x5xi32>
 // CHECK:         %[[IDX5:.*]] = fir.convert %[[CST5]] : (index) -> index
 // CHECK:         %[[UB5:.*]] = arith.subi %[[IDX5]], %{{.*}} : index
 // CHECK:         fir.do_loop %[[INDUC0:.*]] = %{{.*}} to %[[UB5]] step %{{.*}} {
@@ -186,8 +186,8 @@ func @conversion_with_temporary_multidim(%0: !fir.ref<!fir.array<10x5xi32>>) {
 // CHECK:             %[[IDX1:.*]] = arith.addi %[[INDUC1]], %{{.*}} : index
 // CHECK:             %[[IDX2:.*]] = arith.addi %[[INDUC0]], %{{.*}} : index
 // CHECK:             %[[COOR0:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %[[IDX1:.*]], %[[IDX2:.*]] : (!fir.ref<!fir.array<10x5xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
-// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
 // CHECK:             %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}}, %{{.*}} : (!fir.heap<!fir.array<10x5xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
+// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
 // CHECK:             fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:         %{{.*}} = fir.do_loop %[[INDUC0:.*]] = %{{.*}} to %{{.*}} step %{{.*}} unordered iter_args(%{{.*}} = %{{.*}}) -> (!fir.array<10x5xi32>) {
 // CHECK:           %{{.*}} = fir.do_loop %[[INDUC1:.*]] = %{{.*}} to %{{.*}} step %{{.*}} unordered iter_args(%{{.*}} = %{{.*}}) -> (!fir.array<10x5xi32>) {
@@ -208,8 +208,8 @@ func @conversion_with_temporary_multidim(%0: !fir.ref<!fir.array<10x5xi32>>) {
 // CHECK:             %[[IDX1:.*]] = arith.addi %[[INDUC1]], %{{.*}} : index
 // CHECK:             %[[IDX2:.*]] = arith.addi %[[INDUC0]], %{{.*}} : index  
 // CHECK:             %[[COOR0:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %[[IDX1]], %[[IDX2]] : (!fir.heap<!fir.array<10x5xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
-// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
 // CHECK:             %[[COOR1:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}}, %{{.*}} : (!fir.ref<!fir.array<10x5xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
+// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
 // CHECK:             fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:         fir.freemem %[[TEMP]] : <!fir.array<10x5xi32>>
 
@@ -283,12 +283,12 @@ func private @user_defined_assignment(!fir.ref<f32>, !fir.ref<f32>)
 // CHECK-SAME:                               %[[ARR0:.*]]: !fir.ref<!fir.array<100xf32>>) {
 // CHECK:           %[[VAR0:.*]] = fir.alloca f32
 // Allocate the temporary array.
-// CHECK:           %[[TEMP:.*]] = fir.allocmem !fir.array<100xf32>, %{{.*}}
+// CHECK:           %[[TEMP:.*]] = fir.allocmem !fir.array<100xf32>
 // Copy original array to temp.
 // CHECK:           fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 // CHECK:             %[[COOR0:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>, index) -> !fir.ref<f32>
-// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<f32>
 // CHECK:             %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<100xf32>>, !fir.shape<1>, index) -> !fir.ref<f32>
+// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<f32>
 // CHECK:             fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<f32>
 // CHECK:           }
 // CHECK:           %[[VAL_21:.*]] = fir.undefined !fir.array<100xf32>
@@ -304,8 +304,8 @@ func private @user_defined_assignment(!fir.ref<f32>, !fir.ref<f32>)
 // Copy back result to original array from temp.
 // CHECK:           fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 // CHECK:             %[[COOR0:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<100xf32>>, !fir.shape<1>, index) -> !fir.ref<f32>
-// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<f32>
 // CHECK:             %[[COOR1:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>, index) -> !fir.ref<f32>
+// CHECK:             %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<f32>
 // CHECK:             fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<f32>
 // CHECK:           }
 // Free the temporary array.
@@ -374,8 +374,9 @@ func @array_of_types() {
 // Test fir.array_load/boxed array
 func @conversion_with_temporary_boxed_array(%arr0 : !fir.box<!fir.array<10xi32>>) {
   %c10 = arith.constant 10 : index
-  %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %1:3 = fir.box_dims %arr0, %c10 : (!fir.box<!fir.array<10xi32>>, index) -> (index, index, index)
-  %2 = fir.array_load %arr0(%1) : (!fir.box<!fir.array<10xi32>>, !fir.shape<1>) -> !fir.array<10xi32>
+  %shift = fir.shift %1#0 : (index) -> !fir.shift<1>
+  %2 = fir.array_load %arr0(%shift) : (!fir.box<!fir.array<10xi32>>, !fir.shift<1>) -> !fir.array<10xi32>
   %c10_i64 = arith.constant 10 : i64
   %3 = fir.convert %c10_i64 : (i64) -> index
   %c1_i64 = arith.constant 1 : i64
@@ -398,12 +399,12 @@ func @conversion_with_temporary_boxed_array(%arr0 : !fir.box<!fir.array<10xi32>>
 // CHECK-LABEL: func @conversion_with_temporary_boxed_array(
 // CHECK-SAME:                                              %[[ARR0:.*]]: !fir.box<!fir.array<10xi32>>)
 // Allocation of temporary array.
-// CHECK:         %[[TEMP:.*]] = fir.allocmem !fir.array<10xi32>, %{{.*}}
+// CHECK:         %[[TEMP:.*]] = fir.allocmem !fir.array<10xi32>
 // Copy of original array to temp.
 // CHECK:         fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
-// CHECK:           %[[COOR0:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.box<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[COOR0:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
 // CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
-// CHECK:           %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
 // CHECK:           fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:         }
 // Perform the assignment i = i(10:1:-1) using the temporary array.
@@ -412,15 +413,15 @@ func @conversion_with_temporary_boxed_array(%arr0 : !fir.box<!fir.array<10xi32>>
 // CHECK-NOT:       %{{.*}} = fir.update
 // CHECK:           %[[COOR0:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) [%{{.*}}] %{{.*}} : (!fir.box<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<i32>
 // CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0]] : !fir.ref<i32>
-// CHECK:           %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[COOR1:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
 // CHECK:           fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:           fir.result %{{.*}} : !fir.array<10xi32>
 // CHECK:         }
 // Copy the result back to the original array.
 // CHECK:         fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
-// CHECK:           %[[COOR0:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[COOR0:.*]] = fir.array_coor %[[TEMP]](%{{.*}}) %{{.*}} : (!fir.heap<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+// CHECK:           %[[COOR1:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
 // CHECK:           %[[LOAD0:.*]] = fir.load %[[COOR0:.*]] : !fir.ref<i32>
-// CHECK:           %[[COOR1:.*]] = fir.array_coor %[[ARR0]](%{{.*}}) %{{.*}} : (!fir.box<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
 // CHECK:           fir.store %[[LOAD0]] to %[[COOR1]] : !fir.ref<i32>
 // CHECK:         }
 // Free temporary array.