[mlir] OpenMP-to-LLVM: properly set outer alloca insertion point

[ftynse]

Summary

Previously, the OpenMP to LLVM IR conversion was setting the alloca insertion
point to the same position as the main compuation when converting OpenMP
parallel operations. This is problematic if, for example, the parallel
operation is placed inside a loop and would keep allocating on stack on each
iteration leading to stack overflow.

[kiranchandramohan]

Thanks @ftynse for this patch.

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

171

I was planning to a few questions in this area. Will this getParent work correctly with nested parallelism and the if is in the nested parallel operation?
Should we be marking the omp.parallel operation with some of the function attributes like Function-Like, AutomaticAllocationScope?

llvm.func fn() {
  omp.parallel {
    omp.parallel if() {
    }
  }
}

mlir/test/Target/LLVMIR/openmp-llvm.mlir

155

Nit: funciton -> function

158–159

Is this enough to check that it is in the entry? Or should there be a check for a subsequent block name or number?

[kiranchandramohan]

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

171

And if the above does work correctly, consider the following.

llvm.func fn() {
  omp.parallel {
    omp.master { //<- The region of this operation will not be outlined.
      omp.parallel if() {
      }
    }
  }
}

[kumasento]

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

171

Not familiar with this field, but I was thinking about a case similar to (or not) above: what would happen if getParent() doesn't give a function?

173

It seems to me that if bodyGenStatus was failed earlier than restoreIP, the old logic will just abort but this new version will still do the restore. Would this cause any issue, or was the old logic actually problematic?

[jdoerfert]

I skimmed this and if I get it right this won't work reliably.
The OMPIRBuilder provides an alloca insertion point as part of the body code generation callback, that is the one to use.
The user needs to keep a stack for all alloca insertion points and push/pop them as the body callbacks are entered/left.
The same has to happen for other outlined operators, I'm imagining fir.concurrent.do or something like that.

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

92

^^ allocaIP is not used but should be.

[ftynse]

In D101307#2721716, @jdoerfert wrote:

The OMPIRBuilder provides an alloca insertion point as part of the body code generation callback, that is the one to use.

Let me clarify, the alloca insertion point given to the body generation callback should be used for further calls _inside_ the body, right? The outermost, e.g. createParallel, still needs an alloca insertion point, which should normally be the function entry block.

The user needs to keep a stack for all alloca insertion points and push/pop them as the body callbacks are entered/left.

In MLIR, this will be the actual call stack of nested calls to convertOmpParallel, but yes, it needs to carry the insertion point as extra argument.

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

171

@kiranchandramohan have we added the support for nested omp.parallel yet? Where should we place allocations in the nested case, in the AllocaIP of the body of the surrounding parallel operation? I can just condition this new behavior to only work for the outermost omp.parallel because that's the case I am hitting right now...

Should we be marking the omp.parallel operation with some of the function attributes like Function-Like, AutomaticAllocationScope?

Why? It's definitely not FunctionLike (it isn't a symbol, there are no arguments and results, etc.). AutomaticAllocationScope is useless at this point, we don't have any special handling for such operations. @shabalin was trying to add more generic allocation scoping, but was blocked on discussions.

Not familiar with this field, but I was thinking about a case similar to (or not) above: what would happen if getParent() doesn't give a function?

@kumasento, it's LLVM IR, the parent of a block is always a function.

173

It was wrong before. bodyGenStatus can only be modified inside createParallel that actually calls bodyGenCB, but the code was checking it _before_ and ignoring it after the modification.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

158–159

We are checking that these are placed before the icmp which is the first instruction in the function. Should be enough, and there is no block number of the entry block and no block separation after this.

[kiranchandramohan]

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

171

Yes, nested parallel operations work. We have tests also. But the allocaIP issue is not solved. https://github.com/llvm/llvm-project/blob/3d974ac9fc489ac3fec194f324be55e42d1ea4fa/mlir/test/Target/LLVMIR/openmp-llvm.mlir#L230

I can just condition this new behavior to only work for the outermost omp.parallel because that's the case I am hitting right now...

