chore: separate test_sgemm_gpu/cpu tests

huggingface · Aug 16, 2024 · fea7140 · fea7140
1 parent 51f1c0b
commit fea7140
Showing 1 changed file with 19 additions and 1 deletion.
diff --git a/crates/ratchet-core/src/ops/matmul/mod.rs b/crates/ratchet-core/src/ops/matmul/mod.rs
@@ -831,7 +831,25 @@ def matmul(a, b{}):
     }
 
     #[proptest(cases = 64)]
-    fn test_sgemm(prob: SGEMMProblem) {
+    fn test_sgemm_gpu(prob: SGEMMProblem) {
+        let device = Device::request_device(DeviceRequest::GPU).unwrap();
+        let SGEMMProblem {
+            B,
+            M,
+            K,
+            N,
+            has_bias,
+            ref transpose,
+        } = prob;
+        println!(
+            "Running sgemm: B={} M={} N={} K={} has_bias={} transpose={:?}",
+            B, M, N, K, has_bias, transpose
+        );
+        run_matmul_trial(&device, prob).unwrap();
+    }
+
+    #[proptest(cases = 64)]
+    fn test_sgemm_cpu(prob: SGEMMProblem) {
         let device = Device::request_device(DeviceRequest::CPU).unwrap();
         let SGEMMProblem {
             B,