Skip to content

Commit 839275d

Browse files
authored
[MLIR][OpenMP] Add missing OpenMP to LLVM conversion patterns (#104440)
This patch adds conversion patterns to LLVM for the following OpenMP dialect operations: - `omp.critical.declare` - `omp.cancel` - `omp.cancellation_point` - `omp.distribute` - `omp.teams` - `omp.ordered` - `omp.taskloop` Also, arbitrary sorting of operations when passing them as template argument lists when configuring that pass is replaced by alphabetical sorting.
1 parent 6fd4608 commit 839275d

File tree

2 files changed

+132
-31
lines changed

2 files changed

+132
-31
lines changed

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

Lines changed: 38 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -223,22 +223,21 @@ void MultiRegionOpConversion<omp::PrivateClauseOp>::forwardOpAttrs(
223223
void mlir::configureOpenMPToLLVMConversionLegality(
224224
ConversionTarget &target, LLVMTypeConverter &typeConverter) {
225225
target.addDynamicallyLegalOp<
226-
mlir::omp::AtomicReadOp, mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
227-
mlir::omp::ThreadprivateOp, mlir::omp::YieldOp,
228-
mlir::omp::TargetEnterDataOp, mlir::omp::TargetExitDataOp,
229-
mlir::omp::TargetUpdateOp, mlir::omp::MapBoundsOp, mlir::omp::MapInfoOp>(
230-
[&](Operation *op) {
231-
return typeConverter.isLegal(op->getOperandTypes()) &&
232-
typeConverter.isLegal(op->getResultTypes());
233-
});
226+
omp::AtomicReadOp, omp::AtomicWriteOp, omp::CancellationPointOp,
227+
omp::CancelOp, omp::CriticalDeclareOp, omp::FlushOp, omp::MapBoundsOp,
228+
omp::MapInfoOp, omp::OrderedOp, omp::TargetEnterDataOp,
229+
omp::TargetExitDataOp, omp::TargetUpdateOp, omp::ThreadprivateOp,
230+
omp::YieldOp>([&](Operation *op) {
231+
return typeConverter.isLegal(op->getOperandTypes()) &&
232+
typeConverter.isLegal(op->getResultTypes());
233+
});
234234
target.addDynamicallyLegalOp<
235-
mlir::omp::AtomicUpdateOp, mlir::omp::CriticalOp, mlir::omp::TargetOp,
236-
mlir::omp::TargetDataOp, mlir::omp::LoopNestOp,
237-
mlir::omp::OrderedRegionOp, mlir::omp::ParallelOp, mlir::omp::WsloopOp,
238-
mlir::omp::SimdOp, mlir::omp::MasterOp, mlir::omp::SectionOp,
239-
mlir::omp::SectionsOp, mlir::omp::SingleOp, mlir::omp::TaskgroupOp,
240-
mlir::omp::TaskOp, mlir::omp::DeclareReductionOp,
241-
mlir::omp::PrivateClauseOp>([&](Operation *op) {
235+
omp::AtomicUpdateOp, omp::CriticalOp, omp::DeclareReductionOp,
236+
omp::DistributeOp, omp::LoopNestOp, omp::MasterOp, omp::OrderedRegionOp,
237+
omp::ParallelOp, omp::PrivateClauseOp, omp::SectionOp, omp::SectionsOp,
238+
omp::SimdOp, omp::SingleOp, omp::TargetDataOp, omp::TargetOp,
239+
omp::TaskgroupOp, omp::TaskloopOp, omp::TaskOp, omp::TeamsOp,
240+
omp::WsloopOp>([&](Operation *op) {
242241
return std::all_of(op->getRegions().begin(), op->getRegions().end(),
243242
[&](Region &region) {
244243
return typeConverter.isLegal(&region);
@@ -260,23 +259,31 @@ void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
260259
AtomicReadOpConversion, MapInfoOpConversion,
261260
MultiRegionOpConversion<omp::DeclareReductionOp>,
262261
MultiRegionOpConversion<omp::PrivateClauseOp>,
263-
RegionOpConversion<omp::CriticalOp>, RegionOpConversion<omp::LoopNestOp>,
264-
RegionOpConversion<omp::MasterOp>,
265-
RegionOpConversion<omp::OrderedRegionOp>,
266-
RegionOpConversion<omp::ParallelOp>, RegionOpConversion<omp::WsloopOp>,
267-
RegionOpConversion<omp::SectionsOp>, RegionOpConversion<omp::SectionOp>,
268-
RegionOpConversion<omp::SimdOp>, RegionOpConversion<omp::SingleOp>,
269-
RegionOpConversion<omp::TaskgroupOp>, RegionOpConversion<omp::TaskOp>,
270-
RegionOpConversion<omp::TargetDataOp>, RegionOpConversion<omp::TargetOp>,
271-
RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
272-
RegionOpWithVarOperandsConversion<omp::AtomicUpdateOp>,
273-
RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
274-
RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
275-
RegionLessOpConversion<omp::YieldOp>,
262+
RegionLessOpConversion<omp::CancellationPointOp>,
263+
RegionLessOpConversion<omp::CancelOp>,
264+
RegionLessOpConversion<omp::CriticalDeclareOp>,
265+
RegionLessOpConversion<omp::OrderedOp>,
276266
RegionLessOpConversion<omp::TargetEnterDataOp>,
277267
RegionLessOpConversion<omp::TargetExitDataOp>,
278268
RegionLessOpConversion<omp::TargetUpdateOp>,
279-
RegionLessOpWithVarOperandsConversion<omp::MapBoundsOp>>(converter);
269+
RegionLessOpConversion<omp::YieldOp>,
270+
RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
271+
RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
272+
RegionLessOpWithVarOperandsConversion<omp::MapBoundsOp>,
273+
RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
274+
RegionOpConversion<omp::AtomicCaptureOp>,
275+
RegionOpConversion<omp::CriticalOp>,
276+
RegionOpConversion<omp::DistributeOp>,
277+
RegionOpConversion<omp::LoopNestOp>, RegionOpConversion<omp::MaskedOp>,
278+
RegionOpConversion<omp::MasterOp>,
279+
RegionOpConversion<omp::OrderedRegionOp>,
280+
RegionOpConversion<omp::ParallelOp>, RegionOpConversion<omp::SectionOp>,
281+
RegionOpConversion<omp::SectionsOp>, RegionOpConversion<omp::SimdOp>,
282+
RegionOpConversion<omp::SingleOp>, RegionOpConversion<omp::TargetDataOp>,
283+
RegionOpConversion<omp::TargetOp>, RegionOpConversion<omp::TaskgroupOp>,
284+
RegionOpConversion<omp::TaskloopOp>, RegionOpConversion<omp::TaskOp>,
285+
RegionOpConversion<omp::TeamsOp>, RegionOpConversion<omp::WsloopOp>,
286+
RegionOpWithVarOperandsConversion<omp::AtomicUpdateOp>>(converter);
280287
}
281288

282289
namespace {
@@ -301,8 +308,8 @@ void ConvertOpenMPToLLVMPass::runOnOperation() {
301308
populateOpenMPToLLVMConversionPatterns(converter, patterns);
302309

303310
LLVMConversionTarget target(getContext());
304-
target.addLegalOp<omp::TerminatorOp, omp::TaskyieldOp, omp::FlushOp,
305-
omp::BarrierOp, omp::TaskwaitOp>();
311+
target.addLegalOp<omp::BarrierOp, omp::FlushOp, omp::TaskwaitOp,
312+
omp::TaskyieldOp, omp::TerminatorOp>();
306313
configureOpenMPToLLVMConversionLegality(target, converter);
307314
if (failed(applyPartialConversion(module, target, std::move(patterns))))
308315
signalPassFailure();

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,20 @@ func.func @critical_block_arg() {
1818

1919
// -----
2020

21+
// CHECK: omp.critical.declare @[[MUTEX:.*]] hint(contended, speculative)
22+
omp.critical.declare @mutex hint(contended, speculative)
23+
24+
// CHECK: llvm.func @critical_declare
25+
func.func @critical_declare() {
26+
// CHECK: omp.critical(@[[MUTEX]])
27+
omp.critical(@mutex) {
28+
omp.terminator
29+
}
30+
return
31+
}
32+
33+
// -----
34+
2135
// CHECK-LABEL: llvm.func @master_block_arg
2236
func.func @master_block_arg() {
2337
// CHECK: omp.master
@@ -523,3 +537,83 @@ omp.private {type = firstprivate} @y.privatizer : index alloc {
523537
// CHECK: omp.yield(%arg0 : i64)
524538
omp.yield(%arg0 : index)
525539
}
540+
541+
// -----
542+
543+
// CHECK-LABEL: llvm.func @omp_cancel_cancellation_point()
544+
func.func @omp_cancel_cancellation_point() -> () {
545+
omp.parallel {
546+
// CHECK: omp.cancel cancellation_construct_type(parallel)
547+
omp.cancel cancellation_construct_type(parallel)
548+
// CHECK: omp.cancellation_point cancellation_construct_type(parallel)
549+
omp.cancellation_point cancellation_construct_type(parallel)
550+
omp.terminator
551+
}
552+
return
553+
}
554+
555+
// -----
556+
557+
// CHECK-LABEL: llvm.func @omp_distribute(
558+
// CHECK-SAME: %[[ARG0:.*]]: i64)
559+
func.func @omp_distribute(%arg0 : index) -> () {
560+
// CHECK: omp.distribute dist_schedule_static dist_schedule_chunk_size(%[[ARG0]] : i64) {
561+
omp.distribute dist_schedule_static dist_schedule_chunk_size(%arg0 : index) {
562+
omp.loop_nest (%iv) : index = (%arg0) to (%arg0) step (%arg0) {
563+
omp.yield
564+
}
565+
omp.terminator
566+
}
567+
return
568+
}
569+
570+
// -----
571+
572+
// CHECK-LABEL: llvm.func @omp_teams(
573+
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr, %[[ARG1:.*]]: !llvm.ptr, %[[ARG2:.*]]: i64)
574+
func.func @omp_teams(%arg0 : memref<i32>) -> () {
575+
// CHECK: omp.teams allocate(%{{.*}} : !llvm.struct<(ptr, ptr, i64)> -> %{{.*}} : !llvm.struct<(ptr, ptr, i64)>)
576+
omp.teams allocate(%arg0 : memref<i32> -> %arg0 : memref<i32>) {
577+
omp.terminator
578+
}
579+
return
580+
}
581+
582+
// -----
583+
584+
// CHECK-LABEL: llvm.func @omp_ordered(
585+
// CHECK-SAME: %[[ARG0:.*]]: i64)
586+
func.func @omp_ordered(%arg0 : index) -> () {
587+
omp.wsloop ordered(1) {
588+
omp.loop_nest (%iv) : index = (%arg0) to (%arg0) step (%arg0) {
589+
// CHECK: omp.ordered depend_vec(%[[ARG0]] : i64) {doacross_num_loops = 1 : i64}
590+
omp.ordered depend_vec(%arg0 : index) {doacross_num_loops = 1 : i64}
591+
omp.yield
592+
}
593+
omp.terminator
594+
}
595+
return
596+
}
597+
598+
// -----
599+
600+
// CHECK-LABEL: @omp_taskloop(
601+
// CHECK-SAME: %[[ARG0:.*]]: i64, %[[ARG1:.*]]: !llvm.ptr, %[[ARG2:.*]]: !llvm.ptr, %[[ARG3:.*]]: i64)
602+
func.func @omp_taskloop(%arg0: index, %arg1 : memref<i32>) {
603+
// CHECK: omp.parallel {
604+
omp.parallel {
605+
// CHECK: omp.taskloop allocate(%{{.*}} : !llvm.struct<(ptr, ptr, i64)> -> %{{.*}} : !llvm.struct<(ptr, ptr, i64)>) {
606+
omp.taskloop allocate(%arg1 : memref<i32> -> %arg1 : memref<i32>) {
607+
// CHECK: omp.loop_nest (%[[IV:.*]]) : i64 = (%[[ARG0]]) to (%[[ARG0]]) step (%[[ARG0]]) {
608+
omp.loop_nest (%iv) : index = (%arg0) to (%arg0) step (%arg0) {
609+
// CHECK-DAG: %[[CAST_IV:.*]] = builtin.unrealized_conversion_cast %[[IV]] : i64 to index
610+
// CHECK: "test.payload"(%[[CAST_IV]]) : (index) -> ()
611+
"test.payload"(%iv) : (index) -> ()
612+
omp.yield
613+
}
614+
omp.terminator
615+
}
616+
omp.terminator
617+
}
618+
return
619+
}

0 commit comments

Comments
 (0)