[mlir] Add translation of omp.wsloop to LLVM IR

Introduce a translation of OpenMP workshare loop construct to LLVM IR. This is
a minimalist version to enable the pipeline and currently only supports static
loop schedule (default in the specification) on non-collapsed loops. Other
features will be added on per-need basis.

Reviewed By: kiranchandramohan

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

Author: ftynse

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -185,6 +185,11 @@ def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments]> {
   ];
 
   let regions = (region AnyRegion:$region);
+
+  let extraClassDeclaration = [{
+    /// Returns the number of loops in the workshape loop nest.
+    unsigned getNumLoops() { return lowerBound().size(); }
+  }];
 }
 
 def YieldOp : OpenMP_Op<"yield", [NoSideEffect, ReturnLike, Terminator,

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

@@ -100,6 +100,9 @@ class ModuleTranslation {
                            llvm::BasicBlock &continuationIP,
                            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(LLVMType type);
 

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -536,6 +536,126 @@ LogicalResult ModuleTranslation::convertOmpMaster(Operation &opInst,
   return success();
 }
 
+/// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
+LogicalResult ModuleTranslation::convertOmpWsLoop(Operation &opInst,
+                                                  llvm::IRBuilder<> &builder) {
+  auto loop = cast<omp::WsLoopOp>(opInst);
+  // TODO: this should be in the op verifier instead.
+  if (loop.lowerBound().empty())
+    return failure();
+
+  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);
+
+  // Set up the source location value for OpenMP runtime.
+  llvm::DISubprogram *subprogram =
+      builder.GetInsertBlock()->getParent()->getSubprogram();
+  const llvm::DILocation *diLoc =
+      debugTranslation->translateLoc(opInst.getLoc(), subprogram);
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder.saveIP(),
+                                                    llvm::DebugLoc(diLoc));
+
+  // Generator of the canonical loop body. Produces an SESE region of basic
+  // blocks.
+  // TODO: support error propagation in OpenMPIRBuilder and use it instead of
+  // relying on captured variables.
+  LogicalResult bodyGenStatus = success();
+  auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
+    llvm::IRBuilder<>::InsertPointGuard guard(builder);
+
+    // 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);
+  };
+
+  // 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();
+  llvm::CanonicalLoopInfo *loopInfo = ompBuilder->createCanonicalLoop(
+      ompLoc, bodyGen, lowerBound, upperBound, step, /*IsSigned=*/true,
+      /*InclusiveStop=*/false);
+  if (failed(bodyGenStatus))
+    return failure();
+
+  // TODO: get the alloca insertion point from the parallel operation builder.
+  // If we insert the at the top of the current function, they will be passed as
+  // extra arguments into the function the parallel operation builder outlines.
+  // Put them at the start of the current block for now.
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP(
+      insertBlock, insertBlock->getFirstInsertionPt());
+  loopInfo = ompBuilder->createStaticWorkshareLoop(
+      ompLoc, loopInfo, allocaIP,
+      !loop.nowait().hasValue() || loop.nowait().getValue(), chunk);
+
+  // Continue building IR after the loop.
+  builder.restoreIP(loopInfo->getAfterIP());
+  return success();
+}
+
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
 /// (including OpenMP runtime calls).
 LogicalResult
@@ -577,6 +697,13 @@ ModuleTranslation::convertOmpOperation(Operation &opInst,
       .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");
+        return success();
+      })
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();

mlir/test/Target/openmp-llvm.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: define void @test_stand_alone_directives()
 llvm.func @test_stand_alone_directives() {
@@ -291,3 +291,35 @@ llvm.func @test_omp_master() -> () {
   }
   llvm.return
 }
+
+// -----
+
+// CHECK: %struct.ident_t = type
+// CHECK: @[[$parallel_loc:.*]] = private unnamed_addr constant {{.*}} c";LLVMDialectModule;wsloop_simple;{{[0-9]+}};{{[0-9]+}};;\00"
+// CHECK: @[[$parallel_loc_struct:.*]] = private unnamed_addr constant %struct.ident_t {{.*}} @[[$parallel_loc]], {{.*}}
+
+// CHECK: @[[$wsloop_loc:.*]] = private unnamed_addr constant {{.*}} c";LLVMDialectModule;wsloop_simple;{{[0-9]+}};{{[0-9]+}};;\00"
+// CHECK: @[[$wsloop_loc_struct:.*]] = private unnamed_addr constant %struct.ident_t {{.*}} @[[$wsloop_loc]], {{.*}}
+
+// CHECK-LABEL: @wsloop_simple
+llvm.func @wsloop_simple(%arg0: !llvm.ptr<float>) {
+  %0 = llvm.mlir.constant(42 : index) : !llvm.i64
+  %1 = llvm.mlir.constant(10 : index) : !llvm.i64
+  %2 = llvm.mlir.constant(1 : index) : !llvm.i64
+  omp.parallel {
+    "omp.wsloop"(%1, %0, %2) ( {
+    ^bb0(%arg1: !llvm.i64):
+      // The form of the emitted IR is controlled by OpenMPIRBuilder and
+      // tested there. Just check that the right functions are called.
+      // CHECK: call i32 @__kmpc_global_thread_num
+      // CHECK: call void @__kmpc_for_static_init_{{.*}}(%struct.ident_t* @[[$wsloop_loc_struct]],
+      %3 = llvm.mlir.constant(2.000000e+00 : f32) : !llvm.float
+      %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<float>, !llvm.i64) -> !llvm.ptr<float>
+      llvm.store %3, %4 : !llvm.ptr<float>
+      omp.yield
+      // CHECK: call void @__kmpc_for_static_fini(%struct.ident_t* @[[$wsloop_loc_struct]],
+    }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (!llvm.i64, !llvm.i64, !llvm.i64) -> ()
+    omp.terminator
+  }
+  llvm.return
+}