[Mlir-commits] [mlir] [MLIR][OpenMP] Skip host omp ops when compiling for the target device (PR #85239)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Mar 14 08:07:16 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-openmp
Author: Jan Leyonberg (jsjodin)
<details>
<summary>Changes</summary>
This patch separates the lowering dispatch for host and target devices. For the target device, if the current operation is not a top-level operation (e.g. omp.target) or is inside a target device code region it will be ignored, since it belongs to the host code.
This is an alternative approach to #<!-- -->84611, the new test in this PR was taken from there.
---
Full diff: https://github.com/llvm/llvm-project/pull/85239.diff
6 Files Affected:
- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+182-111)
- (modified) mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir (+3-3)
- (modified) mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir (+1-1)
- (modified) mlir/test/Target/LLVMIR/omptarget-wsloop-collapsed.mlir (+1-1)
- (modified) mlir/test/Target/LLVMIR/omptarget-wsloop.mlir (+2-2)
- (added) mlir/test/Target/LLVMIR/openmp-tark-target-device.mlir (+27)
``````````diff
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index bef227f2c583ed..9ad0b99078005b 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2922,6 +2922,178 @@ convertDeclareTargetAttr(Operation *op, mlir::omp::DeclareTargetAttr attribute,
return success();
}
+static bool isInternalTargetDeviceOp(Operation *op) {
+ // Assumes no reverse offloading
+ if (op->getParentOfType<omp::TargetOp>())
+ return true;
+
+ auto parentFn = op->getParentOfType<LLVM::LLVMFuncOp>();
+ if (auto declareTargetIface =
+ llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
+ parentFn.getOperation()))
+ if (declareTargetIface.isDeclareTarget() &&
+ declareTargetIface.getDeclareTargetDeviceType() !=
+ mlir::omp::DeclareTargetDeviceType::host)
+ return true;
+
+ return false;
+}
+
+/// Given an OpenMP MLIR operation, create the corresponding LLVM IR
+/// (including OpenMP runtime calls).
+static LogicalResult
+convertCommonOperation(Operation *op, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+ return llvm::TypeSwitch<Operation *, LogicalResult>(op)
+ .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();
+ })
+ .Case([&](omp::FlushOp) {
+ // No support in Openmp runtime function (__kmpc_flush) to accept
+ // the argument list.
+ // OpenMP standard states the following:
+ // "An implementation may implement a flush with a list by ignoring
+ // the list, and treating it the same as a flush without a list."
+ //
+ // The argument list is discarded so that, flush with a list is treated
+ // same as a flush without a list.
+ ompBuilder->createFlush(builder.saveIP());
+ return success();
+ })
+ .Case([&](omp::ParallelOp op) {
+ return convertOmpParallel(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::ReductionOp reductionOp) {
+ return convertOmpReductionOp(reductionOp, builder, moduleTranslation);
+ })
+ .Case([&](omp::MasterOp) {
+ return convertOmpMaster(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::CriticalOp) {
+ return convertOmpCritical(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::OrderedRegionOp) {
+ return convertOmpOrderedRegion(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::OrderedOp) {
+ return convertOmpOrdered(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::WsLoopOp) {
+ return convertOmpWsLoop(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::SimdLoopOp) {
+ return convertOmpSimdLoop(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::AtomicReadOp) {
+ return convertOmpAtomicRead(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::AtomicWriteOp) {
+ return convertOmpAtomicWrite(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::AtomicUpdateOp op) {
+ return convertOmpAtomicUpdate(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::AtomicCaptureOp op) {
+ return convertOmpAtomicCapture(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::SectionsOp) {
+ return convertOmpSections(*op, builder, moduleTranslation);
+ })
+ .Case([&](omp::SingleOp op) {
+ return convertOmpSingle(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::TeamsOp op) {
+ return convertOmpTeams(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::TaskOp op) {
+ return convertOmpTaskOp(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::TaskGroupOp op) {
+ return convertOmpTaskgroupOp(op, builder, moduleTranslation);
+ })
+ .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
+ omp::CriticalDeclareOp>([](auto op) {
+ // `yield` and `terminator` can be just omitted. The block structure
+ // was created in the region that handles their parent operation.
+ // `reduction.declare` will be used by reductions and is not
+ // converted directly, skip it.
+ // `critical.declare` is only used to declare names of critical
+ // sections which will be used by `critical` ops and hence can be
+ // ignored for lowering. The OpenMP IRBuilder will create unique
+ // name for critical section names.
+ return success();
+ })
+ .Case([&](omp::ThreadprivateOp) {
+ return convertOmpThreadprivate(*op, builder, moduleTranslation);
+ })
+ .Case<omp::DataOp, omp::EnterDataOp, omp::ExitDataOp, omp::UpdateDataOp>(
+ [&](auto op) {
+ return convertOmpTargetData(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::TargetOp) {
+ return convertOmpTarget(*op, builder, moduleTranslation);
+ })
+ .Case<omp::MapInfoOp, omp::DataBoundsOp, omp::PrivateClauseOp>(
+ [&](auto op) {
+ // No-op, should be handled by relevant owning operations e.g.
+ // TargetOp, EnterDataOp, ExitDataOp, DataOp etc. and then
+ // discarded
+ return success();
+ })
+ .Default([&](Operation *inst) {
+ return inst->emitError("unsupported OpenMP operation: ")
+ << inst->getName();
+ });
+}
+
+static LogicalResult
+convertInternalTargetOp(Operation *op, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ return convertCommonOperation(op, builder, moduleTranslation);
+}
+
+static LogicalResult
+convertTopLevelTargetOp(Operation *op, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ return llvm::TypeSwitch<Operation *, LogicalResult>(op)
+ .Case<omp::DataOp, omp::EnterDataOp, omp::ExitDataOp, omp::UpdateDataOp>(
+ [&](auto op) {
+ return convertOmpTargetData(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::TargetOp) {
+ return convertOmpTarget(*op, builder, moduleTranslation);
+ })
+ // Skip omp ops that are not legal top level ops for the target device
+ .Case<omp::BarrierOp, omp::TaskwaitOp, omp::TaskyieldOp, omp::FlushOp,
+ omp::ParallelOp, omp::ReductionOp, omp::MasterOp, omp::CriticalOp,
+ omp::OrderedRegionOp, omp::OrderedOp, omp::WsLoopOp,
+ omp::SimdLoopOp, omp::AtomicReadOp, omp::AtomicWriteOp,
+ omp::AtomicUpdateOp, omp::AtomicCaptureOp, omp::SectionsOp,
+ omp::SingleOp, omp::TeamsOp, omp::TaskOp, omp::TaskGroupOp,
+ omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
+ omp::ThreadprivateOp, omp::DistributeOp, omp::MapInfoOp,
+ omp::DataBoundsOp, omp::CriticalDeclareOp>(
+ [&](auto op) {
+ return success();
+ })
+ .Default([&](Operation *inst) {
+ return inst->emitError("unsupported OpenMP operation: ")
+ << inst->getName();
+ });
+}
+
namespace {
/// Implementation of the dialect interface that converts operations belonging
@@ -2937,8 +3109,8 @@ class OpenMPDialectLLVMIRTranslationInterface
convertOperation(Operation *op, llvm::IRBuilderBase &builder,
LLVM::ModuleTranslation &moduleTranslation) const final;
- /// Given an OpenMP MLIR attribute, create the corresponding LLVM-IR, runtime
- /// calls, or operation amendments
+ /// Given an OpenMP MLIR attribute, create the corresponding LLVM-IR,
+ /// runtime calls, or operation amendments
LogicalResult
amendOperation(Operation *op, ArrayRef<llvm::Instruction *> instructions,
NamedAttribute attribute,
@@ -3043,116 +3215,15 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
LLVM::ModuleTranslation &moduleTranslation) const {
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+ if (ompBuilder->Config.isTargetDevice()) {
+ if (isInternalTargetDeviceOp(op)) {
+ return convertInternalTargetOp(op, builder, moduleTranslation);
+ } else {
+ return convertTopLevelTargetOp(op, builder, moduleTranslation);
+ }
+ }
- return llvm::TypeSwitch<Operation *, LogicalResult>(op)
- .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();
- })
- .Case([&](omp::FlushOp) {
- // No support in Openmp runtime function (__kmpc_flush) to accept
- // the argument list.
- // OpenMP standard states the following:
- // "An implementation may implement a flush with a list by ignoring
- // the list, and treating it the same as a flush without a list."
- //
- // The argument list is discarded so that, flush with a list is treated
- // same as a flush without a list.
- ompBuilder->createFlush(builder.saveIP());
- return success();
- })
- .Case([&](omp::ParallelOp op) {
- return convertOmpParallel(op, builder, moduleTranslation);
- })
- .Case([&](omp::ReductionOp reductionOp) {
- return convertOmpReductionOp(reductionOp, builder, moduleTranslation);
- })
- .Case([&](omp::MasterOp) {
- return convertOmpMaster(*op, builder, moduleTranslation);
- })
- .Case([&](omp::CriticalOp) {
- return convertOmpCritical(*op, builder, moduleTranslation);
- })
- .Case([&](omp::OrderedRegionOp) {
- return convertOmpOrderedRegion(*op, builder, moduleTranslation);
- })
- .Case([&](omp::OrderedOp) {
- return convertOmpOrdered(*op, builder, moduleTranslation);
- })
- .Case([&](omp::WsLoopOp) {
- return convertOmpWsLoop(*op, builder, moduleTranslation);
- })
- .Case([&](omp::SimdLoopOp) {
- return convertOmpSimdLoop(*op, builder, moduleTranslation);
- })
- .Case([&](omp::AtomicReadOp) {
- return convertOmpAtomicRead(*op, builder, moduleTranslation);
- })
- .Case([&](omp::AtomicWriteOp) {
- return convertOmpAtomicWrite(*op, builder, moduleTranslation);
- })
- .Case([&](omp::AtomicUpdateOp op) {
- return convertOmpAtomicUpdate(op, builder, moduleTranslation);
- })
- .Case([&](omp::AtomicCaptureOp op) {
- return convertOmpAtomicCapture(op, builder, moduleTranslation);
- })
- .Case([&](omp::SectionsOp) {
- return convertOmpSections(*op, builder, moduleTranslation);
- })
- .Case([&](omp::SingleOp op) {
- return convertOmpSingle(op, builder, moduleTranslation);
- })
- .Case([&](omp::TeamsOp op) {
- return convertOmpTeams(op, builder, moduleTranslation);
- })
- .Case([&](omp::TaskOp op) {
- return convertOmpTaskOp(op, builder, moduleTranslation);
- })
- .Case([&](omp::TaskGroupOp op) {
- return convertOmpTaskgroupOp(op, builder, moduleTranslation);
- })
- .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
- omp::CriticalDeclareOp>([](auto op) {
- // `yield` and `terminator` can be just omitted. The block structure
- // was created in the region that handles their parent operation.
- // `reduction.declare` will be used by reductions and is not
- // converted directly, skip it.
- // `critical.declare` is only used to declare names of critical
- // sections which will be used by `critical` ops and hence can be
- // ignored for lowering. The OpenMP IRBuilder will create unique
- // name for critical section names.
- return success();
- })
- .Case([&](omp::ThreadprivateOp) {
- return convertOmpThreadprivate(*op, builder, moduleTranslation);
- })
- .Case<omp::DataOp, omp::EnterDataOp, omp::ExitDataOp, omp::UpdateDataOp>(
- [&](auto op) {
- return convertOmpTargetData(op, builder, moduleTranslation);
- })
- .Case([&](omp::TargetOp) {
- return convertOmpTarget(*op, builder, moduleTranslation);
- })
- .Case<omp::MapInfoOp, omp::DataBoundsOp, omp::PrivateClauseOp>(
- [&](auto op) {
- // No-op, should be handled by relevant owning operations e.g.
- // TargetOp, EnterDataOp, ExitDataOp, DataOp etc. and then
- // discarded
- return success();
- })
- .Default([&](Operation *inst) {
- return inst->emitError("unsupported OpenMP operation: ")
- << inst->getName();
- });
+ return convertCommonOperation(op, builder, moduleTranslation);
}
void mlir::registerOpenMPDialectTranslation(DialectRegistry ®istry) {
diff --git a/mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir b/mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir
index 8ab50f05f07167..b0fe642238f14f 100644
--- a/mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir
@@ -4,10 +4,10 @@
// for nested omp do loop inside omp target region
module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>, llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8", llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true } {
- llvm.func @target_parallel_wsloop(%arg0: !llvm.ptr) attributes {
+ llvm.func @target_parallel_wsloop(%arg0: !llvm.ptr) attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>,
target_cpu = "gfx90a",
- target_features = #llvm.target_features<["+gfx9-insts", "+wavefrontsize64"]>
- } {
+ target_features = #llvm.target_features<["+gfx9-insts", "+wavefrontsize64"]>}
+ {
omp.parallel {
%loop_ub = llvm.mlir.constant(9 : i32) : i32
%loop_lb = llvm.mlir.constant(0 : i32) : i32
diff --git a/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir b/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir
index 96cced7a1d584b..c5f89eb2c3274c 100644
--- a/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir
@@ -5,7 +5,7 @@
module attributes {omp.is_target_device = true} {
llvm.func @foo(i32)
- llvm.func @omp_target_teams_shared_simple(%arg0 : i32) {
+ llvm.func @omp_target_teams_shared_simple(%arg0 : i32) attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} {
omp.teams {
llvm.call @foo(%arg0) : (i32) -> ()
omp.terminator
diff --git a/mlir/test/Target/LLVMIR/omptarget-wsloop-collapsed.mlir b/mlir/test/Target/LLVMIR/omptarget-wsloop-collapsed.mlir
index e246c551886cfa..0d77423abcb4f1 100644
--- a/mlir/test/Target/LLVMIR/omptarget-wsloop-collapsed.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-wsloop-collapsed.mlir
@@ -4,7 +4,7 @@
// for nested omp do loop with collapse clause inside omp target region
module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>, llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8", llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true } {
- llvm.func @target_collapsed_wsloop(%arg0: !llvm.ptr) {
+ llvm.func @target_collapsed_wsloop(%arg0: !llvm.ptr) attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} {
%loop_ub = llvm.mlir.constant(99 : i32) : i32
%loop_lb = llvm.mlir.constant(0 : i32) : i32
%loop_step = llvm.mlir.constant(1 : index) : i32
diff --git a/mlir/test/Target/LLVMIR/omptarget-wsloop.mlir b/mlir/test/Target/LLVMIR/omptarget-wsloop.mlir
index 220eb85b3483ec..0f3f503dfa5377 100644
--- a/mlir/test/Target/LLVMIR/omptarget-wsloop.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-wsloop.mlir
@@ -4,7 +4,7 @@
// for nested omp do loop inside omp target region
module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>, llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8", llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true } {
- llvm.func @target_wsloop(%arg0: !llvm.ptr ){
+ llvm.func @target_wsloop(%arg0: !llvm.ptr ) attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} {
%loop_ub = llvm.mlir.constant(9 : i32) : i32
%loop_lb = llvm.mlir.constant(0 : i32) : i32
%loop_step = llvm.mlir.constant(1 : i32) : i32
@@ -16,7 +16,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
llvm.return
}
- llvm.func @target_empty_wsloop(){
+ llvm.func @target_empty_wsloop() attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} {
%loop_ub = llvm.mlir.constant(9 : i32) : i32
%loop_lb = llvm.mlir.constant(0 : i32) : i32
%loop_step = llvm.mlir.constant(1 : i32) : i32
diff --git a/mlir/test/Target/LLVMIR/openmp-tark-target-device.mlir b/mlir/test/Target/LLVMIR/openmp-tark-target-device.mlir
new file mode 100644
index 00000000000000..c167604bcd12c0
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-tark-target-device.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This tests the fix for https://github.com/llvm/llvm-project/issues/84606
+// We are only interested in ensuring that the -mlir-to-llmvir pass doesn't crash.
+// CHECK: {{.*}} = add i32 {{.*}}, 5
+module attributes {omp.is_target_device = true } {
+ llvm.func @_QQmain() attributes {fir.bindc_name = "main", omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>} {
+ %0 = llvm.mlir.constant(0 : i32) : i32
+ %1 = llvm.mlir.constant(1 : i64) : i64
+ %2 = llvm.alloca %1 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr<5>
+ %3 = llvm.addrspacecast %2 : !llvm.ptr<5> to !llvm.ptr
+ omp.task {
+ llvm.store %0, %3 : i32, !llvm.ptr
+ omp.terminator
+ }
+ %4 = omp.map_info var_ptr(%3 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = "a"}
+ omp.target map_entries(%4 -> %arg0 : !llvm.ptr) {
+ ^bb0(%arg0: !llvm.ptr):
+ %5 = llvm.mlir.constant(5 : i32) : i32
+ %6 = llvm.load %arg0 : !llvm.ptr -> i32
+ %7 = llvm.add %6, %5 : i32
+ llvm.store %7, %arg0 : i32, !llvm.ptr
+ omp.terminator
+ }
+ llvm.return
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/85239
More information about the Mlir-commits
mailing list