Fix: invalid tile size (#1043)

Author: nathanielsimard

crates/cubecl-convolution/src/components/selection.rs

@@ -1,7 +1,9 @@
 use cubecl_core::{Runtime, client::ComputeClient};
 use cubecl_matmul::components::stage::PartitionBuffering;
 
-use cubecl_matmul::components::{MatmulElems, MatmulSelection, TilingScheme, adjust_dtypes};
+use cubecl_matmul::components::{
+    MatmulAvailabilityError, MatmulElems, MatmulSelection, TilingScheme, adjust_dtypes,
+};
 use cubecl_matmul::{
     components::tile::TileMatmulFamily,
     kernels::layered::{NUM_SM_APPROX, NUM_TENSOR_CORES_APPROX, find_instruction_size},
@@ -79,14 +81,14 @@ pub fn convolution_matmul_selection<TMM: TileMatmulFamily, R: Runtime>(
     problem: &ConvolutionProblem,
     plane_dim: u32,
     dtypes: &mut MatmulElems,
-) -> MatmulSelection {
+) -> Result<MatmulSelection, MatmulAvailabilityError> {
     adjust_dtypes::<R>(client, dtypes, TMM::requires_accelerator());
 
     // rough heuristic based on previous bench results where 512 channels with a 3x3 kernel seemed
     // to be the rough cutoff for the k=4 size.
     let stage_k = if problem.k >= 4096 { 4 } else { 2 };
 
-    let tile_size = find_instruction_size::<R, TMM>(client, dtypes, problem.m, problem.n);
+    let tile_size = find_instruction_size::<R, TMM>(client, dtypes, problem.m, problem.n)?;
 
     let hardware = &client.properties().hardware;
     let num_sm = hardware
@@ -111,7 +113,7 @@ pub fn convolution_matmul_selection<TMM: TileMatmulFamily, R: Runtime>(
         .build()
         .unwrap();
 
-    MatmulSelection::builder(tiling_scheme, plane_dim)
+    Ok(MatmulSelection::builder(tiling_scheme, plane_dim)
         .partition_buffering(PartitionBuffering::Single)
-        .build()
+        .build())
 }

crates/cubecl-convolution/src/kernels/layered/algorithm/multi_stage_tma.rs

@@ -70,6 +70,6 @@ impl<
             problem,
             plane_dim,
             matmul_elems,
-        ))
+        )?)
     }
 }

crates/cubecl-convolution/src/kernels/layered/algorithm/simple.rs

@@ -70,7 +70,7 @@ impl<
     ) -> Result<MatmulSelection, MatmulSetupError> {
         Ok(convolution_matmul_selection::<TMM, R>(
             client, problem, plane_dim, dtypes,
-        ))
+        )?)
     }
 }
 

crates/cubecl-convolution/src/kernels/layered/algorithm/simple_tma.rs

@@ -80,7 +80,7 @@ impl<
     ) -> Result<MatmulSelection, MatmulSetupError> {
         Ok(convolution_matmul_selection::<TMM, R>(
             client, problem, plane_dim, dtypes,
-        ))
+        )?)
     }
 }
 

crates/cubecl-matmul/src/components/error.rs

@@ -41,6 +41,9 @@ pub enum MatmulAvailabilityError {
         size: Option<TileSize>,
     },
 
+    /// Impossible to find a supported tile size for the problem.
+    TileSizeNotFound,
+
     /// The layout of the matmul is unsupported
     LayoutUnsupported {
         lhs: MatrixLayout,
@@ -172,6 +175,9 @@ impl Debug for MatmulAvailabilityError {
             MatmulAvailabilityError::PlaneOpsUnavailable => {
                 writeln!(f, "Plane-wide operations like plane_sum are not available.")
             }
+            MatmulAvailabilityError::TileSizeNotFound => {
+                writeln!(f, "No tile size is available for the problem.")
+            }
         }
     }
 }

crates/cubecl-matmul/src/kernels/layered/selector/plane.rs

@@ -50,7 +50,7 @@ pub fn plane_matmul_selection<TMM: TileMatmulFamily, R: Runtime>(
         ));
     }
 
-    let tile_size = find_instruction_size::<R, TMM>(client, dtypes, problem.m, problem.n);
+    let tile_size = find_instruction_size::<R, TMM>(client, dtypes, problem.m, problem.n)?;
 
     if options.tiny_selection_enabled && is_tiny(problem, &tile_size) {
         return Ok(selection_tiny::<R>(client, problem, tile_size, plane_dim));
@@ -258,7 +258,7 @@ pub fn find_instruction_size<R: Runtime, TMM: TileMatmulFamily>(
     elems: &MatmulElems,
     m: usize,
     n: usize,
-) -> TileSize {
+) -> Result<TileSize, MatmulAvailabilityError> {
     let supported = |m: u32, n: u32, k: u32| {
         TMM::is_supported::<R>(
             client,
@@ -273,7 +273,7 @@ pub fn find_instruction_size<R: Runtime, TMM: TileMatmulFamily>(
         )
     };
 
-    if m >= 4 * n && supported(32, 8, 16) {
+    let val = if m >= 4 * n && supported(32, 8, 16) {
         (32, 8, 16).into()
     } else if n >= 4 * n && supported(8, 32, 16) {
         (8, 32, 16).into()
@@ -282,16 +282,21 @@ pub fn find_instruction_size<R: Runtime, TMM: TileMatmulFamily>(
     } else if supported(8, 8, 8) {
         (8, 8, 8).into()
     } else {
-        TMM::supported_sizes::<R>(
+        match TMM::supported_sizes::<R>(
             client,
             elems.lhs_register,
             elems.rhs_register,
             elems.acc_register,
         )
         .first()
         .copied()
-        .unwrap_or_else(|| (16, 16, 8).into())
-    }
+        {
+            Some(val) => val,
+            None => return Err(MatmulAvailabilityError::TileSizeNotFound),
+        }
+    };
+
+    Ok(val)
 }
 
 fn selection_tiny<R: Runtime>(