[mlir][OpenMP] omp.sections and omp.section lowering to LLVM IR

[shraiysh]

Summary

This patch adds lowering from omp.sections and omp.section (simple lowering along with the nowait clause) to LLVM IR.
Tests for the same are also added.

[ftynse]

MLIR part LGTM with comments addressed. Please have the OpenMP part reviewed by @jdoerfert or @Meinersbur, or factor it out into a separate patch.

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

567

You probably want to return failure() and stop the conversion process.

574

Nit: please expand auto here. Also, let's not perpetuate the very outdated inst/instruction mlir terminology.

579

Nit: expand auto here too please. MLIR prefers explicit types unless they are too long (iterators) or impossible to spell (lambdas), or if the type is obvious from the RHS (essentially if the RHS ends with a cast).

580

Nit: lowerCamelCase in MLIR please, here and below.

[shraiysh]

Thanks for the review @ftynse. Addressed comments. Waiting for review on OpenMP Part.

[NimishMishra]

Thanks for the patch. LGTM

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

559

Thanks for the patch @shraiysh. I have a question. lastprivate and allocate clauses are supported in PFT to MLIR lowering. Would there be any cause for any problem when we integrate the entire toolchain for sections construct?

[shraiysh]

Thanks for the review @NimishMishra.

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

559

As long as they are not supported in atleast one part of the toolchain, we cannot actually support those clauses and giving an error is better than generating a binary unexpected behavior. So, till there is support for these clauses here flang will not be able to support them.

[kiranchandramohan]

LGTM. I had a look and also compared the behaviour with the clang compiler (without IR builder) and it matches.

Since the change in the OpenMP IRBuilder is minor and to make progress I am approving this although @ftynse specifically asked for a review from @Meinersbur or @jdoerfert. @shraiysh if there are post-commit reviews, please address them.

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

592–593

Nit: Can you expand the comment to say that this situation is only possible if there is only a terminator operation inside the sections operation? May be also assert for that situation.

[shraiysh]

Rebase with main. Addressed comment. Sure @kiranchandramohan. Will address if there are any post-commit reviews.

[shraiysh]

Rebase with main. (Unrelated build failure)

[shraiysh]

Rebase with main

[ftynse]

Since the change in the OpenMP IRBuilder is minor and to make progress I am approving this although @ftynse specifically asked for a review from @Meinersbur or @jdoerfert.

Anyone with sufficient knowledge of OpenMPIRBuilder really. I just know those two have it.

[shraiysh]

@ftynse @kiranchandramohan there are unrelated build failures (and this is happening for other patches too). Looks like there is an issue with the build job. Should I land this change without waiting for the build jobs to be rectified?

[kiranchandramohan]

In D115030#3194308, @shraiysh wrote:

@ftynse @kiranchandramohan there are unrelated build failures (and this is happening for other patches too). Looks like there is an issue with the build job. Should I land this change without waiting for the build jobs to be rectified?

Please go ahead @shraiysh .

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSections(
   Value *LB = ConstantInt::get(I32Ty, 0);
   Value *UB = ConstantInt::get(I32Ty, SectionCBs.size());
   Value *ST = ConstantInt::get(I32Ty, 1);
   llvm::CanonicalLoopInfo *LoopInfo = createCanonicalLoop(
       Loc, LoopBodyGenCB, LB, UB, ST, true, false, AllocaIP, "section_loop");
   Builder.SetInsertPoint(AllocaIP.getBlock()->getTerminator());
   AllocaIP = Builder.saveIP();
   InsertPointTy AfterIP =
-      applyStaticWorkshareLoop(Loc.DL, LoopInfo, AllocaIP, true);
+      applyStaticWorkshareLoop(Loc.DL, LoopInfo, AllocaIP, !IsNowait);
   BasicBlock *LoopAfterBB = AfterIP.getBlock();
   Instruction *SplitPos = LoopAfterBB->getTerminator();
   if (!isa_and_nonnull<BranchInst>(SplitPos))
     SplitPos = new UnreachableInst(Builder.getContext(), LoopAfterBB);
   // ExitBB after LoopAfterBB because LoopAfterBB is used for FinalizationCB,
   // which requires a BB with branch
   BasicBlock *ExitBB =
       LoopAfterBB->splitBasicBlock(SplitPos, "omp_sections.end");

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

for (auto *&CaseBB : CaseBBs) {
     EXPECT_EQ(CaseBB->getSingleSuccessor(), ForExitBB);
   }
 
   ASSERT_EQ(NumBodiesGenerated, 2U);
   ASSERT_EQ(NumFiniCBCalls, 1U);
   EXPECT_FALSE(verifyModule(*M, &errs()));
 }
 
