Implemented direct version of strided-batched GEMM kernel

2 files changed, 121 insertions, 37 deletions

diff --git a/src/kernels/level3/xgemm_direct_batched.opencl b/src/kernels/level3/xgemm_direct_batched.opencl
index d946a056..d15ed31e 100644
--- a/src/kernels/level3/xgemm_direct_batched.opencl
+++ b/src/kernels/level3/xgemm_direct_batched.opencl
@@ -17,15 +17,16 @@
 R"(
 
 // =================================================================================================
+#if defined(ROUTINE_GEMMBATCHED)
 
 // Direct version of the batched GEMM kernel with [A, B] = [non-transposed, non-transposed]
 __kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
 void XgemmDirectBatchedNN(const int kSizeM, const int kSizeN, const int kSizeK,
-                                   const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
-                                   const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
-                                   const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
-                                   __global real* cgm, const __constant int* c_offsets, const int c_ld,
-                                   const int c_transpose, const int a_conjugate, const int b_conjugate) {
+                          const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
+                          const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
+                          const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
+                          __global real* cgm, const __constant int* c_offsets, const int c_ld,
+                          const int c_transpose, const int a_conjugate, const int b_conjugate) {
   const int batch = get_group_id(2);
   const real_arg arg_alpha = arg_alphas[batch];
   const real_arg arg_beta = arg_betas[batch];
@@ -42,11 +43,11 @@ void XgemmDirectBatchedNN(const int kSizeM, const int kSizeN, const int kSizeK,
 // Direct version of the batched GEMM kernel with [A, B] = [non-transposed, transposed]
 __kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
 void XgemmDirectBatchedNT(const int kSizeM, const int kSizeN, const int kSizeK,
-                                   const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
-                                   const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
-                                   const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
-                                   __global real* cgm, const __constant int* c_offsets, const int c_ld,
-                                   const int c_transpose, const int a_conjugate, const int b_conjugate) {
+                          const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
+                          const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
+                          const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
+                          __global real* cgm, const __constant int* c_offsets, const int c_ld,
+                          const int c_transpose, const int a_conjugate, const int b_conjugate) {
   const int batch = get_group_id(2);
   const real_arg arg_alpha = arg_alphas[batch];
   const real_arg arg_beta = arg_betas[batch];
@@ -63,11 +64,11 @@ void XgemmDirectBatchedNT(const int kSizeM, const int kSizeN, const int kSizeK,
 // Direct version of the batched GEMM kernel with [A, B] = [transposed, non-transposed]
 __kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
 void XgemmDirectBatchedTN(const int kSizeM, const int kSizeN, const int kSizeK,
-                                   const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
-                                   const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
-                                   const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
-                                   __global real* cgm, const __constant int* c_offsets, const int c_ld,
-                                   const int c_transpose, const int a_conjugate, const int b_conjugate) {
+                          const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
+                          const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
+                          const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
+                          __global real* cgm, const __constant int* c_offsets, const int c_ld,
+                          const int c_transpose, const int a_conjugate, const int b_conjugate) {
   const int batch = get_group_id(2);
   const real_arg arg_alpha = arg_alphas[batch];
   const real_arg arg_beta = arg_betas[batch];
@@ -84,11 +85,11 @@ void XgemmDirectBatchedTN(const int kSizeM, const int kSizeN, const int kSizeK,
 // Direct version of the batched GEMM kernel with [A, B] = [transposed, transposed]
 __kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
 void XgemmDirectBatchedTT(const int kSizeM, const int kSizeN, const int kSizeK,
-                                   const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
-                                   const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
-                                   const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
-                                   __global real* cgm, const __constant int* c_offsets, const int c_ld,
-                                   const int c_transpose, const int a_conjugate, const int b_conjugate) {
+                          const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
+                          const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
+                          const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
+                          __global real* cgm, const __constant int* c_offsets, const int c_ld,
+                          const int c_transpose, const int a_conjugate, const int b_conjugate) {
   const int batch = get_group_id(2);
   const real_arg arg_alpha = arg_alphas[batch];
   const real_arg arg_beta = arg_betas[batch];
@@ -102,6 +103,87 @@ void XgemmDirectBatchedTT(const int kSizeM, const int kSizeN, const int kSizeK,
               alm, blm, 1, 1, c_transpose, a_conjugate, b_conjugate);
 }
 
+#endif
+// =================================================================================================
+#if defined(ROUTINE_GEMMSTRIDEDBATCHED)
+
+// Direct version of the strided-batched GEMM kernel with [A, B] = [non-transposed, non-transposed]
+__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
+void XgemmDirectStridedBatchedNN(const int kSizeM, const int kSizeN, const int kSizeK,
+                                 const real_arg arg_alpha, const real_arg arg_beta,
+                                 const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
+                                 const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
+                                 __global real* cgm, const int c_offset, const int c_ld, const int c_stride,
+                                 const int c_transpose, const int a_conjugate, const int b_conjugate) {
+  const int batch = get_group_id(2);
+  const int a_offset_batch = a_offset + a_stride * batch;
+  const int b_offset_batch = b_offset + b_stride * batch;
+  const int c_offset_batch = c_offset + c_stride * batch;
+  __local real alm[WGD * (WGD + PADA)];
+  __local real blm[WGD * (WGD + PADB)];
+  XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
+              agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
+              alm, blm, 0, 0, c_transpose, a_conjugate, b_conjugate);
+}
+
+// Direct version of the strided-batched GEMM kernel with [A, B] = [non-transposed, transposed]
+__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
+void XgemmDirectStridedBatchedNT(const int kSizeM, const int kSizeN, const int kSizeK,
+                                 const real_arg arg_alpha, const real_arg arg_beta,
+                                 const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
+                                 const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
+                                 __global real* cgm, const int c_offset, const int c_ld, const int c_stride,
+                                 const int c_transpose, const int a_conjugate, const int b_conjugate) {
+  const int batch = get_group_id(2);
+  const int a_offset_batch = a_offset + a_stride * batch;
+  const int b_offset_batch = b_offset + b_stride * batch;
+  const int c_offset_batch = c_offset + c_stride * batch;
+  __local real alm[WGD * (WGD + PADA)];
+  __local real blm[WGD * (WGD + PADB)];
+  XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
+              agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
+              alm, blm, 0, 1, c_transpose, a_conjugate, b_conjugate);
+}
+
+// Direct version of the strided-batched GEMM kernel with [A, B] = [transposed, non-transposed]
+__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
+void XgemmDirectStridedBatchedTN(const int kSizeM, const int kSizeN, const int kSizeK,
+                                 const real_arg arg_alpha, const real_arg arg_beta,
+                                 const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
+                                 const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
+                                 __global real* cgm, const int c_offset, const int c_ld, const int c_stride,
+                                 const int c_transpose, const int a_conjugate, const int b_conjugate) {
+  const int batch = get_group_id(2);
+  const int a_offset_batch = a_offset + a_stride * batch;
+  const int b_offset_batch = b_offset + b_stride * batch;
+  const int c_offset_batch = c_offset + c_stride * batch;
+  __local real alm[WGD * (WGD + PADA)];
+  __local real blm[WGD * (WGD + PADB)];
+  XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
+              agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
+              alm, blm, 1, 0, c_transpose, a_conjugate, b_conjugate);
+}
+
+// Direct version of the strided-batched GEMM kernel with [A, B] = [transposed, transposed]
+__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
+void XgemmDirectStridedBatchedTT(const int kSizeM, const int kSizeN, const int kSizeK,
+                                 const real_arg arg_alpha, const real_arg arg_beta,
+                                 const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
+                                 const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
+                                 __global real* cgm, const int c_offset, const int c_ld, const int c_stride,
+                                 const int c_transpose, const int a_conjugate, const int b_conjugate) {
+  const int batch = get_group_id(2);
+  const int a_offset_batch = a_offset + a_stride * batch;
+  const int b_offset_batch = b_offset + b_stride * batch;
+  const int c_offset_batch = c_offset + c_stride * batch;
+  __local real alm[WGD * (WGD + PADA)];
+  __local real blm[WGD * (WGD + PADB)];
+  XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
+              agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
+              alm, blm, 1, 1, c_transpose, a_conjugate, b_conjugate);
+}
+
+#endif
 // =================================================================================================
 
 // End of the C++11 raw string literal
diff --git a/src/routines/levelx/xgemmstridedbatched.cpp b/src/routines/levelx/xgemmstridedbatched.cpp
index 3ea52980..ddf7d878 100644
--- a/src/routines/levelx/xgemmstridedbatched.cpp
+++ b/src/routines/levelx/xgemmstridedbatched.cpp
@@ -112,6 +112,7 @@ void XgemmStridedBatched<T>::BatchedGemmIndirect(const size_t m, const size_t n,
                                                  const size_t b_one, const size_t b_two,
                                                  const size_t c_one, const size_t c_two,
                                                  const size_t batch_count) {
+  /* TODO
   // Calculates the ceiled versions of m, n, and k
   const auto m_ceiled = Ceil(Ceil(m, db_["MWG"]), db_["VWM"]);
   const auto n_ceiled = Ceil(Ceil(n, db_["NWG"]), db_["VWN"]);
@@ -123,7 +124,6 @@ void XgemmStridedBatched<T>::BatchedGemmIndirect(const size_t m, const size_t n,
   Xgemm<T>::CalculateInternalDimensions(m, n, k, db_["MWG"], db_["NWG"], db_["KWG"],
                                         a_one_i, a_two_i, b_one_i, b_two_i, c_one_i, c_two_i);
 
-  /* TODO
   // Sets the "internal" offsets, i.e. the perfect offsets
   auto a_offsets_i = 0;//std::vector<int>(batch_count);
   auto b_offsets_i = 0;//std::vector<int>(batch_count);
@@ -244,30 +244,33 @@ void XgemmStridedBatched<T>::BatchedGemmDirect(const size_t m, const size_t n, c
                                                const bool a_do_transpose, const bool b_do_transpose, const bool c_do_transpose,
                                                const bool a_conjugate, const bool b_conjugate,
                                                const size_t batch_count) {
-/* TODO
+
   // Retrieves the proper XgemmDirect kernel from the compiled binary
-  const auto name = (a_do_transpose) ? (b_do_transpose ? "XgemmDirectBatchedTT" : "XgemmDirectBatchedTN") :
-                    (b_do_transpose ? "XgemmDirectBatchedNT" : "XgemmDirectBatchedNN");
+  const auto name = (a_do_transpose) ? (b_do_transpose ? "XgemmDirectStridedBatchedTT" : "XgemmDirectStridedBatchedTN") :
+                    (b_do_transpose ? "XgemmDirectStridedBatchedNT" : "XgemmDirectStridedBatchedNN");
   auto kernel = Kernel(program_, name);
 
   // Sets the kernel arguments
   kernel.SetArgument(0, static_cast<int>(m));
   kernel.SetArgument(1, static_cast<int>(n));
   kernel.SetArgument(2, static_cast<int>(k));
-  kernel.SetArgument(3, alpha);
-  kernel.SetArgument(4, beta);
+  kernel.SetArgument(3, GetRealArg(alpha));
+  kernel.SetArgument(4, GetRealArg(beta));
   kernel.SetArgument(5, a_buffer());
-  kernel.SetArgument(6, a_offset);
+  kernel.SetArgument(6, static_cast<int>(a_offset));
   kernel.SetArgument(7, static_cast<int>(a_ld));
-  kernel.SetArgument(8, b_buffer());
-  kernel.SetArgument(9, b_offset);
-  kernel.SetArgument(10, static_cast<int>(b_ld));
-  kernel.SetArgument(11, c_buffer());
-  kernel.SetArgument(12, c_offset);
-  kernel.SetArgument(13, static_cast<int>(c_ld));
-  kernel.SetArgument(14, static_cast<int>(c_do_transpose));
-  kernel.SetArgument(15, static_cast<int>(a_conjugate));
-  kernel.SetArgument(16, static_cast<int>(b_conjugate));
+  kernel.SetArgument(8, static_cast<int>(a_stride));
+  kernel.SetArgument(9, b_buffer());
+  kernel.SetArgument(10, static_cast<int>(b_offset));
+  kernel.SetArgument(11, static_cast<int>(b_ld));
+  kernel.SetArgument(12, static_cast<int>(b_stride));
+  kernel.SetArgument(13, c_buffer());
+  kernel.SetArgument(14, static_cast<int>(c_offset));
+  kernel.SetArgument(15, static_cast<int>(c_ld));
+  kernel.SetArgument(16, static_cast<int>(c_stride));
+  kernel.SetArgument(17, static_cast<int>(c_do_transpose));
+  kernel.SetArgument(18, static_cast<int>(a_conjugate));
+  kernel.SetArgument(19, static_cast<int>(b_conjugate));
 
   // Computes the global and local thread sizes
   const auto m_ceiled = Ceil(m, db_["WGD"]);
@@ -281,7 +284,6 @@ void XgemmStridedBatched<T>::BatchedGemmDirect(const size_t m, const size_t n, c
 
   // Launches the kernel
   RunKernel(kernel, queue_, device_, global, local, event_);
-  */
 }
 
 // =================================================================================================