Matmul: better error message (minor PR) (#313)

tracel-ai · Nov 27, 2024 · 5296f55 · 5296f55
1 parent 6de4938
commit 5296f55
Show file tree

Hide file tree

Showing 10 changed files with 30 additions and 15 deletions.
diff --git a/crates/cubecl-linalg/src/matmul/components/base.rs b/crates/cubecl-linalg/src/matmul/components/base.rs
@@ -15,7 +15,7 @@ pub trait MatmulKernel<I: Numeric, O: Numeric> {
     /// Checks if the client can handle the features used in this computation
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str>;
+    ) -> Result<(), String>;
 
     /// Create config for this matmul, given launch information
     fn make_config(

diff --git a/crates/cubecl-linalg/src/matmul/components/batch/one_to_many.rs b/crates/cubecl-linalg/src/matmul/components/batch/one_to_many.rs
@@ -85,7 +85,7 @@ impl<EG: Numeric, ES: Numeric, GMM: global::Matmul<EG, ES>, S: SpanMatmul, C: Cu
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         GMM::check_availability::<R>(client)
     }
 

diff --git a/crates/cubecl-linalg/src/matmul/components/batch/one_to_one.rs b/crates/cubecl-linalg/src/matmul/components/batch/one_to_one.rs
@@ -65,7 +65,7 @@ impl<EG: Numeric, ES: Numeric, GMM: global::Matmul<EG, ES>, C: CubeDispatch> Mat
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         GMM::check_availability::<R>(client)
     }
 

diff --git a/crates/cubecl-linalg/src/matmul/components/global/homogeneous/base.rs b/crates/cubecl-linalg/src/matmul/components/global/homogeneous/base.rs
@@ -152,7 +152,7 @@ where
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         SMM::check_availability::<R>(client)
     }
 

diff --git a/crates/cubecl-linalg/src/matmul/components/global/producer_consumer/base.rs b/crates/cubecl-linalg/src/matmul/components/global/producer_consumer/base.rs
@@ -182,7 +182,7 @@ where
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         SMM::check_availability::<R>(client)
     }
 

diff --git a/crates/cubecl-linalg/src/matmul/components/stage/multi_buffer/base.rs b/crates/cubecl-linalg/src/matmul/components/stage/multi_buffer/base.rs
@@ -170,7 +170,7 @@ where
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         TMM::check_availability::<R>(client)
     }
 

diff --git a/crates/cubecl-linalg/src/matmul/components/stage/single_buffer/base.rs b/crates/cubecl-linalg/src/matmul/components/stage/single_buffer/base.rs
@@ -160,7 +160,7 @@ where
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         TMM::check_availability::<R>(client)
     }
 

diff --git a/crates/cubecl-linalg/src/matmul/components/tile/accelerated.rs b/crates/cubecl-linalg/src/matmul/components/tile/accelerated.rs
@@ -100,7 +100,7 @@ macro_rules! instruction {
 
             fn check_availability<R: Runtime>(
                 client: &ComputeClient<R::Server, R::Channel>,
-            ) -> Result<(), &str> {
+            ) -> Result<(), String> {
                 check_availability::<I, O, R>(Self::M, Self::N, Self::K, client)
             }
 
@@ -232,7 +232,7 @@ fn check_availability<I: Numeric, O: Numeric, R: Runtime>(
     n: u32,
     k: u32,
     client: &ComputeClient<R::Server, R::Channel>,
-) -> Result<(), &str> {
+) -> Result<(), String> {
     if !client.properties().feature_enabled(Feature::Cmma {
         a: I::as_elem(),
         b: I::as_elem(),
@@ -241,7 +241,14 @@ fn check_availability<I: Numeric, O: Numeric, R: Runtime>(
         k: k as u8,
         n: n as u8,
     }) {
-        return Err("Cmma not supported.");
+        return Err(format!(
+            "Cmma on inputs {:?} and outputs {:?} with shape m={:?}, n={:?}, k={:?} not supported.",
+            I::as_elem(),
+            O::as_elem(),
+            m,
+            n,
+            k
+        ));
     }
 
     if !(client
@@ -251,7 +258,11 @@ fn check_availability<I: Numeric, O: Numeric, R: Runtime>(
             .properties()
             .feature_enabled(Feature::Type(O::as_elem())))
     {
-        return Err("Types not supported.");
+        return Err(format!(
+            "Types {:?} and/or {:?} not supported.",
+            I::as_elem(),
+            O::as_elem()
+        ));
     }
 
     Ok(())

diff --git a/crates/cubecl-linalg/src/matmul/components/tile/plane.rs b/crates/cubecl-linalg/src/matmul/components/tile/plane.rs
@@ -361,9 +361,9 @@ impl<I: Numeric, O: Numeric, const M: u32, const N: u32, const K: u32> MatmulKer
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         if !client.properties().feature_enabled(Feature::Plane) {
-            return Err("Planes not supported.");
+            return Err("Plane operations not supported.".to_string());
         }
 
         if !(client
@@ -373,7 +373,11 @@ impl<I: Numeric, O: Numeric, const M: u32, const N: u32, const K: u32> MatmulKer
                 .properties()
                 .feature_enabled(Feature::Type(O::as_elem())))
         {
-            return Err("Types not supported.");
+            return Err(format!(
+                "Types {:?} and/or {:?} not supported.",
+                I::as_elem(),
+                O::as_elem()
+            ));
         }
 
         Ok(())

diff --git a/crates/cubecl-linalg/src/matmul/kernels/matmul/algorithm/base.rs b/crates/cubecl-linalg/src/matmul/kernels/matmul/algorithm/base.rs
@@ -54,7 +54,7 @@ pub trait Algorithm<EG: Numeric> {
 
     fn check_availability<R: Runtime>(
         client: &ComputeClient<R::Server, R::Channel>,
-    ) -> Result<(), &str> {
+    ) -> Result<(), String> {
         Self::BatchMatmul::check_availability::<R>(client)
     }