Skip to content

Commit d6c6bde

Browse files
JDPailleuxJDPailleux
authored
[flang] Implement !DIR$ UNROLL_AND_JAM [N] (#125046)
This patch implements support for the UNROLL_AND_JAM directive to enable or disable unrolling and jamming on a `DO LOOP`. It must be placed immediately before a `DO LOOP` and applies only to the loop that follows. N is an integer that specifying the unrolling factor. This is done by adding an attribute to the branch into the loop in LLVM to indicate that the loop should unrolled and jammed.
1 parent 2b340c1 commit d6c6bde

File tree

11 files changed

+156
-5
lines changed

11 files changed

+156
-5
lines changed

‎flang/docs/Directives.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,11 @@ A list of non-standard directives supported by Flang
4545
times if possible. When `n` is omitted, the compiler should attempt to fully
4646
unroll the loop. Some compilers accept an optional `=` before the `n` when `n`
4747
is present in the directive. Flang does not.
48+
* `!dir$ unroll_and_jam [N]` control how many times a loop should be unrolled and
49+
jammed. It must be placed immediately before a loop that follows. `N` is an optional
50+
integer that specifying the unrolling factor. When `N` is `0` or `1`, the loop
51+
should not be unrolled at all. If `N` is omitted the optimizer will
52+
selects the number of times to unroll the loop.
4853

4954
# Directive Details
5055

‎flang/include/flang/Parser/dump-parse-tree.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ class ParseTreeDumper {
209209
NODE(CompilerDirective, Unrecognized)
210210
NODE(CompilerDirective, VectorAlways)
211211
NODE(CompilerDirective, Unroll)
212+
NODE(CompilerDirective, UnrollAndJam)
212213
NODE(parser, ComplexLiteralConstant)
213214
NODE(parser, ComplexPart)
214215
NODE(parser, ComponentArraySpec)

‎flang/include/flang/Parser/parse-tree.h‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3349,6 +3349,8 @@ struct StmtFunctionStmt {
33493349
// !DIR$ IGNORE_TKR [ [(tkrdmac...)] name ]...
33503350
// !DIR$ LOOP COUNT (n1[, n2]...)
33513351
// !DIR$ name[=value] [, name[=value]]... = can be :
3352+
// !DIR$ UNROLL [N]
3353+
// !DIR$ UNROLL_AND_JAM [N]
33523354
// !DIR$ <anything else>
33533355
struct CompilerDirective {
33543356
UNION_CLASS_BOILERPLATE(CompilerDirective);
@@ -3371,10 +3373,13 @@ struct CompilerDirective {
33713373
struct Unroll {
33723374
WRAPPER_CLASS_BOILERPLATE(Unroll, std::optional<std::uint64_t>);
33733375
};
3376+
struct UnrollAndJam {
3377+
WRAPPER_CLASS_BOILERPLATE(UnrollAndJam, std::optional<std::uint64_t>);
3378+
};
33743379
EMPTY_CLASS(Unrecognized);
33753380
CharBlock source;
33763381
std::variant<std::list<IgnoreTKR>, LoopCount, std::list<AssumeAligned>,
3377-
VectorAlways, std::list<NameValue>, Unroll, Unrecognized>
3382+
VectorAlways, std::list<NameValue>, Unroll, UnrollAndJam, Unrecognized>
33783383
u;
33793384
};
33803385

‎flang/lib/Lower/Bridge.cpp‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2205,11 +2205,39 @@ class FirConverter : public Fortran::lower::AbstractConverter {
22052205
/*full=*/fullUnrollAttr, {}, {}, {});
22062206
}
22072207

2208+
// Enabling unroll and jamming directive without a value.
2209+
// For directives with a value, if the value is greater than 1,
2210+
// force unrolling with the given factor. Otherwise, disable unrolling and
2211+
// jamming.
2212+
mlir::LLVM::LoopUnrollAndJamAttr
2213+
genLoopUnrollAndJamAttr(std::optional<std::uint64_t> count) {
2214+
mlir::BoolAttr falseAttr =
2215+
mlir::BoolAttr::get(builder->getContext(), false);
2216+
mlir::BoolAttr trueAttr = mlir::BoolAttr::get(builder->getContext(), true);
2217+
mlir::IntegerAttr countAttr;
2218+
bool shouldUnroll = true;
2219+
if (count.has_value()) {
2220+
auto unrollingFactor = count.value();
2221+
if (unrollingFactor == 0 || unrollingFactor == 1) {
2222+
shouldUnroll = false;
2223+
} else {
2224+
countAttr =
2225+
builder->getIntegerAttr(builder->getI64Type(), unrollingFactor);
2226+
}
2227+
}
2228+
2229+
mlir::BoolAttr disableAttr = shouldUnroll ? falseAttr : trueAttr;
2230+
return mlir::LLVM::LoopUnrollAndJamAttr::get(
2231+
builder->getContext(), /*disable=*/disableAttr, /*count*/ countAttr, {},
2232+
{}, {}, {}, {});
2233+
}
2234+
22082235
void addLoopAnnotationAttr(
22092236
IncrementLoopInfo &info,
22102237
llvm::SmallVectorImpl<const Fortran::parser::CompilerDirective *> &dirs) {
22112238
mlir::LLVM::LoopVectorizeAttr va;
22122239
mlir::LLVM::LoopUnrollAttr ua;
2240+
mlir::LLVM::LoopUnrollAndJamAttr uja;
22132241
bool has_attrs = false;
22142242
for (const auto *dir : dirs) {
22152243
Fortran::common::visit(
@@ -2226,12 +2254,16 @@ class FirConverter : public Fortran::lower::AbstractConverter {
22262254
ua = genLoopUnrollAttr(u.v);
22272255
has_attrs = true;
22282256
},
2257+
[&](const Fortran::parser::CompilerDirective::UnrollAndJam &u) {
2258+
uja = genLoopUnrollAndJamAttr(u.v);
2259+
has_attrs = true;
2260+
},
22292261
[&](const auto &) {}},
22302262
dir->u);
22312263
}
22322264
mlir::LLVM::LoopAnnotationAttr la = mlir::LLVM::LoopAnnotationAttr::get(
2233-
builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua, {}, {},
2234-
{}, {}, {}, {}, {}, {}, {}, {}, {});
2265+
builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua,
2266+
/*unroll_and_jam*/ uja, {}, {}, {}, {}, {}, {}, {}, {}, {}, {});
22352267
if (has_attrs)
22362268
info.doLoop.setLoopAnnotationAttr(la);
22372269
}
@@ -2887,6 +2919,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
28872919
[&](const Fortran::parser::CompilerDirective::Unroll &) {
28882920
attachDirectiveToLoop(dir, &eval);
28892921
},
2922+
[&](const Fortran::parser::CompilerDirective::UnrollAndJam &) {
2923+
attachDirectiveToLoop(dir, &eval);
2924+
},
28902925
[&](const auto &) {}},
28912926
dir.u);
28922927
}

