Skip to content

Commit e811cb0

Browse files
JDPailleuxJDPailleux
authored
[flang] Implement !DIR$ UNROLL [N] (#123331)
This patch implements support for the UNROLL directive to control how many times a loop should be unrolled. 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. The code pushed to support the directive `VECTOR ALWAYS` has been modified to take account of the fact that several directives can be used before a `DO LOOP`.
1 parent 51c7338 commit e811cb0

File tree

11 files changed

+129
-17
lines changed

11 files changed

+129
-17
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ class ParseTreeDumper {
208208
NODE(CompilerDirective, NameValue)
209209
NODE(CompilerDirective, Unrecognized)
210210
NODE(CompilerDirective, VectorAlways)
211+
NODE(CompilerDirective, Unroll)
211212
NODE(parser, ComplexLiteralConstant)
212213
NODE(parser, ComplexPart)
213214
NODE(parser, ComponentArraySpec)

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3368,10 +3368,13 @@ struct CompilerDirective {
33683368
TUPLE_CLASS_BOILERPLATE(NameValue);
33693369
std::tuple<Name, std::optional<std::uint64_t>> t;
33703370
};
3371+
struct Unroll {
3372+
WRAPPER_CLASS_BOILERPLATE(Unroll, std::optional<std::uint64_t>);
3373+
};
33713374
EMPTY_CLASS(Unrecognized);
33723375
CharBlock source;
33733376
std::variant<std::list<IgnoreTKR>, LoopCount, std::list<AssumeAligned>,
3374-
VectorAlways, std::list<NameValue>, Unrecognized>
3377+
VectorAlways, std::list<NameValue>, Unroll, Unrecognized>
33753378
u;
33763379
};
33773380

‎flang/lib/Lower/Bridge.cpp

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2170,14 +2170,42 @@ class FirConverter : public Fortran::lower::AbstractConverter {
21702170
return builder->createIntegerConstant(loc, controlType, 1); // step
21712171
}
21722172

2173-
void addLoopAnnotationAttr(IncrementLoopInfo &info) {
2173+
void addLoopAnnotationAttr(
2174+
IncrementLoopInfo &info,
2175+
llvm::SmallVectorImpl<const Fortran::parser::CompilerDirective *> &dirs) {
21742176
mlir::BoolAttr f = mlir::BoolAttr::get(builder->getContext(), false);
2175-
mlir::LLVM::LoopVectorizeAttr va = mlir::LLVM::LoopVectorizeAttr::get(
2176-
builder->getContext(), /*disable=*/f, {}, {}, {}, {}, {}, {});
2177+
mlir::BoolAttr t = mlir::BoolAttr::get(builder->getContext(), true);
2178+
mlir::LLVM::LoopVectorizeAttr va;
2179+
mlir::LLVM::LoopUnrollAttr ua;
2180+
bool has_attrs = false;
2181+
for (const auto *dir : dirs) {
2182+
Fortran::common::visit(
2183+
Fortran::common::visitors{
2184+
[&](const Fortran::parser::CompilerDirective::VectorAlways &) {
2185+
va = mlir::LLVM::LoopVectorizeAttr::get(builder->getContext(),
2186+
/*disable=*/f, {}, {},
2187+
{}, {}, {}, {});
2188+
has_attrs = true;
2189+
},
2190+
[&](const Fortran::parser::CompilerDirective::Unroll &u) {
2191+
mlir::IntegerAttr countAttr;
2192+
if (u.v.has_value()) {
2193+
countAttr = builder->getIntegerAttr(builder->getI64Type(),
2194+
u.v.value());
2195+
}
2196+
ua = mlir::LLVM::LoopUnrollAttr::get(
2197+
builder->getContext(), /*disable=*/f, /*count*/ countAttr,
2198+
{}, /*full*/ u.v.has_value() ? f : t, {}, {}, {});
2199+
has_attrs = true;
2200+
},
2201+
[&](const auto &) {}},
2202+
dir->u);
2203+
}
21772204
mlir::LLVM::LoopAnnotationAttr la = mlir::LLVM::LoopAnnotationAttr::get(
2178-
builder->getContext(), {}, /*vectorize=*/va, {}, {}, {}, {}, {}, {}, {},
2179-
{}, {}, {}, {}, {}, {});
2180-
info.doLoop.setLoopAnnotationAttr(la);
2205+
builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua, {}, {},
2206+
{}, {}, {}, {}, {}, {}, {}, {}, {});
2207+
if (has_attrs)
2208+
info.doLoop.setLoopAnnotationAttr(la);
21812209
}
21822210

21832211
/// Generate FIR to begin a structured or unstructured increment loop nest.
@@ -2276,14 +2304,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
22762304
if (info.hasLocalitySpecs())
22772305
handleLocalitySpecs(info);
22782306

2279-
for (const auto *dir : dirs) {
2280-
Fortran::common::visit(
2281-
Fortran::common::visitors{
2282-
[&](const Fortran::parser::CompilerDirective::VectorAlways
2283-
&d) { addLoopAnnotationAttr(info); },
2284-
[&](const auto &) {}},
2285-
dir->u);
2286-
}
2307+
addLoopAnnotationAttr(info, dirs);
22872308
continue;
22882309
}
22892310

