[flang] Definitions of fir.pack/unpack_array operations. (#130698)

As defined in #127147.

Author: vzakhari

flang/include/flang/Optimizer/Dialect/FIRAttr.td

@@ -156,4 +156,26 @@ def fir_LocationKindAttr : EnumAttr<FIROpsDialect, fir_LocationKind, "loc_kind">
 def LocationKindArrayAttr : ArrayOfAttr<FIROpsDialect, "LocationKindArray",
     "loc_kind_array", "LocationKindAttr">;
 
+/// Optimization heuristics for fir.pack_array operation.
+def fir_PackArrayHeuristics
+    : I32BitEnumAttr<"PackArrayHeuristics", "",
+                     [
+                         /// fir.pack_array cannot be optimized based on the
+                         /// array usage pattern.
+                         I32BitEnumAttrCaseNone<"None", "none">,
+                         /// fir.pack_array can be optimized away, if the array
+                         /// is not used in a loop.
+                         I32BitEnumAttrCaseBit<"LoopOnly", 0, "loop_only">,
+]> {
+  let separator = ", ";
+  let cppNamespace = "::fir";
+  let genSpecializedAttr = 0;
+}
+
+def fir_PackArrayHeuristicsAttr
+    : EnumAttr<FIROpsDialect, fir_PackArrayHeuristics,
+               "pack_array_heuristics"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
 #endif // FIR_DIALECT_FIR_ATTRS

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

@@ -3320,4 +3320,100 @@ def fir_DummyScopeOp : fir_Op<"dummy_scope",
   let assemblyFormat = "attr-dict `:` type(results)";
 }
 
+def fir_PackArrayOp
+    : fir_Op<"pack_array", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+                            AllTypesMatch<["array", "result"]>]> {
+  let summary = "Pack non-contiguous array into a temporary";
+
+  let description = [{
+    The operation creates a new !fir.box/class<!fir.array<>> value
+    to represent either the original array or a newly allocated
+    temporary array, maybe identical to the original array by value.
+
+    Arguments:
+      - array is the original array.
+        It must have !fir.box/class<!fir.array<>> type.
+      - stack/heap attribute indicates where the temporary array
+        needs to be allocated.
+      - innermost/whole attribute identifies the contiguity mode.
+        innermost means that the repacking has to be done iff the original
+        array is not contiguous in the leading dimension.
+        whole means that the repacking has to be done iff the original
+        array is not contiguous in any dimension.
+        innermost is disallowed for 1D arrays in favor of whole.
+      - no_copy attribute indicates that the original array
+        is not copied into the temporary.
+      - typeparams specify the length parameters of the original array.
+        Even though the array is fully represented with a box, the explicit
+        length parameters might be specified to simplify computing
+        the size of the array's element in compilation time (e.g. constant
+        length parameters might be propagated after MLIR inlining).
+      - optional constraints attributes:
+        * max_size is an unsigned integer attribute specifying the maximum
+          byte size of an array that is eligible for repacking.
+        * max_element_size is an unsigned integer attribute specifying
+          the maximum byte element-size of an array that is eligible
+          for repacking.
+        * min_stride is an unsigned integer attribute specifying
+          the minimum byte stride of the innermost dimension of an array
+          that is eligible for repacking.
+      - heuristics attribute specifies conditions when the array repacking
+        may be optimized.
+  }];
+
+  let arguments = (ins AnyBoxedArray:$array, UnitAttr:$stack,
+      UnitAttr:$innermost, UnitAttr:$no_copy, OptionalAttr<UI64Attr>:$max_size,
+      OptionalAttr<UI64Attr>:$max_element_size,
+      OptionalAttr<UI64Attr>:$min_stride,
+      DefaultValuedAttr<fir_PackArrayHeuristicsAttr,
+                        "::fir::PackArrayHeuristics::None">:$heuristics,
+      Variadic<AnyIntegerType>:$typeparams);
+
+  let results = (outs AnyBoxedArray:$result);
+  let assemblyFormat = [{
+    $array (`stack` $stack^):(`heap`)?
+    (`innermost` $innermost^):(`whole`)?
+    (`no_copy` $no_copy^)?
+    (`constraints` custom<PackArrayConstraints>($max_size, $max_element_size, $min_stride)^)?
+    (`heuristics` $heuristics^)?
+    (`typeparams` $typeparams^)?
+    attr-dict `:` functional-type(operands, results)
+  }];
+
+  let hasVerifier = 1;
+}
+
+def fir_UnpackArrayOp
+    : fir_Op<"unpack_array", [SameTypeOperands,
+                              DeclareOpInterfaceMethods<
+                                  MemoryEffectsOpInterface>]> {
+  let summary = "Unpack values from temporary array into original array";
+
+  let description = [{
+    The operation is either a no-op or deallocates the temporary array,
+    and maybe copies the temporary array into the original array.
+
+    Arguments:
+      - temp is a fir.box/fir.class value produced by fir.pack_array.
+        It describes either the original array or the temporary array.
+      - original is the original array descriptor.
+      - stack/heap attribute indicates where the temporary array
+        was allocated.
+      - no_copy attribute indicates that the temporary array
+        is not copied into the original temporary array.
+  }];
+
+  let arguments = (ins AnyBoxedArray:$temp, AnyBoxedArray:$original,
+      UnitAttr:$stack, UnitAttr:$no_copy);
+
+  let assemblyFormat = [{
+    $temp `to` $original
+    (`stack` $stack^):(`heap`)?
+    (`no_copy` $no_copy^)?
+    attr-dict `:` type($original)
+  }];
+
+  let hasVerifier = 1;
+}
+
 #endif

