[fir] Update fir.array_update, fir.array_fetch and fir.array_merge_store

[clementval]

Summary

Add typeparams to fir.array_update, fir.array_fetch and
fir.array_merge_store operations. Add optional slice operands to
fir.array_merge_store op.
Move verifiers to cpp file.

Co-authored-by: Eric Schweitz <[e-mail]>

[clementval]

Add missing new file

[kiranchandramohan]

A few Nits.
-> The commit message does not mention moving the verifier to cpp file.
-> Are there additional verify tests?

flang/include/flang/Optimizer/Support/Utils.h

27 ↗

(On Diff #375858)

Nit: Newline

[clementval]

Add couple of invalid tests

[clementval]

In D110701#3030749, @kiranchandramohan wrote:

A few Nits.
-> The commit message does not mention moving the verifier to cpp file.
-> Are there additional verify tests?

Thanks for the comments. I updated the message and added couple of invalid tests for the verifiers,

[kiranchandramohan]

LGTM.

[clementval]

Remove duplicated test

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

def fir_ArrayLoadOp : fir_Op<"array_load", [AttrSizedOperandSegments]> {
 
   let verifier = [{ return ::verify(*this); }];
 
   let extraClassDeclaration = [{
     std::vector<mlir::Value> getExtents();
   }];
 }
 
-def fir_ArrayFetchOp : fir_Op<"array_fetch", [NoSideEffect]> {
+def fir_ArrayFetchOp : fir_Op<"array_fetch", [AttrSizedOperandSegments,
+    NoSideEffect]> {
 
   let summary = "Fetch the value of an element of an array value";
 
   let description = [{
     Fetch the value of an element in an array value.
 
     ```fortran
       real :: a(n,m)
```mlir
       %1 = fir.array_fetch %v, %i, %j : (!fir.array<?x?xf32>, index, index) -> f32
     ```
 
     It is only possible to use `array_fetch` on an `array_load` result value.
   }];
 
   let arguments = (ins
     fir_SequenceType:$sequence,
-    Variadic<AnyCoordinateType>:$indices
+    Variadic<AnyCoordinateType>:$indices,
+    Variadic<AnyIntegerType>:$typeparams
   );
 
   let results = (outs AnyType:$element);
 
   let assemblyFormat = [{
-    $sequence `,` $indices attr-dict `:` functional-type(operands, results)
+    $sequence `,` $indices (`typeparams` $typeparams^)? attr-dict `:`
+      functional-type(operands, results)
   }];
 
-  let verifier = [{
-    auto arrTy = sequence().getType().cast<fir::SequenceType>();
-    if (indices().size() != arrTy.getDimension())
-      return emitOpError("number of indices != dimension of array");
-    if (element().getType() != arrTy.getEleTy())
-      return emitOpError("return type does not match array");
-    if (!isa<fir::ArrayLoadOp>(sequence().getDefiningOp()))
-      return emitOpError("argument #0 must be result of fir.array_load");
-    return mlir::success();
-  }];
+  let verifier = "return ::verify(*this);";
 }
 
-def fir_ArrayUpdateOp : fir_Op<"array_update", [NoSideEffect]> {
+def fir_ArrayUpdateOp : fir_Op<"array_update", [AttrSizedOperandSegments,
+    NoSideEffect]> {
 
   let summary = "Update the value of an element of an array value";
 
   let description = [{
     Updates the value of an element in an array value. A new array value is
     returned where all element values of the input array are identical except
     for the selected element which is the value passed in the update.
 
let description = [{
     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,
     AnyType:$merge,
-    Variadic<AnyCoordinateType>:$indices
+    Variadic<AnyCoordinateType>:$indices,
+    Variadic<AnyIntegerType>:$typeparams
   );
 
   let results = (outs fir_SequenceType);
 
   let assemblyFormat = [{
-    $sequence `,` $merge `,` $indices attr-dict `:` functional-type(operands, results)
+    $sequence `,` $merge `,` $indices (`typeparams` $typeparams^)? attr-dict
+      `:` functional-type(operands, results)
   }];
 
-  let verifier = [{
-    auto arrTy = sequence().getType().cast<fir::SequenceType>();
-    if (merge().getType() != arrTy.getEleTy())
-      return emitOpError("merged value does not have element type");
-    if (indices().size() != arrTy.getDimension())
-      return emitOpError("number of indices != dimension of array");
-    return mlir::success();
-  }];
+  let verifier = "return ::verify(*this);";
 }
 
-def fir_ArrayMergeStoreOp : fir_Op<"array_merge_store", [
-    TypesMatchWith<"type of 'original' matches element type of 'memref'",
-                     "memref", "original",
-                     "fir::dyn_cast_ptrOrBoxEleTy($_self)">,
-    TypesMatchWith<"type of 'sequence' matches element type of 'memref'",
-                     "memref", "sequence",
-                     "fir::dyn_cast_ptrOrBoxEleTy($_self)">]> {
+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.
 
     ```fortran
       real :: a(n,m)
let description = [{
 
     This operation merges the original loaded array value, `%v`, with the
     chained updates, `%r`, and stores the result to the array at address, `%a`.
   }];
 
   let arguments = (ins
     fir_SequenceType:$original,
     fir_SequenceType:$sequence,
-    Arg<AnyRefOrBox, "", [MemWrite]>:$memref
+    Arg<AnyRefOrBox, "", [MemWrite]>:$memref,
+    Optional<fir_SliceType>:$slice,
+    Variadic<AnyIntegerType>:$typeparams
   );
 
-  let assemblyFormat = "$original `,` $sequence `to` $memref attr-dict `:` type($memref)";
-
-  let verifier = [{
-    if (!isa<ArrayLoadOp>(original().getDefiningOp()))
-       return emitOpError("operand #0 must be result of a fir.array_load op");
-    return mlir::success();
+  let assemblyFormat = [{
+    $original `,` $sequence `to` $memref (`[` $slice^ `]`)? (`typeparams`
+      $typeparams^)? attr-dict `:` type(operands)
   }];
+
+  let verifier = "return ::verify(*this);";
 }
 
 //===----------------------------------------------------------------------===//
 // Record and array type operations
 //===----------------------------------------------------------------------===//
 
 def fir_ArrayCoorOp : fir_Op<"array_coor",
     [NoSideEffect, AttrSizedOperandSegments]> {

flang/include/flang/Optimizer/Dialect/FIRType.h

 class hash_code;
 } // namespace llvm
 
 namespace mlir {
 class DialectAsmParser;
 class DialectAsmPrinter;
 class ComplexType;
 class FloatType;
+class ValueRange;
 } // namespace mlir
 
 namespace fir {
 
 class FIROpsDialect;
 
 using KindTy = unsigned;
 

 /// Index, or FIR Int). Aborts execution if condition is false.
 void verifyIntegralType(mlir::Type type);
 
 /// Is `t` a FIR or MLIR Complex type?
 inline bool isa_complex(mlir::Type t) {
   return t.isa<fir::ComplexType>() || t.isa<mlir::ComplexType>();
 }
 
+/// Is `t` a CHARACTER type? Does not check the length.
+inline bool isa_char(mlir::Type t) { return t.isa<fir::CharacterType>(); }
+
 /// Is `t` a CHARACTER type with a LEN other than 1?
 inline bool isa_char_string(mlir::Type t) {
   if (auto ct = t.dyn_cast_or_null<fir::CharacterType>())
     return ct.getLen() != fir::CharacterType::singleton();
   return false;
 }
 
 /// Is `t` a box type for which it is not possible to deduce the box size?
 /// It is not possible to deduce the size of a box that describes an entity
 /// of unknown rank or type.
 bool isa_unknown_size_box(mlir::Type t);
 
+/// If `t` is a SequenceType return its element type, otherwise return `t`.
+inline mlir::Type unwrapSequenceType(mlir::Type t) {
+  if (auto seqTy = t.dyn_cast<fir::SequenceType>())
+    return seqTy.getEleTy();
+  return t;
+}
+
 #ifndef NDEBUG
 // !fir.ptr<X> and !fir.heap<X> where X is !fir.ptr, !fir.heap, or !fir.ref
 // is undefined and disallowed.
 inline bool singleIndirectionLevel(mlir::Type ty) {
   return !fir::isa_ref_type(ty);
 }
 #endif
 
+/// Apply the components specified by `path` to `rootTy` to determine the type
+/// of the resulting component element. `rootTy` should be an aggregate type.
+/// Returns null on error.
+mlir::Type applyPathToType(mlir::Type rootTy, mlir::ValueRange path);
+
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_DIALECT_FIRTYPE_H

flang/lib/Optimizer/Dialect/FIROps.cpp

 //===-- FIROps.cpp --------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Support/Utils.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/TypeSwitch.h"
static mlir::LogicalResult verify(fir::ArrayCoorOp op) {
 
   return mlir::success();
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayLoadOp
 //===----------------------------------------------------------------------===//
 
+static mlir::Type adjustedElementType(mlir::Type t) {
+  if (auto ty = t.dyn_cast<fir::ReferenceType>()) {
+    auto eleTy = ty.getEleTy();
+    if (fir::isa_char(eleTy))
+      return eleTy;
+    if (fir::isa_derived(eleTy))
+      return eleTy;
+    if (eleTy.isa<fir::SequenceType>())
+      return eleTy;
+  }
+  return t;
+}
+
 std::vector<mlir::Value> fir::ArrayLoadOp::getExtents() {
   if (auto sh = shape())
     if (auto *op = sh.getDefiningOp()) {
       if (auto shOp = dyn_cast<fir::ShapeOp>(op))
         return shOp.getExtents();
       return cast<fir::ShapeShiftOp>(op).getExtents();
     }
   return {};
if (auto sliceOp = op.slice())
     if (auto sliceTy = sliceOp.getType().dyn_cast<fir::SliceType>())
       if (sliceTy.getRank() != arrDim)
         return op.emitOpError("rank of dimension in slice mismatched");
 
   return mlir::success();
 }
 
 //===----------------------------------------------------------------------===//
+// ArrayMergeStoreOp
+//===----------------------------------------------------------------------===//
+
+static mlir::LogicalResult verify(fir::ArrayMergeStoreOp op) {
+  if (!isa<ArrayLoadOp>(op.original().getDefiningOp()))
+    return op.emitOpError("operand #0 must be result of a fir.array_load op");
+  if (auto sl = op.slice()) {
+    if (auto *slOp = sl.getDefiningOp()) {
+      auto sliceOp = mlir::cast<fir::SliceOp>(slOp);
+      if (!sliceOp.fields().empty()) {
+        // This is an intra-object merge, where the slice is projecting the
+        // subfields that are to be overwritten by the merge operation.
+        auto eleTy = fir::dyn_cast_ptrOrBoxEleTy(op.memref().getType());
+        if (auto seqTy = eleTy.dyn_cast<fir::SequenceType>()) {
+          auto projTy =
+              fir::applyPathToType(seqTy.getEleTy(), sliceOp.fields());
+          if (fir::unwrapSequenceType(op.original().getType()) != projTy)
+            return op.emitOpError(
+                "type of origin does not match sliced memref type");
+          if (fir::unwrapSequenceType(op.sequence().getType()) != projTy)
+            return op.emitOpError(
+                "type of sequence does not match sliced memref type");
+          return mlir::success();
+        }
+        return op.emitOpError("referenced type is not an array");
+      }
+    }
+    return mlir::success();
+  }
+  auto eleTy = fir::dyn_cast_ptrOrBoxEleTy(op.memref().getType());
+  if (op.original().getType() != eleTy)
+    return op.emitOpError("type of origin does not match memref element type");
+  if (op.sequence().getType() != eleTy)
+    return op.emitOpError(
+        "type of sequence does not match memref element type");
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
+// ArrayFetchOp
+//===----------------------------------------------------------------------===//
+
+// Template function used for both array_fetch and array_update verification.
+template <typename A>
+mlir::Type validArraySubobject(A op) {
+  auto ty = op.sequence().getType();
+  return fir::applyPathToType(ty, op.indices());
+}
+
+static mlir::LogicalResult verify(fir::ArrayFetchOp op) {
+  auto arrTy = op.sequence().getType().cast<fir::SequenceType>();
+  auto indSize = op.indices().size();
+  if (indSize < arrTy.getDimension())
+    return op.emitOpError("number of indices != dimension of array");
+  if (indSize == arrTy.getDimension() &&
+      ::adjustedElementType(op.element().getType()) != arrTy.getEleTy())
+    return op.emitOpError("return type does not match array");
+  auto ty = validArraySubobject(op);
+  if (!ty || ty != ::adjustedElementType(op.getType()))
+    return op.emitOpError("return type and/or indices do not type check");
+  if (!isa<fir::ArrayLoadOp>(op.sequence().getDefiningOp()))
+    return op.emitOpError("argument #0 must be result of fir.array_load");
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
+// ArrayUpdateOp
+//===----------------------------------------------------------------------===//
+
+static mlir::LogicalResult verify(fir::ArrayUpdateOp op) {
+  auto arrTy = op.sequence().getType().cast<fir::SequenceType>();
+  auto indSize = op.indices().size();
+  if (indSize < arrTy.getDimension())
+    return op.emitOpError("number of indices != dimension of array");
+  if (indSize == arrTy.getDimension() &&
+      ::adjustedElementType(op.merge().getType()) != arrTy.getEleTy())
+    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();
+}
+
+//===----------------------------------------------------------------------===//
 // 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))
if (auto addressOfOp = mlir::dyn_cast<fir::AddrOfOp>(definingOp)) {
           return globalOp->hasAttr(attributeName);
     }
   }
   // TODO: Construct associated entities attributes. Decide where the fir
   // attributes must be placed/looked for in this case.
   return false;
 }
 
+mlir::Type fir::applyPathToType(mlir::Type eleTy, mlir::ValueRange path) {
+  for (auto i = path.begin(), end = path.end(); eleTy && i < end;) {
+    eleTy = llvm::TypeSwitch<mlir::Type, mlir::Type>(eleTy)
+                .Case<fir::RecordType>([&](fir::RecordType ty) {
+                  if (auto *op = (*i++).getDefiningOp()) {
+                    if (auto off = mlir::dyn_cast<fir::FieldIndexOp>(op))
+                      return ty.getType(off.getFieldName());
+                    if (auto off = mlir::dyn_cast<mlir::ConstantOp>(op))
+                      return ty.getType(fir::toInt(off));
+                  }
+                  return mlir::Type{};
+                })
+                .Case<fir::SequenceType>([&](fir::SequenceType ty) {
+                  bool valid = true;
+                  const auto rank = ty.getDimension();
+                  for (std::remove_const_t<decltype(rank)> ii = 0;
+                       valid && ii < rank; ++ii)
+                    valid = i < end && fir::isa_integer((*i++).getType());
+                  return valid ? ty.getEleTy() : mlir::Type{};
+                })
+                .Case<mlir::TupleType>([&](mlir::TupleType ty) {
+                  if (auto *op = (*i++).getDefiningOp())
+                    if (auto off = mlir::dyn_cast<mlir::ConstantOp>(op))
+                      return ty.getType(fir::toInt(off));
+                  return mlir::Type{};
+                })
+                .Case<fir::ComplexType>([&](fir::ComplexType ty) {
+                  if (fir::isa_integer((*i++).getType()))
+                    return ty.getElementType();
+                  return mlir::Type{};
+                })
+                .Case<mlir::ComplexType>([&](mlir::ComplexType ty) {
+                  if (fir::isa_integer((*i++).getType()))
+                    return ty.getElementType();
+                  return mlir::Type{};
+                })
+                .Default([&](const auto &) { return mlir::Type{}; });
+  }
+  return eleTy;
+}
+
 // Tablegen operators
 
 #define GET_OP_CLASSES
 #include "flang/Optimizer/Dialect/FIROps.cpp.inc"

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) {
   %arr2 = fir.zero_bits !fir.array<10xi32>
     %arr3 = fir.insert_on_range %arr2, %c1_i32, [2 : index, 9 : index] : (!fir.array<10xi32>, i32) -> !fir.array<10xi32>
   fir.call @noret1(%arr3) : (!fir.array<10xi32>) -> ()
 
   // CHECK: [[SHAPE:%.*]] = fir.shape_shift [[INDXM:%.*]], [[INDXN:%.*]], [[INDXO:%.*]], [[INDXP:%.*]] : (index, index, index, index) -> !fir.shapeshift<2>
   // CHECK: [[AV1:%.*]] = fir.array_load [[ARR1]]([[SHAPE]]) : (!fir.ref<!fir.array<?x?xf32>>, !fir.shapeshift<2>) -> !fir.array<?x?xf32>
   // CHECK: [[FVAL:%.*]] = fir.array_fetch [[AV1]], [[I10]], [[J20]] : (!fir.array<?x?xf32>, index, index) -> f32
   // 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.ref<!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.ref<!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>>
 
   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