[flang] [llvm] [mlir] [flang][OpenMP] Enable tiling (PR #143715)

[Kareem Ergawy via llvm-commits]

================
@@ -112,15 +111,19 @@ class CanonicalizationOfOmp {
     // in the same iteration
     //
     // Original:
-    //   ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct
-    //     OmpBeginLoopDirective
+    //   ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct t->
----------------
ergawy wrote:

I found the implementation of this function quite complex to understand tbh. 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 flatterning the logic a little bit:
```c++
  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.

https://github.com/llvm/llvm-project/pull/143715