Reduce TestMatrix calls for xgemmstridedbatched.

Replace the looped test by a single one with the offset of the last batch.

2 files changed, 32 insertions, 5 deletions

diff --git a/src/routines/levelx/xgemmstridedbatched.cpp b/src/routines/levelx/xgemmstridedbatched.cpp
index 5d8945ce..8fb83141 100644
--- a/src/routines/levelx/xgemmstridedbatched.cpp
+++ b/src/routines/levelx/xgemmstridedbatched.cpp
@@ -78,11 +78,9 @@ void XgemmStridedBatched<T>::DoGemmStridedBatched(const Layout layout, const Tra
                              gemm_kernel_id);
 
   // Tests the matrices for validity
-  for (auto batch = size_t{0}; batch < batch_count; ++batch) {
-    TestMatrixA(a_one, a_two, a_buffer, a_offset + a_stride * batch, a_ld);
-    TestMatrixB(b_one, b_two, b_buffer, b_offset + b_stride * batch, b_ld);
-    TestMatrixC(c_one, c_two, c_buffer, c_offset + c_stride * batch, c_ld);
-  }
+  TestStridedBatchedMatrixA(a_one, a_two, a_buffer, a_offset, a_stride, batch_count, a_ld);
+  TestStridedBatchedMatrixB(b_one, b_two, b_buffer, b_offset, b_stride, batch_count, b_ld);
+  TestStridedBatchedMatrixC(c_one, c_two, c_buffer, c_offset, c_stride, batch_count, c_ld);
 
   // Selects which version of the batched GEMM to run
   if (do_gemm_direct) { // single generic kernel
diff --git a/src/utilities/buffer_test.hpp b/src/utilities/buffer_test.hpp
index 9cecce97..4a2a2c95 100644
--- a/src/utilities/buffer_test.hpp
+++ b/src/utilities/buffer_test.hpp
@@ -134,6 +134,35 @@ void TestBatchedMatrixC(const size_t one, const size_t two, const Buffer<T>& buf
 }
 
 // =================================================================================================
+
+// Tests matrix 'A' for validity in a strided batched setting
+template <typename T>
+void TestStridedBatchedMatrixA(const size_t one, const size_t two, const Buffer<T>& buffer,
+                               const size_t offset, const size_t stride, const size_t batch_count,
+                               const size_t ld, const bool test_lead_dim = true) {
+  const auto last_batch_offset = (batch_count - 1) * stride;
+  TestMatrixA(one, two, buffer, offset + last_batch_offset, ld, test_lead_dim);
+}
+
+// Tests matrix 'B' for validity in a strided batched setting
+template <typename T>
+void TestStridedBatchedMatrixB(const size_t one, const size_t two, const Buffer<T>& buffer,
+                               const size_t offset, const size_t stride, const size_t batch_count,
+                               const size_t ld, const bool test_lead_dim = true) {
+  const auto last_batch_offset = (batch_count - 1) * stride;
+  TestMatrixB(one, two, buffer, offset + last_batch_offset, ld, test_lead_dim);
+}
+
+// Tests matrix 'C' for validity in a strided batched setting
+template <typename T>
+void TestStridedBatchedMatrixC(const size_t one, const size_t two, const Buffer<T>& buffer,
+                               const size_t offset, const size_t stride, const size_t batch_count,
+                               const size_t ld) {
+  const auto last_batch_offset = (batch_count - 1) * stride;
+  TestMatrixC(one, two, buffer, offset + last_batch_offset, ld);
+}
+
+// =================================================================================================
 } // namespace clblast
 
 // CLBLAST_BUFFER_TEST_H_