diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -185,6 +185,11 @@ ]; 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, diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h --- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h +++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h @@ -100,6 +100,9 @@ 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); diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -523,6 +523,122 @@ return success(); } +/// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder. +LogicalResult ModuleTranslation::convertOmpWsLoop(Operation &opInst, + llvm::IRBuilder<> &builder) { + auto loop = cast(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 ®ion = 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(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 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(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. + llvm::BasicBlock *insertBlock = builder.GetInsertBlock(); + llvm::CanonicalLoopInfo *loopInfo = ompBuilder->createCanonicalLoop( + ompLoc, bodyGen, lowerBound, upperBound, step, /*IsSigned=*/true, + /*InclusiveStop=*/true); + 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 @@ -564,6 +680,13 @@ .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(); diff --git a/mlir/test/Target/openmp-llvm.mlir b/mlir/test/Target/openmp-llvm.mlir --- a/mlir/test/Target/openmp-llvm.mlir +++ b/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.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) { + %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, !llvm.i64) -> !llvm.ptr + llvm.store %3, %4 : !llvm.ptr + 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 +}