Implemented the in-direct version of the strided-batched GEMM kernel

1 files changed, 66 insertions, 0 deletions

diff --git a/src/routines/common.hpp b/src/routines/common.hpp
index 06d001d9..6cbe1e1b 100644
--- a/src/routines/common.hpp
+++ b/src/routines/common.hpp
@@ -239,6 +239,72 @@ void PadCopyTransposeMatrixBatched(Queue &queue, const Device &device,
   }
 }
 
+// Batched version of the above
+template <typename T>
+void PadCopyTransposeMatrixStridedBatched(Queue &queue, const Device &device,
+                                          const Databases &db,
+                                          EventPointer event, const std::vector<Event> &waitForEvents,
+                                          const size_t src_one, const size_t src_two,
+                                          const size_t src_ld, const size_t src_offset,
+                                          const size_t src_stride, const Buffer<T> &src,
+                                          const size_t dest_one, const size_t dest_two,
+                                          const size_t dest_ld, const size_t dest_offset,
+                                          const size_t dest_stride, const Buffer<T> &dest,
+                                          const Program &program, const bool do_pad,
+                                          const bool do_transpose, const bool do_conjugate,
+                                          const size_t batch_count) {
+
+  // Determines the right kernel
+  auto kernel_name = std::string{};
+  if (do_transpose) {
+    kernel_name = (do_pad) ? "TransposePadMatrixStridedBatched" : "TransposeMatrixStridedBatched";
+  }
+  else {
+    kernel_name = (do_pad) ? "CopyPadMatrixStridedBatched" : "CopyMatrixStridedBatched";
+  }
+
+  // Retrieves the kernel from the compiled binary
+  auto kernel = Kernel(program, kernel_name);
+
+  // Sets the kernel arguments
+  kernel.SetArgument(0, static_cast<int>(src_one));
+  kernel.SetArgument(1, static_cast<int>(src_two));
+  kernel.SetArgument(2, static_cast<int>(src_ld));
+  kernel.SetArgument(3, static_cast<int>(src_offset));
+  kernel.SetArgument(4, static_cast<int>(src_stride));
+  kernel.SetArgument(5, src());
+  kernel.SetArgument(6, static_cast<int>(dest_one));
+  kernel.SetArgument(7, static_cast<int>(dest_two));
+  kernel.SetArgument(8, static_cast<int>(dest_ld));
+  kernel.SetArgument(9, static_cast<int>(dest_offset));
+  kernel.SetArgument(10, static_cast<int>(dest_stride));
+  kernel.SetArgument(11, dest());
+  if (do_pad) {
+    kernel.SetArgument(12, static_cast<int>(do_conjugate));
+  }
+
+  // Launches the kernel and returns the error code. Uses global and local thread sizes based on
+  // parameters in the database.
+  if (do_transpose) {
+    const auto global = std::vector<size_t>{
+        Ceil(CeilDiv(dest_one, db["PADTRA_WPT"]), db["PADTRA_TILE"]),
+        Ceil(CeilDiv(dest_two, db["PADTRA_WPT"]), db["PADTRA_TILE"]),
+        batch_count
+    };
+    const auto local = std::vector<size_t>{db["PADTRA_TILE"], db["PADTRA_TILE"], 1};
+    RunKernel(kernel, queue, device, global, local, event, waitForEvents);
+  }
+  else {
+    const auto global = std::vector<size_t>{
+        Ceil(CeilDiv(dest_one, db["PAD_WPTX"]), db["PAD_DIMX"]),
+        Ceil(CeilDiv(dest_two, db["PAD_WPTY"]), db["PAD_DIMY"]),
+        batch_count
+    };
+    const auto local = std::vector<size_t>{db["PAD_DIMX"], db["PAD_DIMY"], 1};
+    RunKernel(kernel, queue, device, global, local, event, waitForEvents);
+  }
+}
+
 // =================================================================================================
 } // namespace clblast