1 //===- OpenMPToLLVMIRTranslation.cpp - Translate OpenMP dialect to LLVM IR-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements a translation between the MLIR OpenMP dialect and LLVM
10 // IR.
11 //
12 //===----------------------------------------------------------------------===//
13 #include "mlir/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.h"
14 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
15 #include "mlir/IR/Operation.h"
16 #include "mlir/Support/LLVM.h"
17 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
18 
19 #include "llvm/ADT/SetVector.h"
20 #include "llvm/ADT/TypeSwitch.h"
21 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
22 #include "llvm/IR/IRBuilder.h"
23 
24 using namespace mlir;
25 
26 namespace {
27 /// ModuleTranslation stack frame for OpenMP operations. This keeps track of the
28 /// insertion points for allocas.
29 class OpenMPAllocaStackFrame
30     : public LLVM::ModuleTranslation::StackFrameBase<OpenMPAllocaStackFrame> {
31 public:
32   explicit OpenMPAllocaStackFrame(llvm::OpenMPIRBuilder::InsertPointTy allocaIP)
33       : allocaInsertPoint(allocaIP) {}
34   llvm::OpenMPIRBuilder::InsertPointTy allocaInsertPoint;
35 };
36 } // namespace
37 
38 /// Find the insertion point for allocas given the current insertion point for
39 /// normal operations in the builder.
40 static llvm::OpenMPIRBuilder::InsertPointTy
41 findAllocaInsertPoint(llvm::IRBuilderBase &builder,
42                       const LLVM::ModuleTranslation &moduleTranslation) {
43   // If there is an alloca insertion point on stack, i.e. we are in a nested
44   // operation and a specific point was provided by some surrounding operation,
45   // use it.
46   llvm::OpenMPIRBuilder::InsertPointTy allocaInsertPoint;
47   WalkResult walkResult = moduleTranslation.stackWalk<OpenMPAllocaStackFrame>(
48       [&](const OpenMPAllocaStackFrame &frame) {
49         allocaInsertPoint = frame.allocaInsertPoint;
50         return WalkResult::interrupt();
51       });
52   if (walkResult.wasInterrupted())
53     return allocaInsertPoint;
54 
55   // Otherwise, insert to the entry block of the surrounding function.
56   llvm::BasicBlock &funcEntryBlock =
57       builder.GetInsertBlock()->getParent()->getEntryBlock();
58   return llvm::OpenMPIRBuilder::InsertPointTy(
59       &funcEntryBlock, funcEntryBlock.getFirstInsertionPt());
60 }
61 
62 /// Converts the given region that appears within an OpenMP dialect operation to
63 /// LLVM IR, creating a branch from the `sourceBlock` to the entry block of the
64 /// region, and a branch from any block with an successor-less OpenMP terminator
65 /// to `continuationBlock`.
66 static void convertOmpOpRegions(Region &region, StringRef blockName,
67                                 llvm::BasicBlock &sourceBlock,
68                                 llvm::BasicBlock &continuationBlock,
69                                 llvm::IRBuilderBase &builder,
70                                 LLVM::ModuleTranslation &moduleTranslation,
71                                 LogicalResult &bodyGenStatus) {
72   llvm::LLVMContext &llvmContext = builder.getContext();
73   for (Block &bb : region) {
74     llvm::BasicBlock *llvmBB = llvm::BasicBlock::Create(
75         llvmContext, blockName, builder.GetInsertBlock()->getParent());
76     moduleTranslation.mapBlock(&bb, llvmBB);
77   }
78 
79   llvm::Instruction *sourceTerminator = sourceBlock.getTerminator();
80 
81   // Convert blocks one by one in topological order to ensure
82   // defs are converted before uses.
83   SetVector<Block *> blocks =
84       LLVM::detail::getTopologicallySortedBlocks(region);
85   for (Block *bb : blocks) {
86     llvm::BasicBlock *llvmBB = moduleTranslation.lookupBlock(bb);
87     // Retarget the branch of the entry block to the entry block of the
88     // converted region (regions are single-entry).
89     if (bb->isEntryBlock()) {
90       assert(sourceTerminator->getNumSuccessors() == 1 &&
91              "provided entry block has multiple successors");
92       assert(sourceTerminator->getSuccessor(0) == &continuationBlock &&
93              "ContinuationBlock is not the successor of the entry block");
94       sourceTerminator->setSuccessor(0, llvmBB);
95     }
96 
97     llvm::IRBuilderBase::InsertPointGuard guard(builder);
98     if (failed(
99             moduleTranslation.convertBlock(*bb, bb->isEntryBlock(), builder))) {
100       bodyGenStatus = failure();
101       return;
102     }
103 
104     // Special handling for `omp.yield` and `omp.terminator` (we may have more
105     // than one): they return the control to the parent OpenMP dialect operation
106     // so replace them with the branch to the continuation block. We handle this
107     // here to avoid relying inter-function communication through the
108     // ModuleTranslation class to set up the correct insertion point. This is
109     // also consistent with MLIR's idiom of handling special region terminators
110     // in the same code that handles the region-owning operation.
111     if (isa<omp::TerminatorOp, omp::YieldOp>(bb->getTerminator()))
112       builder.CreateBr(&continuationBlock);
113   }
114   // Finally, after all blocks have been traversed and values mapped,
115   // connect the PHI nodes to the results of preceding blocks.
116   LLVM::detail::connectPHINodes(region, moduleTranslation);
117 }
118 
119 /// Converts the OpenMP parallel operation to LLVM IR.
120 static LogicalResult
121 convertOmpParallel(Operation &opInst, llvm::IRBuilderBase &builder,
122                    LLVM::ModuleTranslation &moduleTranslation) {
123   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
124   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
125   // relying on captured variables.
126   LogicalResult bodyGenStatus = success();
127 
128   auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
129                        llvm::BasicBlock &continuationBlock) {
130     // Save the alloca insertion point on ModuleTranslation stack for use in
131     // nested regions.
132     LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
133         moduleTranslation, allocaIP);
134 
135     // ParallelOp has only one region associated with it.
136     auto &region = cast<omp::ParallelOp>(opInst).getRegion();
137     convertOmpOpRegions(region, "omp.par.region", *codeGenIP.getBlock(),
138                         continuationBlock, builder, moduleTranslation,
139                         bodyGenStatus);
140   };
141 
142   // TODO: Perform appropriate actions according to the data-sharing
143   // attribute (shared, private, firstprivate, ...) of variables.
144   // Currently defaults to shared.
145   auto privCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
146                     llvm::Value &, llvm::Value &vPtr,
147                     llvm::Value *&replacementValue) -> InsertPointTy {
148     replacementValue = &vPtr;
149 
150     return codeGenIP;
151   };
152 
153   // TODO: Perform finalization actions for variables. This has to be
154   // called for variables which have destructors/finalizers.
155   auto finiCB = [&](InsertPointTy codeGenIP) {};
156 
157   llvm::Value *ifCond = nullptr;
158   if (auto ifExprVar = cast<omp::ParallelOp>(opInst).if_expr_var())
159     ifCond = moduleTranslation.lookupValue(ifExprVar);
160   llvm::Value *numThreads = nullptr;
161   if (auto numThreadsVar = cast<omp::ParallelOp>(opInst).num_threads_var())
162     numThreads = moduleTranslation.lookupValue(numThreadsVar);
163   llvm::omp::ProcBindKind pbKind = llvm::omp::OMP_PROC_BIND_default;
164   if (auto bind = cast<omp::ParallelOp>(opInst).proc_bind_val())
165     pbKind = llvm::omp::getProcBindKind(bind.getValue());
166   // TODO: Is the Parallel construct cancellable?
167   bool isCancellable = false;
168 
169   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
170       builder.saveIP(), builder.getCurrentDebugLocation());
171   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createParallel(
172       ompLoc, findAllocaInsertPoint(builder, moduleTranslation), bodyGenCB,
173       privCB, finiCB, ifCond, numThreads, pbKind, isCancellable));
174 
175   return bodyGenStatus;
176 }
177 
178 /// Converts an OpenMP 'master' operation into LLVM IR using OpenMPIRBuilder.
179 static LogicalResult
180 convertOmpMaster(Operation &opInst, llvm::IRBuilderBase &builder,
181                  LLVM::ModuleTranslation &moduleTranslation) {
182   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
183   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
184   // relying on captured variables.
185   LogicalResult bodyGenStatus = success();
186 
187   auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
188                        llvm::BasicBlock &continuationBlock) {
189     // MasterOp has only one region associated with it.
190     auto &region = cast<omp::MasterOp>(opInst).getRegion();
191     convertOmpOpRegions(region, "omp.master.region", *codeGenIP.getBlock(),
192                         continuationBlock, builder, moduleTranslation,
193                         bodyGenStatus);
194   };
195 
196   // TODO: Perform finalization actions for variables. This has to be
197   // called for variables which have destructors/finalizers.
198   auto finiCB = [&](InsertPointTy codeGenIP) {};
199 
200   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
201       builder.saveIP(), builder.getCurrentDebugLocation());
202   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createMaster(
203       ompLoc, bodyGenCB, finiCB));
204   return success();
205 }
206 
207 /// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
208 static LogicalResult
209 convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
210                  LLVM::ModuleTranslation &moduleTranslation) {
211   auto loop = cast<omp::WsLoopOp>(opInst);
212   // TODO: this should be in the op verifier instead.
213   if (loop.lowerBound().empty())
214     return failure();
215 
216   if (loop.getNumLoops() != 1)
217     return opInst.emitOpError("collapsed loops not yet supported");
218 
219   // Static is the default.
220   omp::ClauseScheduleKind schedule = omp::ClauseScheduleKind::Static;
221   if (loop.schedule_val().hasValue())
222     schedule =
223         *omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue());
224 
225   // Find the loop configuration.
226   llvm::Value *lowerBound = moduleTranslation.lookupValue(loop.lowerBound()[0]);
227   llvm::Value *upperBound = moduleTranslation.lookupValue(loop.upperBound()[0]);
228   llvm::Value *step = moduleTranslation.lookupValue(loop.step()[0]);
229   llvm::Type *ivType = step->getType();
230   llvm::Value *chunk =
231       loop.schedule_chunk_var()
232           ? moduleTranslation.lookupValue(loop.schedule_chunk_var())
233           : llvm::ConstantInt::get(ivType, 1);
234 
235   // Set up the source location value for OpenMP runtime.
236   llvm::DISubprogram *subprogram =
237       builder.GetInsertBlock()->getParent()->getSubprogram();
238   const llvm::DILocation *diLoc =
239       moduleTranslation.translateLoc(opInst.getLoc(), subprogram);
240   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder.saveIP(),
241                                                     llvm::DebugLoc(diLoc));
242 
243   // Generator of the canonical loop body. Produces an SESE region of basic
244   // blocks.
245   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
246   // relying on captured variables.
247   LogicalResult bodyGenStatus = success();
248   auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
249     llvm::IRBuilder<>::InsertPointGuard guard(builder);
250 
251     // Make sure further conversions know about the induction variable.
252     moduleTranslation.mapValue(loop.getRegion().front().getArgument(0), iv);
253 
254     llvm::BasicBlock *entryBlock = ip.getBlock();
255     llvm::BasicBlock *exitBlock =
256         entryBlock->splitBasicBlock(ip.getPoint(), "omp.wsloop.exit");
257 
258     // Convert the body of the loop.
259     convertOmpOpRegions(loop.region(), "omp.wsloop.region", *entryBlock,
260                         *exitBlock, builder, moduleTranslation, bodyGenStatus);
261   };
262 
263   // Delegate actual loop construction to the OpenMP IRBuilder.
264   // TODO: this currently assumes WsLoop is semantically similar to SCF loop,
265   // i.e. it has a positive step, uses signed integer semantics. Reconsider
266   // this code when WsLoop clearly supports more cases.
267   llvm::CanonicalLoopInfo *loopInfo =
268       moduleTranslation.getOpenMPBuilder()->createCanonicalLoop(
269           ompLoc, bodyGen, lowerBound, upperBound, step, /*IsSigned=*/true,
270           /*InclusiveStop=*/loop.inclusive());
271   if (failed(bodyGenStatus))
272     return failure();
273 
274   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
275       findAllocaInsertPoint(builder, moduleTranslation);
276   llvm::OpenMPIRBuilder::InsertPointTy afterIP;
277   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
278   if (schedule == omp::ClauseScheduleKind::Static) {
279     loopInfo = ompBuilder->createStaticWorkshareLoop(ompLoc, loopInfo, allocaIP,
280                                                      !loop.nowait(), chunk);
281     afterIP = loopInfo->getAfterIP();
282   } else {
283     llvm::omp::OMPScheduleType schedType;
284     switch (schedule) {
285     case omp::ClauseScheduleKind::Dynamic:
286       schedType = llvm::omp::OMPScheduleType::DynamicChunked;
287       break;
288     case omp::ClauseScheduleKind::Guided:
289       schedType = llvm::omp::OMPScheduleType::GuidedChunked;
290       break;
291     case omp::ClauseScheduleKind::Auto:
292       schedType = llvm::omp::OMPScheduleType::Auto;
293       break;
294     case omp::ClauseScheduleKind::Runtime:
295       schedType = llvm::omp::OMPScheduleType::Runtime;
296       break;
297     default:
298       llvm_unreachable("Unknown schedule value");
299       break;
300     }
301 
302     afterIP = ompBuilder->createDynamicWorkshareLoop(
303         ompLoc, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
304   }
305 
306   // Continue building IR after the loop.
307   builder.restoreIP(afterIP);
308   return success();
309 }
310 
311 namespace {
312 
313 /// Implementation of the dialect interface that converts operations belonging
314 /// to the OpenMP dialect to LLVM IR.
315 class OpenMPDialectLLVMIRTranslationInterface
316     : public LLVMTranslationDialectInterface {
317 public:
318   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
319 
320   /// Translates the given operation to LLVM IR using the provided IR builder
321   /// and saving the state in `moduleTranslation`.
322   LogicalResult
323   convertOperation(Operation *op, llvm::IRBuilderBase &builder,
324                    LLVM::ModuleTranslation &moduleTranslation) const final;
325 };
326 
327 } // end namespace
328 
329 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
330 /// (including OpenMP runtime calls).
331 LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
332     Operation *op, llvm::IRBuilderBase &builder,
333     LLVM::ModuleTranslation &moduleTranslation) const {
334 
335   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
336 
337   return llvm::TypeSwitch<Operation *, LogicalResult>(op)
338       .Case([&](omp::BarrierOp) {
339         ompBuilder->createBarrier(builder.saveIP(), llvm::omp::OMPD_barrier);
340         return success();
341       })
342       .Case([&](omp::TaskwaitOp) {
343         ompBuilder->createTaskwait(builder.saveIP());
344         return success();
345       })
346       .Case([&](omp::TaskyieldOp) {
347         ompBuilder->createTaskyield(builder.saveIP());
348         return success();
349       })
350       .Case([&](omp::FlushOp) {
351         // No support in Openmp runtime function (__kmpc_flush) to accept
352         // the argument list.
353         // OpenMP standard states the following:
354         //  "An implementation may implement a flush with a list by ignoring
355         //   the list, and treating it the same as a flush without a list."
356         //
357         // The argument list is discarded so that, flush with a list is treated
358         // same as a flush without a list.
359         ompBuilder->createFlush(builder.saveIP());
360         return success();
361       })
362       .Case([&](omp::ParallelOp) {
363         return convertOmpParallel(*op, builder, moduleTranslation);
364       })
365       .Case([&](omp::MasterOp) {
366         return convertOmpMaster(*op, builder, moduleTranslation);
367       })
368       .Case([&](omp::WsLoopOp) {
369         return convertOmpWsLoop(*op, builder, moduleTranslation);
370       })
371       .Case<omp::YieldOp, omp::TerminatorOp>([](auto op) {
372         // `yield` and `terminator` can be just omitted. The block structure was
373         // created in the function that handles their parent operation.
374         assert(op->getNumOperands() == 0 &&
375                "unexpected OpenMP terminator with operands");
376         return success();
377       })
378       .Default([&](Operation *inst) {
379         return inst->emitError("unsupported OpenMP operation: ")
380                << inst->getName();
381       });
382 }
383 
384 void mlir::registerOpenMPDialectTranslation(DialectRegistry &registry) {
385   registry.insert<omp::OpenMPDialect>();
386   registry.addDialectInterface<omp::OpenMPDialect,
387                                OpenMPDialectLLVMIRTranslationInterface>();
388 }
389 
390 void mlir::registerOpenMPDialectTranslation(MLIRContext &context) {
391   DialectRegistry registry;
392   registerOpenMPDialectTranslation(registry);
393   context.appendDialectRegistry(registry);
394 }
395