[MLIR,OpenMP] Added support for lowering MasterOp to LLVMIR

[SouraVX]

Summary

Some Ops in OMP dialect have regions associated with them i.e
ParallelOp MasterOp. Lowering of these regions involves interfacing
with OMPIRBuilder using callbacks, yet there still exist opportunities
for sharing common code in between.

This patch factors out common code into a separate function and adds
support for lowering MasterOp using that. Lowering of ParallelOp is
also modified appropriately.

[kiranchandramohan]

Thanks @SouraVX for working on this. Looks OK. A few comments to address.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

451–452

Should valueMapping and blockMapping be passed? Aren't these available here since convertOmpOpRegions is a member function? Or did you want this as a static function?

496–500

Can this loop also go into convertOmpOpRegions?

509

lint warning?

mlir/test/Target/openmp-llvm.mlir

216–218

1. Can you add this inside a parallel here or in another test?
2. Nit: I think other tests here use the pretty print syntax.

[kiranktp]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

495

convertOmpMaster and convertOmpParallel function body is 99% same except the casting type (omp::MasterOp/omp::ParallelOp).
Will it be good to have a template function (eg: convertOmpOpWithRegions(..)) and call it with explicit type.
convertOmpOpWithRegions<omp::MasterOp>(opInst, builder);
convertOmpOpWithRegions<omp::ParallelOp>(opInst, builder);
This way we can reduce code duplication.
You can choose to not follow this.

[SouraVX]

Thanks @kiranktp ,@kiranchandramohan for reviewing this!
When adding a test case for a MasterOp inside the region of ParallelOp(typical use case.)
Test snippet:

llvm.func @test_omp_master() -> () {
  omp.parallel {
  "omp.master" ()({
    omp.terminator
  }):()->()
    omp.terminator
}
  llvm.return
}

This produces following crash/assertion failure in OMPIRBuilder:

mlir-translate: llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:717: llvm::IRBuilderBase::InsertPoint llvm::OpenMPIRBuilder::CreateParallel(const llvm::OpenMPIRBuilder::LocationDescription&, llvm::OpenMPIRBuilder::InsertPointTy, llvm::OpenMPIRBuilder::BodyGenCallbackTy, llvm::OpenMPIRBuilder::PrivatizeCallbackTy, llvm::OpenMPIRBuilder::FinalizeCallbackTy, llvm::Value*, llvm::Value*, llvm::omp::ProcBindKind, bool): Assertion `Outputs.empty() && "OpenMP outlining should not produce live-out values!"' failed.

digging in further, I realized this Outputs is getting populated with(below instruction) and finally asserting at Outputs.empty().

%omp_global_thread_num2 = call i32 @__kmpc_global_thread_num(%struct.ident_t* @3), !dbg !9

I'll dig further, BTW if you guys have any inputs WRT this. Please share!

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

451–452

Yes, initially planned for static function. But later realized convertBlock (used here) is private within ModuleTranslation.

And I didn't find enough motivation to expose that as public, if time demands we'll do that.
I'll take care of the NIT's when revising.

495

Thought of this approach too, templatizing this won't help since Op is passed around as generic Operation so it seems to me we can't get away with this.
I even thought pulling out the Callback itself, but that will also incur if/else ladder and so forth.

496–500

It can be, but keeping it separate seems more reasonable.
Since we have to create block with Op Name(this can be done there also but that will incur one more parameter passing and maybe changing the function name itself to convey the intent).

509

Thanks for catching this! Sure, it's only needed for converting ParallelOp.

mlir/test/Target/openmp-llvm.mlir

216–218

Nit: I think other tests here use the pretty print syntax.

MasterOp doesn't has custom assembly registered.

error: custom op 'omp.master' has no custom assembly form

Can you add this inside a parallel here or in another test?

Thanks for this, I discovered one crash.
Details follows:

[kiranchandramohan]

This produces following crash/assertion failure in OMPIRBuilder:

mlir-translate: llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:717: llvm::IRBuilderBase::InsertPoint llvm::OpenMPIRBuilder::CreateParallel(const llvm::OpenMPIRBuilder::LocationDescription&, llvm::OpenMPIRBuilder::InsertPointTy, llvm::OpenMPIRBuilder::BodyGenCallbackTy, llvm::OpenMPIRBuilder::PrivatizeCallbackTy, llvm::OpenMPIRBuilder::FinalizeCallbackTy, llvm::Value*, llvm::Value*, llvm::omp::ProcBindKind, bool): Assertion `Outputs.empty() && "OpenMP outlining should not produce live-out values!"' failed.

I don't have any immediate inputs regarding this crash. Maybe @jdoerfert or @fghanim.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

496–500

No strong opinion here. Let us see whether anyone else has an opinion here before making any change.
But what is the change in name of the function you are hinting here? Creating the llvm blocks is part of the conversion right?

mlir/test/Target/openmp-llvm.mlir

216–218

