From b5fbdc19245face83a35b9c999b0f1e886a21ceb Mon Sep 17 00:00:00 2001 From: Jean-Didier Pailleux Date: Thu, 30 Jan 2025 13:44:55 +0100 Subject: [PATCH] [flang] Implement !DIR$ UNROLL_AND_JAM [N] --- flang/docs/Directives.md | 5 ++ flang/include/flang/Parser/dump-parse-tree.h | 1 + flang/include/flang/Parser/parse-tree.h | 7 ++- flang/lib/Lower/Bridge.cpp | 39 ++++++++++++++- flang/lib/Parser/Fortran-parsers.cpp | 3 ++ flang/lib/Parser/unparse.cpp | 4 ++ .../lib/Semantics/canonicalize-directives.cpp | 6 ++- flang/lib/Semantics/resolve-names.cpp | 3 +- flang/test/Integration/unroll_and_jam.f90 | 48 +++++++++++++++++++ flang/test/Lower/unroll_and_jam.f90 | 34 +++++++++++++ flang/test/Parser/compiler-directives.f90 | 11 +++++ 11 files changed, 156 insertions(+), 5 deletions(-) create mode 100644 flang/test/Integration/unroll_and_jam.f90 create mode 100644 flang/test/Lower/unroll_and_jam.f90 diff --git a/flang/docs/Directives.md b/flang/docs/Directives.md index c6c2e29a420ea..5e76d4331f6de 100644 --- a/flang/docs/Directives.md +++ b/flang/docs/Directives.md @@ -45,6 +45,11 @@ A list of non-standard directives supported by Flang times if possible. When `n` is omitted, the compiler should attempt to fully unroll the loop. Some compilers accept an optional `=` before the `n` when `n` is present in the directive. Flang does not. +* `!dir$ unroll_and_jam [N]` control how many times a loop should be unrolled and + jammed. It must be placed immediately before a loop that follows. `N` is an optional + integer that specifying the unrolling factor. When `N` is `0` or `1`, the loop + should not be unrolled at all. If `N` is omitted the optimizer will + selects the number of times to unroll the loop. # Directive Details diff --git a/flang/include/flang/Parser/dump-parse-tree.h b/flang/include/flang/Parser/dump-parse-tree.h index 21ee1d0517840..75c11301285b3 100644 --- a/flang/include/flang/Parser/dump-parse-tree.h +++ b/flang/include/flang/Parser/dump-parse-tree.h @@ -209,6 +209,7 @@ class ParseTreeDumper { NODE(CompilerDirective, Unrecognized) NODE(CompilerDirective, VectorAlways) NODE(CompilerDirective, Unroll) + NODE(CompilerDirective, UnrollAndJam) NODE(parser, ComplexLiteralConstant) NODE(parser, ComplexPart) NODE(parser, ComponentArraySpec) diff --git a/flang/include/flang/Parser/parse-tree.h b/flang/include/flang/Parser/parse-tree.h index c3a02fca5ade8..a3922a6b61891 100644 --- a/flang/include/flang/Parser/parse-tree.h +++ b/flang/include/flang/Parser/parse-tree.h @@ -3349,6 +3349,8 @@ struct StmtFunctionStmt { // !DIR$ IGNORE_TKR [ [(tkrdmac...)] name ]... // !DIR$ LOOP COUNT (n1[, n2]...) // !DIR$ name[=value] [, name[=value]]... = can be : +// !DIR$ UNROLL [N] +// !DIR$ UNROLL_AND_JAM [N] // !DIR$ struct CompilerDirective { UNION_CLASS_BOILERPLATE(CompilerDirective); @@ -3371,10 +3373,13 @@ struct CompilerDirective { struct Unroll { WRAPPER_CLASS_BOILERPLATE(Unroll, std::optional); }; + struct UnrollAndJam { + WRAPPER_CLASS_BOILERPLATE(UnrollAndJam, std::optional); + }; EMPTY_CLASS(Unrecognized); CharBlock source; std::variant, LoopCount, std::list, - VectorAlways, std::list, Unroll, Unrecognized> + VectorAlways, std::list, Unroll, UnrollAndJam, Unrecognized> u; }; diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp index 36e58e456dea3..781f484f3059a 100644 --- a/flang/lib/Lower/Bridge.cpp +++ b/flang/lib/Lower/Bridge.cpp @@ -2200,11 +2200,39 @@ class FirConverter : public Fortran::lower::AbstractConverter { /*full=*/fullUnrollAttr, {}, {}, {}); } + // Enabling unroll and jamming directive without a value. + // For directives with a value, if the value is greater than 1, + // force unrolling with the given factor. Otherwise, disable unrolling and + // jamming. + mlir::LLVM::LoopUnrollAndJamAttr + genLoopUnrollAndJamAttr(std::optional count) { + mlir::BoolAttr falseAttr = + mlir::BoolAttr::get(builder->getContext(), false); + mlir::BoolAttr trueAttr = mlir::BoolAttr::get(builder->getContext(), true); + mlir::IntegerAttr countAttr; + bool shouldUnroll = true; + if (count.has_value()) { + auto unrollingFactor = count.value(); + if (unrollingFactor == 0 || unrollingFactor == 1) { + shouldUnroll = false; + } else { + countAttr = + builder->getIntegerAttr(builder->getI64Type(), unrollingFactor); + } + } + + mlir::BoolAttr disableAttr = shouldUnroll ? falseAttr : trueAttr; + return mlir::LLVM::LoopUnrollAndJamAttr::get( + builder->getContext(), /*disable=*/disableAttr, /*count*/ countAttr, {}, + {}, {}, {}, {}); + } + void addLoopAnnotationAttr( IncrementLoopInfo &info, llvm::SmallVectorImpl &dirs) { mlir::LLVM::LoopVectorizeAttr va; mlir::LLVM::LoopUnrollAttr ua; + mlir::LLVM::LoopUnrollAndJamAttr uja; bool has_attrs = false; for (const auto *dir : dirs) { Fortran::common::visit( @@ -2221,12 +2249,16 @@ class FirConverter : public Fortran::lower::AbstractConverter { ua = genLoopUnrollAttr(u.v); has_attrs = true; }, + [&](const Fortran::parser::CompilerDirective::UnrollAndJam &u) { + uja = genLoopUnrollAndJamAttr(u.v); + has_attrs = true; + }, [&](const auto &) {}}, dir->u); } mlir::LLVM::LoopAnnotationAttr la = mlir::LLVM::LoopAnnotationAttr::get( - builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua, {}, {}, - {}, {}, {}, {}, {}, {}, {}, {}, {}); + builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua, + /*unroll_and_jam*/ uja, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}); if (has_attrs) info.doLoop.setLoopAnnotationAttr(la); } @@ -2882,6 +2914,9 @@ class FirConverter : public Fortran::lower::AbstractConverter { [&](const Fortran::parser::CompilerDirective::Unroll &) { attachDirectiveToLoop(dir, &eval); }, + [&](const Fortran::parser::CompilerDirective::UnrollAndJam &) { + attachDirectiveToLoop(dir, &eval); + }, [&](const auto &) {}}, dir.u); } diff --git a/flang/lib/Parser/Fortran-parsers.cpp b/flang/lib/Parser/Fortran-parsers.cpp index b5bcb53a12761..cfe9ecb29b0b7 100644 --- a/flang/lib/Parser/Fortran-parsers.cpp +++ b/flang/lib/Parser/Fortran-parsers.cpp @@ -1308,11 +1308,14 @@ constexpr auto vectorAlways{ "VECTOR ALWAYS" >> construct()}; constexpr auto unroll{ "UNROLL" >> construct(maybe(digitString64))}; +constexpr auto unrollAndJam{"UNROLL_AND_JAM" >> + construct(maybe(digitString64))}; TYPE_PARSER(beginDirective >> "DIR$ "_tok >> sourced((construct(ignore_tkr) || construct(loopCount) || construct(assumeAligned) || construct(vectorAlways) || + construct(unrollAndJam) || construct(unroll) || construct( many(construct( diff --git a/flang/lib/Parser/unparse.cpp b/flang/lib/Parser/unparse.cpp index cd91fbe4ea5eb..34d3deb58276c 100644 --- a/flang/lib/Parser/unparse.cpp +++ b/flang/lib/Parser/unparse.cpp @@ -1851,6 +1851,10 @@ class UnparseVisitor { Word("!DIR$ UNROLL"); Walk(" ", unroll.v); }, + [&](const CompilerDirective::UnrollAndJam &unrollAndJam) { + Word("!DIR$ UNROLL_AND_JAM"); + Walk(" ", unrollAndJam.v); + }, [&](const CompilerDirective::Unrecognized &) { Word("!DIR$ "); Word(x.source.ToString()); diff --git a/flang/lib/Semantics/canonicalize-directives.cpp b/flang/lib/Semantics/canonicalize-directives.cpp index b27a27618808b..1a0a0d145b3e2 100644 --- a/flang/lib/Semantics/canonicalize-directives.cpp +++ b/flang/lib/Semantics/canonicalize-directives.cpp @@ -56,7 +56,8 @@ bool CanonicalizeDirectives( static bool IsExecutionDirective(const parser::CompilerDirective &dir) { return std::holds_alternative( dir.u) || - std::holds_alternative(dir.u); + std::holds_alternative(dir.u) || + std::holds_alternative(dir.u); } void CanonicalizationOfDirectives::Post(parser::SpecificationPart &spec) { @@ -115,6 +116,9 @@ void CanonicalizationOfDirectives::Post(parser::Block &block) { [&](parser::CompilerDirective::Unroll &) { CheckLoopDirective(*dir, block, it); }, + [&](parser::CompilerDirective::UnrollAndJam &) { + CheckLoopDirective(*dir, block, it); + }, [&](auto &) {}}, dir->u); } diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp index e64abe6b50e78..df2f7fdc31c0a 100644 --- a/flang/lib/Semantics/resolve-names.cpp +++ b/flang/lib/Semantics/resolve-names.cpp @@ -9535,7 +9535,8 @@ void ResolveNamesVisitor::Post(const parser::AssignedGotoStmt &x) { void ResolveNamesVisitor::Post(const parser::CompilerDirective &x) { if (std::holds_alternative(x.u) || - std::holds_alternative(x.u)) { + std::holds_alternative(x.u) || + std::holds_alternative(x.u)) { return; } if (const auto *tkr{ diff --git a/flang/test/Integration/unroll_and_jam.f90 b/flang/test/Integration/unroll_and_jam.f90 new file mode 100644 index 0000000000000..771b7fb411855 --- /dev/null +++ b/flang/test/Integration/unroll_and_jam.f90 @@ -0,0 +1,48 @@ +! RUN: %flang_fc1 -emit-llvm -o - %s | FileCheck %s + +! CHECK-LABEL: unroll_and_jam_dir +subroutine unroll_and_jam_dir + integer :: a(10) + !dir$ unroll_and_jam 4 + ! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION:.*]] + do i=1,10 + a(i)=i + end do +end subroutine unroll_and_jam_dir + +! CHECK-LABEL: unroll_and_jam_dir_0 +subroutine unroll_and_jam_dir_0 + integer :: a(10) + !dir$ unroll_and_jam 0 + ! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE:.*]] + do i=1,10 + a(i)=i + end do +end subroutine unroll_and_jam_dir_0 + +! CHECK-LABEL: unroll_and_jam_dir_1 +subroutine unroll_and_jam_dir_1 + integer :: a(10) + !dir$ unroll_and_jam 1 + ! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]] + do i=1,10 + a(i)=i + end do +end subroutine unroll_and_jam_dir_1 + +! CHECK-LABEL: unroll_and_jam_dir_no_factor +subroutine unroll_and_jam_dir_no_factor + integer :: a(10) + !dir$ unroll_and_jam + ! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_NO_FACTOR:.*]] + do i=1,10 + a(i)=i + end do +end subroutine unroll_and_jam_dir_no_factor + +! CHECK: ![[ANNOTATION]] = distinct !{![[ANNOTATION]], ![[UNROLL_AND_JAM:.*]], ![[UNROLL_AND_JAM_COUNT:.*]]} +! CHECK: ![[UNROLL_AND_JAM]] = !{!"llvm.loop.unroll_and_jam.enable"} +! CHECK: ![[UNROLL_AND_JAM_COUNT]] = !{!"llvm.loop.unroll_and_jam.count", i32 4} +! CHECK: ![[ANNOTATION_DISABLE]] = distinct !{![[ANNOTATION_DISABLE]], ![[UNROLL_AND_JAM2:.*]]} +! CHECK: ![[UNROLL_AND_JAM2]] = !{!"llvm.loop.unroll_and_jam.disable"} +! CHECK: ![[ANNOTATION_NO_FACTOR]] = distinct !{![[ANNOTATION_NO_FACTOR]], ![[UNROLL_AND_JAM]]} diff --git a/flang/test/Lower/unroll_and_jam.f90 b/flang/test/Lower/unroll_and_jam.f90 new file mode 100644 index 0000000000000..afc5a7b6b271e --- /dev/null +++ b/flang/test/Lower/unroll_and_jam.f90 @@ -0,0 +1,34 @@ +! RUN: %flang_fc1 -emit-hlfir -o - %s | FileCheck %s + +! CHECK: #loop_unroll_and_jam = #llvm.loop_unroll_and_jam +! CHECK: #loop_unroll_and_jam1 = #llvm.loop_unroll_and_jam +! CHECK: #loop_annotation = #llvm.loop_annotation +! CHECK: #loop_annotation1 = #llvm.loop_annotation + +! CHECK-LABEL: unroll_and_jam_dir +subroutine unroll_and_jam_dir + integer :: a(10) + !dir$ unroll_and_jam + !CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation} + do i=1,10 + a(i)=i + end do + + !dir$ unroll_and_jam 2 + !CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation1} + do i=1,10 + a(i)=i + end do +end subroutine unroll_and_jam_dir + + +! CHECK-LABEL: intermediate_directive +subroutine intermediate_directive + integer :: a(10) + !dir$ unroll_and_jam + !dir$ unknown + !CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation} + do i=1,10 + a(i)=i + end do +end subroutine intermediate_directive diff --git a/flang/test/Parser/compiler-directives.f90 b/flang/test/Parser/compiler-directives.f90 index f372a9f533a35..d1e386a01dd4d 100644 --- a/flang/test/Parser/compiler-directives.f90 +++ b/flang/test/Parser/compiler-directives.f90 @@ -46,3 +46,14 @@ subroutine unroll do i=1,10 enddo end subroutine + +subroutine unroll_and_jam + !dir$ unroll_and_jam + ! CHECK: !DIR$ UNROLL_AND_JAM + do i=1,10 + enddo + !dir$ unroll_and_jam 2 + ! CHECK: !DIR$ UNROLL_AND_JAM 2 + do i=1,10 + enddo +end subroutine