[Mlir-commits] [mlir] 32a884c - [mlir] Add translation of omp.wsloop to LLVM IR
Alex Zinenko
llvmlistbot at llvm.org
Wed Dec 23 02:52:38 PST 2020
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
+}
More information about the Mlir-commits
mailing list