From d84dedc7d3346321f3801237fc0ccd4ce772492d Mon Sep 17 00:00:00 2001
From: AndreyChurbanov <andrey.churbanov@intel.com>
Date: Fri, 11 Feb 2022 17:34:26 +0300
Subject: [PATCH] [OpenMP] libomp: fix bug in implementation of distribute
 construct.

Fixed mistaken iterations distribution between different target regions.

Differential Revision: https://reviews.llvm.org/D118393
---
 openmp/runtime/src/kmp_sched.cpp              |  9 +++-
 .../runtime/test/teams/teams-distr-on-host.c  | 46 +++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 openmp/runtime/test/teams/teams-distr-on-host.c

diff --git a/openmp/runtime/src/kmp_sched.cpp b/openmp/runtime/src/kmp_sched.cpp
index 09e497e02914..46d3d26679c6 100644
--- a/openmp/runtime/src/kmp_sched.cpp
+++ b/openmp/runtime/src/kmp_sched.cpp
@@ -194,8 +194,13 @@ static void __kmp_for_static_init(ident_t *loc, kmp_int32 global_tid,
     // we are in DISTRIBUTE construct
     schedtype += kmp_sch_static -
                  kmp_distribute_static; // AC: convert to usual schedule type
-    tid = th->th.th_team->t.t_master_tid;
-    team = th->th.th_team->t.t_parent;
+    if (th->th.th_team->t.t_serialized > 1) {
+      tid = 0;
+      team = th->th.th_team;
+    } else {
+      tid = th->th.th_team->t.t_master_tid;
+      team = th->th.th_team->t.t_parent;
+    }
   } else {
     tid = __kmp_tid_from_gtid(global_tid);
     team = th->th.th_team;
diff --git a/openmp/runtime/test/teams/teams-distr-on-host.c b/openmp/runtime/test/teams/teams-distr-on-host.c
new file mode 100644
index 000000000000..663d1d3a487f
--- /dev/null
+++ b/openmp/runtime/test/teams/teams-distr-on-host.c
@@ -0,0 +1,46 @@
+// The test supposes no offload, pure host execution.
+// It checks that the bug in implementation of distribute construct is fixed.
+
+// RUN: %libomp-compile-and-run
+// UNSUPPORTED: icc
+
+#include <stdio.h>
+#include <omp.h>
+
+int main()
+{
+  const int size = 4;
+  int wrong_counts = 0;
+  omp_set_num_threads(2);
+  #pragma omp parallel reduction(+:wrong_counts)
+  {
+    int i;
+    int A[size];
+    int th = omp_get_thread_num();
+    for(i = 0; i < size; i++)
+      A[i] = 0;
+
+    #pragma omp target teams distribute map(tofrom: A[:size]) private(i)
+    for(i = 0; i < size; i++)
+    {
+      A[i] = i;
+      printf("th %d, team %d, i %d\n", th, omp_get_team_num(), i);
+    }
+    #pragma omp critical
+    {
+      printf("tid = %d\n", th);
+      for(i = 0; i < size; i++)
+      {
+        if (A[i] != i) wrong_counts++;
+        printf("  %d", A[i]);
+      }
+      printf("\n");
+    }
+  }
+  if (wrong_counts) {
+    printf("failed\n");
+  } else {
+    printf("passed\n");
+  }
+  return wrong_counts;
+}