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

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.

Reviewed By: ftynse, kiranchandramohan

Differential Revision: https://reviews.llvm.org/D115030

Author: shraiysh

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -996,7 +996,7 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSections(
   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))

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -3637,6 +3637,34 @@ TEST_F(OpenMPIRBuilderTest, CreateSections) {
   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();

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

@@ -547,6 +547,77 @@ convertOmpOrderedRegion(Operation &opInst, llvm::IRBuilderBase &builder,
   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,
@@ -962,6 +1033,9 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
       .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

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

@@ -736,3 +736,145 @@ llvm.func @omp_atomic_read(%arg0 : !llvm.ptr<i32>) -> () {
   %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
+}