Improve tests for cubecl-reduce (#299)

tracel-ai · Nov 25, 2024 · 781f1c4 · 781f1c4
1 parent 2c09d4d
commit 781f1c4
Show file tree

Hide file tree

Showing 5 changed files with 193 additions and 219 deletions.
diff --git a/crates/cubecl-cuda/src/lib.rs b/crates/cubecl-cuda/src/lib.rs
@@ -20,5 +20,5 @@ mod tests {
     cubecl_linalg::testgen_plane_mma!([f16, bf16, f32], f32);
     cubecl_linalg::testgen_tiling2d!([f16, bf16, f32]);
     cubecl_linalg::testgen_cmma_old!([f16, bf16, f32 /*, f64*/]);
-    cubecl_reduce::testgen_reduce!();
+    cubecl_reduce::testgen_reduce!([f16, bf16, f32, f64]);
 }
diff --git a/crates/cubecl-hip/Cargo.toml b/crates/cubecl-hip/Cargo.toml
@@ -35,6 +35,7 @@ bytemuck = { workspace = true }
 derive-new = { workspace = true }
 half = { workspace = true }
 log = { workspace = true }
+paste = { workspace = true }
 
 [dev-dependencies]
 cubecl-core = { path = "../cubecl-core", version = "0.4.0", features = [

diff --git a/crates/cubecl-hip/src/lib.rs b/crates/cubecl-hip/src/lib.rs
@@ -23,9 +23,10 @@ pub(crate) type HipWmmaCompiler = cubecl_cpp::hip::wmma::WmmaIntrinsicCompiler;
 #[cfg(target_os = "linux")]
 #[cfg(test)]
 mod tests {
+    use half::{bf16, f16};
     pub type TestRuntime = crate::HipRuntime;
 
     cubecl_core::testgen_all!();
     cubecl_linalg::testgen_cmma_matmul!();
-    cubecl_reduce::testgen_reduce!();
+    cubecl_reduce::testgen_reduce!([f16, bf16, f32, f64]);
 }