diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td index 495d72863db56..6e6af7126bf37 100644 --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td @@ -868,7 +868,7 @@ def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> { let hasVerifier = 1; } -def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp"]>, +def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", "ElementalAddrOp"]>, SingleBlockImplicitTerminator<"fir::FirEndOp">]> { let summary = "Yield a value or variable inside a forall, where or region assignment"; @@ -898,4 +898,68 @@ def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp" let assemblyFormat = "$entity attr-dict `:` type($entity) custom($cleanup)"; } +def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"RegionAssignOp">, RecursiveMemoryEffects, RecursivelySpeculatable]> { + let summary = "Yield the address of a vector subscripted variable inside an hlfir.region_assign"; + let description = [{ + Special terminator node for the left-hand side region of an hlfir.region_assign + to a vector subscripted entity. + + It represents how the address of an element of such entity is computed given + one based indices. + + It is very similar to hlfir.elemental, except that it does not produce an SSA + value because there is no hlfir type to describe a vector subscripted entity + (the codegen of such type would be problematic). Hence, it is tightly linked + to an hlfir.region_assign by its terminator property. + + An optional cleanup region may be provided if any of the subscript expressions + of the designator require a cleanup. + This allows documenting cleanups that cannot be generated after the vector + subscripted designator usage (that has not been materizaled yet). The cleanups + will be evaluated after the assignment once the related + hlfir.region_assign is lowered. + + Example: "X(VECTOR) = Y" + + ``` + hlfir.region_assign { + hlfir.yield %y : !fir.ref> + } to { + hlfir.elemental_addr %vector_shape : !fir.shape<1> { + ^bb0(%i: index): + %0 = hlfir.designate %vector (%i) : (!fir.ref>, index) -> !fir.ref + %1 = fir.load %0 : !fir.ref + %x_element_addr = hlfir.designate %x (%1) : (!fir.ref>, i32) -> !fir.ref + hlfir.yield %x_element_addr : !fir.ref + } + } + ``` + }]; + + let arguments = (ins + fir_ShapeType:$shape, + Variadic:$typeparams); + + let regions = (region SizedRegion<1>:$body, + MaxSizedRegion<1>:$cleanup); + + let builders = [ + OpBuilder<(ins "mlir::Value":$shape)> + ]; + + let assemblyFormat = [{ + $shape (`typeparams` $typeparams^)? + attr-dict `:` type(operands) $body + custom($cleanup)}]; + + let extraClassDeclaration = [{ + mlir::Region::BlockArgListType getIndices() { + return getBody().getArguments(); + } + }]; + + let hasVerifier = 1; +} + + #endif // FORTRAN_DIALECT_HLFIR_OPS diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp index 51d18def03478..5b10ecee75f41 100644 --- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp +++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp @@ -1032,8 +1032,8 @@ mlir::LogicalResult hlfir::RegionAssignOp::verify() { if (!mlir::isa_and_nonnull(getTerminator(getRhsRegion()))) return emitOpError( "right-hand side region must be terminated by an hlfir.yield"); - // TODO: allow hlfir.elemental_addr. - if (!mlir::isa_and_nonnull(getTerminator(getLhsRegion()))) + if (!mlir::isa_and_nonnull( + getTerminator(getLhsRegion()))) return emitOpError("left-hand side region must be terminated by an " "hlfir.yield or hlfir.elemental_addr"); return mlir::success(); @@ -1052,7 +1052,6 @@ static mlir::ParseResult parseYieldOpCleanup(mlir::OpAsmParser &parser, hlfir::YieldOp::ensureTerminator(cleanup, parser.getBuilder(), parser.getBuilder().getUnknownLoc()); } - return mlir::success(); } @@ -1066,5 +1065,41 @@ static void printYieldOpCleanup(mlir::OpAsmPrinter &p, YieldOp yieldOp, } } +//===----------------------------------------------------------------------===// +// ElementalAddrOp +//===----------------------------------------------------------------------===// + +void hlfir::ElementalAddrOp::build(mlir::OpBuilder &builder, + mlir::OperationState &odsState, + mlir::Value shape) { + odsState.addOperands(shape); + mlir::Region *bodyRegion = odsState.addRegion(); + bodyRegion->push_back(new mlir::Block{}); + if (auto shapeType = shape.getType().dyn_cast()) { + unsigned dim = shapeType.getRank(); + mlir::Type indexType = builder.getIndexType(); + for (unsigned d = 0; d < dim; ++d) + bodyRegion->front().addArgument(indexType, odsState.location); + } + // Push cleanUp region. + odsState.addRegion(); +} + +mlir::LogicalResult hlfir::ElementalAddrOp::verify() { + hlfir::YieldOp yieldOp = + mlir::dyn_cast_or_null(getTerminator(getBody())); + if (!yieldOp) + return emitOpError("body region must be terminated by an hlfir.yield"); + mlir::Type elementAddrType = yieldOp.getEntity().getType(); + if (!hlfir::isFortranVariableType(elementAddrType) || + hlfir::getFortranElementOrSequenceType(elementAddrType) + .isa()) + return emitOpError("body must compute the address of a scalar entity"); + unsigned shapeRank = getShape().getType().cast().getRank(); + if (shapeRank != getIndices().size()) + return emitOpError("body number of indices must match shape rank"); + return mlir::success(); +} + #define GET_OP_CLASSES #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc" diff --git a/flang/test/HLFIR/element-addr.fir b/flang/test/HLFIR/element-addr.fir new file mode 100644 index 0000000000000..c6f2b06b9f688 --- /dev/null +++ b/flang/test/HLFIR/element-addr.fir @@ -0,0 +1,83 @@ +// Test hlfir.element_addr operation parse, verify (no errors), and unparse. +// RUN: fir-opt %s | fir-opt | FileCheck %s + +func.func @test_element_addr(%x: !fir.ref>, %vector: !fir.ref>, %y: !fir.ref>) { + %c20 = arith.constant 20 : index + %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1> + hlfir.region_assign { + hlfir.yield %y : !fir.ref> + } to { + hlfir.elemental_addr %vector_shape : !fir.shape<1> { + ^bb0(%i: index): + %0 = hlfir.designate %vector (%i) : (!fir.ref>, index) -> !fir.ref + %1 = fir.load %0 : !fir.ref + %x_element_addr = hlfir.designate %x (%1) : (!fir.ref>, i32) -> !fir.ref + hlfir.yield %x_element_addr : !fir.ref + } + } + return +} +// CHECK-LABEL: func.func @test_element_addr( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref>, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.ref>, +// CHECK-SAME: %[[VAL_2:.*]]: !fir.ref>) { +// CHECK: %[[VAL_3:.*]] = arith.constant 20 : index +// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1> +// CHECK: hlfir.region_assign { +// CHECK: hlfir.yield %[[VAL_2]] : !fir.ref> +// CHECK: } to { +// CHECK: hlfir.elemental_addr %[[VAL_4]] : !fir.shape<1> { +// CHECK: ^bb0(%[[VAL_5:.*]]: index): +// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_5]]) : (!fir.ref>, index) -> !fir.ref +// CHECK: %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref +// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_7]]) : (!fir.ref>, i32) -> !fir.ref +// CHECK: hlfir.yield %[[VAL_8]] : !fir.ref +// CHECK: } +// CHECK: } + + + +func.func @test_element_addr_cleanup(%x: !fir.box>>, %y: !fir.box>>) { + hlfir.region_assign { + hlfir.yield %y : !fir.box>> + } to { + %c0 = arith.constant 0 : index + %len = fir.box_elesize %x : (!fir.box>>) -> index + %vector = fir.call @returns_alloc() : () -> !fir.box>> + %vector_dim:3 = fir.box_dims %x, %c0 : (!fir.box>>, index) -> (index, index, index) + %vector_shape = fir.shape %vector_dim#1 : (index) -> !fir.shape<1> + hlfir.elemental_addr %vector_shape typeparams %len : !fir.shape<1>, index { + ^bb0(%i: index): + %0 = hlfir.designate %vector (%i) : (!fir.box>>, index) -> !fir.ref + %1 = fir.load %0 : !fir.ref + %x_element_addr = hlfir.designate %x (%1) typeparams %len : (!fir.box>>, i32, index) -> !fir.boxchar<1> + hlfir.yield %x_element_addr : !fir.boxchar<1> + } cleanup { + %addr = fir.box_addr %vector : (!fir.box>>) -> !fir.heap> + fir.freemem %addr : !fir.heap> + } + } + return +} +// CHECK-LABEL: func.func @test_element_addr_cleanup( +// CHECK-SAME: %[[VAL_0:[^:]*]]: !fir.box>>, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.box>>) { +// CHECK: hlfir.region_assign { +// CHECK: hlfir.yield %[[VAL_1]] : !fir.box>> +// CHECK: } to { +// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_3:.*]] = fir.box_elesize %[[VAL_0]] : (!fir.box>>) -> index +// CHECK: %[[VAL_4:.*]] = fir.call @returns_alloc() : () -> !fir.box>> +// CHECK: %[[VAL_5:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box>>, index) -> (index, index, index) +// CHECK: %[[VAL_6:.*]] = fir.shape %[[VAL_5]]#1 : (index) -> !fir.shape<1> +// CHECK: hlfir.elemental_addr %[[VAL_6]] typeparams %[[VAL_3]] : !fir.shape<1>, index { +// CHECK: ^bb0(%[[VAL_7:.*]]: index): +// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_4]] (%[[VAL_7]]) : (!fir.box>>, index) -> !fir.ref +// CHECK: %[[VAL_9:.*]] = fir.load %[[VAL_8]] : !fir.ref +// CHECK: %[[VAL_10:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_9]]) typeparams %[[VAL_3]] : (!fir.box>>, i32, index) -> !fir.boxchar<1> +// CHECK: hlfir.yield %[[VAL_10]] : !fir.boxchar<1> +// CHECK: } cleanup { +// CHECK: %[[VAL_11:.*]] = fir.box_addr %[[VAL_4]] : (!fir.box>>) -> !fir.heap> +// CHECK: fir.freemem %[[VAL_11]] : !fir.heap> +// CHECK: } +// CHECK: } diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir index a1e532563b4da..5eb772b4aad5e 100644 --- a/flang/test/HLFIR/invalid.fir +++ b/flang/test/HLFIR/invalid.fir @@ -541,3 +541,69 @@ func.func @bad_region_assign_2(%x: !fir.box>) { } return } + +// ----- +func.func @bad_element_addr_1(%x: !fir.ref>) { + %c20 = arith.constant 20 : index + %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1> + hlfir.region_assign { + hlfir.yield %x : !fir.ref> + } to { + // expected-error@+1 {{'hlfir.elemental_addr' op body must compute the address of a scalar entity}} + hlfir.elemental_addr %vector_shape : !fir.shape<1> { + ^bb0(%i: index): + %c42 = arith.constant 42.0 : f32 + hlfir.yield %c42 : f32 + } + } + return +} + +// ----- +func.func @bad_element_addr_2(%x: !fir.ref>) { + %c20 = arith.constant 20 : index + %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1> + hlfir.region_assign { + hlfir.yield %x : !fir.ref> + } to { + // expected-error@+1 {{'hlfir.elemental_addr' op body must compute the address of a scalar entity}} + hlfir.elemental_addr %vector_shape : !fir.shape<1> { + ^bb0(%i: index): + hlfir.yield %x : !fir.ref> + } + } + return +} + +// ----- +func.func @bad_element_addr_3(%x: !fir.ref>) { + %c20 = arith.constant 20 : index + %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1> + hlfir.region_assign { + hlfir.yield %x : !fir.ref> + } to { + // expected-error@+1 {{'hlfir.elemental_addr' op body region must be terminated by an hlfir.yield}} + hlfir.elemental_addr %vector_shape : !fir.shape<1> { + ^bb0(%i: index): + %c42 = arith.constant 42.0 : f32 + } + } + return +} + +// ----- +func.func @bad_element_addr_4(%x: !fir.ref>, %y: !fir.ref>) { + %c20 = arith.constant 20 : index + %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1> + hlfir.region_assign { + hlfir.yield %x : !fir.ref> + } to { + // expected-error@+1 {{'hlfir.elemental_addr' op body number of indices must match shape rank}} + hlfir.elemental_addr %vector_shape : !fir.shape<1> { + ^bb0(%i: index, %j: index): + %elt = hlfir.designate %y(%i, %j) : (!fir.ref>, index, index) -> !fir.ref + hlfir.yield %elt : !fir.ref + } + } + return +}