Use more convenient function to get OpenMPLoopConstruct. Fix comments.

Author: jsjodin

flang/lib/Semantics/canonicalize-omp.cpp

@@ -137,23 +137,19 @@ class CanonicalizationOfOmp {
       // 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 *innerOmpLoop{GetOmpIf<parser::OpenMPLoopConstruct>(*nextIt)}) {
+        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);
         }
       }
 

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -2991,7 +2991,6 @@ convertOmpLoopNest(Operation &opInst, llvm::IRBuilderBase &builder,
     loopInfos.push_back(*loopResult);
   }
 
-  //  llvm::OpenMPIRBuilder::InsertPointTy afterIP = builder.saveIP();
   llvm::OpenMPIRBuilder::InsertPointTy afterIP =
       loopInfos.front()->getAfterIP();
 
@@ -3012,10 +3011,10 @@ convertOmpLoopNest(Operation &opInst, llvm::IRBuilderBase &builder,
     std::vector<llvm::CanonicalLoopInfo *> NewLoops =
         ompBuilder->tileLoops(ompLoc.DL, loopInfos, TileSizes);
 
-    // Collapse loops. Store the insertion point because LoopInfos may get
-    // invalidated.
-    auto AfterBB = NewLoops.front()->getAfter();
-    auto AfterAfterBB = AfterBB->getSingleSuccessor();
+    // Update afterIP to get the correct insertion point after
+    // tiling.
+    llvm::BasicBlock *AfterBB = NewLoops.front()->getAfter();
+    llvm::BasicBlock  *AfterAfterBB = AfterBB->getSingleSuccessor();
     afterIP = {AfterAfterBB, AfterAfterBB->begin()};
     NewTopLoopInfo = NewLoops[0];