Author: Alex Zinenko Date: 2020-12-23T11:52:28+01:00 New Revision: 32a884c9c52c1216d57835e557233b238d601726
URL: https://github.com/llvm/llvm-project/commit/32a884c9c52c1216d57835e557233b238d601726 DIFF: https://github.com/llvm/llvm-project/commit/32a884c9c52c1216d57835e557233b238d601726.diff LOG: [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 Added: Modified: mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h mlir/lib/Target/LLVMIR/ModuleTranslation.cpp mlir/test/Target/openmp-llvm.mlir Removed: ################################################################################ diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td index f915afcf32c9..6c6230f0c2e8 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/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, diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h index d3d289414b38..5259ed7fe182 100644 --- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h +++ b/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); diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp index ae0745b0be28..0b2cf7de270f 100644 --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/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 ®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<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(); diff --git a/mlir/test/Target/openmp-llvm.mlir b/mlir/test/Target/openmp-llvm.mlir index 0651c6f5df40..c5ef16f4393d 100644 --- 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.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 +} _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits