Merge pull request #1975 from bstatcomp/fix_kg_local_variables

make kernel generator declare local variables at top level of kernel function.

Author: t4c1

stan/math/opencl/kernel_generator/colwise_reduction.hpp

@@ -101,9 +101,10 @@ class colwise_reduction
                                const bool view_handled,
                                const std::string& var_name_arg) const {
     kernel_parts res;
-    res.initialization = type_str<Scalar>() + " " + var_name_ + " = " + init_
-        + ";\n" "__local " + type_str<Scalar>() + " "
-        + var_name_ + "_local[LOCAL_SIZE_];\n";
+    res.declarations = "__local " + type_str<Scalar>() + " " + var_name_
+                       + "_local[LOCAL_SIZE_];\n";
+    res.initialization
+        = type_str<Scalar>() + " " + var_name_ + " = " + init_ + ";\n";
     res.body = var_name_ + " = " + var_name_arg + ";\n";
     res.reduction =
           var_name_ + "_local[lid_i] = " + var_name_ + ";\n"

stan/math/opencl/kernel_generator/multi_result_kernel.hpp

@@ -343,6 +343,7 @@ class results_cl {
           "const int i0 = lsize_i * wg_id_i;\n"
           "const int i = i0 + lid_i;\n"
           "const int j0 = lsize_i * wg_id_j;\n"
+          + parts.declarations +
           "for(int lid_j = 0; lid_j < min(cols - j0, lsize_i); lid_j++){\n"
           "const int j = j0 + lid_j;\n"
           + parts.initialization +
@@ -360,6 +361,7 @@ class results_cl {
           "){\n"
           "int i = get_global_id(0);\n"
           "int j = get_global_id(1);\n"
+          + parts.declarations
           + parts.initialization
           + parts.body
           + parts.reduction +

stan/math/opencl/kernel_generator/operation_cl.hpp

@@ -33,6 +33,7 @@ namespace math {
 struct kernel_parts {
   std::string includes;  // any function definitions - as if they were includet
                          // at the start of kernel source
+  std::string declarations;    // declarations of any local variables
   std::string initialization;  // the code for initializations done by all
                                // threads, even if they have no work
   std::string body_prefix;     // the code that should be placed at the start of
@@ -43,14 +44,18 @@ struct kernel_parts {
   std::string args;       // kernel arguments
 
   kernel_parts operator+(const kernel_parts& other) {
-    return {
-        includes + other.includes,       initialization + other.initialization,
-        body_prefix + other.body_prefix, body + other.body,
-        reduction + other.reduction,     args + other.args};
+    return {includes + other.includes,
+            declarations += other.declarations,
+            initialization + other.initialization,
+            body_prefix + other.body_prefix,
+            body + other.body,
+            reduction + other.reduction,
+            args + other.args};
   }
 
   kernel_parts operator+=(const kernel_parts& other) {
     includes += other.includes;
+    declarations += other.declarations;
     initialization += other.initialization;
     body_prefix += other.body_prefix;
     body += other.body;

test/unit/math/opencl/kernel_generator/reference_kernels/colwise_sum.cl

@@ -10,10 +10,10 @@ const int blocks_cols = (cols + lsize_i - 1) / lsize_i;
 const int i0 = lsize_i * wg_id_i;
 const int i = i0 + lid_i;
 const int j0 = lsize_i * wg_id_j;
+__local double var1_local[LOCAL_SIZE_];
 for(int lid_j = 0; lid_j < min(cols - j0, lsize_i); lid_j++){
 const int j = j0 + lid_j;
 double var1 = 0;
-__local double var1_local[LOCAL_SIZE_];
 if(i < rows){
 double var2 = 0; if (!((!contains_nonzero(var2_view, LOWER) && j < i) || (!contains_nonzero(var2_view, UPPER) && j > i))) {var2 = var2_global[i + var2_rows * j];}
 var1 = var2;