[flang][OpenMP] Lower OpenMP `taskgroup` construct

[SouraVX]

Summary

Lower Fortran OpenMP taskgroup to FIR + OpenMP Dialect.

[SouraVX]

NIT corrected.

[kiranchandramohan]

Welcome back @SouraVX! Have one minor comment.

flang/test/Lower/OpenMP/taskgroup.f90

3

This test (fir-to-llvm-ir) should go here. https://github.com/llvm/llvm-project/blob/main/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

[SouraVX]

Thanks @kiranchandramohan for reviewing this.
Changes:

• Addressed review comments.

[kiranchandramohan]

LG.

flang/test/Lower/OpenMP/taskgroup.f90

2

You can remove the prefixes since we are not using them separately.

[shraiysh]

Thanks for the patch @SouraVX. LGTM.

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

288

Here we are missing the following around the function call inside.

CHECK: omp.taskgroup

CHECK: omp.terminator

Please let me know if thats not needed here.

[peixin]

flang/lib/Lower/OpenMP.cpp

843

The allocate clause operands are added here, so can you add one test case for this?

[SouraVX]

Thanks for the review @shraiysh & @peixin.
Changes:

• Updated taskgroup.f90 test to include allocate clause. task_reduction clause not added, for support missing/todo while lowering.
• Updated convert-to-llvm-openmp-and-fir.fir test based on feedback. (Some how that test still passed previously as reflected in build-status of the patch also).

Not completely orthogonal to this patch, I tried to generate llvmir for the test case having allocate clause:

subroutine omp_taskgroup
use omp_lib
integer :: allocated_x
!$omp taskgroup allocate(omp_high_bw_mem_alloc: allocated_x)
!$omp task
   call work()
!$omp end task
end subroutine

$ bbc taskgroup.f90 -o -|tco
Relevant pieces of Err:
loc("taskgroup.mlir":4:3): error: null operand found
"llvm.func"() ({
  %0 = "llvm.mlir.constant"() {value = 1 : i32} : () -> i32
  %1 = "llvm.mlir.constant"() {value = 1 : i64} : () -> i64
  %2 = "llvm.alloca"(%1) {bindc_name = "allocated_x", in_type = i32, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_taskgroupEallocated_x"} : (i64) -> !llvm.ptr<i32>
  "omp.taskgroup"(<<NULL VALUE>>, %0) ({
    "omp.task"() ({
      "llvm.call"() {callee = @_QPwork, fastmathFlags = #llvm.fastmath<>} : () -> ()
      "omp.terminator"() : () -> ()
    }) {operand_segment_sizes = array<i32: 0, 0, 0, 0, 0, 0>} : () -> ()
    "omp.terminator"() : () -> ()
  }) {operand_segment_sizes = array<i32: 0, 1, 1>} : (<<NULL TYPE>>, i32) -> ()
  "llvm.return"() : () -> ()
}) {CConv = #llvm.cconv<ccc>, function_type = !llvm.func<void ()>, linkage = #llvm.linkage<external>, sym_name = "_QPomp_taskgroup"} : () -> ()

It seems like something is broken WRT allocate clause lowering/codegen in taskgroup.

[SouraVX]

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

277

Not added,allocate clause for reasons prescribed in above comment.

[peixin]

LGTM since this patch is doing the lowering part as the title and summary mentioned.

It would be good to file a ticket about the error of allocate clause when lowering to LLVMIR.

flang/lib/Lower/OpenMP.cpp

841

Nit

[SouraVX]

FYI: https://github.com/llvm/llvm-project/issues/57786

flang/lib/Lower/OpenMP.cpp

createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation,
                                          eval);
   } else if (blockDirective.v == llvm::omp::OMPD_task) {
     auto taskOp = firOpBuilder.create<mlir::omp::TaskOp>(
         currentLocation, ifClauseOperand, finalClauseOperand, untiedAttr,
         mergeableAttr, /*in_reduction_vars=*/ValueRange(),
         /*in_reductions=*/nullptr, priorityClauseOperand, allocateOperands,
         allocatorOperands);
     createBodyOfOp(taskOp, converter, currentLocation, eval, &opClauseList);
+  } else if (blockDirective.v == llvm::omp::OMPD_taskgroup) {
+      // TODO: Add task_reduction support
+    auto taskGroupOp = firOpBuilder.create<mlir::omp::TaskGroupOp>(
+        currentLocation, /*task_reduction_vars=*/ValueRange(),
+        /*task_reductions=*/nullptr, allocateOperands, allocatorOperands);
+    createBodyOfOp(taskGroupOp, converter, currentLocation, eval,
+                   &opClauseList);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
 }
 
 /// This function returns the identity value of the operator \p reductionOpName.
 /// For example:
 ///    0 + x = x,

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

 // CHECK:             llvm.cond_br %[[VAL_6]], ^bb2, ^bb3
 // CHECK:           ^bb2:
 // CHECK:             llvm.br ^bb1(%[[VAL_7:.*]], %[[VAL_8:.*]] : i64, i64)
 // CHECK:           ^bb3:
 // CHECK:             omp.yield
 // CHECK:           }
 // CHECK:           llvm.return
 // CHECK:         }
+
+// -----
+
+func.func @_QPomp_taskgroup() {
+  omp.taskgroup {
+    omp.task   {
+      fir.call @_QPwork() : () -> ()
+      omp.terminator
+    }
+    omp.terminator
+  }
+  return
+}
+func.func private @_QPwork()
+
+// CHECK-LABEL: llvm.func @_QPomp_taskgroup() {
+// CHECK:          omp.taskgroup   {
+// CHECK:            omp.task   {
+// CHECK:              llvm.call @_QPwork() : () -> ()
+// CHECK:              omp.terminator
+// CHECK:            }
+// CHECK:            omp.terminator
+// CHECK:          }
+// CHECK:          llvm.return
+// CHECK:        }
+// CHECK:        llvm.func @_QPwork() attributes {sym_visibility = "private"}
+// CHECK:      }

flang/test/Lower/OpenMP/taskgroup.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+
+!CHECK-LABEL: @_QPomp_taskgroup
+subroutine omp_taskgroup
+use omp_lib
+integer :: allocated_x
+!CHECK-DAG: %{{.*}} = fir.alloca i32 {bindc_name = "allocated_x", uniq_name = "_QFomp_taskgroupEallocated_x"}
+!CHECK-DAG: %{{.*}} = arith.constant 1 : i32
+
+!CHECK: omp.taskgroup  allocate(%{{.*}} : i32 -> %0 : !fir.ref<i32>)
+!$omp taskgroup allocate(omp_high_bw_mem_alloc: allocated_x)
+!$omp task
+!CHECK: fir.call @_QPwork() : () -> ()
+   call work()
+!CHECK: omp.terminator
+!$omp end task
+!CHECK: omp.terminator
+!$omp end taskgroup
+end subroutine