Feat/execution error (#1075)

Author: nathanielsimard

crates/cubecl-attention/src/base.rs

@@ -156,11 +156,8 @@ pub fn launch_attention<R: Runtime, A: Algorithm>(
         .hypercube_config()
         .cube_count_plan(&problem, &selection);
 
-    unsafe {
-        <BlackboxAcceleratedAlgorithm as Algorithm>::BatchAttention::launch_unchecked::<
-            TensorArgs,
-            R,
-        >(
+    let result = unsafe {
+        <BlackboxAcceleratedAlgorithm as Algorithm>::BatchAttention::launch_unchecked::<TensorArgs, R>(
             client,
             config.cube_dim(),
             cube_count_plan.resolve(),
@@ -176,8 +173,11 @@ pub fn launch_attention<R: Runtime, A: Algorithm>(
             cube_count_plan.as_args(),
             config,
             attention_elems,
-        );
-    }
+        )
+    };
 
-    Ok(())
+    match result {
+        Ok(_) => Ok(()),
+        Err(err) => Err(AttentionSetupError::Execution(err)),
+    }
 }

crates/cubecl-attention/src/components/batch/base.rs

@@ -35,7 +35,7 @@ pub trait BatchAttentionFamily: Send + Sync + 'static {
         cube_count_input: CubeCountInputArgs<'a, R>,
         config: Self::Config,
         dtypes: &AttentionElems,
-    );
+    ) -> Result<(), LaunchError>;
 
     /// Constructs the configuration based on the Attention problem, selection, and line sizes.
     ///

crates/cubecl-attention/src/components/batch/simple/setup.rs

@@ -1,6 +1,6 @@
 use std::marker::PhantomData;
 
-use cubecl_core::client::ComputeClient;
+use cubecl_core::{client::ComputeClient, server::LaunchError};
 
 use crate::components::{
     AttentionElems, AttentionLineSizes, AttentionPrecision, AttentionProblem, AttentionSelection,
@@ -49,7 +49,7 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for SimpleBatchAttentionFam
         cube_count_input: crate::components::batch::CubeCountInputArgs<'a, R>,
         config: Self::Config,
         dtypes: &AttentionElems,
-    ) {
+    ) -> Result<(), LaunchError> {
         unsafe {
             attention::launch_unchecked::<AA, Self, R>(
                 client,
@@ -60,7 +60,7 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for SimpleBatchAttentionFam
                 cube_count_input,
                 config,
                 dtypes.into(),
-            );
+            )
         }
     }
 }

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

@@ -1,4 +1,4 @@
-use cubecl_core::{CubeCount, CubeDim, LineSizeError};
+use cubecl_core::{CubeCount, CubeDim, LineSizeError, server::LaunchError};
 use cubecl_matmul::components::MatmulSetupError;
 use std::fmt::{Debug, Display};
 
@@ -15,6 +15,9 @@ pub enum AttentionSetupError {
 
     /// Error in underlying matmul
     MatmulSetup(MatmulSetupError),
+
+    /// An error that happened during execution.
+    Execution(LaunchError),
 }
 
 /// A specific feature required for attention is not available in the current runtime or hardware.
@@ -75,6 +78,9 @@ impl Debug for AttentionSetupError {
             AttentionSetupError::MatmulSetup(matmul_setup_error) => {
                 writeln!(f, "{matmul_setup_error:?}")
             }
+            AttentionSetupError::Execution(error) => {
+                writeln!(f, "{error:?}")
+            }
         }
     }
 }

crates/cubecl-attention/src/tests/attention_test_launcher.rs

@@ -95,7 +95,7 @@ pub fn test_attention_algorithm<A, P, R>(
         .hypercube_config()
         .cube_count_plan(&problem, &selection);
 
-    unsafe {
+    let result = unsafe {
         A::BatchAttention::launch_unchecked::<TensorArgs, R>(
             &client,
             config.cube_dim(),
@@ -138,8 +138,16 @@ pub fn test_attention_algorithm<A, P, R>(
             cube_count_plan.as_args(),
             config,
             &attention_elems,
-        );
-    }
+        )
+    };
+
+    match result {
+        Ok(_) => {}
+        Err(err) => {
+            println!("Skipping the test with an execution error {err:?}");
+            return;
+        }
+    };
 
     P::assert_result(
         &query.original_data.unwrap(),

crates/cubecl-attention/src/tests/test_utils.rs

@@ -276,7 +276,7 @@ macro_rules! sample_float {
                     let dtype = Self::as_type_native_unchecked();
                     let output = TensorHandle::empty(client, shape.to_vec(), dtype);
 
-                    cubecl_random::random_uniform(&client, f32::from_int(-1), f32::from_int(1), output.as_ref(), dtype);
+                    cubecl_random::random_uniform(&client, f32::from_int(-1), f32::from_int(1), output.as_ref(), dtype).unwrap();
 
                     output
                 }
@@ -306,7 +306,8 @@ impl Sampleable for flex32 {
             f32::from_int(1),
             output.as_ref(),
             dtype,
-        );
+        )
+        .unwrap();
 
         output
     }