flang/include/flang/Optimizer/Dialect/FIRTypes.td

@@ -658,5 +658,12 @@ def ArrayOrBoxOrRecord : TypeConstraint<Or<[fir_SequenceType.predicate,
     IsBaseBoxTypePred, fir_RecordType.predicate]>,
     "fir.box, fir.array or fir.type">;
 
+// Returns true iff the type is an array box or a reference to such type.
+def IsArrayBoxPred : CPred<"::fir::getBoxRank($_self) != 0">;
+
+// Any boxed array type (not a reference to a boxed array type).
+def AnyBoxedArray
+    : TypeConstraint<And<[BoxOrClassType.predicate, IsArrayBoxPred]>,
+                     "any boxed array">;
 
 #endif // FIR_DIALECT_FIR_TYPES

flang/lib/Optimizer/Dialect/FIRAttr.cpp

@@ -300,5 +300,5 @@ void FIROpsDialect::registerAttributes() {
                 FortranProcedureFlagsEnumAttr, FortranVariableFlagsAttr,
                 LowerBoundAttr, PointIntervalAttr, RealAttr, ReduceAttr,
                 SubclassAttr, UpperBoundAttr, LocationKindAttr,
-                LocationKindArrayAttr>();
+                LocationKindArrayAttr, PackArrayHeuristicsAttr>();
 }

flang/lib/Optimizer/Dialect/FIROps.cpp

