[Openmp-commits] [openmp] [OpenMP] Remove optimization skipping reduction struct initialization (PR #65697)

[Rodrigo Ceccato de Freitas via Openmp-commits]

https://github.com/rodrigo-ceccato updated https://github.com/llvm/llvm-project/pull/65697:

>From 225afe47cbf00132889129cdfcf085bef7bf1415 Mon Sep 17 00:00:00 2001
From: Rodrigo Ceccato <rodrigoceccatodefreitas at gmail.com>
Date: Tue, 5 Sep 2023 21:18:11 +0000
Subject: [PATCH] [OpenMP] Ensure initialization of reduction struct in
 taskloop

This commit ensures the initialization of the reduction struct in task loops,
even when the number of threads in a team is 1, provided that hidden helper
threads are active. Previously, an optimization skipped this initialization
in cases where there was only a single thread in the team, which led to a bug.

The bug behavior resulted in a segmentation fault when hidden helper threads
were active. This occurred because the reduction struct was not initialized
when the task group was created by a team with only one thread, causing the
hidden helper thread to encounter issues when accessing the reduction data.

In addition to this fix, a LIT test for issue #57522 is included to validate
the resolution and prevent any regressions.

Fixes: #57522
---
 .../offloading/task_in_reduction_target.c     | 31 +++++++++++++++++++
 openmp/runtime/src/kmp_tasking.cpp            |  8 +++--
 2 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 openmp/libomptarget/test/offloading/task_in_reduction_target.c

diff --git a/openmp/libomptarget/test/offloading/task_in_reduction_target.c b/openmp/libomptarget/test/offloading/task_in_reduction_target.c
new file mode 100644
index 000000000000000..b546d73d544a584
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/task_in_reduction_target.c
@@ -0,0 +1,31 @@
+// RUN: %libomptarget-compile-generic && \
+// RUN: %libomptarget-run-generic
+
+#include <omp.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+int main(int argc, char *argv[]) {
+
+  int num_devices = omp_get_num_devices();
+
+  // No target devices, just return
+  if (num_devices == 0) {
+    printf("PASS\n");
+    return 0;
+  }
+
+  double sum = 999;
+  double A = 311;
+
+#pragma omp taskgroup task_reduction(+ : sum)
+  {
+#pragma omp target map(to : A) in_reduction(+ : sum) device(0) nowait
+    { sum += A; }
+  }
+
+  printf("PASS\n");
+  return EXIT_SUCCESS;
+}
+
+// CHECK: PASS
diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index fefa609927e8933..e8eb6b02650377c 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -2512,7 +2512,7 @@ void *__kmp_task_reduction_init(int gtid, int num, T *data) {
   KMP_ASSERT(tg != NULL);
   KMP_ASSERT(data != NULL);
   KMP_ASSERT(num > 0);
-  if (nth == 1) {
+  if (nth == 1 && !__kmp_enable_hidden_helper) {
     KA_TRACE(10, ("__kmpc_task_reduction_init: T#%d, tg %p, exiting nth=1\n",
                   gtid, tg));
     return (void *)tg;
@@ -2699,6 +2699,7 @@ void *__kmpc_task_reduction_get_th_data(int gtid, void *tskgrp, void *data) {
         return p_priv[tid];
       }
     }
+    KMP_ASSERT(tg->parent);
     tg = tg->parent;
     arr = (kmp_taskred_data_t *)(tg->reduce_data);
     num = tg->reduce_num_data;
@@ -2711,7 +2712,10 @@ void *__kmpc_task_reduction_get_th_data(int gtid, void *tskgrp, void *data) {
 // Called from __kmpc_end_taskgroup()
 static void __kmp_task_reduction_fini(kmp_info_t *th, kmp_taskgroup_t *tg) {
   kmp_int32 nth = th->th.th_team_nproc;
-  KMP_DEBUG_ASSERT(nth > 1); // should not be called if nth == 1
+  KMP_DEBUG_ASSERT(
+      nth > 1 ||
+      __kmp_enable_hidden_helper); // should not be called if nth == 1 unless we
+                                   // are using hidden helper threads
   kmp_taskred_data_t *arr = (kmp_taskred_data_t *)tg->reduce_data;
   kmp_int32 num = tg->reduce_num_data;
   for (int i = 0; i < num; ++i) {