This is OK with me for now.

Where should we place allocations in the nested case, in the AllocaIP of the body of the surrounding parallel operation?
Why? It's definitely not FunctionLike (it isn't a symbol, there are no arguments and results, etc.). AutomaticAllocationScope is useless at this point, we don't have any special handling for such operations. @shabalin was trying to add more generic allocation scoping, but was blocked on discussions.

In the OpenMP dialect (and possibly others like FIR) we have to distinguish operations with regions which will be outlined (like parallel, task) and which will not (like master, single) be outlined. The operations which will be outlined have some properties similar to functions, like having an entry-block for allocas and finally existing as a function in LLVM IR. Ideally, I want a function (getAllocaBlock) which if I call will traverse up through the nested operations and get me the entry-block of the operation which will be outlined or if there is no operation like that then the function entry-block.

[kiranchandramohan]

In D101307#2722384, @ftynse wrote:

In D101307#2721716, @jdoerfert wrote:

The OMPIRBuilder provides an alloca insertion point as part of the body code generation callback, that is the one to use.

Let me clarify, the alloca insertion point given to the body generation callback should be used for further calls _inside_ the body, right? The outermost, e.g. createParallel, still needs an alloca insertion point, which should normally be the function entry block.

The user needs to keep a stack for all alloca insertion points and push/pop them as the body callbacks are entered/left.

In MLIR, this will be the actual call stack of nested calls to convertOmpParallel, but yes, it needs to carry the insertion point as extra argument.

If I remember correctly there are two allocaIPs, outerAllocaIP which is passed to the createParallel function of the OpenMPIRBuilder and the innerAllocaIP which is the allocaIP of the outlined region which the OpenMPIRBuilder creates. The BodyGenCallback function is called with the innerAllocaIP as argument.
You are correct that for the outermost OpenMP Operation it is the function entry block that should be the allocaIP. But for any nested Operations it should be the allocaIP provided by the BodyGenCallback that should be the AllocaIP.

We alos need some changes so that there is only a single allocaBlock. I believe we need to either,

1. Set the block corresponding to the entry-block of the omp.parallel operation as the innerAllocaIP provided by the OpenMPIRBuilder in the bodyGenCallBack. https://github.com/llvm/llvm-project/blob/ab5823867c4aee7f3e02ddfaa217905c87471bf9/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L40
2. Reserve the entry-block of the omp.parallel operation in MLIR for allocas and pass it to the OpenMPIRBuilder in the createParallel Function. And change the OpenMPIRBuilder to work with this as the inner-alloca entry block of the outlined region.

[jdoerfert]

In D101307#2722384, @ftynse wrote:

In D101307#2721716, @jdoerfert wrote:

The OMPIRBuilder provides an alloca insertion point as part of the body code generation callback, that is the one to use.

Let me clarify, the alloca insertion point given to the body generation callback should be used for further calls _inside_ the body, right? The outermost, e.g. createParallel, still needs an alloca insertion point, which should normally be the function entry block.

Right. I was expecting lowering should keep a stack, each new function pushes the entry block, each time a operator gives you a new alloca-IP you push/pop it appropriately.

The user needs to keep a stack for all alloca insertion points and push/pop them as the body callbacks are entered/left.

In MLIR, this will be the actual call stack of nested calls to convertOmpParallel, but yes, it needs to carry the insertion point as extra argument.

I assume some extra state is preferable as there might be various places, called in arbitrary nesting, that need to push/pop a alloca-IP. If you prefer an
argument passed through all of lowering instead of a "global stack", that should work too. However, alloca-IP probably will influence, and be influenced, by
non-OpenMP operators too.

[kiranchandramohan]

I can just condition this new behavior to only work for the outermost omp.parallel because that's the case I am hitting right now...

This is OK with me for now.

Just re-iterating that special casing this only for the outermost omp.parallel is fine with me.

[ftynse]

Handle nested scopes.

[ftynse]

In D101307#2723445, @jdoerfert wrote:

