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 @@ -1258,4 +1258,49 @@ let hasVerifier = 1; } +def hlfir_ForallIndexOp : hlfir_Op<"forall_index", [fir_FortranVariableOpInterface, + hlfir_OrderedAssignmentTreeOpInterface, Pure]> { + let summary = "represent a Fortran forall index declaration"; + let description = [{ + This operation allows placing an hlfir.forall index in memory with + the related Fortran index-value name and type. + + So far, lowering needs to manipulate symbols as memory entities. + This operation allows fulfilling this requirements without allowing + bare alloca/declare/store inside the body of hlfir.forall, which would + make their analysis more complex. + + Given Forall index-value cannot be modified it also allows defining + a canonicalization of all its loads into a fir.convert of the + hlfir.forall index, which helps simplifying the data dependency analysis + of hlfir.forall. + }]; + + let arguments = (ins AnyIntegerType:$index, + Builtin_StringAttr:$name); + + let results = (outs AnyFortranVariable); + + let assemblyFormat = [{ + $name $index attr-dict `:` functional-type(operands, results) + }]; + + let extraClassDeclaration = [{ + /// Implement FortranVariableInterface interface. + std::optional getFortranAttrs() const { + return std::nullopt; + } + mlir::Value getShape() const {return mlir::Value{};} + mlir::OperandRange getExplicitTypeParams() const { + // Return an empty range. + return {(*this)->getOperands().begin(), (*this)->getOperands().begin()}; + } + /// Implement OrderedAssignmentTreeOpInterface interface. + void getLeafRegions(llvm::SmallVectorImpl& regions) {} + mlir::Region* getSubTreeRegion() { return nullptr; } + }]; + + let hasCanonicalizeMethod = 1; +} + #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 @@ -1244,6 +1244,28 @@ return verifyWhereAndElseWhereBody(*this); } +//===----------------------------------------------------------------------===// +// ForallIndexOp +//===----------------------------------------------------------------------===// + +mlir::LogicalResult +hlfir::ForallIndexOp::canonicalize(hlfir::ForallIndexOp indexOp, + mlir::PatternRewriter &rewriter) { + auto insertPt = rewriter.saveInsertionPoint(); + for (mlir::Operation *user : indexOp->getResult(0).getUsers()) + if (auto loadOp = mlir::dyn_cast_or_null(user)) { + rewriter.setInsertionPoint(loadOp); + rewriter.replaceOpWithNewOp( + user, loadOp.getResult().getType(), indexOp.getIndex()); + } + rewriter.restoreInsertionPoint(insertPt); + if (indexOp.use_empty()) { + rewriter.eraseOp(indexOp); + return mlir::success(); + } + return mlir::failure(); +} + #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-index.fir b/flang/test/HLFIR/forall-index.fir new file mode 100644 --- /dev/null +++ b/flang/test/HLFIR/forall-index.fir @@ -0,0 +1,44 @@ +// Test hlfir.forall_index operation parse, verify (no errors), unparse, +// and canonicalization. +// RUN: fir-opt %s | fir-opt | FileCheck %s +// RUN: fir-opt -canonicalize %s | FileCheck %s --check-prefix=CANONICALIZATION + +func.func @forall_index(%x: !fir.ref>, %y: !fir.ref>) { + %c1 = arith.constant 1 : index + %c10 = arith.constant 10 : index + hlfir.forall lb { + hlfir.yield %c1 : index + } ub { + hlfir.yield %c10 : index + } (%arg0: i64) { + %i = hlfir.forall_index "i" %arg0 : (i64) -> !fir.ref + hlfir.region_assign { + %ival = fir.load %i : !fir.ref + %yi = hlfir.designate %y(%ival) : (!fir.ref>, i64) -> !fir.ref + hlfir.yield %yi : !fir.ref + } to { + %ival = fir.load %i : !fir.ref + %xi = hlfir.designate %x(%ival) : (!fir.ref>, i64) -> !fir.ref + hlfir.yield %xi : !fir.ref + } + hlfir.region_assign { + %res = fir.call @taking_address(%i) : (!fir.ref) -> f32 + hlfir.yield %res : f32 + } to { + %ival = fir.load %i : !fir.ref + %xi = hlfir.designate %x(%ival) : (!fir.ref>, i64) -> !fir.ref + hlfir.yield %xi : !fir.ref + } + } + return +} +// CHECK: hlfir.forall lb { +// CHECK: } ub { +// CHECK: } (%[[VAL_4:.*]]: i64) { +// CHECK: hlfir.forall_index "i" %[[VAL_4]] : (i64) -> !fir.ref + +// CANONICALIZATION: %[[VAL_5:.*]] = hlfir.forall_index "i" %[[VAL_4:.*]] : (i64) -> !fir.ref +// CANONICALIZATION: hlfir.designate %{{.*}} (%[[VAL_4]]) : (!fir.ref>, i64) -> !fir.ref +// CANONICALIZATION: hlfir.designate %{{.*}} (%[[VAL_4]]) : (!fir.ref>, i64) -> !fir.ref +// CANONICALIZATION: fir.call @taking_address(%[[VAL_5]]) : (!fir.ref) -> f32 +// CANONICALIZATION: hlfir.designate %{{.*}} (%[[VAL_4]]) : (!fir.ref>, i64) -> !fir.ref