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 'critical' operation into LLVM IR using OpenMPIRBuilder.
208 static LogicalResult
209 convertOmpCritical(Operation &opInst, llvm::IRBuilderBase &builder,
210                    LLVM::ModuleTranslation &moduleTranslation) {
211   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
212   auto criticalOp = cast<omp::CriticalOp>(opInst);
213   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
214   // relying on captured variables.
215   LogicalResult bodyGenStatus = success();
216 
217   auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
218                        llvm::BasicBlock &continuationBlock) {
219     // CriticalOp has only one region associated with it.
220     auto &region = cast<omp::CriticalOp>(opInst).getRegion();
221     convertOmpOpRegions(region, "omp.critical.region", *codeGenIP.getBlock(),
222                         continuationBlock, builder, moduleTranslation,
223                         bodyGenStatus);
224   };
225 
226   // TODO: Perform finalization actions for variables. This has to be
227   // called for variables which have destructors/finalizers.
228   auto finiCB = [&](InsertPointTy codeGenIP) {};
229 
230   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
231       builder.saveIP(), builder.getCurrentDebugLocation());
232   llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
233   llvm::Constant *hint = nullptr;
234   if (criticalOp.hint().hasValue()) {
235     hint =
236         llvm::ConstantInt::get(llvm::Type::getInt32Ty(llvmContext),
237                                static_cast<int>(criticalOp.hint().getValue()));
238   } else {
239     hint = llvm::ConstantInt::get(llvm::Type::getInt32Ty(llvmContext), 0);
240   }
241   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createCritical(
242       ompLoc, bodyGenCB, finiCB, criticalOp.name().getValueOr(""), hint));
243   return success();
244 }
245 
246 /// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
247 static LogicalResult
248 convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
249                  LLVM::ModuleTranslation &moduleTranslation) {
250   auto loop = cast<omp::WsLoopOp>(opInst);
251   // TODO: this should be in the op verifier instead.
252   if (loop.lowerBound().empty())
253     return failure();
254 
255   if (loop.getNumLoops() != 1)
256     return opInst.emitOpError("collapsed loops not yet supported");
257 
258   // Static is the default.
259   omp::ClauseScheduleKind schedule = omp::ClauseScheduleKind::Static;
260   if (loop.schedule_val().hasValue())
261     schedule =
262         *omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue());
263 
264   // Find the loop configuration.
265   llvm::Value *lowerBound = moduleTranslation.lookupValue(loop.lowerBound()[0]);
266   llvm::Value *upperBound = moduleTranslation.lookupValue(loop.upperBound()[0]);
267   llvm::Value *step = moduleTranslation.lookupValue(loop.step()[0]);
268   llvm::Type *ivType = step->getType();
269   llvm::Value *chunk =
270       loop.schedule_chunk_var()
271           ? moduleTranslation.lookupValue(loop.schedule_chunk_var())
272           : llvm::ConstantInt::get(ivType, 1);
273 
274   // Set up the source location value for OpenMP runtime.
275   llvm::DISubprogram *subprogram =
276       builder.GetInsertBlock()->getParent()->getSubprogram();
277   const llvm::DILocation *diLoc =
278       moduleTranslation.translateLoc(opInst.getLoc(), subprogram);
279   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder.saveIP(),
280                                                     llvm::DebugLoc(diLoc));
281 
282   // Generator of the canonical loop body. Produces an SESE region of basic
283   // blocks.
284   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
285   // relying on captured variables.
286   LogicalResult bodyGenStatus = success();
287   auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
288     llvm::IRBuilder<>::InsertPointGuard guard(builder);
289 
290     // Make sure further conversions know about the induction variable.
291     moduleTranslation.mapValue(loop.getRegion().front().getArgument(0), iv);
292 
293     llvm::BasicBlock *entryBlock = ip.getBlock();
294     llvm::BasicBlock *exitBlock =
295         entryBlock->splitBasicBlock(ip.getPoint(), "omp.wsloop.exit");
296 
297     // Convert the body of the loop.
298     convertOmpOpRegions(loop.region(), "omp.wsloop.region", *entryBlock,
299                         *exitBlock, builder, moduleTranslation, bodyGenStatus);
300   };
301 
302   // Delegate actual loop construction to the OpenMP IRBuilder.
303   // TODO: this currently assumes WsLoop is semantically similar to SCF loop,
304   // i.e. it has a positive step, uses signed integer semantics. Reconsider
305   // this code when WsLoop clearly supports more cases.
306   llvm::CanonicalLoopInfo *loopInfo =
307       moduleTranslation.getOpenMPBuilder()->createCanonicalLoop(
308           ompLoc, bodyGen, lowerBound, upperBound, step, /*IsSigned=*/true,
309           /*InclusiveStop=*/loop.inclusive());
310   if (failed(bodyGenStatus))
311     return failure();
312 
313   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
314       findAllocaInsertPoint(builder, moduleTranslation);
315   llvm::OpenMPIRBuilder::InsertPointTy afterIP;
316   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
317   if (schedule == omp::ClauseScheduleKind::Static) {
318     loopInfo = ompBuilder->createStaticWorkshareLoop(ompLoc, loopInfo, allocaIP,
319                                                      !loop.nowait(), chunk);
320     afterIP = loopInfo->getAfterIP();
321   } else {
322     llvm::omp::OMPScheduleType schedType;
323     switch (schedule) {
324     case omp::ClauseScheduleKind::Dynamic:
325       schedType = llvm::omp::OMPScheduleType::DynamicChunked;
326       break;
327     case omp::ClauseScheduleKind::Guided:
328       schedType = llvm::omp::OMPScheduleType::GuidedChunked;
329       break;
330     case omp::ClauseScheduleKind::Auto:
331       schedType = llvm::omp::OMPScheduleType::Auto;
332       break;
333     case omp::ClauseScheduleKind::Runtime:
334       schedType = llvm::omp::OMPScheduleType::Runtime;
335       break;
336     default:
337       llvm_unreachable("Unknown schedule value");
338       break;
339     }
340 
341     afterIP = ompBuilder->createDynamicWorkshareLoop(
342         ompLoc, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
343   }
344 
345   // Continue building IR after the loop.
346   builder.restoreIP(afterIP);
347   return success();
348 }
349 
350 namespace {
351 
352 /// Implementation of the dialect interface that converts operations belonging
353 /// to the OpenMP dialect to LLVM IR.
354 class OpenMPDialectLLVMIRTranslationInterface
355     : public LLVMTranslationDialectInterface {
356 public:
357   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
358 
359   /// Translates the given operation to LLVM IR using the provided IR builder
360   /// and saving the state in `moduleTranslation`.
361   LogicalResult
362   convertOperation(Operation *op, llvm::IRBuilderBase &builder,
363                    LLVM::ModuleTranslation &moduleTranslation) const final;
364 };
365 
366 } // end namespace
367 
368 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
369 /// (including OpenMP runtime calls).
370 LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
371     Operation *op, llvm::IRBuilderBase &builder,
372     LLVM::ModuleTranslation &moduleTranslation) const {
373 
374   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
375 
376   return llvm::TypeSwitch<Operation *, LogicalResult>(op)
377       .Case([&](omp::BarrierOp) {
378         ompBuilder->createBarrier(builder.saveIP(), llvm::omp::OMPD_barrier);
379         return success();
380       })
381       .Case([&](omp::TaskwaitOp) {
382         ompBuilder->createTaskwait(builder.saveIP());
383         return success();
384       })
385       .Case([&](omp::TaskyieldOp) {
386         ompBuilder->createTaskyield(builder.saveIP());
387         return success();
388       })
389       .Case([&](omp::FlushOp) {
390         // No support in Openmp runtime function (__kmpc_flush) to accept
391         // the argument list.
392         // OpenMP standard states the following:
393         //  "An implementation may implement a flush with a list by ignoring
394         //   the list, and treating it the same as a flush without a list."
395         //
396         // The argument list is discarded so that, flush with a list is treated
397         // same as a flush without a list.
398         ompBuilder->createFlush(builder.saveIP());
399         return success();
400       })
401       .Case([&](omp::ParallelOp) {
402         return convertOmpParallel(*op, builder, moduleTranslation);
403       })
404       .Case([&](omp::MasterOp) {
405         return convertOmpMaster(*op, builder, moduleTranslation);
406       })
407       .Case([&](omp::CriticalOp) {
408         return convertOmpCritical(*op, builder, moduleTranslation);
409       })
410       .Case([&](omp::WsLoopOp) {
411         return convertOmpWsLoop(*op, builder, moduleTranslation);
412       })
413       .Case<omp::YieldOp, omp::TerminatorOp>([](auto op) {
414         // `yield` and `terminator` can be just omitted. The block structure was
415         // created in the function that handles their parent operation.
416         assert(op->getNumOperands() == 0 &&
417                "unexpected OpenMP terminator with operands");
418         return success();
419       })
420       .Default([&](Operation *inst) {
421         return inst->emitError("unsupported OpenMP operation: ")
422                << inst->getName();
423       });
424 }
425 
426 void mlir::registerOpenMPDialectTranslation(DialectRegistry &registry) {
427   registry.insert<omp::OpenMPDialect>();
428   registry.addDialectInterface<omp::OpenMPDialect,
429                                OpenMPDialectLLVMIRTranslationInterface>();
430 }
431 
432 void mlir::registerOpenMPDialectTranslation(MLIRContext &context) {
433   DialectRegistry registry;
434   registerOpenMPDialectTranslation(registry);
435   context.appendDialectRegistry(registry);
436 }
437