From 7ecee63e719aea7fe3a6764e535cb865669c2fd5 Mon Sep 17 00:00:00 2001 From: Kiran Kumar T P <kirankumartp@gmail.com> Date: Fri, 10 Apr 2020 07:41:04 +0000 Subject: [PATCH] [MLIR] Support for taskwait and taskyield operations, and translating the same to LLVM IR This patch adds support for taskwait and taskyield operations in OpenMP dialect and translation of the these constructs to LLVM IR. The OpenMP IRBuilder is used for this translation. The patch includes code changes and a testcase modifications. Differential Revision: https://reviews.llvm.org/D77634 --- mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 23 ++++++++++- .../mlir/Target/LLVMIR/ModuleTranslation.h | 2 + mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 41 ++++++++++++++----- mlir/test/Dialect/OpenMP/ops.mlir | 12 ++++++ mlir/test/Target/openmp-llvm.mlir | 15 +++++-- 5 files changed, 77 insertions(+), 16 deletions(-) diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td index 8aa8635597b0..f283d3a7bf3d 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -30,8 +30,27 @@ def BarrierOp : OpenMP_Op<"barrier"> { the construct appears. }]; - let parser = [{ return success(); }]; - let printer = [{ p << getOperationName(); }]; + let assemblyFormat = "attr-dict"; +} + +def TaskwaitOp : OpenMP_Op<"taskwait"> { + let summary = "taskwait construct"; + let description = [{ + The taskwait construct specifies a wait on the completion of child tasks + of the current task. + }]; + + let assemblyFormat = "attr-dict"; +} + +def TaskyieldOp : OpenMP_Op<"taskyield"> { + let summary = "taskyield construct"; + let description = [{ + The taskyield construct specifies that the current task can be suspended + in favor of execution of a different task. + }]; + + let assemblyFormat = "attr-dict"; } #endif // OPENMP_OPS diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h index f6b299dbd1ae..7ba4a7d21adf 100644 --- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h +++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h @@ -83,6 +83,8 @@ protected: virtual LogicalResult convertOperation(Operation &op, llvm::IRBuilder<> &builder); + virtual LogicalResult convertOmpOperation(Operation &op, + llvm::IRBuilder<> &builder); static std::unique_ptr<llvm::Module> prepareLLVMModule(Operation *m); /// A helper to look up remapped operands in the value remapping table. diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp index 84b64b910eef..125b88f7789a 100644 --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -14,6 +14,7 @@ #include "mlir/Target/LLVMIR/ModuleTranslation.h" #include "DebugTranslation.h" +#include "mlir/ADT/TypeSwitch.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/OpenMP/OpenMPDialect.h" #include "mlir/IR/Attributes.h" @@ -306,6 +307,34 @@ ModuleTranslation::ModuleTranslation(Operation *module, } ModuleTranslation::~ModuleTranslation() {} +/// Given an OpenMP MLIR operation, create the corresponding LLVM IR +/// (including OpenMP runtime calls). +LogicalResult +ModuleTranslation::convertOmpOperation(Operation &opInst, + llvm::IRBuilder<> &builder) { + if (!ompBuilder) { + ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule); + ompBuilder->initialize(); + } + return mlir::TypeSwitch<Operation *, LogicalResult>(&opInst) + .Case([&](omp::BarrierOp) { + ompBuilder->CreateBarrier(builder.saveIP(), llvm::omp::OMPD_barrier); + return success(); + }) + .Case([&](omp::TaskwaitOp) { + ompBuilder->CreateTaskwait(builder.saveIP()); + return success(); + }) + .Case([&](omp::TaskyieldOp) { + ompBuilder->CreateTaskyield(builder.saveIP()); + return success(); + }) + .Default([&](Operation *inst) { + return inst->emitError("unsupported OpenMP operation: ") + << inst->getName(); + }); +} + /// Given a single MLIR operation, create the corresponding LLVM IR operation /// using the `builder`. LLVM IR Builder does not have a generic interface so /// this has to be a long chain of `if`s calling different functions with a @@ -415,17 +444,7 @@ LogicalResult ModuleTranslation::convertOperation(Operation &opInst, } if (opInst.getDialect() == ompDialect) { - if (!ompBuilder) { - ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule); - ompBuilder->initialize(); - } - - if (isa<omp::BarrierOp>(opInst)) { - ompBuilder->CreateBarrier(builder.saveIP(), llvm::omp::OMPD_barrier); - return success(); - } - return opInst.emitError("unsupported OpenMP operation: ") - << opInst.getName(); + return convertOmpOperation(opInst, builder); } return opInst.emitError("unsupported or non-LLVM operation: ") diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 08bc5d796bb1..e26ad4236266 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -5,3 +5,15 @@ func @omp_barrier() -> () { omp.barrier return } + +func @omp_taskwait() -> () { + // CHECK: omp.taskwait + omp.taskwait + return +} + +func @omp_taskyield() -> () { + // CHECK: omp.taskyield + omp.taskyield + return +} diff --git a/mlir/test/Target/openmp-llvm.mlir b/mlir/test/Target/openmp-llvm.mlir index fd9c6bc583c2..0bd1601d7868 100644 --- a/mlir/test/Target/openmp-llvm.mlir +++ b/mlir/test/Target/openmp-llvm.mlir @@ -1,10 +1,19 @@ // RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s // CHECK-LABEL: define void @empty() -// CHECK: [[OMP_THREAD:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}}) -// CHECK-NEXT: call void @__kmpc_barrier(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD]]) -// CHECK-NEXT: ret void llvm.func @empty() { + // CHECK: [[OMP_THREAD:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}}) + // CHECK-NEXT: call void @__kmpc_barrier(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD]]) omp.barrier + + // CHECK: [[OMP_THREAD1:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}}) + // CHECK-NEXT: [[RET_VAL:%.*]] = call i32 @__kmpc_omp_taskwait(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD1]]) + omp.taskwait + + // CHECK: [[OMP_THREAD2:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}}) + // CHECK-NEXT: [[RET_VAL:%.*]] = call i32 @__kmpc_omp_taskyield(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD2]], i32 0) + omp.taskyield + +// CHECK-NEXT: ret void llvm.return }