[OpenMP][MLIR] Fix for nested parallel regions

[kiranchandramohan]

Summary

[MLIR][OpenMP] Fix for nested parallel regions

Fix contains the following changes,

1. Don't set the insertion point in the body callback.
2. Save the continuation IP in a stack and set the branch to continuation IP at the terminator.

Note: This is required for supporting the master Operation (https://reviews.llvm.org/D87247).

[kiranchandramohan]

@SouraVX In this patch I have removed the OpenMPIRBuilder alloca changes. This contains only the changes to in ModuleTranslation and it seems to work fine. Can you check at your end as well as with the master Operation?

[SouraVX]

In D88720#2307968, @kiranchandramohan wrote:

@SouraVX In this patch I have removed the OpenMPIRBuilder alloca changes. This contains only the changes to in ModuleTranslation and it seems to work fine. Can you check at your end as well as with the master Operation?

Thanks! for the patch @kiranchandramohan :). I tried/applied this patch locally, it is working as great/expected(i.e MasterOp is getting lowered correctly(verifed the well formedness of the IR by generating objects)).

However I haven't reviewed the code/patch fully(since I was OOO). Will look in as soon as I can.

[ftynse]

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

133

Nit: SmallVector is re-exported, no need to prefix it with llvm namespace.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

505

I'd much rather have IRBuilder used here for consistency with the rest of the code, even if it would require setting the insertion point twice.

[SouraVX]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

460

Is this TODO still valid ? @jdoerfert .
Patch fixes the problem in hand, i.e it's the reset of CodeGenIP that was RC for this.

[kiranchandramohan]

Using the builder to create branch instruction for a terminator.
Added one more assertion.
Added on more test.

[kiranchandramohan]

@SouraVX do you have more comments?

[kiranchandramohan]

Fixed the comment to remove the llvm prefix for smallvector.

@SouraVX I believe some changes are required for handling the allocaIP, but as you say it is not related to the nested parallel region issue.

[SouraVX]

Thanks for the patch. You may want to wait for @jdoerfert for final sign off!

[kiranchandramohan]

ping @jdoerfert. Is this patch OK? This fixes the nested parallel regions issue and will unblock the master op lowering (https://reviews.llvm.org/D87247).

[jdoerfert]

LGTM

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

private:
   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.
+  llvm::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,
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
 
   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);
 
-    builder.SetInsertPoint(codeGenIPBB);
     // ParallelOp has only `1` 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;
     }
 
     // Then, 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];
-      if (bb->isEntryBlock())
+      if (bb->isEntryBlock()) {
+        assert(codeGenIPBBTI->getNumSuccessors() == 1 &&
+               "OpenMPIRBuilder provided entry block has multiple successors");
         codeGenIPBBTI->setSuccessor(0, curLLVMBB);
+      }
 
       // TODO: Error not returned up the hierarchy
       if (failed(convertBlock(*bb, /*ignoreArguments=*/indexedBB.index() == 0)))
         return;
-
-      // If this block has the terminator then add a jump to
-      // continuation bb
-      for (auto &op : *bb) {
-        if (isa<omp::TerminatorOp>(op)) {
-          builder.SetInsertPoint(curLLVMBB);
-          builder.CreateBr(&continuationIP);
-        }
-      }
     }
+
+    ompContinuationIPStack.pop_back();
+
     // Finally, after all blocks have been traversed and values mapped,
     // connect the PHI nodes to the results of preceding blocks.
     connectPHINodes(region, valueMapping, blockMapping);
   };
 
   // TODO: Perform appropriate actions according to the data-sharing
   // attribute (shared, private, firstprivate, ...) of variables.
   // Currently defaults to shared.
ModuleTranslation::convertOmpParallel(Operation &opInst,
   llvm::Value *numThreads = nullptr;
   if (auto numThreadsVar = cast<omp::ParallelOp>(opInst).num_threads_var())
     numThreads = valueMapping.lookup(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());
   builder.restoreIP(
       ompBuilder->CreateParallel(builder, allocaIP, bodyGenCB, privCB, finiCB,
                                  ifCond, numThreads, pbKind, isCancellable));
   return success();
 }
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) { return success(); })
+      .Case([&](omp::TerminatorOp) {
+        llvm::BranchInst::Create(ompContinuationIPStack.back(),
+                                 builder.GetInsertBlock());
+        return success();
+      })
       .Case(
           [&](omp::ParallelOp) { return convertOmpParallel(opInst, builder); })
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();
       });
 }
 

mlir/test/Target/openmp-llvm.mlir

llvm.func @test_omp_parallel_3() -> () {
   }
 
   llvm.return
 }
 
 // CHECK: define internal void @[[OMP_OUTLINED_FN_3_3]]
 // CHECK: define internal void @[[OMP_OUTLINED_FN_3_2]]
 // CHECK: define internal void @[[OMP_OUTLINED_FN_3_1]]
+
+// CHECK-LABEL: define void @test_omp_parallel_4()
+llvm.func @test_omp_parallel_4() -> () {
+// CHECK: call void {{.*}}@__kmpc_fork_call{{.*}} @[[OMP_OUTLINED_FN_4_1:.*]] to
+// CHECK: define internal void @[[OMP_OUTLINED_FN_4_1]]
+// CHECK: call void @__kmpc_barrier
+// CHECK: call void {{.*}}@__kmpc_fork_call{{.*}} @[[OMP_OUTLINED_FN_4_1_1:.*]] to
+// CHECK: call void @__kmpc_barrier
+  omp.parallel {
+    omp.barrier
+
+// CHECK: define internal void @[[OMP_OUTLINED_FN_4_1_1]]
+// CHECK: call void @__kmpc_barrier
+    omp.parallel {
+      omp.barrier
+      omp.terminator
+    }
+
+    omp.barrier
+    omp.terminator
+  }
+  llvm.return
+}