Insert kmpc_parallel_51 call loop distribution

Pragma target teams distribute parallel for calls
kmpc_parallel_51 function in different way in comparison to
pragma target parallel.

Pragma target parallel calls kmpc_parallel_51 function
just after initialization procedure.

Pragma target teams distribute parallel for calls
kmpc_parallel_51 function in the body of the outer
loop.

Signed-off-by: Dominik Adamski <[email protected]>

Author: DominikAdamski

tools/flang2/flang2exe/expand.cpp

@@ -238,6 +238,9 @@ expand(void)
   auto it = process_expanded_map.find(gbl.currsub);
   int process_expanded = 0;
 
+  //we are at the beginning of pragma expansion
+  //make sure that mploop_counter equals to zero
+  reset_mploop_counter();
   // we reset flag because we do not know if we generate initialization
   // function for SPMD kernel (the function with kmpc_parallel_51 call)
   // or the proper kernel code (the function which is passed as an argument

tools/flang2/flang2exe/kmpcutil.cpp

@@ -1778,7 +1778,11 @@ bool check_if_skip_symbol(SPTR sym)
 }
 
 int
-ll_make_kmpc_parallel_51(int global_tid_sptr, std::vector<int> &symbols, SPTR helper_func)
+ll_make_kmpc_parallel_51(int global_tid_sptr,
+                         std::vector<int> &symbols,
+                         SPTR helper_func,
+                         SPTR lower,
+                         SPTR upper)
 {
   static int id;
   int n_symbols = get_n_symbols(ompaccel_tinfo_get(gbl.currsub));
@@ -1787,6 +1791,10 @@ ll_make_kmpc_parallel_51(int global_tid_sptr, std::vector<int> &symbols, SPTR he
   DTYPE void_ptr_ptr_t = get_type(2, TY_PTR, void_ptr_t);
   DTYPE arr_dtype;
   int args[9];
+
+  if (lower && upper)
+    n_symbols += 2;
+
   SPTR captured_vars = make_array_sptr(const_cast<char*>("captured_vars_addrs"),
                                                          void_ptr_t,
                                                          n_symbols);
@@ -1798,7 +1806,25 @@ ll_make_kmpc_parallel_51(int global_tid_sptr, std::vector<int> &symbols, SPTR he
                             ad_icon(0),
                             FALSE);
   int j = 0;
-  for (int i = 0; i < n_symbols; ++i) {
+  int i = 0;
+  /* Store lower and upper bounds for loop distribution */
+  if (lower && upper) {
+    ilix = mk_ompaccel_ldsptr(lower);
+    ilix = mk_ompaccel_store(ilix,
+                             DT_INT8,
+                             nme_args,
+                             ad_acon(captured_vars, i * TARGET_PTRSIZE));
+    chk_block(ilix);
+    i++;
+    ilix = mk_ompaccel_ldsptr(upper);
+    ilix = mk_ompaccel_store(ilix,
+                             DT_INT8,
+                             nme_args,
+                             ad_acon(captured_vars, i * TARGET_PTRSIZE));
+    chk_block(ilix);
+    i++;
+  }
+  for (; i < n_symbols; ++i) {
     if (check_if_skip_symbol(ompaccel_tinfo_get(gbl.currsub)->symbols[i].device_sym))
       continue;
     else if (PASSBYVALG(ompaccel_tinfo_get(gbl.currsub)->symbols[i].device_sym) &&
@@ -1825,7 +1851,6 @@ ll_make_kmpc_parallel_51(int global_tid_sptr, std::vector<int> &symbols, SPTR he
     }
     chk_block(ilix);
   }
-  
 
   arg_types[0] = DT_CPTR;        /* ident */
   arg_types[1] = DT_INT;         /* global_tid */
@@ -1842,7 +1867,10 @@ ll_make_kmpc_parallel_51(int global_tid_sptr, std::vector<int> &symbols, SPTR he
   args[6] = ad_icon(1);                /* if_expr */
   args[5] = ad_icon(-1);               /* num_threads */
   args[4] = ad_icon(-1);               /* proc_bind */
-  args[3] = ad_acon(helper_func, 0);
+  if (helper_func)
+    args[3] = ad_acon(helper_func, 0);
+  else
+    args[3] = gen_null_arg();
   args[2] = gen_null_arg();            /* wrapper_fn */
   args[1] = ad_acon(captured_vars, 0); /* args */
   args[0] = ad_icon(n_symbols);        /* n_args */

tools/flang2/flang2exe/kmpcutil.h

@@ -509,7 +509,11 @@ int ll_make_kmpc_target_init(OMP_TARGET_MODE);
 /**
   \brief Generate kmpc_parallel_51 function call
 */
-int ll_make_kmpc_parallel_51(int global_tid_sptr, std::vector<int> &, SPTR);
+int ll_make_kmpc_parallel_51(int global_tid_sptr,
+                             std::vector<int> &,
+                             SPTR,
+                             SPTR lower = (SPTR)0,
+                             SPTR upper = (SPTR)0);
 
 #ifdef OMP_OFFLOAD_AMD
 /**

tools/flang2/flang2exe/ompaccel.cpp

@@ -66,6 +66,9 @@
 // Should be in sync with clang::GPU::AMDGPUGpuGridValues in clang
 int warp_size_log2;
 int warp_size_log2_mask;
+// count if we expand the second MPLOOP instruction
+// inside single OpenMP pragma
+int mploop_counter;
 // AOCC End
 #include "../../flang1/flang1exe/global.h"
 
@@ -2758,7 +2761,20 @@ exp_ompaccel_mploop(ILM *ilmp, int curilm)
       ili = ll_make_kmpc_for_static_init(&loop_args);
     // AOCC end
     } else {
-      ili = ll_make_kmpc_for_static_init_simple_spmd(&loop_args, sched);
+      mploop_counter++;
+      if (mploop_counter != 2)
+        ili = ll_make_kmpc_for_static_init_simple_spmd(&loop_args, sched);
+      else {
+	std::vector<int> allocated_symbols;
+        int ilix = ll_make_kmpc_global_thread_num();
+        ilix = ll_make_kmpc_parallel_51(ilix,
+			                allocated_symbols,
+                                        (SPTR)0, /*TODO: replace with wrapper fn ptr */
+                                        loop_args.lower,
+                                        loop_args.upper);
+        iltb.callfg = 1;
+        chk_block(ilix);
+      }
     }
     break;
   default:
@@ -3732,6 +3748,11 @@ bool is_SPMD_mode(OMP_TARGET_MODE mode) {
   return false;
 }
 
+void reset_mploop_counter()
+{
+  mploop_counter = 0;
+}
+
 // AOCC End
 #endif
 /* Expander - OpenMP Accelerator Model */

tools/flang2/flang2exe/ompaccel.h

@@ -590,4 +590,8 @@ void ompaccel_set_target_declare();
  */
 bool is_SPMD_mode(OMP_TARGET_MODE mode);
 // AOCC End
+/**
+ \brief Reset counts of MPLOOP instruction
+ */
+void reset_mploop_counter();
 #endif