[Mlir-commits] [mlir] [MLIR][OpenMP] Normalize lowering of omp.loop_nest (PR #127217)

[Sergio Afonso]

================
@@ -2348,57 +2364,31 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
     llvm::Expected<llvm::CanonicalLoopInfo *> loopResult =
         ompBuilder->createCanonicalLoop(
             loc, bodyGen, lowerBound, upperBound, step,
-            /*IsSigned=*/true, /*InclusiveStop=*/true, computeIP);
+            /*IsSigned=*/true, loopOp.getLoopInclusive(), computeIP);
 
     if (failed(handleError(loopResult, *loopOp)))
       return failure();
 
     loopInfos.push_back(*loopResult);
   }
 
-  // Collapse loops.
-  llvm::IRBuilderBase::InsertPoint afterIP = loopInfos.front()->getAfterIP();
-  llvm::CanonicalLoopInfo *loopInfo =
-      ompBuilder->collapseLoops(ompLoc.DL, loopInfos, {});
-
-  llvm::ConstantInt *simdlen = nullptr;
-  if (std::optional<uint64_t> simdlenVar = simdOp.getSimdlen())
-    simdlen = builder.getInt64(simdlenVar.value());
+  // Collapse loops. Store the insertion point because LoopInfos may get
+  // invalidated.
+  llvm::OpenMPIRBuilder::InsertPointTy afterIP =
+      loopInfos.front()->getAfterIP();
 
-  llvm::ConstantInt *safelen = nullptr;
-  if (std::optional<uint64_t> safelenVar = simdOp.getSafelen())
-    safelen = builder.getInt64(safelenVar.value());
-
-  llvm::MapVector<llvm::Value *, llvm::Value *> alignedVars;
-  llvm::omp::OrderKind order = convertOrderKind(simdOp.getOrder());
-  llvm::BasicBlock *sourceBlock = builder.GetInsertBlock();
-  std::optional<ArrayAttr> alignmentValues = simdOp.getAlignments();
-  mlir::OperandRange operands = simdOp.getAlignedVars();
-  for (size_t i = 0; i < operands.size(); ++i) {
-    llvm::Value *alignment = nullptr;
-    llvm::Value *llvmVal = moduleTranslation.lookupValue(operands[i]);
-    llvm::Type *ty = llvmVal->getType();
-    if (auto intAttr = llvm::dyn_cast<IntegerAttr>((*alignmentValues)[i])) {
-      alignment = builder.getInt64(intAttr.getInt());
-      assert(ty->isPointerTy() && "Invalid type for aligned variable");
-      assert(alignment && "Invalid alignment value");
-      auto curInsert = builder.saveIP();
-      builder.SetInsertPoint(sourceBlock->getTerminator());
-      llvmVal = builder.CreateLoad(ty, llvmVal);
-      builder.restoreIP(curInsert);
-      alignedVars[llvmVal] = alignment;
-    }
-  }
-  ompBuilder->applySimd(loopInfo, alignedVars,
-                        simdOp.getIfExpr()
-                            ? moduleTranslation.lookupValue(simdOp.getIfExpr())
-                            : nullptr,
-                        order, simdlen, safelen);
+  // Add a stack frame holding information about the resulting loop after
+  // applying transformations, to be further transformed by parent loop
+  // wrappers.
+  moduleTranslation.stackPush<OpenMPLoopInfoStackFrame>(
----------------
skatrak wrote:

This is a tricky issue. The problem with `SaveStack` is that it pushes and pops from within the same function. In this case, we could move those two calls to `convertHostOrTargetOperation`, but they would have to be only done conditionally for the outermost loop wrapper. That means we couldn't use the `SaveStack` class either way, but at least the push and pop calls would be located in the same function (which is certainly easier to follow).

However, if we do that, another issue that we encounter is that the current design of `StackFrame` handling expects them to be created once and never modified. Here, we want to set the `loopInfo` pointer based on the value produced while translating the inner `omp.loop_nest` and then use it during translation of the associated parent loop wrappers. So, we would have to create a stack frame with a `null` value initially and then update/create some way of accessing it in a non-const way (removing `const` from `ModuleTranslation::stackWalk`, for example).

This all boils down to the core issue of trying to pass information from inside (the nested loop) to outside (its wrappers), and the stack frame system is designed to do the opposite. It doesn't quite fit either way we try to do it, but so far it's the best I've been able to come up with (the alternative being a global variable). The implementation I proposed here is the one that doesn't introduce changes to `ModuleTranslation` for something that's a very unusual use of the stack frame feature.

> The pop is conditional in `convertHostOrTargetOperation`, the pairing is not obviously to me.

For each `omp.loop_nest` there will be 1 or more loop wrappers. If we push a stack frame when we create the `omp.loop_nest` to be used up the chain, the place where we can finally pop that stack frame is when we finish processing the outermost loop wrapper. That's how the pairing is done.

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