Skip to content

[flang][OpenMP] Enable tiling #143715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 23 commits into
base: main
Choose a base branch
from
Open

[flang][OpenMP] Enable tiling #143715

Show file tree
Hide file tree
Changes from 18 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
943ecef
Initial implementation of tiling.
jsjodin Mar 14, 2025
bcaacf9
Fix tests and limit the nesting of construct to only tiling.
jsjodin Jun 10, 2025
8c2cd48
Enable stand-alone tiling, but it gives a warning and converting to s…
jsjodin Jun 10, 2025
5c17a0d
Add minimal test, remove debug print.
jsjodin Jun 11, 2025
7e7437b
Fix formatting
jsjodin Jun 13, 2025
3367c5e
Fix formatting
jsjodin Jun 14, 2025
753653d
Fix test.
jsjodin Jun 19, 2025
2d09a69
Add more mlir tests. Set collapse value when lowering from SCF to Ope…
jsjodin Jun 20, 2025
db09c91
Fix formatting
jsjodin Jun 20, 2025
ab74c1a
Use llvm::SmallVector instead of std::stack
jsjodin Jun 20, 2025
88348a3
Improve test a bit to make sure IVs are used as expected.
jsjodin Jun 21, 2025
d362617
Fix comments to clarify canonicalization.
jsjodin Jun 21, 2025
60aa4b0
Special handling of tile directive when dealing with start end end lo…
jsjodin Jun 21, 2025
ad32cc3
Inline functions.
jsjodin Jun 21, 2025
6da33a4
Remove debug code.
jsjodin Jun 23, 2025
991e042
Reuse loop op lowering, add comment.
jsjodin Jun 23, 2025
b4e2109
Fix formatting.
jsjodin Jun 23, 2025
5a3d8d2
Remove curly braces.
jsjodin Jun 23, 2025
3029793
Avoid attaching the sizes clause to the parent construct, instead fin…
jsjodin Jun 25, 2025
bdead72
Fix formatting
jsjodin Jun 25, 2025
f1c260d
Fix unparse and add a test for nested loop constructs.
jsjodin Jun 26, 2025
7701e07
Use more convenient function to get OpenMPLoopConstruct. Fix comments.
jsjodin Jun 26, 2025
7c3b3e5
Fix formatting.
jsjodin Jun 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion flang/include/flang/Lower/OpenMP.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,6 @@ void genOpenMPDeclarativeConstruct(AbstractConverter &,
void genOpenMPSymbolProperties(AbstractConverter &converter,
const pft::Variable &var);

int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
void genThreadprivateOp(AbstractConverter &, const pft::Variable &);
void genDeclareTargetIntGlobal(AbstractConverter &, const pft::Variable &);
bool isOpenMPTargetConstruct(const parser::OpenMPConstruct &);
Expand Down
4 changes: 3 additions & 1 deletion flang/include/flang/Parser/parse-tree.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5024,8 +5024,10 @@ struct OpenMPBlockConstruct {
struct OpenMPLoopConstruct {
TUPLE_CLASS_BOILERPLATE(OpenMPLoopConstruct);
OpenMPLoopConstruct(OmpBeginLoopDirective &&a)
: t({std::move(a), std::nullopt, std::nullopt}) {}
: t({std::move(a), std::nullopt, std::nullopt, std::nullopt}) {}
std::tuple<OmpBeginLoopDirective, std::optional<DoConstruct>,
// Inner loop construct used to handle tiling for now.
std::optional<common::Indirection<OpenMPLoopConstruct>>,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: can you add a comment here explaining that this is used mainly for tiling constructs?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this strictly only support nesting of two loop constructs or do you plan for it to work like a linked list: with each loop construct pointing to the next inner-most construct?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It can work as a linked list for multiple nested loop constructs. Note that we also might want to keep the DoConstruct inside the innermost loop construct in the future, instead of the outer one. The choice to keep it with the outer loop construct was to minimize the changes in the rest of the lowering since we can only handle tiling. Moving the DoConstruct in the canonicalization to the tiling loop construct (inner most) can be a separate PR if that is a better representation. Allowing multiple nested ops would be the final step, at that point we would need to use the CLIs, canonical loops and loop transformation ops.

Copy link
Contributor

@Stylie777 Stylie777 Jun 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Allowing multiple nested ops would be the final step, at that point we would need to use the CLIs, canonical loops and loop transformation ops.

As part of my work for #110008 I have implemented nested ops for tile and unroll in the semantics checks (PR coming soon, just doing finishing touches). It does work slightly differently in that std::optional<DoConstruct> has become std::optional<std::variant<DoConstruct, common::Indirection<OpenMPLoopConstruct>>> in the Parse Tree.

This approach stores either one or the other, and then Flang can determine which is being used. I personally think this is a better approach as then you are not opening the possibility for there being a DoConstruct and OpenMPLoopConstruct. It works in a similar way where it will link together the LoopConstructs until you reach a DoConstruct.

Its currently a standalone change, as this is not merged yet, but it would be good to agree on an approach so there is minimal changes required when one of these is merged.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree that having the nesting done in this way is a better approach, I forgot that I actually tried out putting the DoConstruct inside the innermost OpenMPLoopConstruct, and the code generation was fine, my comment above was incorrect. The parse tree after canonicalization looks like this:

| | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct
| | | OmpBeginLoopDirective
| | | | OmpLoopDirective -> llvm::omp::Directive = do
| | | | OmpClauseList -> 
| | | OpenMPLoopConstruct
| | | | OmpBeginLoopDirective
| | | | | OmpLoopDirective -> llvm::omp::Directive = tile
| | | | | OmpClauseList -> OmpClause -> Sizes -> Scalar -> Integer -> Expr = '2_4'
| | | | | | LiteralConstant -> IntLiteralConstant = '2'
| | | | DoConstruct
| | | | | NonLabelDoStmt
| | | | | | LoopControl -> LoopBounds
| | | | | | | Scalar -> Name = 'x'
| | | | | | | Scalar -> Expr = '1_4'
| | | | | | | | LiteralConstant -> IntLiteralConstant = '1'
| | | | | | | Scalar -> Expr = '100_4'
| | | | | | | | LiteralConstant -> IntLiteralConstant = '100'

So this should be pretty compatible with what you are working on.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, the changes I have worked on give a similar parsing outcome. For reference, this will be my version of OpenMPLoopConstruct

struct OpenMPLoopConstruct {
  TUPLE_CLASS_BOILERPLATE(OpenMPLoopConstruct);
  OpenMPLoopConstruct(OmpBeginLoopDirective &&a)
      : t({std::move(a), std::nullopt, std::nullopt}) {}
  std::tuple<OmpBeginLoopDirective, std::optional<std::variant<DoConstruct, common::Indirection<OpenMPLoopConstruct>>>,
      std::optional<OmpEndLoopDirective>>
      t;
};

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just pushed a change for a test and to fix unparse.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just for reference, my PR for the Semantics covering unroll and tile are here: #145917

There does seem to be some crossover with this patch, so some rebasing will be needed depending on which is merged first.

std::optional<OmpEndLoopDirective>>
t;
};
Expand Down
75 changes: 60 additions & 15 deletions flang/lib/Lower/OpenMP/OpenMP.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,7 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
return;

const parser::OmpClauseList *beginClauseList = nullptr;
const parser::OmpClauseList *middleClauseList = nullptr;
const parser::OmpClauseList *endClauseList = nullptr;
common::visit(
common::visitors{
Expand All @@ -473,6 +474,23 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
beginClauseList =
&std::get<parser::OmpClauseList>(beginDirective.t);

// FIXME(JAN): For now we check if there is an inner
// OpenMPLoopConstruct, and extract the size clause from there
const auto &innerOptional = std::get<std::optional<
common::Indirection<parser::OpenMPLoopConstruct>>>(
ompConstruct.t);
if (innerOptional.has_value()) {
const auto &innerLoopDirective = innerOptional.value().value();
const auto &innerBegin =
std::get<parser::OmpBeginLoopDirective>(
innerLoopDirective.t);
const auto &innerDirective =
std::get<parser::OmpLoopDirective>(innerBegin.t);
if (innerDirective.v == llvm::omp::Directive::OMPD_tile) {
middleClauseList =
&std::get<parser::OmpClauseList>(innerBegin.t);
}
}
if (auto &endDirective =
std::get<std::optional<parser::OmpEndLoopDirective>>(
ompConstruct.t))
Expand All @@ -485,6 +503,9 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
assert(beginClauseList && "expected begin directive");
clauses.append(makeClauses(*beginClauseList, semaCtx));

if (middleClauseList)
clauses.append(makeClauses(*middleClauseList, semaCtx));

if (endClauseList)
clauses.append(makeClauses(*endClauseList, semaCtx));
};
Expand Down Expand Up @@ -960,6 +981,7 @@ static void genLoopVars(
storeOp =
createAndSetPrivatizedLoopVar(converter, loc, indexVal, argSymbol);
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Extra whitespace

firOpBuilder.setInsertionPointAfter(storeOp);
}

Expand Down Expand Up @@ -1712,6 +1734,23 @@ genLoopNestClauses(lower::AbstractConverter &converter,
cp.processCollapse(loc, eval, clauseOps, iv);

clauseOps.loopInclusive = converter.getFirOpBuilder().getUnitAttr();

fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
for (auto &clause : clauses) {
if (clause.id == llvm::omp::Clause::OMPC_collapse) {
const auto &collapse = std::get<clause::Collapse>(clause.u);
int64_t collapseValue = evaluate::ToInt64(collapse.v).value();
clauseOps.numCollapse = firOpBuilder.getI64IntegerAttr(collapseValue);
} else if (clause.id == llvm::omp::Clause::OMPC_sizes) {
const auto &sizes = std::get<clause::Sizes>(clause.u);
llvm::SmallVector<int64_t> sizeValues;
for (auto &size : sizes.v) {
int64_t sizeValue = evaluate::ToInt64(size).value();
sizeValues.push_back(sizeValue);
}
clauseOps.tileSizes = sizeValues;
}
}
}

static void genLoopClauses(
Expand Down Expand Up @@ -2085,9 +2124,9 @@ static mlir::omp::LoopNestOp genLoopNestOp(
return llvm::SmallVector<const semantics::Symbol *>(iv);
};

auto *nestedEval =
getCollapsedLoopEval(eval, getCollapseValue(item->clauses));

uint64_t nestValue = getCollapseValue(item->clauses);
nestValue = nestValue < iv.size() ? iv.size() : nestValue;
auto *nestedEval = getCollapsedLoopEval(eval, nestValue);
return genOpWithBody<mlir::omp::LoopNestOp>(
OpWithBodyGenInfo(converter, symTable, semaCtx, loc, *nestedEval,
directive)
Expand Down Expand Up @@ -3610,6 +3649,8 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
item);
break;
case llvm::omp::Directive::OMPD_tile:
newOp = genLoopOp(converter, symTable, semaCtx, eval, loc, queue, item);
break;
case llvm::omp::Directive::OMPD_unroll: {
unsigned version = semaCtx.langOptions().OpenMPVersion;
TODO(loc, "Unhandled loop directive (" +
Expand Down Expand Up @@ -4186,6 +4227,22 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
std::get<parser::OmpBeginLoopDirective>(loopConstruct.t);
List<Clause> clauses = makeClauses(
std::get<parser::OmpClauseList>(beginLoopDirective.t), semaCtx);

const auto &innerOptional =
std::get<std::optional<common::Indirection<parser::OpenMPLoopConstruct>>>(
loopConstruct.t);
if (innerOptional.has_value()) {
const auto &innerLoopDirective = innerOptional.value().value();
const auto &innerBegin =
std::get<parser::OmpBeginLoopDirective>(innerLoopDirective.t);
const auto &innerDirective =
std::get<parser::OmpLoopDirective>(innerBegin.t);
if (innerDirective.v == llvm::omp::Directive::OMPD_tile) {
clauses.append(
makeClauses(std::get<parser::OmpClauseList>(innerBegin.t), semaCtx));
}
}

if (auto &endLoopDirective =
std::get<std::optional<parser::OmpEndLoopDirective>>(
loopConstruct.t)) {
Expand Down Expand Up @@ -4292,18 +4349,6 @@ void Fortran::lower::genOpenMPSymbolProperties(
lower::genDeclareTargetIntGlobal(converter, var);
}

int64_t
Fortran::lower::getCollapseValue(const parser::OmpClauseList &clauseList) {
for (const parser::OmpClause &clause : clauseList.v) {
if (const auto &collapseClause =
std::get_if<parser::OmpClause::Collapse>(&clause.u)) {
const auto *expr = semantics::GetExpr(collapseClause->v);
return evaluate::ToInt64(*expr).value();
}
}
return 1;
}

void Fortran::lower::genThreadprivateOp(lower::AbstractConverter &converter,
const lower::pft::Variable &var) {
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
Expand Down
32 changes: 24 additions & 8 deletions flang/lib/Lower/OpenMP/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,21 @@ namespace lower {
namespace omp {

int64_t getCollapseValue(const List<Clause> &clauses) {
auto iter = llvm::find_if(clauses, [](const Clause &clause) {
return clause.id == llvm::omp::Clause::OMPC_collapse;
});
if (iter != clauses.end()) {
const auto &collapse = std::get<clause::Collapse>(iter->u);
return evaluate::ToInt64(collapse.v).value();
int64_t collapseValue = 1;
int64_t numTileSizes = 0;
for (auto &clause : clauses) {
if (clause.id == llvm::omp::Clause::OMPC_collapse) {
const auto &collapse = std::get<clause::Collapse>(clause.u);
collapseValue = evaluate::ToInt64(collapse.v).value();
} else if (clause.id == llvm::omp::Clause::OMPC_sizes) {
const auto &sizes = std::get<clause::Sizes>(clause.u);
numTileSizes = sizes.v.size();
}
Comment on lines +65 to +71
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This mixes the directive that collapse goes on (DO, DISTRIBUTE, ..) and those sizes goes on (TILE). During resolving merge conflicts of #145917 and #144785 the approach of handling tile like a compound construct already caused problems. E.g. it could be tile twice with different sizes, or unroll partial followed by tile vs the other way around, and there is no way to disambiguate them.

So it will need to be tangled apart eventually, and I would appreciate if we could avoid the assumption that there is just one tile which is always first.

}
return 1;

collapseValue = collapseValue - numTileSizes;
int64_t result = collapseValue > numTileSizes ? collapseValue : numTileSizes;
return result;
Comment on lines +74 to +76
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you explain to me why we need this calculation? To see the effect of this, I tried replacing these few lines with simply return collapseValue; and ran all tests but no tests failed. So it seems this part is not tested. A test can also help explaining the purpose of the change.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The general thing this computes is the number of loops that need to be considered in the source code. If you have collapse(4) on a loop nest with 2 loops that would be incorrect since we can max collapse 2 loops. However tiling creates new loops, so collapse(4) would theoretically be legal if tiling is done first e.g. tile(5,10) since that will result in 4 loops. This is not really testable though since collapse requires independent loops, which is only true for the 2 outer loops after tiling is done. There is a check for this, and an error message is given if the collapse value is larger than the number of loops that are tiled to prevent incorrect code. We could just use numTileSizes if that is present, but if collapse could handle dependent loops in the future the above calculation should be the correct one.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Support for collapsing grid- and tile-loops have only been added in OpenMP 6.0. Before that, there was a restriction:
image

The value of the collapse clause should not really be modified implicitly. Is there is check in FortranSemantics that ensures collapse is not larger than the number of grid loops?

}

void genObjectList(const ObjectList &objects,
Expand Down Expand Up @@ -611,6 +618,7 @@ bool collectLoopRelatedInfo(
lower::pft::Evaluation &eval, const omp::List<omp::Clause> &clauses,
mlir::omp::LoopRelatedClauseOps &result,
llvm::SmallVectorImpl<const semantics::Symbol *> &iv) {

bool found = false;
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();

Expand All @@ -626,7 +634,16 @@ bool collectLoopRelatedInfo(
collapseValue = evaluate::ToInt64(clause->v).value();
found = true;
}
std::int64_t sizesLengthValue = 0l;
if (auto *clause =
ClauseFinder::findUniqueClause<omp::clause::Sizes>(clauses)) {
sizesLengthValue = clause->v.size();
found = true;
}

collapseValue = collapseValue - sizesLengthValue;
collapseValue =
collapseValue < sizesLengthValue ? sizesLengthValue : collapseValue;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This as well does not seem to affect on tests added/modified by the PR.

Let me know if I missed something.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No you didn't miss anything, and because of the limitations with tile/collapse only one case (tile vs no tile) is possible. We could go with a simpler calculation if we want to.

It was also a but of an illustration for when we allow multiple loop transformations it could be pretty tricky to know what is legal and not depending on how they are combined.

std::size_t loopVarTypeSize = 0;
do {
lower::pft::Evaluation *doLoop =
Expand Down Expand Up @@ -659,7 +676,6 @@ bool collectLoopRelatedInfo(
} while (collapseValue > 0);

convertLoopBounds(converter, currentLocation, result, loopVarTypeSize);

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nit] unrelated change

return found;
}
} // namespace omp
Expand Down
66 changes: 54 additions & 12 deletions flang/lib/Semantics/canonicalize-omp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

#include "canonicalize-omp.h"
#include "flang/Parser/parse-tree-visitor.h"

// After Loop Canonicalization, rewrite OpenMP parse tree to make OpenMP
// Constructs more structured which provide explicit scopes for later
// structural checks and semantic analysis.
Expand Down Expand Up @@ -112,15 +111,19 @@ class CanonicalizationOfOmp {
// in the same iteration
//
// Original:
// ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct
// OmpBeginLoopDirective
// ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct t->
Copy link
Member

@ergawy ergawy Jun 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found the implementation of this function quite complex to understand tbh, for me at least. Specially with the 2 nested while loops (the old one and the one introduced by the PR).

While trying to understand what is going on, I tried to simpify it by flattening the logic a little bit:

  void RewriteOpenMPLoopConstruct(parser::OpenMPLoopConstruct &x,
      parser::Block &block, parser::Block::iterator it) {
    // Check the sequence of DoConstruct and OmpEndLoopDirective
    // in the same iteration
    //
    // Original:
    //   ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct t->
    //     OmpBeginLoopDirective t-> OmpLoopDirective
    //   [ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct u->
    //     OmpBeginLoopDirective t-> OmpLoopDirective t-> Tile v-> OMP_tile]
    //   ExecutableConstruct -> DoConstruct
    //   [ExecutableConstruct -> OmpEndLoopDirective] (note: tile)
    //   ExecutableConstruct -> OmpEndLoopDirective (if available)
    //
    // After rewriting:
    //   ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct t->
    //     [OpenMPLoopConstruct t -> OmpBeginLoopDirective -> OmpLoopDirective
    //                               OmpEndLoopDirective] (note: tile)
    //     OmpBeginLoopDirective t -> OmpLoopDirective -> DoConstruct
    //     OmpEndLoopDirective (if available)
    parser::Block::iterator nextIt;
    auto &beginDir{std::get<parser::OmpBeginLoopDirective>(x.t)};
    auto &dir{std::get<parser::OmpLoopDirective>(beginDir.t)};

    // Ignore compiler directives.
    nextIt = std::find_if_not(++it, block.end(), [this](auto &y) {
      return GetConstructIf<parser::CompilerDirective>(y);
    });

    if (nextIt == block.end()) {
      return;
    }

    auto &optionalInnerTile = std::get<
        std::optional<common::Indirection<parser::OpenMPLoopConstruct>>>(x.t);

    if (auto *innerConstruct{
            GetConstructIf<parser::OpenMPConstruct>(*nextIt)}) {
      if (auto *innerOmpLoop{
              std::get_if<parser::OpenMPLoopConstruct>(&innerConstruct->u)}) {
        auto &innerBeginDir{
            std::get<parser::OmpBeginLoopDirective>(innerOmpLoop->t)};
        auto &innerDir{std::get<parser::OmpLoopDirective>(innerBeginDir.t)};
        if (innerDir.v == llvm::omp::Directive::OMPD_tile) {
          optionalInnerTile = std::move(*innerOmpLoop);
          // Retrieveing the address so that DoConstruct or inner loop can be
          // set later.
          nextIt = block.erase(nextIt);
        }
      }
    }

    if (auto *doCons{GetConstructIf<parser::DoConstruct>(*nextIt)}) {
      if (doCons->GetLoopControl()) {
        parser::OpenMPLoopConstruct *innermostLoop =
            optionalInnerTile ? &optionalInnerTile.value().value() : &x;
        std::get<std::optional<parser::DoConstruct>>(innermostLoop->t) =
            std::move(*doCons);
        nextIt = block.erase(nextIt);

      } else {
        messages_.Say(dir.source,
            "DO loop after the %s directive must have loop control"_err_en_US,
            parser::ToUpperCaseLetters(dir.source.ToString()));
      }
    } else {
      messages_.Say(dir.source,
          "A DO loop must follow the %s directive"_err_en_US,
          parser::ToUpperCaseLetters(dir.source.ToString()));
    }

    auto tryToMatchEndLoopDirective = [&, this](
                                          parser::OpenMPLoopConstruct *loop) {
      if (nextIt == block.end()) {
        return false;
      }

      if (auto *endDir{GetConstructIf<parser::OmpEndLoopDirective>(*nextIt)}) {
        auto &endOmpDirective{std::get<parser::OmpLoopDirective>(endDir->t)};
        auto &loopBegin{std::get<parser::OmpBeginLoopDirective>(loop->t)};
        auto &loopDir{std::get<parser::OmpLoopDirective>(loopBegin.t)};

        // If the directive is a tile we try to match the corresponding
        // end tile if it exsists. If it is not a tile directive we
        // always assign the end loop directive and fall back on the
        // existing directive structure checks.
        if (loopDir.v != llvm::omp::Directive::OMPD_tile ||
            loopDir.v == endOmpDirective.v) {
          std::get<std::optional<parser::OmpEndLoopDirective>>(loop->t) =
              std::move(*endDir);
          nextIt = block.erase(nextIt);
        }

        return true;
      }

      // If there is a mismatch bail out.
      return false;
    };

    if (optionalInnerTile) {
      if (tryToMatchEndLoopDirective(&optionalInnerTile.value().value())) {
        tryToMatchEndLoopDirective(&x);
      }
    } else {
      tryToMatchEndLoopDirective(&x);
    }
  }

This is a, hopefully, more simplified version that matches the original logic. Feel free to ignore if I got the logic wrong.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would like to keep the existing one for now, so other reviewers don't have to go through it again. We can make this a separate cleanup PR afterwards. I'd move the lambda function outside of this function definition btw since it is not passed anywhere.

// OmpBeginLoopDirective t-> OmpLoopDirective
// [ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct u->
// OmpBeginLoopDirective t-> OmpLoopDirective t-> Tile v-> OMP_tile]
// ExecutableConstruct -> DoConstruct
// [ExecutableConstruct -> OmpEndLoopDirective] (note: tile)
// ExecutableConstruct -> OmpEndLoopDirective (if available)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you mean to delete this line?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is the optional OmpEndLoopDirective associated with the optional BeginLoopDirective for the tile construct, so it is meant to be there. I will update the "After rewriting" comment to hopefully clarify what is going on.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done updating comments.

//
// After rewriting:
// ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct
// OmpBeginLoopDirective
// DoConstruct
// ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct t->
// [OpenMPLoopConstruct t -> OmpBeginLoopDirective -> OmpLoopDirective
// OmpEndLoopDirective] (note: tile)
// OmpBeginLoopDirective t -> OmpLoopDirective -> DoConstruct
// OmpEndLoopDirective (if available)
parser::Block::iterator nextIt;
auto &beginDir{std::get<parser::OmpBeginLoopDirective>(x.t)};
Expand All @@ -131,20 +134,59 @@ class CanonicalizationOfOmp {
// Ignore compiler directives.
if (GetConstructIf<parser::CompilerDirective>(*nextIt))
continue;
// Keep track of the loops to handle the end loop directives
llvm::SmallVector<parser::OpenMPLoopConstruct *> loops;
loops.push_back(&x);
if (auto *innerConstruct{
GetConstructIf<parser::OpenMPConstruct>(*nextIt)}) {
if (auto *innerOmpLoop{
std::get_if<parser::OpenMPLoopConstruct>(&innerConstruct->u)}) {
auto &innerBeginDir{
std::get<parser::OmpBeginLoopDirective>(innerOmpLoop->t)};
auto &innerDir{std::get<parser::OmpLoopDirective>(innerBeginDir.t)};
if (innerDir.v == llvm::omp::Directive::OMPD_tile) {
auto &innerLoop = std::get<std::optional<
common::Indirection<parser::OpenMPLoopConstruct>>>(
loops.back()->t);
innerLoop = std::move(*innerOmpLoop);
// Retrieveing the address so that DoConstruct or inner loop can be
// set later.
loops.push_back(&(innerLoop.value().value()));
nextIt = block.erase(nextIt);
}
}
}

if (auto *doCons{GetConstructIf<parser::DoConstruct>(*nextIt)}) {
if (doCons->GetLoopControl()) {
// move DoConstruct
std::get<std::optional<parser::DoConstruct>>(x.t) =
std::get<std::optional<parser::DoConstruct>>(loops.back()->t) =
std::move(*doCons);
nextIt = block.erase(nextIt);
// try to match OmpEndLoopDirective
if (nextIt != block.end()) {
while (nextIt != block.end() && !loops.empty()) {
if (auto *endDir{
GetConstructIf<parser::OmpEndLoopDirective>(*nextIt)}) {
std::get<std::optional<parser::OmpEndLoopDirective>>(x.t) =
std::move(*endDir);
block.erase(nextIt);
auto &endOmpDirective{
std::get<parser::OmpLoopDirective>(endDir->t)};
auto &loopBegin{
std::get<parser::OmpBeginLoopDirective>(loops.back()->t)};
auto &loopDir{std::get<parser::OmpLoopDirective>(loopBegin.t)};

// If the directive is a tile we try to match the corresponding
// end tile if it exsists. If it is not a tile directive we
// always assign the end loop directive and fall back on the
// existing directive structure checks.
if (loopDir.v != llvm::omp::Directive::OMPD_tile ||
loopDir.v == endOmpDirective.v) {
std::get<std::optional<parser::OmpEndLoopDirective>>(
loops.back()->t) = std::move(*endDir);
nextIt = block.erase(nextIt);
}

loops.pop_back();
} else {
// If there is a mismatch bail out.
break;
}
}
} else {
Expand Down
Loading
Loading