@@ -2835,6 +2856,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
28352856
[&](const Fortran::parser::CompilerDirective::VectorAlways &) {
28362857
attachDirectiveToLoop(dir, &eval);
28372858
},
2859+
[&](const Fortran::parser::CompilerDirective::Unroll &) {
2860+
attachDirectiveToLoop(dir, &eval);
2861+
},
28382862
[&](const auto &) {}},
28392863
dir.u);
28402864
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1293,6 +1293,7 @@ TYPE_PARSER(construct<StatOrErrmsg>("STAT =" >> statVariable) ||
12931293
// !DIR$ IGNORE_TKR [ [(tkrdmac...)] name ]...
12941294
// !DIR$ LOOP COUNT (n1[, n2]...)
12951295
// !DIR$ name[=value] [, name[=value]]...
1296+
// !DIR$ UNROLL [n]
12961297
// !DIR$ <anything else>
12971298
constexpr auto ignore_tkr{
12981299
"IGNORE_TKR" >> optionalList(construct<CompilerDirective::IgnoreTKR>(
@@ -1305,11 +1306,14 @@ constexpr auto assumeAligned{"ASSUME_ALIGNED" >>
13051306
indirect(designator), ":"_tok >> digitString64))};
13061307
constexpr auto vectorAlways{
13071308
"VECTOR ALWAYS" >> construct<CompilerDirective::VectorAlways>()};
1309+
constexpr auto unroll{
1310+
"UNROLL" >> construct<CompilerDirective::Unroll>(maybe(digitString64))};
13081311
TYPE_PARSER(beginDirective >> "DIR$ "_tok >>
13091312
sourced((construct<CompilerDirective>(ignore_tkr) ||
13101313
construct<CompilerDirective>(loopCount) ||
13111314
construct<CompilerDirective>(assumeAligned) ||
13121315
construct<CompilerDirective>(vectorAlways) ||
1316+
construct<CompilerDirective>(unroll) ||
13131317
construct<CompilerDirective>(
13141318
many(construct<CompilerDirective::NameValue>(
13151319
name, maybe(("="_tok || ":"_tok) >> digitString64))))) /

‎flang/lib/Parser/unparse.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1847,6 +1847,10 @@ class UnparseVisitor {
18471847
[&](const std::list<CompilerDirective::NameValue> &names) {
18481848
Walk("!DIR$ ", names, " ");
18491849
},
1850+
[&](const CompilerDirective::Unroll &unroll) {
1851+
Word("!DIR$ UNROLL");
1852+
Walk(" ", unroll.v);
1853+
},
18501854
[&](const CompilerDirective::Unrecognized &) {
18511855
Word("!DIR$ ");
18521856
Word(x.source.ToString());

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,9 @@ bool CanonicalizeDirectives(
5454
}
5555

5656
static bool IsExecutionDirective(const parser::CompilerDirective &dir) {
57-
return std::holds_alternative<parser::CompilerDirective::VectorAlways>(dir.u);
57+
return std::holds_alternative<parser::CompilerDirective::VectorAlways>(
58+
dir.u) ||
59+
std::holds_alternative<parser::CompilerDirective::Unroll>(dir.u);
5860
}
5961

6062
void CanonicalizationOfDirectives::Post(parser::SpecificationPart &spec) {
@@ -110,6 +112,9 @@ void CanonicalizationOfDirectives::Post(parser::Block &block) {
110112
common::visitors{[&](parser::CompilerDirective::VectorAlways &) {
111113
CheckLoopDirective(*dir, block, it);
112114
},
115+
[&](parser::CompilerDirective::Unroll &) {
116+
CheckLoopDirective(*dir, block, it);
117+
},
113118
[&](auto &) {}},
114119
dir->u);
115120
}

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

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

94609460
void ResolveNamesVisitor::Post(const parser::CompilerDirective &x) {
9461-
if (std::holds_alternative<parser::CompilerDirective::VectorAlways>(x.u)) {
9461+
if (std::holds_alternative<parser::CompilerDirective::VectorAlways>(x.u) ||
9462+
std::holds_alternative<parser::CompilerDirective::Unroll>(x.u)) {
94629463
return;
94639464
}
94649465
if (const auto *tkr{

‎flang/test/Integration/unroll.f90

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
! RUN: %flang_fc1 -emit-llvm -o - %s | FileCheck %s
2+
3+
! CHECK-LABEL: unroll_dir
4+
subroutine unroll_dir
5+
integer :: a(10)
6+
!dir$ unroll
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_dir
12+
13+
! CHECK: ![[ANNOTATION]] = distinct !{![[ANNOTATION]], ![[UNROLL:.*]], ![[UNROLL_FULL:.*]]}
14+
! CHECK: ![[UNROLL]] = !{!"llvm.loop.unroll.enable"}
15+
! CHECK: ![[UNROLL_FULL]] = !{!"llvm.loop.unroll.full"}
16+

‎flang/test/Lower/unroll.f90

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

‎flang/test/Parser/compiler-directives.f90

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,3 +35,14 @@ subroutine vector_always
3535
do i=1,10
3636
enddo
3737
end subroutine
38+
39+
subroutine unroll
40+
!dir$ unroll
41+
! CHECK: !DIR$ UNROLL
42+
do i=1,10
43+
enddo
44+
!dir$ unroll 2
45+
! CHECK: !DIR$ UNROLL 2
46+
do i=1,10
47+
enddo
48+
end subroutine

‎flang/test/Semantics/loop-directives.f90

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,24 @@
44
subroutine empty
55
! WARNING: A DO loop must follow the VECTOR ALWAYS directive
66
!dir$ vector always
7+
! WARNING: A DO loop must follow the UNROLL directive
8+
!dir$ unroll
79
end subroutine empty
810

911
subroutine non_do
1012
! WARNING: A DO loop must follow the VECTOR ALWAYS directive
1113
!dir$ vector always
14+
! WARNING: A DO loop must follow the UNROLL directive
15+
!dir$ unroll
1216
a = 1
1317
end subroutine non_do
1418

1519
subroutine execution_part
1620
do i=1,10
1721
! WARNING: A DO loop must follow the VECTOR ALWAYS directive
1822
!dir$ vector always
23+
! WARNING: A DO loop must follow the UNROLL directive
24+
!dir$ unroll
1925
end do
2026
end subroutine execution_part
2127

@@ -28,3 +34,13 @@ subroutine test_vector_always_before_acc(a, b, c)
2834
a(i) = b(i) + c(i)
2935
enddo
3036
end subroutine
37+
38+
! OK
39+
subroutine test_unroll_before_acc(a, b, c)
40+
real, dimension(10) :: a,b,c
41+
!dir$ unroll
42+
!$acc loop
43+
do i=1,N
44+
a(i) = b(i) + c(i)
45+
enddo
46+
end subroutine

0 commit comments

Comments
 (0)