@@ -380,11 +380,16 @@ llvm::LogicalResult fir::AllocMemOp::verify() {
 
 // CHARACTERs and derived types with LEN PARAMETERs are dependent types that
 // require runtime values to fully define the type of an object.
-static bool validTypeParams(mlir::Type dynTy, mlir::ValueRange typeParams) {
+static bool validTypeParams(mlir::Type dynTy, mlir::ValueRange typeParams,
+                            bool allowParamsForBox = false) {
   dynTy = fir::unwrapAllRefAndSeqType(dynTy);
-  // A box value will contain type parameter values itself.
-  if (mlir::isa<fir::BoxType>(dynTy))
-    return typeParams.size() == 0;
+  if (mlir::isa<fir::BaseBoxType>(dynTy)) {
+    // A box value will contain type parameter values itself.
+    if (!allowParamsForBox)
+      return typeParams.size() == 0;
+
+    dynTy = fir::getFortranElementType(dynTy);
+  }
   // Derived type must have all type parameters satisfied.
   if (auto recTy = mlir::dyn_cast<fir::RecordType>(dynTy))
     return typeParams.size() == recTy.getNumLenParams();
@@ -4561,6 +4566,111 @@ llvm::LogicalResult fir::DeclareOp::verify() {
   return fortranVar.verifyDeclareLikeOpImpl(getMemref());
 }
 
+//===----------------------------------------------------------------------===//
+// PackArrayOp
+//===----------------------------------------------------------------------===//
+
+llvm::LogicalResult fir::PackArrayOp::verify() {
+  mlir::Type arrayType = getArray().getType();
+  if (!validTypeParams(arrayType, getTypeparams(), /*allowParamsForBox=*/true))
+    return emitOpError("invalid type parameters");
+
+  if (getInnermost() && fir::getBoxRank(arrayType) == 1)
+    return emitOpError(
+        "'innermost' is invalid for 1D arrays, use 'whole' instead");
+  return mlir::success();
+}
+
+void fir::PackArrayOp::getEffects(
+    llvm::SmallVectorImpl<
+        mlir::SideEffects::EffectInstance<mlir::MemoryEffects::Effect>>
+        &effects) {
+  if (getStack())
+    effects.emplace_back(
+        mlir::MemoryEffects::Allocate::get(),
+        mlir::SideEffects::AutomaticAllocationScopeResource::get());
+  else
+    effects.emplace_back(mlir::MemoryEffects::Allocate::get(),
+                         mlir::SideEffects::DefaultResource::get());
+
+  if (!getNoCopy())
+    effects.emplace_back(mlir::MemoryEffects::Read::get(),
+                         mlir::SideEffects::DefaultResource::get());
+}
+
+static mlir::ParseResult
+parsePackArrayConstraints(mlir::OpAsmParser &parser, mlir::IntegerAttr &maxSize,
+                          mlir::IntegerAttr &maxElementSize,
+                          mlir::IntegerAttr &minStride) {
+  mlir::OperationName opName = mlir::OperationName(
+      fir::PackArrayOp::getOperationName(), parser.getContext());
+  struct {
+    llvm::StringRef name;
+    mlir::IntegerAttr &ref;
+  } attributes[] = {
+      {fir::PackArrayOp::getMaxSizeAttrName(opName), maxSize},
+      {fir::PackArrayOp::getMaxElementSizeAttrName(opName), maxElementSize},
+      {fir::PackArrayOp::getMinStrideAttrName(opName), minStride}};
+
+  mlir::NamedAttrList parsedAttrs;
+  if (succeeded(parser.parseOptionalAttrDict(parsedAttrs))) {
+    for (auto parsedAttr : parsedAttrs) {
+      for (auto opAttr : attributes) {
+        if (parsedAttr.getName() == opAttr.name)
+          opAttr.ref = mlir::cast<mlir::IntegerAttr>(parsedAttr.getValue());
+      }
+    }
+    return mlir::success();
+  }
+  return mlir::failure();
+}
+
+static void printPackArrayConstraints(mlir::OpAsmPrinter &p,
+                                      fir::PackArrayOp &op,
+                                      const mlir::IntegerAttr &maxSize,
+                                      const mlir::IntegerAttr &maxElementSize,
+                                      const mlir::IntegerAttr &minStride) {
+  llvm::SmallVector<mlir::NamedAttribute> attributes;
+  if (maxSize)
+    attributes.emplace_back(op.getMaxSizeAttrName(), maxSize);
+  if (maxElementSize)
+    attributes.emplace_back(op.getMaxElementSizeAttrName(), maxElementSize);
+  if (minStride)
+    attributes.emplace_back(op.getMinStrideAttrName(), minStride);
+
+  p.printOptionalAttrDict(attributes);
+}
+
+//===----------------------------------------------------------------------===//
+// UnpackArrayOp
+//===----------------------------------------------------------------------===//
+
+llvm::LogicalResult fir::UnpackArrayOp::verify() {
+  if (auto packOp = getTemp().getDefiningOp<fir::PackArrayOp>())
+    if (getStack() != packOp.getStack())
+      return emitOpError() << "the pack operation uses different memory for "
+                              "the temporary (stack vs heap): "
+                           << *packOp.getOperation() << "\n";
+  return mlir::success();
+}
+
+void fir::UnpackArrayOp::getEffects(
+    llvm::SmallVectorImpl<
+        mlir::SideEffects::EffectInstance<mlir::MemoryEffects::Effect>>
+        &effects) {
+  if (getStack())
+    effects.emplace_back(
+        mlir::MemoryEffects::Free::get(),
+        mlir::SideEffects::AutomaticAllocationScopeResource::get());
+  else
+    effects.emplace_back(mlir::MemoryEffects::Free::get(),
+                         mlir::SideEffects::DefaultResource::get());
+
+  if (!getNoCopy())
+    effects.emplace_back(mlir::MemoryEffects::Write::get(),
+                         mlir::SideEffects::DefaultResource::get());
+}
+
 //===----------------------------------------------------------------------===//
 // FIROpsDialect
 //===----------------------------------------------------------------------===//

flang/test/Fir/fir-ops.fir

@@ -942,3 +942,28 @@ func.func @test_copy(%arg0: !fir.ref<!fir.type<sometype{i:i32}>>, %arg1: !fir.pt
   fir.copy %arg0 to %arg1 no_overlap : !fir.ref<!fir.type<sometype{i:i32}>>, !fir.ptr<!fir.type<sometype{i:i32}>>
   return
 }
+
+// CHECK-LABEL:   func.func @test_pack_unpack_array(
+// CHECK-SAME:                                      %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !fir.ref<!fir.box<none>>,
+// CHECK-SAME:                                      %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !fir.box<!fir.array<?xi32>>) {
+func.func @test_pack_unpack_array(%arg0: !fir.ref<!fir.box<none>>, %arg1: !fir.box<!fir.array<?xi32>>) {
+// CHECK:           %[[VAL_2:.*]] = fir.pack_array %[[VAL_1]] heap whole : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %0 = fir.pack_array %arg1 heap whole : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %1 = fir.convert %0 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
+  fir.store %1 to %arg0 : !fir.ref<!fir.box<none>>
+// CHECK:           %[[VAL_4:.*]] = fir.pack_array %[[VAL_1]] stack whole : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %2 = fir.pack_array %arg1 stack whole : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %3 = fir.convert %2 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
+  fir.store %3 to %arg0 : !fir.ref<!fir.box<none>>
+// CHECK:           %[[VAL_6:.*]] = fir.pack_array %[[VAL_1]] heap whole no_copy : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %4 = fir.pack_array %arg1 heap whole no_copy constraints {} heuristics <none> : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %5 = fir.convert %4 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
+  fir.store %5 to %arg0 : !fir.ref<!fir.box<none>>
+// CHECK:           %[[VAL_8:.*]] = fir.pack_array %[[VAL_1]] stack whole constraints  {max_size = 100 : ui64, max_element_size = 1 : ui64, min_stride = 10 : ui64} heuristics <loop_only> : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %6 = fir.pack_array %arg1 stack whole constraints {max_size = 100 : ui64, max_element_size = 1 : ui64, min_stride = 10 : ui64} heuristics <loop_only> : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+  %7 = fir.convert %6 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
+  fir.store %7 to %arg0 : !fir.ref<!fir.box<none>>
+// CHECK:           fir.unpack_array %[[VAL_8]] to %[[VAL_1]] stack no_copy : !fir.box<!fir.array<?xi32>>
+  fir.unpack_array %6 to %arg1 stack no_copy : !fir.box<!fir.array<?xi32>>
+  return
+}