+TEST_F(OpenMPIRBuilderTest, CreateSectionsNoWait) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  using BodyGenCallbackTy = llvm::OpenMPIRBuilder::StorableBodyGenCallbackTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  F->setName("func");
+  IRBuilder<> Builder(BB);
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+  IRBuilder<>::InsertPoint AllocaIP(&F->getEntryBlock(),
+                                    F->getEntryBlock().getFirstInsertionPt());
+  llvm::SmallVector<BodyGenCallbackTy, 4> SectionCBVector;
+  auto PrivCB = [](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
+                   llvm::Value &, llvm::Value &Val,
+                   llvm::Value *&ReplVal) { return CodeGenIP; };
+  auto FiniCB = [&](InsertPointTy IP) {};
+
+  Builder.restoreIP(OMPBuilder.createSections(Loc, AllocaIP, SectionCBVector,
+                                              PrivCB, FiniCB, false, true));
+  Builder.CreateRetVoid(); // Required at the end of the function
+  for (auto &Inst : instructions(*F)) {
+    EXPECT_FALSE(isa<CallInst>(Inst) &&
+                 cast<CallInst>(&Inst)->getCalledFunction()->getName() ==
+                     "__kmpc_barrier" &&
+                 "call to function __kmpc_barrier found with nowait");
+  }
+}
+
 TEST_F(OpenMPIRBuilderTest, CreateOffloadMaptypes) {
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.initialize();
 
   IRBuilder<> Builder(BB);
 
   SmallVector<uint64_t> Mappings = {0, 1};
   GlobalVariable *OffloadMaptypesGlobal =

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

convertOmpOrderedRegion(Operation &opInst, llvm::IRBuilderBase &builder,
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
       builder.saveIP(), builder.getCurrentDebugLocation());
   builder.restoreIP(
       moduleTranslation.getOpenMPBuilder()->createOrderedThreadsSimd(
           ompLoc, bodyGenCB, finiCB, !orderedRegionOp.simd()));
   return bodyGenStatus;
 }
 
+static LogicalResult
+convertOmpSections(Operation &opInst, llvm::IRBuilderBase &builder,
+                   LLVM::ModuleTranslation &moduleTranslation) {
+  using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
+  using StorableBodyGenCallbackTy =
+      llvm::OpenMPIRBuilder::StorableBodyGenCallbackTy;
+
+  auto sectionsOp = cast<omp::SectionsOp>(opInst);
+
+  // TODO: Support the following clauses: private, firstprivate, lastprivate,
+  // reduction, allocate
+  if (!sectionsOp.private_vars().empty() ||
+      !sectionsOp.firstprivate_vars().empty() ||
+      !sectionsOp.lastprivate_vars().empty() ||
+      !sectionsOp.reduction_vars().empty() || sectionsOp.reductions() ||
+      !sectionsOp.allocate_vars().empty() ||
+      !sectionsOp.allocators_vars().empty())
+    return emitError(sectionsOp.getLoc())
+           << "private, firstprivate, lastprivate, reduction and allocate "
+              "clauses are not supported for sections construct";
+
+  LogicalResult bodyGenStatus = success();
+  SmallVector<StorableBodyGenCallbackTy> sectionCBs;
+
+  for (Operation &op : *sectionsOp.region().begin()) {
+    auto sectionOp = dyn_cast<omp::SectionOp>(op);
+    if (!sectionOp) // omp.terminator
+      continue;
+
+    Region &region = sectionOp.region();
+    auto sectionCB = [&region, &builder, &moduleTranslation, &bodyGenStatus](
+                         InsertPointTy allocaIP, InsertPointTy codeGenIP,
+                         llvm::BasicBlock &finiBB) {
+      builder.restoreIP(codeGenIP);
+      builder.CreateBr(&finiBB);
+      convertOmpOpRegions(region, "omp.section.region", *codeGenIP.getBlock(),
+                          finiBB, builder, moduleTranslation, bodyGenStatus);
+    };
+    sectionCBs.push_back(sectionCB);
+  }
+
+  // No sections within omp.sections operation - skip generation. This situation
+  // is only possible if there is only a terminator operation inside the
+  // sections operation
+  if (sectionCBs.size() == 0)
+    return success();
+
+  assert(isa<omp::SectionOp>(*sectionsOp.region().op_begin()));
+
+  // TODO: Perform appropriate actions according to the data-sharing
+  // attribute (shared, private, firstprivate, ...) of variables.
+  // Currently defaults to shared.
+  auto privCB = [&](InsertPointTy, InsertPointTy codeGenIP, llvm::Value &,
+                    llvm::Value &vPtr,
+                    llvm::Value *&replacementValue) -> InsertPointTy {
+    replacementValue = &vPtr;
+    return codeGenIP;
+  };
+
+  // TODO: Perform finalization actions for variables. This has to be
+  // called for variables which have destructors/finalizers.
+  auto finiCB = [&](InsertPointTy codeGenIP) {};
+
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(
+      builder.saveIP(), builder.getCurrentDebugLocation());
+  builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSections(
+      ompLoc, findAllocaInsertPoint(builder, moduleTranslation), sectionCBs,
+      privCB, finiCB, false, sectionsOp.nowait()));
+  return bodyGenStatus;
+}
+
 /// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
   auto loop = cast<omp::WsLoopOp>(opInst);
   // TODO: this should be in the op verifier instead.
   if (loop.lowerBound().empty())
     return failure();