@@ -328,7 +329,8 @@ impl Sampleable for tf32 {
             f32::from_int(1),
             output.as_ref(),
             dtype,
-        );
+        )
+        .unwrap();
 
         output
     }
@@ -344,7 +346,7 @@ impl Sampleable for bool {
         let dtype = bool::as_type_native_unchecked();
         let output = TensorHandle::empty(client, shape.to_vec(), dtype);
 
-        cubecl_random::random_bernoulli(client, 0.5, output.as_ref(), dtype);
+        cubecl_random::random_bernoulli(client, 0.5, output.as_ref(), dtype).unwrap();
 
         output
     }
@@ -360,7 +362,7 @@ impl Sampleable for u8 {
         let dtype = u8::as_type_native_unchecked();
         let output = TensorHandle::empty(client, shape.to_vec(), dtype);
 
-        cubecl_random::random_bernoulli(client, 0.5, output.as_ref(), dtype);
+        cubecl_random::random_bernoulli(client, 0.5, output.as_ref(), dtype).unwrap();
 
         output
     }

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

@@ -1,11 +1,19 @@
+use core::fmt::Debug;
+use cubecl_core::server::LaunchError;
 use cubecl_matmul::components::{MatmulAvailabilityError, MatmulSetupError};
-use std::fmt::Debug;
 
 #[allow(clippy::large_enum_variant)]
 pub enum ConvSetupError {
     Matmul(MatmulSetupError),
     Groups(usize),
     Unknown,
+    Launch(LaunchError),
+}
+
+impl From<LaunchError> for ConvSetupError {
+    fn from(value: LaunchError) -> Self {
+        Self::Launch(value)
+    }
 }
 
 impl Debug for ConvSetupError {
@@ -21,6 +29,7 @@ impl Debug for ConvSetupError {
                 )
             }
             ConvSetupError::Unknown => write!(f, "Unknown"),
+            ConvSetupError::Launch(err) => write!(f, "Launch error {err:?}"),
         }
     }
 }

crates/cubecl-convolution/src/components/global/entry_point.rs

@@ -38,7 +38,7 @@ pub trait ConvolutionLaunch<Config> {
         problem: &ConvolutionProblem,
         config: Config,
         dtypes: &MatmulElems,
-    );
+    ) -> Result<(), LaunchError>;
 }
 
 #[cube(launch_unchecked)]

crates/cubecl-convolution/src/components/global/multi_stage/tma/launch.rs

@@ -1,4 +1,6 @@
-use cubecl_core::{CubeCount, CubeDim, Runtime, client::ComputeClient, prelude::ScalarArg};
+use cubecl_core::{
+    CubeCount, CubeDim, Runtime, client::ComputeClient, prelude::ScalarArg, server::LaunchError,
+};
 use cubecl_matmul::components::{
     InputRuntimeArg, MatmulElems, OutputRuntimeArg,
     global::{PartitionedStageFamily, args::MatmulArgs},
@@ -36,7 +38,7 @@ impl<
         problem: &ConvolutionProblem,
         config: GlobalConfig<Self>,
         dtypes: &MatmulElems,
-    ) {
+    ) -> Result<(), LaunchError> {
         let padded_channels =
             (problem.channels as u32).next_multiple_of(config.stage_config.elements_in_tile_k());
 
@@ -66,7 +68,7 @@ impl<
                 dtypes.lhs_stage,
                 dtypes.rhs_stage,
                 dtypes.acc_stage,
-            );
+            )
         }
     }
 }

crates/cubecl-convolution/src/components/global/single_stage/simple/launch.rs

@@ -1,4 +1,6 @@
-use cubecl_core::{CubeCount, CubeDim, Runtime, client::ComputeClient, prelude::ScalarArg};
+use cubecl_core::{
+    CubeCount, CubeDim, Runtime, client::ComputeClient, prelude::ScalarArg, server::LaunchError,
+};
 use cubecl_matmul::components::{
     InputRuntimeArg, MatmulElems, OutputRuntimeArg,
     global::{PartitionedStageFamily, args::MatmulArgs},
@@ -36,7 +38,7 @@ impl<
         problem: &ConvolutionProblem,
         config: GlobalConfig<Self>,
         dtypes: &MatmulElems,
-    ) {
+    ) -> Result<(), LaunchError> {
         let shape_channels = FastDivmodArgs::new(client, problem.channels as u32);
 
         let runtime_args = RuntimeArgsLaunch::new(
@@ -63,7 +65,7 @@ impl<
                 dtypes.lhs_stage,
                 dtypes.rhs_stage,
                 dtypes.acc_stage,
-            );
+            )
         }
     }
 }