diff --git a/flang/include/flang/Optimizer/Dialect/FIRAttr.td b/flang/include/flang/Optimizer/Dialect/FIRAttr.td index e3474da6685af..8e86d82f38df4 100644 --- a/flang/include/flang/Optimizer/Dialect/FIRAttr.td +++ b/flang/include/flang/Optimizer/Dialect/FIRAttr.td @@ -156,4 +156,26 @@ def fir_LocationKindAttr : EnumAttr def LocationKindArrayAttr : ArrayOfAttr; +/// 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 { + let assemblyFormat = "`<` $value `>`"; +} + #endif // FIR_DIALECT_FIR_ATTRS diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td index 8325468c4b210..7147a2401baa7 100644 --- a/flang/include/flang/Optimizer/Dialect/FIROps.td +++ b/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, + AllTypesMatch<["array", "result"]>]> { + let summary = "Pack non-contiguous array into a temporary"; + + let description = [{ + The operation creates a new !fir.box/class> 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> 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:$max_size, + OptionalAttr:$max_element_size, + OptionalAttr:$min_stride, + DefaultValuedAttr:$heuristics, + Variadic:$typeparams); + + let results = (outs AnyBoxedArray:$result); + let assemblyFormat = [{ + $array (`stack` $stack^):(`heap`)? + (`innermost` $innermost^):(`whole`)? + (`no_copy` $no_copy^)? + (`constraints` custom($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 diff --git a/flang/include/flang/Optimizer/Dialect/FIRTypes.td b/flang/include/flang/Optimizer/Dialect/FIRTypes.td index 41e765c1cb7b9..fd5bbbe44751f 100644 --- a/flang/include/flang/Optimizer/Dialect/FIRTypes.td +++ b/flang/include/flang/Optimizer/Dialect/FIRTypes.td @@ -658,5 +658,12 @@ def ArrayOrBoxOrRecord : TypeConstraint, "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, + "any boxed array">; #endif // FIR_DIALECT_FIR_TYPES diff --git a/flang/lib/Optimizer/Dialect/FIRAttr.cpp b/flang/lib/Optimizer/Dialect/FIRAttr.cpp index 4c78e223b4178..d190b307dc7f2 100644 --- a/flang/lib/Optimizer/Dialect/FIRAttr.cpp +++ b/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>(); } diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp index 203a72af61b92..033d6453a619a 100644 --- a/flang/lib/Optimizer/Dialect/FIROps.cpp +++ b/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(dynTy)) - return typeParams.size() == 0; + if (mlir::isa(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(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> + &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(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 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()) + 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> + &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 //===----------------------------------------------------------------------===// diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir index 06b0bbbf0bd20..c5bcba9103db1 100644 --- a/flang/test/Fir/fir-ops.fir +++ b/flang/test/Fir/fir-ops.fir @@ -942,3 +942,28 @@ func.func @test_copy(%arg0: !fir.ref>, %arg1: !fir.pt fir.copy %arg0 to %arg1 no_overlap : !fir.ref>, !fir.ptr> return } + +// CHECK-LABEL: func.func @test_pack_unpack_array( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !fir.ref>, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !fir.box>) { +func.func @test_pack_unpack_array(%arg0: !fir.ref>, %arg1: !fir.box>) { +// CHECK: %[[VAL_2:.*]] = fir.pack_array %[[VAL_1]] heap whole : (!fir.box>) -> !fir.box> + %0 = fir.pack_array %arg1 heap whole : (!fir.box>) -> !fir.box> + %1 = fir.convert %0 : (!fir.box>) -> !fir.box + fir.store %1 to %arg0 : !fir.ref> +// CHECK: %[[VAL_4:.*]] = fir.pack_array %[[VAL_1]] stack whole : (!fir.box>) -> !fir.box> + %2 = fir.pack_array %arg1 stack whole : (!fir.box>) -> !fir.box> + %3 = fir.convert %2 : (!fir.box>) -> !fir.box + fir.store %3 to %arg0 : !fir.ref> +// CHECK: %[[VAL_6:.*]] = fir.pack_array %[[VAL_1]] heap whole no_copy : (!fir.box>) -> !fir.box> + %4 = fir.pack_array %arg1 heap whole no_copy constraints {} heuristics : (!fir.box>) -> !fir.box> + %5 = fir.convert %4 : (!fir.box>) -> !fir.box + fir.store %5 to %arg0 : !fir.ref> +// 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 : (!fir.box>) -> !fir.box> + %6 = fir.pack_array %arg1 stack whole constraints {max_size = 100 : ui64, max_element_size = 1 : ui64, min_stride = 10 : ui64} heuristics : (!fir.box>) -> !fir.box> + %7 = fir.convert %6 : (!fir.box>) -> !fir.box + fir.store %7 to %arg0 : !fir.ref> +// CHECK: fir.unpack_array %[[VAL_8]] to %[[VAL_1]] stack no_copy : !fir.box> + fir.unpack_array %6 to %arg1 stack no_copy : !fir.box> + return +} diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir index feb2cd55b3786..1876f8c7d965e 100644 --- a/flang/test/Fir/invalid.fir +++ b/flang/test/Fir/invalid.fir @@ -1055,3 +1055,94 @@ func.func @bad_copy_4(%arg0: !fir.ref, %arg1: !fir.ref) { fir.copy %arg0 to %arg1 no_overlap : !fir.ref, !fir.ref return } + +// ----- + +func.func @bad_pack_array1(%arg0: !fir.ref>>) { + // expected-error@+1{{op operand #0 must be any boxed array, but got '!fir.ref>>'}} + %0 = fir.pack_array %arg0 stack whole : (!fir.ref>>) -> !fir.ref>> + return +} + +// ----- + +func.func @bad_pack_array2(%arg0: !fir.box>) { + // expected-error@+1{{op failed to verify that all of {array, result} have same type}} + %0 = fir.pack_array %arg0 stack whole : (!fir.box>) -> !fir.box> + return +} + +// ----- + +func.func @bad_pack_array3(%arg0: !fir.box>>, %arg1: i32) { + // expected-error@+1{{op invalid type parameters}} + %0 = fir.pack_array %arg0 stack whole typeparams %arg1, %arg1 : (!fir.box>>, i32, i32) -> !fir.box>> + return +} + +// ----- + +func.func @bad_pack_array4(%arg0: !fir.box>) { + // expected-error@+1{{op attribute 'max_size' failed to satisfy constraint: 64-bit unsigned integer attribute}} + %0 = fir.pack_array %arg0 stack whole constraints {max_size = -1 : i64} : (!fir.box>) -> !fir.box> + return +} + +// ----- + +func.func @bad_pack_array5(%arg0: !fir.box>) { + // expected-error@+1{{op attribute 'max_element_size' failed to satisfy constraint: 64-bit unsigned integer attribute}} + %0 = fir.pack_array %arg0 stack whole constraints {max_element_size = -1 : i64} : (!fir.box>) -> !fir.box> + return +} + +// ----- + +func.func @bad_pack_array6(%arg0: !fir.box>) { + // expected-error@+1{{op attribute 'min_stride' failed to satisfy constraint: 64-bit unsigned integer attribute}} + %0 = fir.pack_array %arg0 stack whole constraints {min_stride = -1 : i64} : (!fir.box>) -> !fir.box> + return +} + +// ----- + +func.func @bad_pack_array7(%arg0: !fir.box>) { + // expected-error@+1{{op 'innermost' is invalid for 1D arrays, use 'whole' instead}} + %0 = fir.pack_array %arg0 stack innermost : (!fir.box>) -> !fir.box> + return +} + +// ----- + +func.func @bad_unpack_array1(%arg0: !fir.ref>>, %arg1: !fir.ref>>) { + // expected-error@+1{{op operand #0 must be any boxed array, but got '!fir.ref>>'}} + fir.unpack_array %arg0 to %arg1 stack : !fir.ref>> + return +} + +// ----- + +func.func @bad_unpack_array2(%arg0: !fir.box>) { + %0 = fir.pack_array %arg0 stack whole : (!fir.box>) -> !fir.box> + // expected-error@+1{{op the pack operation uses different memory for the temporary (stack vs heap)}} + fir.unpack_array %0 to %arg0 heap no_copy : !fir.box> + return +} + +// ----- + +func.func @bad_unpack_array3(%arg0: !fir.box>) { + %0 = fir.pack_array %arg0 heap whole no_copy: (!fir.box>) -> !fir.box> + // expected-error@+1{{op the pack operation uses different memory for the temporary (stack vs heap)}} + fir.unpack_array %0 to %arg0 stack : !fir.box> + return +} + +// ----- + +func.func @bad_unpack_array4(%arg0: !fir.box>, %arg1: !fir.box>) { + // expected-note@-1 {{prior use here}} + // expected-error@+1{{use of value '%arg0' expects different type than prior uses: '!fir.box>' vs '!fir.box>'}} + fir.unpack_array %arg0 to %arg1 stack : !fir.box> + return +}