Relaxed requirement on a_ld and b_ld for batched GEMM

3 files changed, 15 insertions, 6 deletions

diff --git a/src/routines/levelx/xgemmbatched.cpp b/src/routines/levelx/xgemmbatched.cpp
index 0fea1922..ee8448d2 100644
--- a/src/routines/levelx/xgemmbatched.cpp
+++ b/src/routines/levelx/xgemmbatched.cpp
@@ -94,8 +94,8 @@ void XgemmBatched<T>::DoGemmBatched(const Layout layout, const Transpose a_trans
 
   // Tests the matrices for validity
   for (auto batch = size_t{0}; batch < batch_count; ++batch) {
-    TestMatrixA(a_one, a_two, a_buffer, a_offsets[batch], a_ld);
-    TestMatrixB(b_one, b_two, b_buffer, b_offsets[batch], b_ld);
+    TestMatrixA(a_one, a_two, a_buffer, a_offsets[batch], a_ld, false); // don't test for invalid LD
+    TestMatrixB(b_one, b_two, b_buffer, b_offsets[batch], b_ld, false); // don't test for invalid LD
     TestMatrixC(c_one, c_two, c_buffer, c_offsets[batch], c_ld);
   }
 
diff --git a/src/utilities/buffer_test.hpp b/src/utilities/buffer_test.hpp
index 652ab8c6..b5693181 100644
--- a/src/utilities/buffer_test.hpp
+++ b/src/utilities/buffer_test.hpp
@@ -23,8 +23,8 @@ namespace clblast {
 // Tests matrix 'A' for validity
 template <typename T>
 void TestMatrixA(const size_t one, const size_t two, const Buffer<T> &buffer,
-                 const size_t offset, const size_t ld) {
-  if (ld < one) { throw BLASError(StatusCode::kInvalidLeadDimA); }
+                 const size_t offset, const size_t ld, const bool test_lead_dim = true) {
+  if (test_lead_dim && ld < one) { throw BLASError(StatusCode::kInvalidLeadDimA); }
   try {
     const auto required_size = (ld * (two - 1) + one + offset) * sizeof(T);
     if (buffer.GetSize() < required_size) { throw BLASError(StatusCode::kInsufficientMemoryA); }
@@ -34,8 +34,8 @@ void TestMatrixA(const size_t one, const size_t two, const Buffer<T> &buffer,
 // Tests matrix 'B' for validity
 template <typename T>
 void TestMatrixB(const size_t one, const size_t two, const Buffer<T> &buffer,
-                 const size_t offset, const size_t ld) {
-  if (ld < one) { throw BLASError(StatusCode::kInvalidLeadDimB); }
+                 const size_t offset, const size_t ld, const bool test_lead_dim = true) {
+  if (test_lead_dim && ld < one) { throw BLASError(StatusCode::kInvalidLeadDimB); }
   try {
     const auto required_size = (ld * (two - 1) + one + offset) * sizeof(T);
     if (buffer.GetSize() < required_size) { throw BLASError(StatusCode::kInsufficientMemoryB); }
diff --git a/test/routines/levelx/xgemmbatched.hpp b/test/routines/levelx/xgemmbatched.hpp
index 56823e47..704d0578 100644
--- a/test/routines/levelx/xgemmbatched.hpp
+++ b/test/routines/levelx/xgemmbatched.hpp
@@ -110,6 +110,15 @@ class TestXgemmBatched {
   static StatusCode RunRoutine(const Arguments<T> &args, Buffers<T> &buffers, Queue &queue) {
     auto queue_plain = queue();
     auto event = cl_event{};
+    // Relaxed requirement on ld_a and ld_b within the library, this is here to match clBLAS
+    auto a_rotated = (args.layout == Layout::kColMajor && args.a_transpose != Transpose::kNo) ||
+                     (args.layout == Layout::kRowMajor && args.a_transpose == Transpose::kNo);
+    auto b_rotated = (args.layout == Layout::kColMajor && args.b_transpose != Transpose::kNo) ||
+                     (args.layout == Layout::kRowMajor && args.b_transpose == Transpose::kNo);
+    auto a_one = (!a_rotated) ? args.m : args.k;
+    auto b_one = (!b_rotated) ? args.k : args.n;
+    if (args.a_ld < a_one) { return StatusCode::kInvalidLeadDimA; }
+    if (args.b_ld < b_one) { return StatusCode::kInvalidLeadDimB; }
     auto status = GemmBatched(args.layout, args.a_transpose, args.b_transpose,
                               args.m, args.n, args.k, args.alphas.data(),
                               buffers.a_mat(), args.a_offsets.data(), args.a_ld,