Implemented a batched version of the AXPY kernel

1 files changed, 7 insertions, 7 deletions

diff --git a/src/kernels/level1/xaxpy.opencl b/src/kernels/level1/xaxpy.opencl
index 0d730c9e..3f5ab2b5 100644
--- a/src/kernels/level1/xaxpy.opencl
+++ b/src/kernels/level1/xaxpy.opencl
@@ -57,17 +57,17 @@ void XaxpyFast(const int n, const real_arg arg_alpha,
 
 // Full version of the kernel with offsets and strided accesses: batched version
 __kernel __attribute__((reqd_work_group_size(WGS, 1, 1)))
-void XaxpyBatched(const int n, const real_arg arg_alpha,
-                  const __global real* restrict xgm, const int x_offset, const int x_inc,
-                  __global real* ygm, const int y_offset, const int y_inc,
-                  const int batch) {
-  const real alpha = GetRealArg(arg_alpha);
+void XaxpyBatched(const int n, const __global real_arg* arg_alphas,
+                  const __global real* restrict xgm, const __global int* restrict x_offsets, const int x_inc,
+                  __global real* ygm, const __global int* restrict y_offsets, const int y_inc) {
+  const int batch = get_group_id(1);
+  const real alpha = GetRealArg(arg_alphas[batch]);
 
   // Loops over the work that needs to be done (allows for an arbitrary number of threads)
   #pragma unroll
   for (int id = get_global_id(0); id<n; id += get_global_size(0)) {
-    real xvalue = xgm[id*x_inc + x_offset];
-    MultiplyAdd(ygm[id*y_inc + y_offset], alpha, xvalue);
+    real xvalue = xgm[id*x_inc + x_offsets[batch]];
+    MultiplyAdd(ygm[id*y_inc + y_offsets[batch]], alpha, xvalue);
   }
 }