In D101307#2722384, @ftynse wrote:

In D101307#2721716, @jdoerfert wrote:

The OMPIRBuilder provides an alloca insertion point as part of the body code generation callback, that is the one to use.

Let me clarify, the alloca insertion point given to the body generation callback should be used for further calls _inside_ the body, right? The outermost, e.g. createParallel, still needs an alloca insertion point, which should normally be the function entry block.

Right. I was expecting lowering should keep a stack, each new function pushes the entry block, each time a operator gives you a new alloca-IP you push/pop it appropriately.

The user needs to keep a stack for all alloca insertion points and push/pop them as the body callbacks are entered/left.

In MLIR, this will be the actual call stack of nested calls to convertOmpParallel, but yes, it needs to carry the insertion point as extra argument.

I assume some extra state is preferable as there might be various places, called in arbitrary nesting, that need to push/pop a alloca-IP. If you prefer an
argument passed through all of lowering instead of a "global stack", that should work too. However, alloca-IP probably will influence, and be influenced, by
non-OpenMP operators too.

Yeah, the trick is to do this in a sufficiently modular way for MLIR. That is, the main translation engine must not care about OpenMP specifically. And the rest of the dialects are not concerned with regions (FWIW, I would have required regions to be flattened and the OpenMP dialect to be translated to the LLVM dialect hadn't there been the argument of code reused with OpenMPIRBuilder) or alloca insertion points, all that should happen before translation. I defined a globally visible stack, it is currently only used by the OpenMP dialect, but if others need it, we can factor out the OpenMP stack frame class to be a common alloca insertion point stack frame class.

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

171

Not familiar with this field, but I was thinking about a case similar to (or not) above: what would happen if getParent() doesn't give a function?

It's LLVM IR, getParent on a Block is always a function.

171

Yes, nested parallel operations work. We have tests also. But the allocaIP issue is not solved.

Now it is.

In the OpenMP dialect (and possibly others like FIR) we have to distinguish operations with regions which will be outlined (like parallel, task) and which will not (like master, single) be outlined. The operations which will be outlined have some properties similar to functions, like having an entry-block for allocas and finally existing as a function in LLVM IR.

This sounds like you want an interface for outlinable OpenMP dialect ops. That is perfectly doable, it just isn't FunctionLike but a different interface.

Ideally, I want a function (getAllocaBlock) which if I call will traverse up through the nested operations and get me the entry-block of the operation which will be outlined or if there is no operation like that then the function entry-block.

This can be done with the interface above. Just go up the region tree until you hit an op that is either OutlineableOpenMPOp or an LLVMFuncOp, and take the first block of the first region. For translation purposes, this isn't necessary though. We can keep the stack as @jdoerfert suggested.

[ftynse]

In D101307#2738760, @kiranchandramohan wrote:

I can just condition this new behavior to only work for the outermost omp.parallel because that's the case I am hitting right now...

This is OK with me for now.

Just re-iterating that special casing this only for the outermost omp.parallel is fine with me.

Sorry, I was busy with other things.

[ftynse]

Nits.

[kiranchandramohan]

LGTM. This looks great. More general than what I thought it would be.

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

194

Nit: spelling operations, ModuleTranslation.

[ftynse]

Address review.

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

public:
   /// the predecessors that may not exist yet.
   LogicalResult convertBlock(Block &bb, bool ignoreArguments,
                              llvm::IRBuilderBase &builder);
 
   /// Gets the named metadata in the LLVM IR module being constructed, creating
   /// it if it does not exist.
   llvm::NamedMDNode *getOrInsertNamedModuleMetadata(StringRef name);
 
+  /// Common CRTP base class for ModuleTranslation stack frames.
+  class StackFrame {
+  public:
+    virtual ~StackFrame() {}
+    TypeID getTypeID() const { return typeID; }
+
+  protected:
+    explicit StackFrame(TypeID typeID) : typeID(typeID) {}
+
+  private:
+    const TypeID typeID;
+    virtual void anchor();
+  };
+
+  /// Concrete CRTP base class for ModuleTranslation stack frames. When
+  /// translating operations with regions, users of ModuleTranslation can store
+  /// state on ModuleTranslation stack before entering the region and inspect
+  /// it when converting operations nested within that region. Users are
+  /// expected to derive this class and put any relevant information into fields
+  /// of the derived class. The usual isa/dyn_cast functionality is available
+  /// for instances of derived classes.
+  template <typename Derived>
+  class StackFrameBase : public StackFrame {
+  public:
+    explicit StackFrameBase() : StackFrame(TypeID::get<Derived>()) {}
+  };
+
+  /// Creates a stack frame of type `T` on ModuleTranslation stack. `T` must
+  /// be derived from `StackFrameBase<T>` and constructible from the provided
+  /// arguments. Doing this before entering the region of the op being
+  /// translated makes the frame available when translating ops within that
+  /// region.
+  template <typename T, typename... Args>
+  void stackPush(Args &&... args) {
+    static_assert(
+        std::is_base_of<StackFrame, T>::value,
+        "can only push instances of StackFrame on ModuleTranslation stack");
+    stack.push_back(std::make_unique<T>(std::forward<Args>(args)...));
+  }
+
+  /// Pops the last element from the ModuleTranslation stack.
+  void stackPop() { stack.pop_back(); }
+
+  /// Calls `callback` for every ModuleTranslation stack frame of type `T`
+  /// starting from the top of the stack.
+  template <typename T>
+  WalkResult
+  stackWalk(llvm::function_ref<WalkResult(const T &)> callback) const {
+    static_assert(std::is_base_of<StackFrame, T>::value,
+                  "expected T derived from StackFrame");
+    if (!callback)
+      return WalkResult::skip();
+    for (const std::unique_ptr<StackFrame> &frame : llvm::reverse(stack)) {
+      if (T *ptr = dyn_cast_or_null<T>(frame.get())) {
+        WalkResult result = callback(*ptr);
+        if (result.wasInterrupted())
+          return result;
+      }
+    }
+    return WalkResult::advance();
+  }
+
+  /// RAII object calling stackPush/stackPop on construction/destruction.
+  template <typename T>
+  struct SaveStack {
+    template <typename... Args>
+    explicit SaveStack(ModuleTranslation &m, Args &&...args)
+        : moduleTranslation(m) {
+      moduleTranslation.stackPush<T>(std::forward<Args>(args)...);
+    }
+    ~SaveStack() { moduleTranslation.stackPop(); }
+
+  private:
+    ModuleTranslation &moduleTranslation;
+  };
+
 private:
   ModuleTranslation(Operation *module,
                     std::unique_ptr<llvm::Module> llvmModule);
   ~ModuleTranslation();
 
   /// Converts individual components.
   LogicalResult convertOperation(Operation &op, llvm::IRBuilderBase &builder);
   LogicalResult convertFunctionSignatures();
private:
   /// This map is populated on module entry and is used to annotate loops (as
   /// identified via their branches) and contained memory accesses.
   DenseMap<Operation *, llvm::MDNode *> accessGroupMetadataMapping;
 
   /// Mapping from an attribute describing loop codegen options to its LLVM
   /// metadata. The metadata is attached to Latch block branches with this
   /// attribute.
   DenseMap<Attribute, llvm::MDNode *> loopOptionsMetadataMapping;
+
+  /// Stack of user-specified state elements, useful when translating operations
+  /// with regions.
+  SmallVector<std::unique_ptr<StackFrame>> stack;
 };
 
 namespace detail {
 /// For all blocks in the region that were converted to LLVM IR using the given
 /// ModuleTranslation, connect the PHI nodes of the corresponding LLVM IR blocks
 /// to the results of preceding blocks.
 void connectPHINodes(Region &region, const ModuleTranslation &state);
 

 llvm::Value *createNvvmIntrinsicCall(llvm::IRBuilderBase &builder,
                                      llvm::Intrinsic::ID intrinsic,
                                      ArrayRef<llvm::Value *> args = {});
 } // namespace detail
 
 } // namespace LLVM
 } // namespace mlir
 
