[Mlir-commits] [mlir] [mlir][OpenMP] Translate reductions on taskloop (PR #199670)

[Sairudra More]

================
@@ -0,0 +1,245 @@
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
+
+// Single scalar reduction on omp.taskloop.context. The lowering must:
+//   1. Emit an implicit __kmpc_taskgroup in the encountering function (since
+//      the user did not write nogroup);
+//   2. Build a kmp_taskred_input_t descriptor array and call
+//      __kmpc_taskred_init, capturing the returned descriptor handle;
+//   3. Force nogroup=1 on the inner __kmpc_taskloop call so that the
+//      OpenMPIRBuilder does not emit a second taskgroup;
+//   4. Inside the outlined task body, call __kmpc_global_thread_num to obtain
+//      the executing thread's gtid, then look up the per-task private storage
+//      via __kmpc_task_reduction_get_th_data(gtid, redDesc, orig);
+//   5. Close the implicit taskgroup with __kmpc_end_taskgroup.
+
+omp.declare_reduction @add_i32 : i32
+init {
+^bb0(%arg0: i32):
+  %c0 = llvm.mlir.constant(0 : i32) : i32
+  omp.yield(%c0 : i32)
+}
+combiner {
+^bb0(%arg0: i32, %arg1: i32):
+  %s = llvm.add %arg0, %arg1 : i32
+  omp.yield(%s : i32)
+}
+
+llvm.func @taskloop_reduction_single(%x : !llvm.ptr, %lb : i32, %ub : i32, %step : i32) {
+  omp.taskloop.context reduction(@add_i32 %x -> %prv : !llvm.ptr) {
+    omp.taskloop.wrapper {
+      omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+        %v = llvm.load %prv : !llvm.ptr -> i32
+        %s = llvm.add %v, %iv : i32
+        llvm.store %s, %prv : i32, !llvm.ptr
+        omp.yield
+      }
+    }
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK: %kmp_taskred_input_t = type { ptr, ptr, i64, ptr, ptr, ptr, i32 }
+
+// Encountering function emits taskgroup + descriptor + taskred_init.
+// CHECK-LABEL: define void @taskloop_reduction_single(
+// CHECK-SAME:    ptr %[[X:[^,]+]],
+// CHECK:         %[[ARR:.+]] = alloca [1 x %kmp_taskred_input_t]
+// CHECK:         call void @__kmpc_taskgroup(
+// CHECK:         %[[ELEM:.+]] = getelementptr inbounds [1 x %kmp_taskred_input_t], ptr %[[ARR]], i32 0, i32 0
+// CHECK:         %[[SHAR:.+]] = getelementptr {{.+}} %kmp_taskred_input_t, ptr %[[ELEM]], i32 0, i32 0
+// CHECK:         store ptr %[[X]], ptr %[[SHAR]]
+// CHECK:         store ptr @__omp_taskloop_taskred_add_i32.red.init
+// CHECK:         store ptr @__omp_taskloop_taskred_add_i32.red.comb
+// CHECK:         %[[DESC:.+]] = call ptr @__kmpc_taskred_init(i32 %{{.+}}, i32 1, ptr %[[ARR]])
+// The returned descriptor is stored into the structArg captured by
+// __kmpc_omp_task_alloc so the outlined task body can load it back.
+// CHECK:         store ptr %[[DESC]], ptr %{{.+}}
+// __kmpc_taskloop must be called with nogroup=1 because we already opened
+// our own taskgroup above.
+// CHECK:         call void @__kmpc_taskloop(ptr {{.+}}, i32 {{.+}}, ptr {{.+}}, i32 1,
----------------
Saieiei wrote:

I updated the test to check the actual behavior: for `taskloop reduction`, the lowering opens the taskgroup explicitly around the taskloop, so the taskloop builder must not emit a second `__kmpc_taskgroup` / `__kmpc_end_taskgroup` pair around it.

https://github.com/llvm/llvm-project/pull/199670