Can you add an assembly format to print/parse this operation in custom syntax as part of this patch? I think it will look a bit odd if this operation only has the generic syntax.
From the latest mlir news discourse post I see that they have added support for region parsing also in assembly format. master looks like a simple op to try this out.
https://github.com/llvm/llvm-project/commit/eaeadce9bd11

[fghanim]

I don't know enough about the OMPIR for me to be helpful, however, I noticed a couple of things that I wanted to ask about for better understanding. :)

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

487

I don't see where you set the location where alloca should be codegened. how do you handle variable declarations inside the omp master scope? when do you codegen the relevant alloca?

509

All llvm passes expect Allocas to be in the entry block of the function. In this case, the soon-to-be-outlined region. is builder's current insertion point in the entry block? Also, is it guaranteed to not be empty?

[fghanim]

@kiranchandramohan This error happens when there are things within the OMP region still alive/are used outside the region, I cannot comment on what causes this here. Usually, this suggests that "something" that should be contained within the parallel region is not being detected as such, and so when the region is outlined, this User remains as part of the original function.

The first thing that comes to mind, make sure continuationBB post dominate everything within OMP parallel

[jdoerfert]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

493

Why do you set the insertion point to something other than codeGenIP?

509

Where is allocaIP used?

[jdoerfert]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

493

Do not insert code or modify anything else. Basically you get this:

block:
  ... 
  %...      <- codeGenIP
  ...
  terminator

And if you want to insert control flow, split it:

block.pre:
  ... 
  %...      <- codeGenIP
  br %block.post

block.post:
  ...
  terminator

No you can remove the branch that was introduced and create instructions and control flow that converge at block.post.

[ftynse]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

464

It's a good time to do this.

507

Nit: a more comprehensible variable name is welcome.

[kiranchandramohan]

mlir/test/Target/openmp-llvm.mlir

216–218

Have created a trivial patch for adding the assembly format.
https://reviews.llvm.org/D87549

[SouraVX]

• Rebased.
• Addressed @kiranchandramohan review comments for test case.
• At https://github.com/llvm/llvm-project/blob/master/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L918 @jdoerfert could you please explain the need for passing empty AllocaIP ? I meant this will be empty? while interfacing MasterOp from flang/mlir. Why is that.
• A more rigorous end-to-endtesting can't be done at this moment, because of divergence and code scattered between LLVM Master and FIR.

[jdoerfert]

At https://github.com/llvm/llvm-project/blob/master/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L918 @jdoerfert could you please explain the need for passing empty AllocaIP ? I meant this will be empty? while interfacing MasterOp from flang/mlir. Why is that.

That looks to me like a mistake which needs to be fixed. And, by glancing at it, it seems we now need to do what I feared in D82470, we need to add the AllocaIP to a lot more APIs so it can be passed down properly. Maybe there is no way around it.

[fghanim]

In D87247#2341467, @SouraVX wrote:

• At https://github.com/llvm/llvm-project/blob/master/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L918 @jdoerfert could you please explain the need for passing empty AllocaIP ? I meant this will be empty? while interfacing MasterOp from flang/mlir. Why is that.

It reuses BodyGenCallbackTy which has an AllocaIP argument as part of the type. But since this is an inlined region it doesn't benefit from it. Why? because inlined-regions should be using the AllocaIP of the innermost outlined region they belong to, for any allocations that happen inside them. Hence, you don't need to update the AllocaIP for inlined regions and so it is left empty. I wrote this with a design like clang's in mind, where an AllocaIP is always maintained in the frontend, and BodyCB is a place for the frontend to generate whatever non-OMP code it wants according to its own rules.

[SouraVX]

Rebase + Ping!

[kiranchandramohan]

I guess the only remaining point to address is to handle return the pass/fail status to MLIR module translation.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

487

AllocaIP issues will be handled in a separate patch.
https://reviews.llvm.org/D88706

493

We do not change the codeGenIP now. This was fixed in the following review.
https://reviews.llvm.org/D88720

509

AllocaIP issues will be handled in a separate patch.
https://reviews.llvm.org/D88706

mlir/test/Target/openmp-llvm.mlir

213–233

Can you nest the omp operations with appropriate formatting?

[fghanim]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

487

Sure. If I may suggest things to consider:
1- Privatization clauses: The OMPBuilder will never handle privatization or provide creators for it - with minor exceptions- and so this is something you guys will have to handle on the MLIR side. and it is something you should consider because you will need to allocate local variables inside the entry block of an outlined region based on how Fortran does it. Again, the OMPBuilder will never do any of that.
2- Variables declared inside a scope: I don't know if fortran has the concept of scopes like C/C++, but it is also something to keep in mind - although this is likely to be handled by FIR.

[SouraVX]

Addressed comments related to formatting test case + Rebase.

[SouraVX]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

470