+namespace llvm {
+template <typename T>
+struct isa_impl<T, ::mlir::LLVM::ModuleTranslation::StackFrame> {
+  static inline bool
+  doit(const ::mlir::LLVM::ModuleTranslation::StackFrame &frame) {
+    return frame.getTypeID() == ::mlir::TypeID::get<T>();
+  }
+};
+} // namespace llvm
+
 #endif // MLIR_TARGET_LLVMIR_MODULETRANSLATION_H

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

 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 
 using namespace mlir;
 
+namespace {
+/// ModuleTranslation stack frame for OpenMP operations. This keeps track of the
+/// insertion points for allocas.
+class OpenMPAllocaStackFrame
+    : public LLVM::ModuleTranslation::StackFrameBase<OpenMPAllocaStackFrame> {
+public:
+  explicit OpenMPAllocaStackFrame(llvm::OpenMPIRBuilder::InsertPointTy allocaIP)
+      : allocaInsertPoint(allocaIP) {}
+  llvm::OpenMPIRBuilder::InsertPointTy allocaInsertPoint;
+};
+} // namespace
+
+/// Find the insertion point for allocas given the current insertion point for
+/// normal operations in the builder.
+static llvm::OpenMPIRBuilder::InsertPointTy
+findAllocaInsertPoint(llvm::IRBuilderBase &builder,
+                      const LLVM::ModuleTranslation &moduleTranslation) {
+  // If there is an alloca insertion point on stack, i.e. we are in a nested
+  // operation and a specific point was provided by some surrounding operation,
+  // use it.
+  llvm::OpenMPIRBuilder::InsertPointTy allocaInsertPoint;
+  WalkResult walkResult = moduleTranslation.stackWalk<OpenMPAllocaStackFrame>(
+      [&](const OpenMPAllocaStackFrame &frame) {
+        allocaInsertPoint = frame.allocaInsertPoint;
+        return WalkResult::interrupt();
+      });
+  if (walkResult.wasInterrupted())
+    return allocaInsertPoint;
+
+  // Otherwise, insert to the entry block of the surrounding function.
+  llvm::BasicBlock &funcEntryBlock =
+      builder.GetInsertBlock()->getParent()->getEntryBlock();
+  return llvm::OpenMPIRBuilder::InsertPointTy(
+      &funcEntryBlock, funcEntryBlock.getFirstInsertionPt());
+}
+
 /// Converts the given region that appears within an OpenMP dialect operation to
 /// LLVM IR, creating a branch from the `sourceBlock` to the entry block of the
 /// region, and a branch from any block with an successor-less OpenMP terminator
 /// to `continuationBlock`.
 static void convertOmpOpRegions(Region &region, StringRef blockName,
                                 llvm::BasicBlock &sourceBlock,
                                 llvm::BasicBlock &continuationBlock,
                                 llvm::IRBuilderBase &builder,

 static LogicalResult
 convertOmpParallel(Operation &opInst, llvm::IRBuilderBase &builder,
                    LLVM::ModuleTranslation &moduleTranslation) {
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
   // relying on captured variables.
   LogicalResult bodyGenStatus = success();
 
   auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
                        llvm::BasicBlock &continuationBlock) {
+    // Save the alloca insertion point on ModuleTranslation stack for use in
+    // nested regions.
+    LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
+        moduleTranslation, allocaIP);
+
     // ParallelOp has only one region associated with it.
     auto &region = cast<omp::ParallelOp>(opInst).getRegion();
     convertOmpOpRegions(region, "omp.par.region", *codeGenIP.getBlock(),
                         continuationBlock, builder, moduleTranslation,
                         bodyGenStatus);
   };
 
   // TODO: Perform appropriate actions according to the data-sharing
