From fe79d0ab958303e442598c7d492905fc2b3a6a75 Mon Sep 17 00:00:00 2001
From: Brodey Newman <brodeynewman@gmail.com>
Date: Thu, 5 Dec 2024 07:02:43 +0000
Subject: [PATCH] chore: test for cublas batched

---
 local.sh | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/local.sh b/local.sh
index ae0b62d..caa9c70 100755
--- a/local.sh
+++ b/local.sh
@@ -146,6 +146,23 @@ test_cudnn() {
   fi
 }
 
+test_cublas_batched() {
+  output=$(LD_PRELOAD="$libscuda_path" ./cublas_batched.o | tail -n 5)
+
+  expected_output=$'=====\nC[1]\n111.00 122.00\n151.00 166.00\n====='
+
+  # trim ugly output from the file
+  output=$(echo "$output" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+  expected_output=$(echo "$expected_output" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+
+  if [[ "$output" == "$expected_output" ]]; then
+    ansi_format "pass" "$pass_message"
+  else
+    ansi_format "fail" "test_cublas_batched failed. Got [$output]."
+    return 1
+  fi
+}
+
 #---- declare test cases ----#
 declare -A test_cuda_avail=(
   ["function"]="test_cuda_available"
@@ -172,8 +189,13 @@ declare -A test_cudnn=(
   ["pass"]="cuDNN correctly applies sigmoid activation on a tensor."
 )
 
+declare -A test_cublas_batched=(
+  ["function"]="test_cublas_batched"
+  ["pass"]="Batched cublas works via test/cublas_batched.cu."
+)
+
 #---- assign them to our associative array ----#
-tests=("test_cuda_avail" "test_tensor_to_cuda" "test_tensor_to_cuda_to_cpu" "test_vector_add" "test_cudnn")
+tests=("test_cuda_avail" "test_tensor_to_cuda" "test_tensor_to_cuda_to_cpu" "test_vector_add" "test_cudnn" "test_cublas_batched")
 
 test() {
   build