diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.h b/flang/include/flang/Optimizer/HLFIR/HLFIROps.h --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.h +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.h @@ -11,6 +11,7 @@ #include "flang/Optimizer/Dialect/FIRAttr.h" #include "flang/Optimizer/Dialect/FIRDialect.h" +#include "flang/Optimizer/Dialect/FIROps.h" #include "flang/Optimizer/Dialect/FIRType.h" #include "flang/Optimizer/Dialect/FortranVariableInterface.h" #include "flang/Optimizer/HLFIR/HLFIRDialect.h" diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td @@ -814,4 +814,88 @@ let builders = [OpBuilder<(ins "mlir::Value":$shape, "unsigned":$dim)>]; } +def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> { + let summary = "represent a Fortran assignment using regions for the LHS and RHS evaluation"; + let description = [{ + This operation can represent Forall and Where assignment when inside an + hlfir.forall or hlfir.where "ordered assignment tree". It can + also represent user defined assignments and assignment to vector + subscripted entities without requiring the materialization of the + right-hand side temporary copy that may be needed to implement Fortran + assignment semantic. + + The right-hand side and left-hand side evaluations are held in their + own regions terminated with hlfir.yield operations (or hlfir.elemental_addr + for a left-hand side with vector subscript). + + An optional region may be added to implement user defined assignment. + This region provides two block arguments with the same type as the + yielded rhs and lhs entities (in that order), or the element type if this + is an elemental user defined assignment. + + If this optional region is not provided, intrinsic assignment is performed. + + Example: "X = Y", where "=" is a user defined elemental assignment "foo" + taking Y by value. + ``` + hlfir.region_assign { + hlfir.yield %y : !fir.box> + } to { + hlfir.yield %x : !fir.box>> + } user_defined_assignment (%rhs_elt: !fir.ref) to (%lhs_elt: !fir.ref>) { + %0 = fir.load %rhs_elt : !fir.ref + fir.call @foo(%lhs_elt, %0) : (!fir.ref>, f32) -> () + } + ``` + + TODO: add optional "realloc" semantics like for hlfir.assign. + }]; + + let regions = (region SizedRegion<1>:$rhs_region, + SizedRegion<1>:$lhs_region, + MaxSizedRegion<1>:$user_defined_assignment); + + let extraClassDeclaration = [{ + mlir::Value getUserAssignmentRhs() { + return getUserDefinedAssignment().getArguments()[0]; + } + mlir::Value getUserAssignmentLhs() { + return getUserDefinedAssignment().getArguments()[1]; + } + }]; + + let hasCustomAssemblyFormat = 1; + let hasVerifier = 1; +} + +def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp"]>, + SingleBlockImplicitTerminator<"fir::FirEndOp">]> { + + let summary = "Yield a value or variable inside a forall, where or region assignment"; + + let description = [{ + Terminator operation that yields an HLFIR value or variable that was computed in + a region and hold the yielded entity cleanup, if any, into its own region. + This allows representing any Fortran expression evaluation in its own region so + that the evaluation can easily be scheduled/moved around in a pass. + + Example: "foo(x)" where foo returns an allocatable array. + ``` + { + // In some region. + %0 = fir.call @foo(x) (!fir.ref) -> !fir.box>> + hlfir.yield %0 : !fir.box>> cleanup { + %1 = fir.box_addr %0 : !fir.box>> -> !fir.heap> + %fir.freemem %1 : !fir.heap> + } + } + ``` + }]; + + let arguments = (ins AnyFortranEntity:$entity); + let regions = (region MaxSizedRegion<1>:$cleanup); + + let assemblyFormat = "$entity attr-dict `:` type($entity) custom($cleanup)"; +} + #endif // FORTRAN_DIALECT_HLFIR_OPS diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp --- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp +++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp @@ -960,5 +960,111 @@ return mlir::success(); } +//===----------------------------------------------------------------------===// +// RegionAssignOp +//===----------------------------------------------------------------------===// + +/// Add a fir.end terminator to a parsed region if it does not already has a +/// terminator. +static void ensureTerminator(mlir::Region ®ion, mlir::Builder &builder, + mlir::Location loc) { + // Borrow YielOp::ensureTerminator MLIR generated implementation to add a + // fir.end if there is no terminator. This has nothing to do with YielOp, + // other than the fact that yieldOp has the + // SingleBlocklicitTerminator<"fir::FirEndOp"> interface that + // cannot be added on other HLFIR operations with several regions which are + // not all terminated the same way. + hlfir::YieldOp::ensureTerminator(region, builder, loc); +} + +mlir::ParseResult hlfir::RegionAssignOp::parse(mlir::OpAsmParser &parser, + mlir::OperationState &result) { + mlir::Region &rhsRegion = *result.addRegion(); + if (parser.parseRegion(rhsRegion)) + return mlir::failure(); + mlir::Region &lhsRegion = *result.addRegion(); + if (parser.parseKeyword("to") || parser.parseRegion(lhsRegion)) + return mlir::failure(); + mlir::Region &userDefinedAssignmentRegion = *result.addRegion(); + if (succeeded(parser.parseOptionalKeyword("user_defined_assign"))) { + mlir::OpAsmParser::Argument rhsArg, lhsArg; + if (parser.parseLParen() || parser.parseArgument(rhsArg) || + parser.parseColon() || parser.parseType(rhsArg.type) || + parser.parseRParen() || parser.parseKeyword("to") || + parser.parseLParen() || parser.parseArgument(lhsArg) || + parser.parseColon() || parser.parseType(lhsArg.type) || + parser.parseRParen()) + return mlir::failure(); + if (parser.parseRegion(userDefinedAssignmentRegion, {rhsArg, lhsArg})) + return mlir::failure(); + ensureTerminator(userDefinedAssignmentRegion, parser.getBuilder(), + result.location); + } + return mlir::success(); +} + +void hlfir::RegionAssignOp::print(mlir::OpAsmPrinter &p) { + p << " "; + p.printRegion(getRhsRegion(), /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/true); + p << " to "; + p.printRegion(getLhsRegion(), /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/true); + if (!getUserDefinedAssignment().empty()) { + p << " user_defined_assign "; + mlir::Value userAssignmentRhs = getUserAssignmentRhs(); + mlir::Value userAssignmentLhs = getUserAssignmentLhs(); + p << " (" << userAssignmentRhs << ": " << userAssignmentRhs.getType() + << ") to ("; + p << userAssignmentLhs << ": " << userAssignmentLhs.getType() << ") "; + p.printRegion(getUserDefinedAssignment(), /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + } +} + +static mlir::Operation *getTerminator(mlir::Region ®ion) { + if (region.empty() || region.back().empty()) + return nullptr; + return ®ion.back().back(); +} + +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()))) + return emitOpError("left-hand side region must be terminated by an " + "hlfir.yield or hlfir.elemental_addr"); + return mlir::success(); +} + +//===----------------------------------------------------------------------===// +// YieldOp +//===----------------------------------------------------------------------===// + +static mlir::ParseResult parseYieldOpCleanup(mlir::OpAsmParser &parser, + mlir::Region &cleanup) { + if (succeeded(parser.parseOptionalKeyword("cleanup"))) { + if (parser.parseRegion(cleanup, /*arguments=*/{}, + /*argTypes=*/{})) + return mlir::failure(); + hlfir::YieldOp::ensureTerminator(cleanup, parser.getBuilder(), + parser.getBuilder().getUnknownLoc()); + } + + return mlir::success(); +} + +template +static void printYieldOpCleanup(mlir::OpAsmPrinter &p, YieldOp yieldOp, + mlir::Region &cleanup) { + if (!cleanup.empty()) { + p << "cleanup "; + p.printRegion(cleanup, /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + } +} + #define GET_OP_CLASSES #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc" diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir --- a/flang/test/HLFIR/invalid.fir +++ b/flang/test/HLFIR/invalid.fir @@ -518,3 +518,26 @@ // expected-error@+1 {{'hlfir.get_extent' op dimension index out of bounds}} %0 = hlfir.get_extent %arg0 {dim = 1 : index} : (!fir.shape<1>) -> index } + +// ----- +func.func @bad_region_assign_1(%x: !fir.box>) { +// expected-error@+1 {{'hlfir.region_assign' op right-hand side region must be terminated by an hlfir.yield}} + hlfir.region_assign { + %c100 = arith.constant 100 : index + } to { + hlfir.yield %x : !fir.box> + } + return +} + +// ----- +func.func @bad_region_assign_2(%x: !fir.box>) { +// expected-error@+1 {{'hlfir.region_assign' op left-hand side region must be terminated by an hlfir.yield or hlfir.elemental_addr}} + hlfir.region_assign { + hlfir.yield %x : !fir.box> + } to { + %c100 = arith.constant 100 : index + } user_defined_assign (%rhs: !fir.ref) to (%lhs: !fir.ref) { + } + return +} diff --git a/flang/test/HLFIR/region-assign.fir b/flang/test/HLFIR/region-assign.fir new file mode 100644 --- /dev/null +++ b/flang/test/HLFIR/region-assign.fir @@ -0,0 +1,87 @@ +// Test hlfir.region_assign and hlfir.yield operation parse, verify (no errors), +// and unparse. +// RUN: fir-opt %s | fir-opt | FileCheck %s + +func.func @region_assign_test(%y : !fir.box>, %x: !fir.box>) { + %c100 = arith.constant 100 : index + %shape = fir.shape %c100 : (index) -> !fir.shape<1> + hlfir.region_assign { + %expr = hlfir.elemental %shape : (!fir.shape<1>) -> !hlfir.expr { + ^bb0(%i : index): + %yelt = hlfir.designate %y(%i) : (!fir.box>, index) -> !fir.ref + %elt = fir.call @some_elemental(%yelt) : (!fir.ref) -> f32 + hlfir.yield_element %elt : f32 + } + hlfir.yield %expr : !hlfir.expr + } to { + hlfir.yield %x : !fir.box> + } + return +} +// CHECK-LABEL: func.func @region_assign_test( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.box>, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.box>) { +// CHECK: %[[VAL_2:.*]] = arith.constant 100 : index +// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1> +// CHECK: hlfir.region_assign { +// CHECK: %[[VAL_4:.*]] = hlfir.elemental %[[VAL_3]] : (!fir.shape<1>) -> !hlfir.expr { +// CHECK: ^bb0(%[[VAL_5:.*]]: index): +// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]]) : (!fir.box>, index) -> !fir.ref +// CHECK: %[[VAL_7:.*]] = fir.call @some_elemental(%[[VAL_6]]) : (!fir.ref) -> f32 +// CHECK: hlfir.yield_element %[[VAL_7]] : f32 +// CHECK: } +// CHECK: hlfir.yield %[[VAL_8:.*]] : !hlfir.expr +// CHECK: } to { +// CHECK: hlfir.yield %[[VAL_1]] : !fir.box> +// CHECK: } + +func.func @region_user_assign_test(%y : !fir.box>, %x: !fir.box>) { + hlfir.region_assign { + hlfir.yield %y : !fir.box> + } to { + hlfir.yield %x : !fir.box> + } user_defined_assign (%rhs : !fir.ref) to (%lhs : !fir.ref) { + %0 = fir.load %rhs : !fir.ref + fir.call @user_assign(%lhs, %0) : (!fir.ref, i64) -> () + } + return +} +// CHECK-LABEL: func.func @region_user_assign_test( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.box>, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.box>) { +// CHECK: hlfir.region_assign { +// CHECK: hlfir.yield %[[VAL_0]] : !fir.box> +// CHECK: } to { +// CHECK: hlfir.yield %[[VAL_1]] : !fir.box> +// CHECK: } user_defined_assign (%[[VAL_2:.*]]: !fir.ref) to (%[[VAL_3:.*]]: !fir.ref) { +// CHECK: %[[VAL_4:.*]] = fir.load %[[VAL_2]] : !fir.ref +// CHECK: fir.call @user_assign(%[[VAL_3]], %[[VAL_4]]) : (!fir.ref, i64) -> () +// CHECK: } + +func.func @yield_cleanup(%x: !fir.box>) { + hlfir.region_assign { + %0 = fir.allocmem !fir.array<100xf32> + fir.call @fillin_some_values(%0) : (!fir.heap>) -> () + hlfir.yield %0 : !fir.heap> cleanup { + fir.freemem %0 : !fir.heap> + } + } to { + hlfir.yield %x : !fir.box> + } + return +} +// CHECK-LABEL: func.func @yield_cleanup( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.box>) { +// CHECK: hlfir.region_assign { +// CHECK: %[[VAL_1:.*]] = fir.allocmem !fir.array<100xf32> +// CHECK: fir.call @fillin_some_values(%[[VAL_1]]) : (!fir.heap>) -> () +// CHECK: hlfir.yield %[[VAL_1]] : !fir.heap> cleanup { +// CHECK: fir.freemem %[[VAL_1]] : !fir.heap> +// CHECK: } +// CHECK: } to { +// CHECK: hlfir.yield %[[VAL_0]] : !fir.box> +// CHECK: } + +func.func private @user_assign(!fir.ref, i64) -> () +func.func private @some_elemental(!fir.ref) -> f32 +func.func private @fillin_some_values(!fir.heap>) -> ()