convertOmpParallel(Operation &opInst, llvm::IRBuilderBase &builder,
   llvm::Value *numThreads = nullptr;
   if (auto numThreadsVar = cast<omp::ParallelOp>(opInst).num_threads_var())
     numThreads = moduleTranslation.lookupValue(numThreadsVar);
   llvm::omp::ProcBindKind pbKind = llvm::omp::OMP_PROC_BIND_default;
   if (auto bind = cast<omp::ParallelOp>(opInst).proc_bind_val())
     pbKind = llvm::omp::getProcBindKind(bind.getValue());
   // TODO: Is the Parallel construct cancellable?
   bool isCancellable = false;
-  // TODO: Determine the actual alloca insertion point, e.g., the function
-  // entry or the alloca insertion point as provided by the body callback
-  // above.
-  llvm::OpenMPIRBuilder::InsertPointTy allocaIP(builder.saveIP());
-  if (failed(bodyGenStatus))
-    return failure();
+
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
       builder.saveIP(), builder.getCurrentDebugLocation());
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createParallel(
-      ompLoc, allocaIP, bodyGenCB, privCB, finiCB, ifCond, numThreads, pbKind,
-      isCancellable));
-  return success();
+      ompLoc, findAllocaInsertPoint(builder, moduleTranslation), bodyGenCB,
+      privCB, finiCB, ifCond, numThreads, pbKind, isCancellable));
+
+  return bodyGenStatus;
 }
 
 /// Converts an OpenMP 'master' operation into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpMaster(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
     convertOmpOpRegions(loop.region(), "omp.wsloop.region", *entryBlock,
                         *exitBlock, builder, moduleTranslation, bodyGenStatus);
   };
 
   // Delegate actual loop construction to the OpenMP IRBuilder.
   // TODO: this currently assumes WsLoop is semantically similar to SCF loop,
   // i.e. it has a positive step, uses signed integer semantics. Reconsider
   // this code when WsLoop clearly supports more cases.
