Added PrepareData function for TRSM to create proper test input

3 files changed, 42 insertions, 3 deletions

diff --git a/test/routines/level2/xtrsv.hpp b/test/routines/level2/xtrsv.hpp
index 811feac5..72ebdf9e 100644
--- a/test/routines/level2/xtrsv.hpp
+++ b/test/routines/level2/xtrsv.hpp
@@ -46,10 +46,10 @@ void PrepareData(const Arguments<T> &args, Buffers<T> &buffers, Queue &queue) {
     auto diagonal = a_mat_cpu[i*args.a_ld + i + args.a_offset];
     diagonal = AbsoluteValue(diagonal) + static_cast<T>(args.n / size_t{4});
     for (auto j = size_t{0}; j < args.n; ++j) {
-      a_mat_cpu[j*args.a_ld + i + args.a_offset] /= T{2.0};
+      a_mat_cpu[j*args.a_ld + i + args.a_offset] /= ConstantTwo<T>();
     }
     a_mat_cpu[i*args.a_ld + i + args.a_offset] = diagonal;
-    x_vec_cpu[i * args.x_inc + args.x_offset] /= T{2.0};
+    x_vec_cpu[i * args.x_inc + args.x_offset] /= ConstantTwo<T>();
   }
 
   // Copies input buffers back to the OpenCL device
diff --git a/test/routines/level3/xtrsm.hpp b/test/routines/level3/xtrsm.hpp
index e59c96ff..246cb930 100644
--- a/test/routines/level3/xtrsm.hpp
+++ b/test/routines/level3/xtrsm.hpp
@@ -29,6 +29,37 @@
 namespace clblast {
 // =================================================================================================
 
+// Prepares the data
+template <typename T>
+void PrepareData(const Arguments<T> &args, Buffers<T> &buffers, Queue &queue) {
+  const auto k = (args.side == Side::kLeft) ? args.m : args.n;
+  if (args.a_ld < k) { return; }
+
+  // Copies input buffers to the host
+  std::vector<T> a_mat_cpu(args.a_size, static_cast<T>(0));
+  std::vector<T> b_mat_cpu(args.b_size, static_cast<T>(0));
+  buffers.a_mat.Read(queue, args.a_size, a_mat_cpu);
+  buffers.b_mat.Read(queue, args.b_size, b_mat_cpu);
+
+  // Generates 'proper' input for the TRSM routine
+  // TODO: Improve this
+  for (auto i = size_t{0}; i < k; ++i) {
+    for (auto j = size_t{0}; j < k; ++j) {
+      auto value = a_mat_cpu[j*args.a_ld + i + args.a_offset];
+      value *= ConstantTwo<T>();
+      if (IsCloseToZero(value)) { value += ConstantOne<T>(); }
+      a_mat_cpu[j*args.a_ld + i + args.a_offset] = value;
+    }
+  }
+
+  // Copies input buffers back to the OpenCL device
+  buffers.a_mat.Write(queue, args.a_size, a_mat_cpu);
+  buffers.b_mat.Write(queue, args.b_size, b_mat_cpu);
+  return;
+}
+
+// =================================================================================================
+
 // See comment at top of file for a description of the class
 template <typename T>
 class TestXtrsm {
@@ -75,6 +106,7 @@ class TestXtrsm {
 
   // Describes how to run the CLBlast routine
   static StatusCode RunRoutine(const Arguments<T> &args, Buffers<T> &buffers, Queue &queue) {
+    PrepareData(args, buffers, queue);
     auto queue_plain = queue();
     auto event = cl_event{};
     auto status = Trsm(args.layout, args.side, args.triangle, args.a_transpose, args.diagonal,
@@ -89,6 +121,7 @@ class TestXtrsm {
   // Describes how to run the clBLAS routine (for correctness/performance comparison)
   #ifdef CLBLAST_REF_CLBLAS
     static StatusCode RunReference1(const Arguments<T> &args, Buffers<T> &buffers, Queue &queue) {
+      PrepareData(args, buffers, queue);
       auto queue_plain = queue();
       auto event = cl_event{};
       auto status = clblasXtrsm(convertToCLBLAS(args.layout),
@@ -108,6 +141,7 @@ class TestXtrsm {
   // Describes how to run the CPU BLAS routine (for correctness/performance comparison)
   #ifdef CLBLAST_REF_CBLAS
     static StatusCode RunReference2(const Arguments<T> &args, Buffers<T> &buffers, Queue &queue) {
+      PrepareData(args, buffers, queue);
       std::vector<T> a_mat_cpu(args.a_size, static_cast<T>(0));
       std::vector<T> b_mat_cpu(args.b_size, static_cast<T>(0));
       buffers.a_mat.Read(queue, args.a_size, a_mat_cpu);
diff --git a/test/routines/levelx/xinvert.hpp b/test/routines/levelx/xinvert.hpp
index 4408e8d5..c6ce4b07 100644
--- a/test/routines/levelx/xinvert.hpp
+++ b/test/routines/levelx/xinvert.hpp
@@ -41,9 +41,14 @@ StatusCode RunReference(const Arguments<T> &args, Buffers<T> &buffers, Queue &qu
   const auto num_blocks = CeilDiv(args.n, block_size);
   const auto a_ld = args.a_ld;
   const auto b_ld = block_size;
-  if ((block_size == 0) || (args.n == 0) || (block_size > args.n)) {
+
+  // Checks for valid arguments
+  if ((block_size == 0) || (args.n == 0)) {
     return StatusCode::kInvalidDimension;
   }
+  if ((block_size % 16 != 0) || (block_size > 128)) {
+    return StatusCode::kUnknownError;
+  }
 
   // Loops over the amount of diagonal blocks of size args.m by args.m each
   for (auto block_id = size_t{0}; block_id < num_blocks; ++block_id) {