‎flang/lib/Parser/Fortran-parsers.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1308,11 +1308,14 @@ constexpr auto vectorAlways{
13081308
"VECTOR ALWAYS" >> construct<CompilerDirective::VectorAlways>()};
13091309
constexpr auto unroll{
13101310
"UNROLL" >> construct<CompilerDirective::Unroll>(maybe(digitString64))};
1311+
constexpr auto unrollAndJam{"UNROLL_AND_JAM" >>
1312+
construct<CompilerDirective::UnrollAndJam>(maybe(digitString64))};
13111313
TYPE_PARSER(beginDirective >> "DIR$ "_tok >>
13121314
sourced((construct<CompilerDirective>(ignore_tkr) ||
13131315
construct<CompilerDirective>(loopCount) ||
13141316
construct<CompilerDirective>(assumeAligned) ||
13151317
construct<CompilerDirective>(vectorAlways) ||
1318+
construct<CompilerDirective>(unrollAndJam) ||
13161319
construct<CompilerDirective>(unroll) ||
13171320
construct<CompilerDirective>(
13181321
many(construct<CompilerDirective::NameValue>(

‎flang/lib/Parser/unparse.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1851,6 +1851,10 @@ class UnparseVisitor {
18511851
Word("!DIR$ UNROLL");
18521852
Walk(" ", unroll.v);
18531853
},
1854+
[&](const CompilerDirective::UnrollAndJam &unrollAndJam) {
1855+
Word("!DIR$ UNROLL_AND_JAM");
1856+
Walk(" ", unrollAndJam.v);
1857+
},
18541858
[&](const CompilerDirective::Unrecognized &) {
18551859
Word("!DIR$ ");
18561860
Word(x.source.ToString());

‎flang/lib/Semantics/canonicalize-directives.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,8 @@ bool CanonicalizeDirectives(
5656
static bool IsExecutionDirective(const parser::CompilerDirective &dir) {
5757
return std::holds_alternative<parser::CompilerDirective::VectorAlways>(
5858
dir.u) ||
59-
std::holds_alternative<parser::CompilerDirective::Unroll>(dir.u);
59+
std::holds_alternative<parser::CompilerDirective::Unroll>(dir.u) ||
60+
std::holds_alternative<parser::CompilerDirective::UnrollAndJam>(dir.u);
6061
}
6162

6263
void CanonicalizationOfDirectives::Post(parser::SpecificationPart &spec) {
@@ -115,6 +116,9 @@ void CanonicalizationOfDirectives::Post(parser::Block &block) {
115116
[&](parser::CompilerDirective::Unroll &) {
116117
CheckLoopDirective(*dir, block, it);
117118
},
119+
[&](parser::CompilerDirective::UnrollAndJam &) {
120+
CheckLoopDirective(*dir, block, it);
121+
},
118122
[&](auto &) {}},
119123
dir->u);
120124
}

‎flang/lib/Semantics/resolve-names.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9552,7 +9552,8 @@ void ResolveNamesVisitor::Post(const parser::AssignedGotoStmt &x) {
95529552

95539553
void ResolveNamesVisitor::Post(const parser::CompilerDirective &x) {
95549554
if (std::holds_alternative<parser::CompilerDirective::VectorAlways>(x.u) ||
9555-
std::holds_alternative<parser::CompilerDirective::Unroll>(x.u)) {
9555+
std::holds_alternative<parser::CompilerDirective::Unroll>(x.u) ||
9556+
std::holds_alternative<parser::CompilerDirective::UnrollAndJam>(x.u)) {
95569557
return;
95579558
}
95589559
if (const auto *tkr{

‎flang/test/Integration/unroll_and_jam.f90‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
! RUN: %flang_fc1 -emit-llvm -o - %s | FileCheck %s
2+
3+
! CHECK-LABEL: unroll_and_jam_dir
4+
subroutine unroll_and_jam_dir
5+
integer :: a(10)
6+
!dir$ unroll_and_jam 4
7+
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
8+
do i=1,10
9+
a(i)=i
10+
end do
11+
end subroutine unroll_and_jam_dir
12+
13+
! CHECK-LABEL: unroll_and_jam_dir_0
14+
subroutine unroll_and_jam_dir_0
15+
integer :: a(10)
16+
!dir$ unroll_and_jam 0
17+
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE:.*]]
18+
do i=1,10
19+
a(i)=i
20+
end do
21+
end subroutine unroll_and_jam_dir_0
22+
23+
! CHECK-LABEL: unroll_and_jam_dir_1
24+
subroutine unroll_and_jam_dir_1
25+
integer :: a(10)
26+
!dir$ unroll_and_jam 1
27+
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]]
28+
do i=1,10
29+
a(i)=i
30+
end do
31+
end subroutine unroll_and_jam_dir_1
32+
33+
! CHECK-LABEL: unroll_and_jam_dir_no_factor
34+
subroutine unroll_and_jam_dir_no_factor
35+
integer :: a(10)
36+
!dir$ unroll_and_jam
37+
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_NO_FACTOR:.*]]
38+
do i=1,10
39+
a(i)=i
40+
end do
41+
end subroutine unroll_and_jam_dir_no_factor
42+
43+
! CHECK: ![[ANNOTATION]] = distinct !{![[ANNOTATION]], ![[UNROLL_AND_JAM:.*]], ![[UNROLL_AND_JAM_COUNT:.*]]}
44+
! CHECK: ![[UNROLL_AND_JAM]] = !{!"llvm.loop.unroll_and_jam.enable"}
45+
! CHECK: ![[UNROLL_AND_JAM_COUNT]] = !{!"llvm.loop.unroll_and_jam.count", i32 4}
46+
! CHECK: ![[ANNOTATION_DISABLE]] = distinct !{![[ANNOTATION_DISABLE]], ![[UNROLL_AND_JAM2:.*]]}
47+
! CHECK: ![[UNROLL_AND_JAM2]] = !{!"llvm.loop.unroll_and_jam.disable"}
48+
! CHECK: ![[ANNOTATION_NO_FACTOR]] = distinct !{![[ANNOTATION_NO_FACTOR]], ![[UNROLL_AND_JAM]]}

‎flang/test/Lower/unroll_and_jam.f90‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
! RUN: %flang_fc1 -emit-hlfir -o - %s | FileCheck %s
2+
3+
! CHECK: #loop_unroll_and_jam = #llvm.loop_unroll_and_jam<disable = false>
4+
! CHECK: #loop_unroll_and_jam1 = #llvm.loop_unroll_and_jam<disable = false, count = 2 : i64>
5+
! CHECK: #loop_annotation = #llvm.loop_annotation<unrollAndJam = #loop_unroll_and_jam>
6+
! CHECK: #loop_annotation1 = #llvm.loop_annotation<unrollAndJam = #loop_unroll_and_jam1>
7+
8+
! CHECK-LABEL: unroll_and_jam_dir
9+
subroutine unroll_and_jam_dir
10+
integer :: a(10)
11+
!dir$ unroll_and_jam
12+
!CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation}
13+
do i=1,10
14+
a(i)=i
15+
end do
16+
17+
!dir$ unroll_and_jam 2
18+
!CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation1}
19+
do i=1,10
20+
a(i)=i
21+
end do
22+
end subroutine unroll_and_jam_dir
23+
24+
25+
! CHECK-LABEL: intermediate_directive
26+
subroutine intermediate_directive
27+
integer :: a(10)
28+
!dir$ unroll_and_jam
29+
!dir$ unknown
30+
!CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation}
31+
do i=1,10
32+
a(i)=i
33+
end do
34+
end subroutine intermediate_directive

0 commit comments

Comments
 (0)