diff --git a/flang/include/flang/Optimizer/HLFIR/CMakeLists.txt b/flang/include/flang/Optimizer/HLFIR/CMakeLists.txt index bf90996edf3c3..823918bf8e9f9 100644 --- a/flang/include/flang/Optimizer/HLFIR/CMakeLists.txt +++ b/flang/include/flang/Optimizer/HLFIR/CMakeLists.txt @@ -7,6 +7,8 @@ mlir_tablegen(HLFIRAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=hlfir) mlir_tablegen(HLFIRAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=hlfir) set(LLVM_TARGET_DEFINITIONS HLFIROps.td) +mlir_tablegen(HLFIROpInterfaces.h.inc -gen-op-interface-decls) +mlir_tablegen(HLFIROpInterfaces.cpp.inc -gen-op-interface-defs) mlir_tablegen(HLFIROps.h.inc -gen-op-decls) mlir_tablegen(HLFIROps.cpp.inc -gen-op-defs) diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.h b/flang/include/flang/Optimizer/HLFIR/HLFIROps.h index fe82b691bc526..26bfc9a805dcc 100644 --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.h +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.h @@ -20,6 +20,7 @@ #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" +#include "flang/Optimizer/HLFIR/HLFIROpInterfaces.h.inc" #define GET_OP_CLASSES #include "flang/Optimizer/HLFIR/HLFIROps.h.inc" diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td index 6e6af7126bf37..792cfe1d005a9 100644 --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td @@ -814,7 +814,69 @@ def hlfir_GetExtentOp : hlfir_Op<"get_extent", [Pure]> { let builders = [OpBuilder<(ins "mlir::Value":$shape, "unsigned":$dim)>]; } -def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> { +def hlfir_OrderedAssignmentTreeOpInterface : OpInterface<"OrderedAssignmentTreeOpInterface"> { + let description = [{ + Interface for the operations representing Forall and Where constructs and + statements as an mlir::Region tree. + + These operations all have in common that they have "leaf" regions that contains + some code that should be evaluated for "all active combinations of Forall + index-name values" before the next OrderedAssignmentTreeOpInterface is + evaluated. + + These operations are ordered in a tree fashion: Some operations, like + hlfir.forall or hlfir.where, contain a list of OrderedAssignmentTreeOpInterface + that should be evaluated after the "Leaf" regions, and before the next + OrderedAssignmentTreeOpInterface. + + Nested OrderedAssignmentTreeOpInterface operations are affected by the + OrderedAssignmentTreeOpInterface operations that contain them (e.g: + hlfir.region_assign may be masked by the value of the mask region of + an hlfir.where that contains it). + + OrderedAssignmentTreeOpInterface operations that contain nested operation + must return a "sub-tree" region that contains the list of nested + OrderedAssignmentTreeOpInterface operations. + + There is no constraints over what IR a leaf region may contain. There is also + no restriction regarding how many leaf regions an + OrderedAssignmentTreeOpInterface operation may contain. + + A "sub-tree" region, if any, must contain only OrderedAssignmentTreeOpInterface + operations and, maybe, a fir.end terminator. + }]; + + let methods = [ + InterfaceMethod< + /*desc=*/"Get the OrderedAssignmentTreeOpInterface leaf regions that contain evaluation code", + /*retTy=*/"void", + /*methodName=*/"getLeafRegions", + /*args=*/(ins "llvm::SmallVectorImpl&":$regions), + /*methodBody=*/[{}] + >, + InterfaceMethod< + /*desc=*/"Get the region, if any, containing the list of sub-tree OrderedAssignmentTreeOpInterface nodes", + /*retTy=*/"mlir::Region*", + /*methodName=*/"getSubTreeRegion", + /*args=*/(ins), + /*methodBody=*/[{}] + >, + ]; + + let extraClassDeclaration = [{ + /// Interface verifier imlementation. + mlir::LogicalResult verifyImpl(); + }]; + + let verify = [{ + return ::mlir::cast<::hlfir::OrderedAssignmentTreeOpInterface>($_op).verifyImpl(); + }]; + + let cppNamespace = "hlfir"; +} + + +def hlfir_RegionAssignOp : hlfir_Op<"region_assign", [hlfir_OrderedAssignmentTreeOpInterface]> { 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 @@ -862,13 +924,22 @@ def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> { mlir::Value getUserAssignmentLhs() { return getUserDefinedAssignment().getArguments()[1]; } + void getLeafRegions(llvm::SmallVectorImpl& regions) { + regions.push_back(&getRhsRegion()); + regions.push_back(&getLhsRegion()); + if (!getUserDefinedAssignment().empty()) + regions.push_back(&getUserDefinedAssignment()); + } + mlir::Region* getSubTreeRegion() { return nullptr; } + }]; let hasCustomAssemblyFormat = 1; let hasVerifier = 1; } -def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", "ElementalAddrOp"]>, +def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", + "ElementalAddrOp", "ForallOp"]>, SingleBlockImplicitTerminator<"fir::FirEndOp">]> { let summary = "Yield a value or variable inside a forall, where or region assignment"; @@ -961,5 +1032,86 @@ def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"R let hasVerifier = 1; } +/// Define ODS constraints to verify that a region ends with a yield of a +/// certain type. +def YieldIntegerOrEmpty : CPred<"yieldsIntegerOrEmpty($_self)">; +def YieldIntegerRegion : RegionConstraint< + And<[SizedRegion<1>.predicate, YieldIntegerOrEmpty]>, + "single block region that yields an integer scalar value">; +def MaybeYieldIntegerRegion : RegionConstraint< + And<[MaxSizedRegion<1>.predicate, YieldIntegerOrEmpty]>, + "optional single block region that yields an integer scalar value">; + +def hlfir_ForallOp : hlfir_Op<"forall", [hlfir_OrderedAssignmentTreeOpInterface]> { + let summary = "represent a Fortran forall"; + let description = [{ + This operation allows representing Fortran forall. It computes + a set of "index-name" values based on lower bound, upper bound, + and step values whose evaluations are represented in their own + regions. + + Operations nested in its body region are evaluated in order. + As opposed to a regular loop, each nested operation is + fully evaluated for all the values in the "active set of + index-name" before the next nested operation. In practice, the + nested operation evaluation may be fused if it is proven that + they do not have data dependency. + + The "index-name" value is represented as the argument of the + body region. + + The lower, upper, and step region (if provided), must be terminated + by hlfir.yield that yields scalar integers. + + The body region must only contain other OrderedAssignmentTreeOpInterface + operations (like hlfir.region_assign, or other hlfir.forall). + + A Fortran forall with several indices is represented as a nest + of hlfir.forall. + + Example: FORALL(I=1:10) X(I) = FOO(I) + ``` + hlfir.forall lb { + hlfir.yield %c1 : index + } ub { + hlfir.yield %c10 : index + } (%i : index) { + hlfir.region_assign { + %res = fir.call @foo(%i) : (index) -> f32 + hlfir.yield %res : f32 + } to { + %xi = hlfir.designate %x(%i) : (!fir.box>, index) -> !fir.ref + hlfir.yield %xi : !fir.ref + } + } + ``` + + }]; + + let regions = (region YieldIntegerRegion:$lb_region, + YieldIntegerRegion:$ub_region, + MaybeYieldIntegerRegion:$step_region, + SizedRegion<1>:$body); + + let extraClassDeclaration = [{ + mlir::Value getForallIndexValue() { + return getBody().getArguments()[0]; + } + void getLeafRegions(llvm::SmallVectorImpl& regions) { + regions.push_back(&getLbRegion()); + regions.push_back(&getUbRegion()); + if (!getStepRegion().empty()) + regions.push_back(&getStepRegion()); + } + mlir::Region* getSubTreeRegion() { return &getBody(); } + }]; + + let assemblyFormat = [{ + attr-dict `lb` $lb_region + `ub` $ub_region + (`step` $step_region^)? + custom($body) + }]; +} #endif // FORTRAN_DIALECT_HLFIR_OPS diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp index 5b10ecee75f41..c03b7c92248fb 100644 --- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp +++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp @@ -1101,5 +1101,56 @@ mlir::LogicalResult hlfir::ElementalAddrOp::verify() { return mlir::success(); } +//===----------------------------------------------------------------------===// +// OrderedAssignmentTreeOpInterface +//===----------------------------------------------------------------------===// + +mlir::LogicalResult hlfir::OrderedAssignmentTreeOpInterface::verifyImpl() { + if (mlir::Region *body = getSubTreeRegion()) + if (!body->empty()) + for (mlir::Operation &op : body->front()) + if (!mlir::isa( + op)) + return emitOpError( + "body region must only contain OrderedAssignmentTreeOpInterface " + "operations or fir.end"); + return mlir::success(); +} + +//===----------------------------------------------------------------------===// +// ForallOp +//===----------------------------------------------------------------------===// + +static mlir::ParseResult parseForallOpBody(mlir::OpAsmParser &parser, + mlir::Region &body) { + mlir::OpAsmParser::Argument bodyArg; + if (parser.parseLParen() || parser.parseArgument(bodyArg) || + parser.parseColon() || parser.parseType(bodyArg.type) || + parser.parseRParen()) + return mlir::failure(); + if (parser.parseRegion(body, {bodyArg})) + return mlir::failure(); + ensureTerminator(body, parser.getBuilder(), + parser.getBuilder().getUnknownLoc()); + return mlir::success(); +} + +static void printForallOpBody(mlir::OpAsmPrinter &p, hlfir::ForallOp forall, + mlir::Region &body) { + mlir::Value forallIndex = forall.getForallIndexValue(); + p << " (" << forallIndex << ": " << forallIndex.getType() << ") "; + p.printRegion(body, /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); +} + +/// Predicate implementation of YieldIntegerOrEmpty. +static bool yieldsIntegerOrEmpty(mlir::Region ®ion) { + if (region.empty()) + return true; + auto yield = mlir::dyn_cast_or_null(getTerminator(region)); + return yield && fir::isa_integer(yield.getEntity().getType()); +} + +#include "flang/Optimizer/HLFIR/HLFIROpInterfaces.cpp.inc" #define GET_OP_CLASSES #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc" diff --git a/flang/test/HLFIR/forall.fir b/flang/test/HLFIR/forall.fir new file mode 100644 index 0000000000000..2ae91deafcacd --- /dev/null +++ b/flang/test/HLFIR/forall.fir @@ -0,0 +1,110 @@ +// Test hlfir.forall operation parse, verify (no errors), and unparse. +// RUN: fir-opt %s | fir-opt | FileCheck %s + +func.func @forall_test(%x: !fir.box>, %x2: !fir.box>) { + %c1 = arith.constant 1 : index + %c10 = arith.constant 10 : index + hlfir.forall lb { + hlfir.yield %c1 : index + } ub { + hlfir.yield %c10 : index + } (%i : index) { + hlfir.region_assign { + %res = fir.call @foo(%i) : (index) -> f32 + hlfir.yield %res : f32 + } to { + %xi = hlfir.designate %x(%i) : (!fir.box>, index) -> !fir.ref + hlfir.yield %xi : !fir.ref + } + hlfir.forall lb { + hlfir.yield %c1 : index + } ub { + hlfir.yield %c10 : index + } (%j : index) { + hlfir.region_assign { + %jf = fir.convert %j : (index) -> f32 + hlfir.yield %jf : f32 + } to { + %xij = hlfir.designate %x2(%i, %j) : (!fir.box>, index, index) -> !fir.ref + hlfir.yield %xij : !fir.ref + } + } + } + return +} +// CHECK-LABEL: func.func @forall_test( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.box>, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.box>) { +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_3:.*]] = arith.constant 10 : index +// CHECK: hlfir.forall lb { +// CHECK: hlfir.yield %[[VAL_2]] : index +// CHECK: } ub { +// CHECK: hlfir.yield %[[VAL_3]] : index +// CHECK: } (%[[VAL_4:.*]]: index) { +// CHECK: hlfir.region_assign { +// CHECK: %[[VAL_5:.*]] = fir.call @foo(%[[VAL_4]]) : (index) -> f32 +// CHECK: hlfir.yield %[[VAL_5]] : f32 +// CHECK: } to { +// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_4]]) : (!fir.box>, index) -> !fir.ref +// CHECK: hlfir.yield %[[VAL_6]] : !fir.ref +// CHECK: } +// CHECK: hlfir.forall lb { +// CHECK: hlfir.yield %[[VAL_2]] : index +// CHECK: } ub { +// CHECK: hlfir.yield %[[VAL_3]] : index +// CHECK: } (%[[VAL_7:.*]]: index) { +// CHECK: hlfir.region_assign { +// CHECK: %[[VAL_8:.*]] = fir.convert %[[VAL_7]] : (index) -> f32 +// CHECK: hlfir.yield %[[VAL_8]] : f32 +// CHECK: } to { +// CHECK: %[[VAL_9:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_4]], %[[VAL_7]]) : (!fir.box>, index, index) -> !fir.ref +// CHECK: hlfir.yield %[[VAL_9]] : !fir.ref +// CHECK: } +// CHECK: } +// CHECK: } + +func.func @forall_test_step(%x : !fir.box>, %y: !fir.box>) { + hlfir.forall lb { + %c1 = arith.constant 1 : index + hlfir.yield %c1 : index + } ub { + %c10 = arith.constant 10 : index + hlfir.yield %c10 : index + } step { + %c2 = arith.constant 2 : index + hlfir.yield %c2 : index + } (%i : index) { + hlfir.region_assign { + %yi = hlfir.designate %y(%i) : (!fir.box>, index) -> !fir.ref + %val = fir.load %yi : !fir.ref + hlfir.yield %val : f32 + } to { + %xi = hlfir.designate %x(%i) : (!fir.box>, index) -> !fir.ref + hlfir.yield %xi : !fir.ref + } + } + return +} +// CHECK-LABEL: func.func @forall_test_step( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.box>, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.box>) { +// CHECK: hlfir.forall lb { +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index +// CHECK: hlfir.yield %[[VAL_2]] : index +// CHECK: } ub { +// CHECK: %[[VAL_3:.*]] = arith.constant 10 : index +// CHECK: hlfir.yield %[[VAL_3]] : index +// CHECK: } step { +// CHECK: %[[VAL_4:.*]] = arith.constant 2 : index +// CHECK: hlfir.yield %[[VAL_4]] : index +// CHECK: } (%[[VAL_5:.*]]: index) { +// CHECK: hlfir.region_assign { +// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_5]]) : (!fir.box>, index) -> !fir.ref +// CHECK: %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref +// CHECK: hlfir.yield %[[VAL_7]] : f32 +// CHECK: } to { +// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]]) : (!fir.box>, index) -> !fir.ref +// CHECK: hlfir.yield %[[VAL_8]] : !fir.ref +// CHECK: } +// CHECK: } diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir index 5eb772b4aad5e..c1bcdaf687c78 100644 --- a/flang/test/HLFIR/invalid.fir +++ b/flang/test/HLFIR/invalid.fir @@ -607,3 +607,38 @@ func.func @bad_element_addr_4(%x: !fir.ref>, %y: !fir.ref>, %y: f32, %bad : !fir.ref>) { + // expected-error@+1 {{'hlfir.forall' op region #0 ('lb_region') failed to verify constraint: single block region that yields an integer scalar value}} + hlfir.forall lb { + hlfir.yield %bad : !fir.ref> + } ub { + %c10 = arith.constant 10 : index + hlfir.yield %c10 : index + } (%i : index) { + hlfir.region_assign { + hlfir.yield %y : f32 + } to { + %xi = hlfir.designate %x(%i) : (!fir.box>, index) -> !fir.ref + hlfir.yield %xi : !fir.ref + } + } + return +} + +// ----- +func.func @bad_forall_2(%x : !fir.box>, %y: f32) { + // expected-error@+1 {{'hlfir.forall' op body region must only contain OrderedAssignmentTreeOpInterface operations or fir.end}} + hlfir.forall lb { + %c1 = arith.constant 1 : index + hlfir.yield %c1 : index + } ub { + %c10 = arith.constant 10 : index + hlfir.yield %c10 : index + } (%i : index) { + %xi = hlfir.designate %x(%i) : (!fir.box>, index) -> !fir.ref + hlfir.assign %y to %xi : f32, !fir.ref + } + return +}