Add OpenMPToLLVM conversion pattern for taskgroup

Fixes part of the issue in https://github.com/llvm/llvm-project/issues/62013

Reviewed By: psoni2628

Differential Revision: https://reviews.llvm.org/D151206

diff --git a/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir b/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir
index 832f2a3..bf03c24 100644 (file)
--- a/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir
+++ b/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir
@@ -436,6 +436,51 @@ func.func private @_QPwork()
 
 // -----
 
+
+func.func @_QQmain() {
+  %c0 = arith.constant 0 : index
+  %c5 = arith.constant 5 : index
+  %c1 = arith.constant 1 : index
+  %0 = fir.alloca i32
+  omp.taskgroup   {
+    %1 = fir.convert %c1 : (index) -> i32
+    cf.br ^bb1(%1, %c5 : i32, index)
+  ^bb1(%2: i32, %3: index):  // 2 preds: ^bb0, ^bb2
+    %4 = arith.cmpi sgt, %3, %c0 : index
+    cf.cond_br %4, ^bb2, ^bb3
+  ^bb2:  // pred: ^bb1
+    fir.store %2 to %0 : !fir.ref<i32>
+    omp.task   {
+      fir.call @_QFPdo_work(%0) : (!fir.ref<i32>) -> ()
+      omp.terminator
+    }
+    %5 = fir.load %0 : !fir.ref<i32>
+    %6 = arith.addi %5, %1 : i32
+    %7 = arith.subi %3, %c1 : index
+    cf.br ^bb1(%6, %7 : i32, index)
+  ^bb3:  // pred: ^bb1
+    fir.store %2 to %0 : !fir.ref<i32>
+    omp.terminator
+  }
+  return
+}
+func.func private @_QFPdo_work(!fir.ref<i32>)
+
+// CHECK-LABEL: llvm.func @_QQmain
+// CHECK:          omp.taskgroup   {
+// CHECK:            omp.task   {
+// CHECK:              llvm.call @_QFPdo_work({{.*}}) : (!llvm.ptr<i32>) -> ()
+// CHECK:              omp.terminator
+// CHECK:            }
+// CHECK:            omp.terminator
+// CHECK:          }
+// CHECK:          llvm.return
+// CHECK:        }
+// CHECK:        llvm.func @_QFPdo_work(!llvm.ptr<i32>)
+// CHECK:      }
+
+// -----
+
 func.func @_QPs() {
   %0 = fir.address_of(@_QFsEc) : !fir.ref<i32>
   omp.atomic.update   %0 : !fir.ref<i32> {

diff --git a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
index 34ec00f..a5fdc4e 100644 (file)
--- a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
+++ b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
@@ -183,12 +183,12 @@ void mlir::configureOpenMPToLLVMConversionLegality(
       mlir::omp::AtomicUpdateOp, mlir::omp::CriticalOp, mlir::omp::TargetOp,
       mlir::omp::DataOp, mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
       mlir::omp::SimdLoopOp, mlir::omp::MasterOp, mlir::omp::SectionOp,
-      mlir::omp::SectionsOp, mlir::omp::SingleOp, mlir::omp::TaskOp>(
-      [&](Operation *op) {
-        return typeConverter.isLegal(&op->getRegion(0)) &&
-               typeConverter.isLegal(op->getOperandTypes()) &&
-               typeConverter.isLegal(op->getResultTypes());
-      });
+      mlir::omp::SectionsOp, mlir::omp::SingleOp, mlir::omp::TaskGroupOp,
+      mlir::omp::TaskOp>([&](Operation *op) {
+    return typeConverter.isLegal(&op->getRegion(0)) &&
+           typeConverter.isLegal(op->getOperandTypes()) &&
+           typeConverter.isLegal(op->getResultTypes());
+  });
   target.addDynamicallyLegalOp<mlir::omp::AtomicReadOp,
                                mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
                                mlir::omp::ThreadprivateOp, mlir::omp::YieldOp,
@@ -218,8 +218,9 @@ void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
       ReductionOpConversion, RegionOpConversion<omp::ParallelOp>,
       RegionOpConversion<omp::WsLoopOp>, RegionOpConversion<omp::SectionsOp>,
       RegionOpConversion<omp::SectionOp>, RegionOpConversion<omp::SimdLoopOp>,
-      RegionOpConversion<omp::SingleOp>, RegionOpConversion<omp::TaskOp>,
-      RegionOpConversion<omp::DataOp>, RegionOpConversion<omp::TargetOp>,
+      RegionOpConversion<omp::SingleOp>, RegionOpConversion<omp::TaskGroupOp>,
+      RegionOpConversion<omp::TaskOp>, RegionOpConversion<omp::DataOp>,
+      RegionOpConversion<omp::TargetOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionOpWithVarOperandsConversion<omp::AtomicUpdateOp>,

diff --git a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
index 6cbe02c..ab91729 100644 (file)
--- a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
+++ b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
@@ -343,3 +343,44 @@ llvm.func @_QPsimple_reduction(%arg0: !llvm.ptr<array<100 x i32>> {fir.bindc_nam
   }
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL:  @_QQmain
+llvm.func @_QQmain() {
+  %0 = llvm.mlir.constant(0 : index) : i64
+  %1 = llvm.mlir.constant(5 : index) : i64
+  %2 = llvm.mlir.constant(1 : index) : i64
+  %3 = llvm.mlir.constant(1 : i64) : i64
+  %4 = llvm.alloca %3 x i32 : (i64) -> !llvm.ptr<i32>
+// CHECK: omp.taskgroup
+  omp.taskgroup   {
+    %5 = llvm.trunc %2 : i64 to i32
+    llvm.br ^bb1(%5, %1 : i32, i64)
+  ^bb1(%6: i32, %7: i64):  // 2 preds: ^bb0, ^bb2
+    %8 = llvm.icmp "sgt" %7, %0 : i64
+    llvm.cond_br %8, ^bb2, ^bb3
+  ^bb2:  // pred: ^bb1
+    llvm.store %6, %4 : !llvm.ptr<i32>
+// CHECK: omp.task
+    omp.task   {
+// CHECK: llvm.call @[[CALL_FUNC:.*]]({{.*}}) :
+      llvm.call @_QFPdo_work(%4) : (!llvm.ptr<i32>) -> ()
+// CHECK: omp.terminator
+      omp.terminator
+    }
+    %9 = llvm.load %4 : !llvm.ptr<i32>
+    %10 = llvm.add %9, %5  : i32
+    %11 = llvm.sub %7, %2  : i64
+    llvm.br ^bb1(%10, %11 : i32, i64)
+  ^bb3:  // pred: ^bb1
+    llvm.store %6, %4 : !llvm.ptr<i32>
+// CHECK: omp.terminator
+    omp.terminator
+  }
+  llvm.return
+}
+// CHECK: @[[CALL_FUNC]]
+llvm.func @_QFPdo_work(%arg0: !llvm.ptr<i32> {fir.bindc_name = "i"}) {
+  llvm.return
+}