Add OpenMPToLLVM conversion pattern for taskgroup

[kiranchandramohan]

Summary

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

[psoni2628]

Do we need to split up the Flang and MLIR parts of this patch?

[kiranchandramohan]

In D151206#4364989, @psoni2628 wrote:

Do we need to split up the Flang and MLIR parts of this patch?

Since this is only a test, i guess it is OK.

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

 // CHECK:          }
 // CHECK:          llvm.return
 // CHECK:        }
 // CHECK:        llvm.func @_QPwork() attributes {sym_visibility = "private"}
 // CHECK:      }
 
 // -----
 
+
+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> {
   ^bb0(%arg0: i32):
     %c1_i32 = arith.constant 1 : i32
     %1 = arith.addi %arg0, %c1_i32 : i32
     omp.yield(%1 : i32)
   }

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

 } // namespace
 
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
   target.addDynamicallyLegalOp<
       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) {
   target.addDynamicallyLegalOp<mlir::omp::AtomicReadOp,
                                mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
                                mlir::omp::ThreadprivateOp, mlir::omp::YieldOp,
                                mlir::omp::EnterDataOp, mlir::omp::ExitDataOp>(
       [&](Operation *op) {
         return typeConverter.isLegal(op->getOperandTypes()) &&
                typeConverter.isLegal(op->getResultTypes());
       });

 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<
       ReductionOpConversion, ReductionDeclareOpConversion,
       RegionOpConversion<omp::CriticalOp>, RegionOpConversion<omp::MasterOp>,
       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>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
       RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
       RegionLessOpConversion<omp::YieldOp>,
       RegionLessOpConversion<omp::EnterDataOp>,
       RegionLessOpConversion<omp::ExitDataOp>>(converter);

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

omp.wsloop   reduction(@eqv_reduction -> %4 : !llvm.ptr<i32>) for  (%arg1) : i32 = (%1) to (%0) inclusive step (%1) {
       %11 = llvm.load %10 : !llvm.ptr<i32>
       omp.reduction %11, %4 : i32, !llvm.ptr<i32>
       omp.yield
     }
     omp.terminator
   }
   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
+}