diff --git a/flang/include/flang/Optimizer/Dialect/FIRAttr.td b/flang/include/flang/Optimizer/Dialect/FIRAttr.td index 989319ff3ddaf..0c34b640a5c9c 100644 --- a/flang/include/flang/Optimizer/Dialect/FIRAttr.td +++ b/flang/include/flang/Optimizer/Dialect/FIRAttr.td @@ -70,4 +70,15 @@ def fir_BoxFieldAttr : I32EnumAttr< // mlir::SideEffects::Resource for modelling operations which add debugging information def DebuggingResource : Resource<"::fir::DebuggingResource">; +def fir_LowerBoundModifierAttribute : I32EnumAttr< + "LowerBoundModifierAttribute", + "Describes how to modify lower bounds", + [ + I32EnumAttrCase<"Preserve", 0, "preserve">, + I32EnumAttrCase<"SetToOnes", 1, "ones">, + I32EnumAttrCase<"SetToZeroes", 2, "zeroes">, + ]> { + let cppNamespace = "::fir"; +} + #endif // FIR_DIALECT_FIR_ATTRS diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td index d9c1149040066..584b7e82bf27a 100644 --- a/flang/include/flang/Optimizer/Dialect/FIROps.td +++ b/flang/include/flang/Optimizer/Dialect/FIROps.td @@ -857,6 +857,43 @@ def fir_ReboxOp : fir_Op<"rebox", [NoMemoryEffect, AttrSizedOperandSegments]> { let hasVerifier = 1; } +def fir_ReboxAssumedRankOp : fir_Op<"rebox_assumed_rank", + [DeclareOpInterfaceMethods]> { + let summary = "create an assumed-rank box given another assumed-rank box"; + + let description = [{ + Limited version of fir.rebox for assumed-rank. Only the lower bounds, + attribute, and element type may change. + + The input may be a box or a reference to a box, in which case the operation + reads the incoming reference. + Since a fir.shift cannot be built without knowing the rank statically, + lower bound changes are encoded via a LowerBoundModifierAttribute. + Attribute and element type change are encoded in the result type. + Changing the element type is only allowed if the input type is a derived + type that extends the output element type. + + Example: + ``` + fir.rebox_assumed_rank %1 lbs zeroes : (!fir.box>) -> !fir.box> + ``` + }]; + + let arguments = (ins + AnyRefOrBoxType:$box, + fir_LowerBoundModifierAttribute:$lbs_modifier + ); + + let results = (outs BoxOrClassType); + + let assemblyFormat = [{ + $box `lbs` $lbs_modifier + attr-dict `:` functional-type(operands, results) + }]; + + let hasVerifier = 1; +} + def fir_EmboxCharOp : fir_Op<"emboxchar", [NoMemoryEffect]> { let summary = "boxes a given CHARACTER reference and its LEN parameter"; diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp index 94113da9a46cf..998e9535582cb 100644 --- a/flang/lib/Optimizer/Dialect/FIROps.cpp +++ b/flang/lib/Optimizer/Dialect/FIROps.cpp @@ -2412,6 +2412,52 @@ mlir::LogicalResult fir::ReboxOp::verify() { return mlir::success(); } +//===----------------------------------------------------------------------===// +// ReboxAssumedRankOp +//===----------------------------------------------------------------------===// + +static bool areCompatibleAssumedRankElementType(mlir::Type inputEleTy, + mlir::Type outEleTy) { + if (inputEleTy == outEleTy) + return true; + // Output is unlimited polymorphic -> output dynamic type is the same as input + // type. + if (mlir::isa(outEleTy)) + return true; + // Output/Input are derived types. Assuming input extends output type, output + // dynamic type is the output static type, unless output is polymorphic. + if (mlir::isa(inputEleTy) && + mlir::isa(outEleTy)) + return true; + if (areCompatibleCharacterTypes(inputEleTy, outEleTy)) + return true; + return false; +} + +mlir::LogicalResult fir::ReboxAssumedRankOp::verify() { + mlir::Type inputType = getBox().getType(); + if (!mlir::isa(inputType) && !fir::isBoxAddress(inputType)) + return emitOpError("input must be a box or box address"); + mlir::Type inputEleTy = + mlir::cast(fir::unwrapRefType(inputType)) + .unwrapInnerType(); + mlir::Type outEleTy = + mlir::cast(getType()).unwrapInnerType(); + if (!areCompatibleAssumedRankElementType(inputEleTy, outEleTy)) + return emitOpError("input and output element types are incompatible"); + return mlir::success(); +} + +void fir::ReboxAssumedRankOp::getEffects( + llvm::SmallVectorImpl< + mlir::SideEffects::EffectInstance> + &effects) { + mlir::Value inputBox = getBox(); + if (fir::isBoxAddress(inputBox.getType())) + effects.emplace_back(mlir::MemoryEffects::Read::get(), inputBox, + mlir::SideEffects::DefaultResource::get()); +} + //===----------------------------------------------------------------------===// // ResultOp //===----------------------------------------------------------------------===// diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir index 962621c4e2e1a..a826dd49ef99d 100644 --- a/flang/test/Fir/fir-ops.fir +++ b/flang/test/Fir/fir-ops.fir @@ -900,3 +900,15 @@ fir.global @t1 {keep_my_attr = "data"} : i32 { } // CHECK-LABEL: fir.global @t1 {keep_my_attr = "data"} : i32 + +func.func @test_rebox_assumed_rank(%arg0: !fir.box> ) { + %1 = fir.rebox_assumed_rank %arg0 lbs ones : (!fir.box>) -> !fir.box> + %2 = fir.rebox_assumed_rank %arg0 lbs zeroes : (!fir.box>) -> !fir.box> + %3 = fir.rebox_assumed_rank %arg0 lbs preserve : (!fir.box>) -> !fir.box> + return +} +// CHECK-LABEL: func.func @test_rebox_assumed_rank( +// CHECK-SAME: %[[A:.*]]: !fir.box>) + // CHECK: fir.rebox_assumed_rank %[[A]] lbs ones : (!fir.box>) -> !fir.box> + // CHECK: fir.rebox_assumed_rank %[[A]] lbs zeroes : (!fir.box>) -> !fir.box> + // CHECK: fir.rebox_assumed_rank %[[A]] lbs preserve : (!fir.box>) -> !fir.box> diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir index 049e108ba992d..f1e1aa433b9b0 100644 --- a/flang/test/Fir/invalid.fir +++ b/flang/test/Fir/invalid.fir @@ -978,3 +978,27 @@ func.func @bad_box_offset(%no_addendum : !fir.ref>) { %addr1 = fir.box_offset %no_addendum derived_type : (!fir.ref>) -> !fir.llvm_ptr>> return } + +// ----- + +func.func @bad_rebox_assumed_rank_1(%arg0: !fir.ref> ) { + // expected-error@+1{{'fir.rebox_assumed_rank' op input must be a box or box address}} + %1 = fir.rebox_assumed_rank %arg0 lbs ones : (!fir.ref>) -> !fir.box> + return +} + +// ----- + +func.func @bad_rebox_assumed_rank_2(%arg0: !fir.box> ) { + // expected-error@+1{{'fir.rebox_assumed_rank' op result #0 must be box or class, but got '!fir.ref>>'}} + %1 = fir.rebox_assumed_rank %arg0 lbs ones : (!fir.box>) -> !fir.ref>> + return +} + +// ----- + +func.func @bad_rebox_assumed_rank_3(%arg0: !fir.box> ) { + // expected-error@+1{{'fir.rebox_assumed_rank' op input and output element types are incompatible}} + %1 = fir.rebox_assumed_rank %arg0 lbs ones : (!fir.box>) -> !fir.box> + return +}