return llvm::TypeSwitch<Operation *, LogicalResult>(op)
         return convertOmpOrdered(*op, builder, moduleTranslation);
       })
       .Case([&](omp::WsLoopOp) {
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
       .Case([&](omp::AtomicReadOp) {
         return convertOmpAtomicRead(*op, builder, moduleTranslation);
       })
+      .Case([&](omp::SectionsOp) {
+        return convertOmpSections(*op, builder, moduleTranslation);
+      })
       .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
             omp::CriticalDeclareOp>([](auto op) {
         // `yield` and `terminator` can be just omitted. The block structure
         // was created in the region that handles their parent operation.
         // `reduction.declare` will be used by reductions and is not
         // converted directly, skip it.
         // `critical.declare` is only used to declare names of critical
         // sections which will be used by `critical` ops and hence can be

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

llvm.func @omp_atomic_read(%arg0 : !llvm.ptr<i32>) -> () {
   // CHECK: store i32 %[[X3]], i32* %{{.*}}, align 4
   %x3 = omp.atomic.read %arg0 memory_order(acquire) : !llvm.ptr<i32> -> i32
 
   // CHECK: %[[X4:.*]] = load atomic i32, i32* %[[ARG0]] monotonic, align 4
   // CHECK: store i32 %[[X4]], i32* %{{.*}}, align 4
   %x4 = omp.atomic.read %arg0 memory_order(relaxed) : !llvm.ptr<i32> -> i32
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: @omp_sections_empty
+llvm.func @omp_sections_empty() -> () {
+  omp.sections {
+    omp.terminator
+  }
+  // CHECK-NEXT: ret void
+  llvm.return
+}
+
+// -----
+
+// Check IR generation for simple empty sections. This only checks the overall
+// shape of the IR, detailed checking is done by the OpenMPIRBuilder.
+
+// CHECK-LABEL: @omp_sections_trivial
+llvm.func @omp_sections_trivial() -> () {
+  // CHECK:   br label %[[PREHEADER:.*]]
+
+  // CHECK: [[PREHEADER]]:
+  // CHECK:   %{{.*}} = call i32 @__kmpc_global_thread_num({{.*}})
+  // CHECK:   call void @__kmpc_for_static_init_4u({{.*}})
+  // CHECK:   br label %[[HEADER:.*]]
+
+  // CHECK: [[HEADER]]:
+  // CHECK:   br label %[[COND:.*]]
+
+  // CHECK: [[COND]]:
+  // CHECK:   br i1 %{{.*}}, label %[[BODY:.*]], label %[[EXIT:.*]]
+  // CHECK: [[BODY]]:
+  // CHECK:   switch i32 %{{.*}}, label %[[INC:.*]] [
+  // CHECK-NEXT:     i32 0, label %[[SECTION1:.*]]
+  // CHECK-NEXT:     i32 1, label %[[SECTION2:.*]]
+  // CHECK-NEXT: ]
+
+  // CHECK: [[INC]]:
+  // CHECK:   %{{.*}} = add {{.*}}, 1
+  // CHECK:   br label %[[HEADER]]
+
+  // CHECK: [[EXIT]]:
+  // CHECK:   call void @__kmpc_for_static_fini({{.*}})
+  // CHECK:   call void @__kmpc_barrier({{.*}})
+  // CHECK:   br label %[[AFTER:.*]]
+
+  // CHECK: [[AFTER]]:
+  // CHECK:   br label %[[END:.*]]
+
+  // CHECK: [[END]]:
+  // CHECK:   ret void
+  omp.sections {
+    omp.section {
+      // CHECK: [[SECTION1]]:
+      // CHECK-NEXT: br label %[[REGION1:[^ ,]*]]
+      // CHECK: [[REGION1]]:
+      // CHECK-NEXT: br label %[[EXIT]]
+      omp.terminator
+    }
+    omp.section {
+      // CHECK: [[SECTION2]]:
+      // CHECK-NEXT: br label %[[REGION2:[^ ,]*]]
+      // CHECK: [[REGION2]]:
+      // CHECK-NEXT: br label %[[EXIT]]
+      omp.terminator
+    }
+    omp.terminator
+  }
+  llvm.return
+}
+
+// -----
+
+// CHECK: declare void @foo()
+llvm.func @foo()
+
+// CHECK: declare void @bar(i32)
+llvm.func @bar(%arg0 : i32)
+
+// CHECK-LABEL: @omp_sections
+llvm.func @omp_sections(%arg0 : i32, %arg1 : i32, %arg2 : !llvm.ptr<i32>) -> () {
+
+  // CHECK: switch i32 %{{.*}}, label %{{.*}} [
+  // CHECK-NEXT:   i32 0, label %[[SECTION1:.*]]
+  // CHECK-NEXT:   i32 1, label %[[SECTION2:.*]]
+  // CHECK-NEXT:   i32 2, label %[[SECTION3:.*]]
+  // CHECK-NEXT: ]
+  omp.sections {
+    omp.section {
+      // CHECK: [[SECTION1]]:
+      // CHECK:   br label %[[REGION1:[^ ,]*]]
+      // CHECK: [[REGION1]]:
+      // CHECK:   call void @foo()
+      // CHECK:   br label %{{.*}}
+      llvm.call @foo() : () -> ()
+      omp.terminator
+    }
+    omp.section {
+      // CHECK: [[SECTION2]]:
+      // CHECK:   br label %[[REGION2:[^ ,]*]]
+      // CHECK: [[REGION2]]:
+      // CHECK:   call void @bar(i32 %{{.*}})
+      // CHECK:   br label %{{.*}}
+      llvm.call @bar(%arg0) : (i32) -> ()
+      omp.terminator
+    }
+    omp.section {
+      // CHECK: [[SECTION3]]:
+      // CHECK:   br label %[[REGION3:[^ ,]*]]
+      // CHECK: [[REGION3]]:
+      // CHECK:   %11 = add i32 %{{.*}}, %{{.*}}
+      %add = llvm.add %arg0, %arg1 : i32
+      // CHECK:   store i32 %{{.*}}, i32* %{{.*}}, align 4
+      // CHECK:   br label %{{.*}}
+      llvm.store %add, %arg2 : !llvm.ptr<i32>
+      omp.terminator
+    }
+    omp.terminator
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @foo()
+
+// CHECK-LABEL: @omp_sections_with_clauses
+llvm.func @omp_sections_with_clauses() -> () {
+  // CHECK-NOT: call void @__kmpc_barrier
+  omp.sections nowait {
+    omp.section {
+      llvm.call @foo() : () -> ()
+      omp.terminator
+    }
+    omp.section {
+      llvm.call @foo() : () -> ()
+      omp.terminator
+    }
+    omp.terminator
+  }
+  llvm.return
+}