-  llvm::BasicBlock *insertBlock = builder.GetInsertBlock();
   llvm::CanonicalLoopInfo *loopInfo =
       moduleTranslation.getOpenMPBuilder()->createCanonicalLoop(
           ompLoc, bodyGen, lowerBound, upperBound, step, /*IsSigned=*/true,
           /*InclusiveStop=*/loop.inclusive());
   if (failed(bodyGenStatus))
     return failure();
 
-  // TODO: get the alloca insertion point from the parallel operation builder.
-  // If we insert the at the top of the current function, they will be passed as
-  // extra arguments into the function the parallel operation builder outlines.
-  // Put them at the start of the current block for now.
-  llvm::OpenMPIRBuilder::InsertPointTy allocaIP(
-      insertBlock, insertBlock->getFirstInsertionPt());
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
   llvm::OpenMPIRBuilder::InsertPointTy afterIP;
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
   if (isStatic) {
     loopInfo = ompBuilder->createStaticWorkshareLoop(ompLoc, loopInfo, allocaIP,
                                                      !loop.nowait(), chunk);
     afterIP = loopInfo->getAfterIP();
   } else {
     afterIP = ompBuilder->createDynamicWorkshareLoop(ompLoc, loopInfo, allocaIP,

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

ModuleTranslation::translateLoc(Location loc, llvm::DILocalScope *scope) {
   return debugTranslation->translateLoc(loc, scope);
 }
 
 llvm::NamedMDNode *
 ModuleTranslation::getOrInsertNamedModuleMetadata(StringRef name) {
   return llvmModule->getOrInsertNamedMetadata(name);
 }
 
+void ModuleTranslation::StackFrame::anchor() {}
+
 static std::unique_ptr<llvm::Module>
 prepareLLVMModule(Operation *m, llvm::LLVMContext &llvmContext,
                   StringRef name) {
   m->getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
   auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
   if (auto dataLayoutAttr =
           m->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName()))
     llvmModule->setDataLayout(dataLayoutAttr.cast<StringAttr>().getValue());

mlir/test/Target/LLVMIR/openmp-llvm.mlir

// CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_3_2]]
   // CHECK: call void @__kmpc_barrier
 
 // CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_3_1]]
   // CHECK: call void @__kmpc_barrier
 
 // CHECK: define void @test_omp_parallel_if_1(i32 %[[IF_VAR_1:.*]])
 llvm.func @test_omp_parallel_if_1(%arg0: i32) -> () {
 
+// Check that the allocas are emitted by the OpenMPIRBuilder at the top of the
+// function, before the condition. Allocas are only emitted by the builder when
+// the `if` clause is present. We match specific SSA value names since LLVM
+// actually produces those names.
+// CHECK: %tid.addr{{.*}} = alloca i32
+// CHECK: %zero.addr{{.*}} = alloca i32
+
 // CHECK: %[[IF_COND_VAR_1:.*]] = icmp slt i32 %[[IF_VAR_1]], 0
   %0 = llvm.mlir.constant(0 : index) : i32
   %1 = llvm.icmp "slt" %arg0, %0 : i32
 
 // CHECK: %[[GTN_IF_1:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[SI_VAR_IF_1:.*]])
 // CHECK: br i1 %[[IF_COND_VAR_1]], label %[[IF_COND_TRUE_BLOCK_1:.*]], label %[[IF_COND_FALSE_BLOCK_1:.*]]
 // CHECK: [[IF_COND_TRUE_BLOCK_1]]:
 // CHECK: br label %[[OUTLINED_CALL_IF_BLOCK_1:.*]]

 // CHECK: [[RETURN_BLOCK_IF_1]]:
 // CHECK: ret void
   llvm.return
 }
 
 // CHECK: define internal void @[[OMP_OUTLINED_FN_IF_1]]
   // CHECK: call void @__kmpc_barrier
 
+// -----
+
+// CHECK-LABEL: @test_nested_alloca_ip
+llvm.func @test_nested_alloca_ip(%arg0: i32) -> () {
+
+  // Check that the allocas are emitted by the OpenMPIRBuilder at the top of
+  // the function, before the condition. Allocas are only emitted by the
+  // builder when the `if` clause is present. We match specific SSA value names
+  // since LLVM actually produces those names and ensure they come before the
+  // "icmp" that is the first operation we emit.
+  // CHECK: %tid.addr{{.*}} = alloca i32
+  // CHECK: %zero.addr{{.*}} = alloca i32
+  // CHECK: icmp slt i32 %{{.*}}, 0
+  %0 = llvm.mlir.constant(0 : index) : i32
+  %1 = llvm.icmp "slt" %arg0, %0 : i32
+
+  omp.parallel if(%1 : i1) {
+    // The "parallel" operation will be outlined, check the the function is
+    // produced. Inside that function, further allocas should be placed before
+    // another "icmp".
+    // CHECK: define
+    // CHECK: %tid.addr{{.*}} = alloca i32
+    // CHECK: %zero.addr{{.*}} = alloca i32
+    // CHECK: icmp slt i32 %{{.*}}, 1
+    %2 = llvm.mlir.constant(1 : index) : i32
+    %3 = llvm.icmp "slt" %arg0, %2 : i32
+
+    omp.parallel if(%3 : i1) {
+      // One more nesting level.
+      // CHECK: define
+      // CHECK: %tid.addr{{.*}} = alloca i32
+      // CHECK: %zero.addr{{.*}} = alloca i32
+      // CHECK: icmp slt i32 %{{.*}}, 2
+
+      %4 = llvm.mlir.constant(2 : index) : i32
+      %5 = llvm.icmp "slt" %arg0, %4 : i32
+
+      omp.parallel if(%5 : i1) {
+        omp.barrier
+        omp.terminator
+      }
+
+      omp.barrier
+      omp.terminator
+    }
+    omp.barrier
+    omp.terminator
+  }
+
+  llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: define void @test_omp_parallel_3()
 llvm.func @test_omp_parallel_3() -> () {
   // CHECK: [[OMP_THREAD_3_1:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}})
   // CHECK: call void @__kmpc_push_proc_bind(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD_3_1]], i32 2)
   // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_3_1:.*]] to {{.*}}
   omp.parallel proc_bind(master) {
     omp.barrier
     omp.terminator