[mlir] Refactor translation of OpenMP dialect ops to LLVM IR

[ftynse]

Summary

The original implementation of the OpenMP dialect to LLVM IR translation has
been relying on a stack of insertion points for delayed insertion of branch
instructions that correspond to terminator ops. This is an intrusive into
ModuleTranslation and makes the translation non-local. A recent addition of the
WsLoop translation exercised another approach where the parent op is
responsible for converting terminators of all blocks in its regions. Use this
approach for other OpenMP dialect operations with regions, remove the stack and
deduplicate the code for converting such regions.

[kiranchandramohan]

Thanks for this patch Alex. LGTM.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

599

Nit: Can this comment be also moved to the refactored code?

[ftynse]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

599

Will do before landing.

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

protected:
   virtual LogicalResult convertOperation(Operation &op,
                                          llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpOperation(Operation &op,
                                             llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpParallel(Operation &op,
                                            llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpMaster(Operation &op,
                                          llvm::IRBuilder<> &builder);
-  void convertOmpOpRegions(Region &region,
+  void convertOmpOpRegions(Region &region, StringRef blockName,
                            DenseMap<Value, llvm::Value *> &valueMapping,
                            DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
-                           llvm::Instruction *codeGenIPBBTI,
-                           llvm::BasicBlock &continuationIP,
+                           llvm::BasicBlock &sourceBlock,
+                           llvm::BasicBlock &continuationBlock,
                            llvm::IRBuilder<> &builder,
                            LogicalResult &bodyGenStatus);
   virtual LogicalResult convertOmpWsLoop(Operation &opInst,
                                          llvm::IRBuilder<> &builder);
 
   /// Converts the type from MLIR LLVM dialect to LLVM.
   llvm::Type *convertType(Type type);
 
   static std::unique_ptr<llvm::Module>
   prepareLLVMModule(Operation *m, llvm::LLVMContext &llvmContext,
                     StringRef name);
 
   /// A helper to look up remapped operands in the value remapping table.
   SmallVector<llvm::Value *, 8> lookupValues(ValueRange values);
 
 private:
   /// Check whether the module contains only supported ops directly in its body.
   static LogicalResult checkSupportedModuleOps(Operation *m);
 
   LogicalResult convertFunctionSignatures();
   LogicalResult convertFunctions();
   LogicalResult convertGlobals();
   LogicalResult convertOneFunction(LLVMFuncOp func);
-  LogicalResult convertBlock(Block &bb, bool ignoreArguments);
+  LogicalResult convertBlock(Block &bb, bool ignoreArguments,
+                             llvm::IRBuilder<> &builder);
 
   llvm::Constant *getLLVMConstant(llvm::Type *llvmType, Attribute attr,
                                   Location loc);
 
   /// Original and translated module.
   Operation *mlirModule;
   std::unique_ptr<llvm::Module> llvmModule;
   /// A converter for translating debug information.
   std::unique_ptr<detail::DebugTranslation> debugTranslation;
 
   /// Builder for LLVM IR generation of OpenMP constructs.
   std::unique_ptr<llvm::OpenMPIRBuilder> ompBuilder;
+
   /// Precomputed pointer to OpenMP dialect. Note this can be nullptr if the
   /// OpenMP dialect hasn't been loaded (it is always loaded if there are OpenMP
   /// operations in the module though).
   const Dialect *ompDialect;
-  /// Stack which stores the target block to which a branch a must be added when
-  /// a terminator is seen. A stack is required to handle nested OpenMP parallel
-  /// regions.
-  SmallVector<llvm::BasicBlock *, 4> ompContinuationIPStack;
 
   /// Mappings between llvm.mlir.global definitions and corresponding globals.
   DenseMap<Operation *, llvm::GlobalValue *> globalsMapping;
 
   /// A stateful object used to translate types.
   TypeToLLVMIRTranslator typeTranslator;
 
 protected:

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

ModuleTranslation::convertOmpParallel(Operation &opInst,
                                       llvm::IRBuilder<> &builder) {
   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 &continuationIP) {
-    llvm::LLVMContext &llvmContext = llvmModule->getContext();
-
-    llvm::BasicBlock *codeGenIPBB = codeGenIP.getBlock();
-    llvm::Instruction *codeGenIPBBTI = codeGenIPBB->getTerminator();
-    ompContinuationIPStack.push_back(&continuationIP);
-
-    // ParallelOp has only `1` region associated with it.
+    // ParallelOp has only one region associated with it.
     auto &region = cast<omp::ParallelOp>(opInst).getRegion();
-    for (auto &bb : region) {
-      auto *llvmBB = llvm::BasicBlock::Create(
-          llvmContext, "omp.par.region", codeGenIP.getBlock()->getParent());
-      blockMapping[&bb] = llvmBB;
-    }
-
-    convertOmpOpRegions(region, valueMapping, blockMapping, codeGenIPBBTI,
-                        continuationIP, builder, bodyGenStatus);
-    ompContinuationIPStack.pop_back();
-
+    convertOmpOpRegions(region, "omp.par.region", valueMapping, blockMapping,
+                        *codeGenIP.getBlock(), continuationIP, builder,
+                        bodyGenStatus);
   };
 
   // TODO: Perform appropriate actions according to the data-sharing
   // attribute (shared, private, firstprivate, ...) of variables.
   // Currently defaults to shared.
   auto privCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
                     llvm::Value &, llvm::Value &vPtr,
                     llvm::Value *&replacementValue) -> InsertPointTy {
if (failed(bodyGenStatus))
     return failure();
   builder.restoreIP(
       ompBuilder->createParallel(builder, allocaIP, bodyGenCB, privCB, finiCB,
                                  ifCond, numThreads, pbKind, isCancellable));
   return success();
 }
 
 void ModuleTranslation::convertOmpOpRegions(
-    Region &region, DenseMap<Value, llvm::Value *> &valueMapping,
+    Region &region, StringRef blockName,
+    DenseMap<Value, llvm::Value *> &valueMapping,
     DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
-    llvm::Instruction *codeGenIPBBTI, llvm::BasicBlock &continuationIP,
+    llvm::BasicBlock &sourceBlock, llvm::BasicBlock &continuationBlock,
     llvm::IRBuilder<> &builder, LogicalResult &bodyGenStatus) {
+  llvm::LLVMContext &llvmContext = builder.getContext();
+  for (Block &bb : region) {
+    llvm::BasicBlock *llvmBB = llvm::BasicBlock::Create(
+        llvmContext, blockName, builder.GetInsertBlock()->getParent());
+    blockMapping[&bb] = llvmBB;
+  }
+
+  llvm::Instruction *sourceTerminator = sourceBlock.getTerminator();
+
   // Convert blocks one by one in topological order to ensure
   // defs are converted before uses.
   llvm::SetVector<Block *> blocks = topologicalSort(region);
-  for (auto indexedBB : llvm::enumerate(blocks)) {
-    Block *bb = indexedBB.value();
-    llvm::BasicBlock *curLLVMBB = blockMapping[bb];
+  for (Block *bb : blocks) {
+    llvm::BasicBlock *llvmBB = blockMapping[bb];
+    // Retarget the branch of the entry block to the entry block of the
+    // converted region (regions are single-entry).
     if (bb->isEntryBlock()) {
-      assert(codeGenIPBBTI->getNumSuccessors() == 1 &&
-             "OpenMPIRBuilder provided entry block has multiple successors");
-      assert(codeGenIPBBTI->getSuccessor(0) == &continuationIP &&
-             "ContinuationIP is not the successor of OpenMPIRBuilder "
-             "provided entry block");
-      codeGenIPBBTI->setSuccessor(0, curLLVMBB);
+      assert(sourceTerminator->getNumSuccessors() == 1 &&
+             "provided entry block has multiple successors");
+      assert(sourceTerminator->getSuccessor(0) == &continuationBlock &&
+             "ContinuationBlock is not the successor of the entry block");
+      sourceTerminator->setSuccessor(0, llvmBB);
     }
 
-    if (failed(convertBlock(*bb, /*ignoreArguments=*/indexedBB.index() == 0))) {
+    llvm::IRBuilder<>::InsertPointGuard guard(builder);
+    if (failed(convertBlock(*bb, bb->isEntryBlock(), builder))) {
       bodyGenStatus = failure();
       return;
     }
+
+    // Special handling for `omp.yield` and `omp.terminator` (we may have more
+    // than one): they return the control to the parent OpenMP dialect operation
+    // so replace them with the branch to the continuation block. We handle this
+    // here to avoid relying inter-function communication through the
+    // ModuleTranslation class to set up the correct insertion point. This is
+    // also consistent with MLIR's idiom of handling special region terminators
+    // in the same code that handles the region-owning operation.
+    if (isa<omp::TerminatorOp, omp::YieldOp>(bb->getTerminator()))
+      builder.CreateBr(&continuationBlock);
   }
   // Finally, after all blocks have been traversed and values mapped,
   // connect the PHI nodes to the results of preceding blocks.
   connectPHINodes(region, valueMapping, blockMapping, branchMapping);
 }
 
 LogicalResult ModuleTranslation::convertOmpMaster(Operation &opInst,
                                                   llvm::IRBuilder<> &builder) {
   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 &continuationIP) {
-    llvm::LLVMContext &llvmContext = llvmModule->getContext();
-
-    llvm::BasicBlock *codeGenIPBB = codeGenIP.getBlock();
-    llvm::Instruction *codeGenIPBBTI = codeGenIPBB->getTerminator();
-    ompContinuationIPStack.push_back(&continuationIP);
-
-    // MasterOp has only `1` region associated with it.
+    // MasterOp has only one region associated with it.
     auto &region = cast<omp::MasterOp>(opInst).getRegion();
-    for (auto &bb : region) {
-      auto *llvmBB = llvm::BasicBlock::Create(
-          llvmContext, "omp.master.region", codeGenIP.getBlock()->getParent());
-      blockMapping[&bb] = llvmBB;
-    }
-    convertOmpOpRegions(region, valueMapping, blockMapping, codeGenIPBBTI,
-                        continuationIP, builder, bodyGenStatus);
-    ompContinuationIPStack.pop_back();
+    convertOmpOpRegions(region, "omp.master.region", valueMapping, blockMapping,
+                        *codeGenIP.getBlock(), continuationIP, builder,
+                        bodyGenStatus);
   };
 
   // TODO: Perform finalization actions for variables. This has to be
   // called for variables which have destructors/finalizers.
   auto finiCB = [&](InsertPointTy codeGenIP) {};
 
   builder.restoreIP(ompBuilder->createMaster(builder, bodyGenCB, finiCB));
   return success();
if (loop.getNumLoops() != 1)
     return opInst.emitOpError("collapsed loops not yet supported");
 
   if (loop.schedule_val().hasValue() &&
       omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue()) !=
           omp::ClauseScheduleKind::Static)
     return opInst.emitOpError(
         "only static (default) loop schedule is currently supported");
 
-  llvm::Function *func = builder.GetInsertBlock()->getParent();
-  llvm::LLVMContext &llvmContext = llvmModule->getContext();
-
   // Find the loop configuration.
   llvm::Value *lowerBound = valueMapping.lookup(loop.lowerBound()[0]);
   llvm::Value *upperBound = valueMapping.lookup(loop.upperBound()[0]);
   llvm::Value *step = valueMapping.lookup(loop.step()[0]);
   llvm::Type *ivType = step->getType();
   llvm::Value *chunk = loop.schedule_chunk_var()
                            ? valueMapping[loop.schedule_chunk_var()]
                            : llvm::ConstantInt::get(ivType, 1);
auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
     // Make sure further conversions know about the induction variable.
     valueMapping[loop.getRegion().front().getArgument(0)] = iv;
 
     llvm::BasicBlock *entryBlock = ip.getBlock();
     llvm::BasicBlock *exitBlock =
         entryBlock->splitBasicBlock(ip.getPoint(), "omp.wsloop.exit");
 
     // Convert the body of the loop.
-    Region &region = loop.region();
-    for (Block &bb : region) {
-      llvm::BasicBlock *llvmBB =
-          llvm::BasicBlock::Create(llvmContext, "omp.wsloop.region", func);
-      blockMapping[&bb] = llvmBB;
-
-      // Retarget the branch of the entry block to the entry block of the
-      // converted region (regions are single-entry).
-      if (bb.isEntryBlock()) {
-        auto *branch = cast<llvm::BranchInst>(entryBlock->getTerminator());
-        branch->setSuccessor(0, llvmBB);
-      }
-    }
-
-    // Block conversion creates a new IRBuilder every time so need not bother
-    // about maintaining the insertion point.
-    llvm::SetVector<Block *> blocks = topologicalSort(region);
-    for (Block *bb : blocks) {
-      if (failed(convertBlock(*bb, bb->isEntryBlock()))) {
-        bodyGenStatus = failure();
-        return;
-      }
-
-      // Special handling for `omp.yield` terminators (we may have more than
-      // one): they return the control to the parent WsLoop operation so replace
-      // them with the branch to the exit block. We handle this here to avoid
-      // relying inter-function communication through the ModuleTranslation
-      // class to set up the correct insertion point. This is also consistent
-      // with MLIR's idiom of handling special region terminators in the same
-      // code that handles the region-owning operation.
-      if (isa<omp::YieldOp>(bb->getTerminator())) {
-        llvm::BasicBlock *llvmBB = blockMapping[bb];
-        builder.SetInsertPoint(llvmBB, llvmBB->end());
-        builder.CreateBr(exitBlock);
-      }
-    }
-
-    connectPHINodes(region, valueMapping, blockMapping, branchMapping);
+    convertOmpOpRegions(loop.region(), "omp.wsloop.region", valueMapping,
+                        blockMapping, *entryBlock, *exitBlock, builder,
+                        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, and its upper
   // bound is not included. Reconsider this code when WsLoop clearly supports
   // more cases.
   llvm::BasicBlock *insertBlock = builder.GetInsertBlock();
return llvm::TypeSwitch<Operation *, LogicalResult>(&opInst)
         //  "An implementation may implement a flush with a list by ignoring
         //   the list, and treating it the same as a flush without a list."
         //
         // The argument list is discarded so that, flush with a list is treated
         // same as a flush without a list.
         ompBuilder->createFlush(builder.saveIP());
         return success();
       })
-      .Case([&](omp::TerminatorOp) {
-        builder.CreateBr(ompContinuationIPStack.back());
-        return success();
-      })
       .Case(
           [&](omp::ParallelOp) { return convertOmpParallel(opInst, builder); })
       .Case([&](omp::MasterOp) { return convertOmpMaster(opInst, builder); })
       .Case([&](omp::WsLoopOp) { return convertOmpWsLoop(opInst, builder); })
-      .Case([&](omp::YieldOp op) {
-        // Yields are loop terminators that can be just omitted. The loop
-        // structure was created in the function that handles WsLoopOp.
-        assert(op.getNumOperands() == 0 && "unexpected yield with operands");
+      .Case<omp::YieldOp, omp::TerminatorOp>([](auto op) {
+        // `yield` and `terminator` can be just omitted. The block structure was
+        // created in the function that handles their parent operation.
+        assert(op->getNumOperands() == 0 &&
+               "unexpected OpenMP terminator with operands");
         return success();
       })
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();
       });
 }
 
if (ompDialect && opInst.getDialect() == ompDialect)
     return convertOmpOperation(opInst, builder);
 
   return opInst.emitError("unsupported or non-LLVM operation: ")
          << opInst.getName();
 }
 
 /// Convert block to LLVM IR.  Unless `ignoreArguments` is set, emit PHI nodes
 /// to define values corresponding to the MLIR block arguments.  These nodes
-/// are not connected to the source basic blocks, which may not exist yet.
-LogicalResult ModuleTranslation::convertBlock(Block &bb, bool ignoreArguments) {
-  llvm::IRBuilder<> builder(blockMapping[&bb]);
+/// are not connected to the source basic blocks, which may not exist yet.  Uses
+/// `builder` to construct the LLVM IR. Expects the LLVM IR basic block to have
+/// been created for `bb` and included in the block mapping.  Inserts new
+/// instructions at the end of the block and leaves `builder` in a state
+/// suitable for further insertion into the end of the block.
+LogicalResult ModuleTranslation::convertBlock(Block &bb, bool ignoreArguments,
+                                              llvm::IRBuilder<> &builder) {
+  builder.SetInsertPoint(blockMapping[&bb]);
   auto *subprogram = builder.GetInsertBlock()->getParent()->getSubprogram();
 
   // Before traversing operations, make block arguments available through
   // value remapping and PHI nodes, but do not add incoming edges for the PHI
   // nodes just yet: those values may be defined by this or following blocks.
   // This step is omitted if "ignoreArguments" is set.  The arguments of the
   // first block have been already made available through the remapping of
   // LLVM function arguments.
for (auto &bb : func) {
     auto *llvmBB = llvm::BasicBlock::Create(llvmContext);
     llvmBB->insertInto(llvmFunc);
     blockMapping[&bb] = llvmBB;
   }
 
   // Then, convert blocks one by one in topological order to ensure defs are
   // converted before uses.
   auto blocks = topologicalSort(func);
-  for (auto indexedBB : llvm::enumerate(blocks)) {
-    auto *bb = indexedBB.value();
-    if (failed(convertBlock(*bb, /*ignoreArguments=*/indexedBB.index() == 0)))
+  for (Block *bb : blocks) {
+    llvm::IRBuilder<> builder(llvmContext);
+    if (failed(convertBlock(*bb, bb->isEntryBlock(), builder)))
       return failure();
   }
 
   // Finally, after all blocks have been traversed and values mapped, connect
   // the PHI nodes to the results of preceding blocks.
   connectPHINodes(func, valueMapping, blockMapping, branchMapping);
   return success();
 }