This is getting trickier than what we initially thought. Even if we return LogicalResult back to caller(which is callback bodyGenCB). Propagating same error upward in hierarchy may not be possible(or degenerate the interface) because this bodyGenCB will be called very late inside OMPIRBuilder using createParallel.
Now createParallel is primary interfacing which returns InsertionPoint.

[SouraVX]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

470

@ftynse are you Okay with this ? One other approach we have is something of sort:

bool ErrorStatus = false/*Initial state*/;
bodyGenCB = [&] (..){
...
if (failed(convertBlock(*bb, /*ignoreArguments=*/indexedBB.index() == 0)))
      ErrorStatus = true;
...
};

builder.restoreIP(ompBuilder->createMaster(builder, bodyGenCB, finiCB));

if (ErrorStatus)
    return failure();

return success();

};

[SouraVX]

Added error propagation support. @ftynse I was waiting for your response, however I noticed same approach in D92055. So I've revised this.

[ftynse]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

454

bodyGenStatus should be a reference here

470

LGTM modulo the fix above.

That being said, I expect all this callback and delayed execution to bite us back one day.

504

Not for this diff, but I am not a fan of this stack approach. The parent op of the region can handle conversion of the terminators of all the blocks instead. This will avoid the need to communicate through this explicit stack between functions (it will be instead on the actual call stack), and better corresponds to how MLIR expects one to handle regions -- the parent op knows more about terminator semantics than terminators themselves.

[SouraVX]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

504

Thanks for the pointers, me and @kiranchandramohan will consider this and how to use this.

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

protected:
   virtual ~ModuleTranslation();
 
   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,
+                           DenseMap<Value, llvm::Value *> &valueMapping,
+                           DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
+                           llvm::Instruction *codeGenIPBBTI,
+                           llvm::BasicBlock &continuationIP,
+                           llvm::IRBuilder<> &builder);
   /// Converts the type from MLIR LLVM dialect to LLVM.
   llvm::Type *convertType(LLVMType 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.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
     // 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())
-        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);
-        }
-      }
-    }
-    // Finally, after all blocks have been traversed and values mapped,
-    // connect the PHI nodes to the results of preceding blocks.
-    connectPHINodes(region, valueMapping, blockMapping);
+    convertOmpOpRegions(region, valueMapping, blockMapping, codeGenIPBBTI,
+                        continuationIP, builder);
   };
 
   // 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 &vPtr,
                     llvm::Value *&replacementValue) -> InsertPointTy {
ModuleTranslation::convertOmpParallel(Operation &opInst,
   // above.
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP(builder.saveIP());
   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,
+    DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
+    llvm::Instruction *codeGenIPBBTI, llvm::BasicBlock &continuationIP,
+    llvm::IRBuilder<> &builder) {
+  // 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())
+      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);
+      }
+    }
+  }
+  // Finally, after all blocks have been traversed and values mapped,
+  // connect the PHI nodes to the results of preceding blocks.
+  connectPHINodes(region, valueMapping, blockMapping);
+}
+
+LogicalResult ModuleTranslation::convertOmpMaster(Operation &opInst,
+                                                  llvm::IRBuilder<> &builder) {
+  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();
+
+    builder.SetInsertPoint(codeGenIPBB);
+    // MasterOp has only `1` 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);
+  };
+
+  // TODO: Perform finalization actions for variables. This has to be
+  // called for variables which have destructors/finalizers.
+  auto finiCB = [&](InsertPointTy codeGenIP) {};
+
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP(builder.saveIP());
+  builder.restoreIP(ompBuilder->CreateMaster(builder, bodyGenCB, finiCB));
+  return success();
+}
+
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
 /// (including OpenMP runtime calls).
 LogicalResult
 ModuleTranslation::convertOmpOperation(Operation &opInst,
                                        llvm::IRBuilder<> &builder) {
   if (!ompBuilder) {
     ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
     ompBuilder->initialize();
return llvm::TypeSwitch<Operation *, LogicalResult>(&opInst)
         // 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::ParallelOp) { return convertOmpParallel(opInst, builder); })
+      .Case([&](omp::MasterOp) { return convertOmpMaster(opInst, builder); })
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();
       });
 }
 
 /// Given a single MLIR operation, create the corresponding LLVM IR operation
 /// using the `builder`.  LLVM IR Builder does not have a generic interface so

mlir/test/Target/openmp-llvm.mlir

llvm.func @test_omp_parallel_3() -> () {
   // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_3_3:.*]] to {{.*}}
   omp.parallel proc_bind(spread) {
     omp.barrier
     omp.terminator
   }
 
   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_master()
+// CHECK: [[OMP_THREAD_3_4:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}})
+// CHECK-NEXT: {{[0-9]+}} = call i32 @__kmpc_master(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD_3_4]])
+// CHECK: omp.master.region
+// CHECK: call void @__kmpc_end_master(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD_3_4]])
+// CHECK: br label %omp_region.end
+llvm.func @test_omp_master() -> () {
+ "omp.master" ()({
+    omp.terminator
+  }):()->()
+  llvm.return
+}