[Feat] Debug improvements (#474)

crates/cubecl-core/Cargo.toml

@@ -44,6 +44,7 @@ num-traits = { workspace = true }
 paste = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+variadics_please = { workspace = true }
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }

crates/cubecl-core/src/frontend/cmma.rs

@@ -53,7 +53,9 @@ use crate::{
     unexpanded,
 };
 
-use super::{CubePrimitive, CubeType, ExpandElementTyped, Init, IntoRuntime, Slice, SliceMut};
+use super::{
+    CubeDebug, CubePrimitive, CubeType, ExpandElementTyped, Init, IntoRuntime, Slice, SliceMut,
+};
 
 use cubecl_ir::{ExpandElement, Scope};
 pub use ir::{MatrixIdent, MatrixLayout};
@@ -100,6 +102,12 @@ impl<C: CubeType> Init for MatrixExpand<C> {
     }
 }
 
+impl<C: CubeType> CubeDebug for MatrixExpand<C> {
+    fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
+        scope.update_variable_name(*self.elem, name);
+    }
+}
+
 impl<C: CubePrimitive> Matrix<C> {
     /// Create a new uninitialized matrix that is going to be used in the
     /// [matrix-multiply and accumulate](execute()) function.
@@ -465,3 +473,5 @@ impl Init for MatrixLayout {
         self
     }
 }
+
+impl CubeDebug for MatrixLayout {}

crates/cubecl-core/src/frontend/container/registry/base.rs

@@ -2,7 +2,7 @@ use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
 
 use cubecl_ir::Scope;
 
-use crate::prelude::{CubeType, ExpandElementTyped, Init, IntoRuntime};
+use crate::prelude::{CubeDebug, CubeType, ExpandElementTyped, Init, IntoRuntime};
 
 /// It is similar to a map, but where the keys are stored at comptime, but the values can be runtime
 /// variables.
@@ -136,6 +136,8 @@ impl<K: PartialOrd + Ord, V> Init for Registry<K, V> {
     }
 }
 
+impl<K: PartialOrd + Ord, V> CubeDebug for Registry<K, V> {}
+
 impl<K: PartialOrd + Ord, V: CubeType> IntoRuntime for Registry<K, V> {
     fn __expand_runtime_method(self, _scope: &mut Scope) -> Registry<K, V::ExpandType> {
         unimplemented!("Comptime registry can't be moved to runtime.");

crates/cubecl-core/src/frontend/container/sequence/base.rs

@@ -1,8 +1,11 @@
 use cubecl_ir::{ExpandElement, Scope};
 use serde::{Deserialize, Serialize};
 
-use crate::frontend::{
-    branch::Iterable, indexation::Index, CubeType, ExpandElementTyped, Init, IntoRuntime,
+use crate::{
+    frontend::{
+        branch::Iterable, indexation::Index, CubeType, ExpandElementTyped, Init, IntoRuntime,
+    },
+    prelude::CubeDebug,
 };
 use std::{cell::RefCell, rc::Rc};
 
@@ -29,6 +32,7 @@ impl<T: CubeType> Init for Sequence<T> {
         self
     }
 }
+impl<T: CubeType> CubeDebug for Sequence<T> {}
 
 impl<T: CubeType + Clone> Sequence<T> {
     pub fn rev(&self) -> Self {
@@ -144,6 +148,7 @@ impl<T: CubeType> Init for SequenceExpand<T> {
         self
     }
 }
+impl<T: CubeType> CubeDebug for SequenceExpand<T> {}
 
 impl<T: CubeType> Clone for SequenceExpand<T> {
     fn clone(&self) -> Self {

crates/cubecl-core/src/frontend/debug.rs

@@ -1,29 +1,25 @@
+use cubecl_ir::CubeFnSource;
+
 use crate::ir::{NonSemantic, Scope, Variable};
 
+use super::CubeDebug;
+
 /// Calls a function and inserts debug symbols if debug is enabled.
 #[track_caller]
 pub fn debug_call_expand<C>(
     scope: &mut Scope,
-    name: &'static str,
     line: u32,
     col: u32,
     call: impl FnOnce(&mut Scope) -> C,
 ) -> C {
-    if scope.debug_enabled {
-        scope.register(NonSemantic::BeginCall {
-            name: name.to_string(),
-            line,
-            col,
-        });
-
-        let ret = call(scope);
-
-        scope.register(NonSemantic::EndCall);
-
-        ret
-    } else {
-        call(scope)
-    }
+    // Save source_loc before the call so it can be restored once the call returns
+    let source_loc = scope.debug.source_loc.take();
+    scope.update_span(line, col);
+    scope.register(NonSemantic::EnterDebugScope);
+    let ret = call(scope);
+    scope.register(NonSemantic::ExitDebugScope);
+    scope.debug.source_loc = source_loc;
+    ret
 }
 
 /// Calls an intrinsic op and inserts debug symbols if debug is enabled.
@@ -34,27 +30,35 @@ pub fn spanned_expand<C>(
     col: u32,
     call: impl FnOnce(&mut Scope) -> C,
 ) -> C {
-    if scope.debug_enabled {
-        scope.register(NonSemantic::Line { line, col });
-        call(scope)
-    } else {
-        call(scope)
-    }
+    scope.update_span(line, col);
+    call(scope)
 }
 
 /// Adds source instruction if debug is enabled
 #[track_caller]
-pub fn debug_source_expand(scope: &mut Scope, name: &str, file: &str, line: u32, col: u32) {
-    if scope.debug_enabled {
-        // Normalize to linux separators
-        let file = file.replace("\\", "/");
-        scope.register(NonSemantic::Source {
-            name: name.into(),
-            file_name: format!("./{file}"),
-            line,
-            col,
-        });
-    }
+pub fn debug_source_expand(
+    scope: &mut Scope,
+    name: &'static str,
+    file: &'static str,
+    source_text: &'static str,
+    line: u32,
+    column: u32,
+) {
+    let file = file.replace("\\", "/");
+    scope.update_source(CubeFnSource {
+        function_name: name.into(),
+        file: file.into(),
+        source_text: source_text.into(),
+        line,
+        column,
+    });
+}
+
+/// Registers name for an expand if possible
+#[track_caller]
+pub fn debug_var_expand<E: CubeDebug>(scope: &mut Scope, name: &'static str, expand: E) -> E {
+    expand.set_debug_name(scope, name);
+    expand
 }
 
 /// Prints a formatted message using the print debug layer in Vulkan, or `printf` in CUDA.

crates/cubecl-core/src/frontend/element/base.rs

@@ -8,6 +8,7 @@ use cubecl_common::{flex32, tf32};
 use cubecl_ir::ExpandElement;
 use half::{bf16, f16};
 use std::marker::PhantomData;
+use variadics_please::all_tuples;
 
 /// Types used in a cube function must implement this trait
 ///
@@ -23,7 +24,7 @@ use std::marker::PhantomData;
 /// the generated code.
 #[diagnostic::on_unimplemented(note = "Consider using `#[derive(CubeType)]` on `{Self}`")]
 pub trait CubeType {
-    type ExpandType: Clone + Init;
+    type ExpandType: Clone + Init + CubeDebug;
 
     /// Wrapper around the init method, necessary to type inference.
     fn init(scope: &mut Scope, expand: Self::ExpandType) -> Self::ExpandType {
@@ -52,6 +53,13 @@ pub trait Init: Sized {
     fn init(self, scope: &mut Scope) -> Self;
 }
 
+pub trait CubeDebug: Sized {
+    /// Set the debug name of this type's expansion. Should do nothing for types that don't appear
+    /// at runtime
+    #[allow(unused)]
+    fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {}
+}
+
 /// Argument used during the compilation of kernels.
 pub trait CompilationArg:
     serde::Serialize
@@ -139,16 +147,6 @@ impl LaunchArgExpand for () {
     }
 }
 
-impl CubeType for () {
-    type ExpandType = ();
-}
-
-impl Init for () {
-    fn init(self, _scope: &mut Scope) -> Self {
-        self
-    }
-}
-
 impl<T: Clone> CubeType for PhantomData<T> {
     type ExpandType = ();
 }
@@ -208,7 +206,7 @@ macro_rules! tuple_cube_type {
 macro_rules! tuple_init {
     ($($P:ident),*) => {
         impl<$($P: Init),*> Init for ($($P,)*) {
-            #[allow(non_snake_case)]
+            #[allow(non_snake_case, unused, clippy::unused_unit)]
             fn init(self, scope: &mut Scope) -> Self {
                 let ($($P,)*) = self;
                 ($(
@@ -218,10 +216,15 @@ macro_rules! tuple_init {
         }
     }
 }
+macro_rules! tuple_debug {
+    ($($P:ident),*) => {
+        impl<$($P: CubeDebug),*> CubeDebug for ($($P,)*) {}
+    }
+}
 macro_rules! tuple_runtime {
     ($($P:ident),*) => {
         impl<$($P: IntoRuntime),*> IntoRuntime for ($($P,)*) {
-            #[allow(non_snake_case)]
+            #[allow(non_snake_case, unused, clippy::unused_unit)]
             fn __expand_runtime_method(self, scope: &mut Scope) -> Self::ExpandType {
                 let ($($P,)*) = self;
                 ($(
@@ -232,26 +235,12 @@ macro_rules! tuple_runtime {
     }
 }
 
-tuple_cube_type!(P1);
-tuple_cube_type!(P1, P2);
-tuple_cube_type!(P1, P2, P3);
-tuple_cube_type!(P1, P2, P3, P4);
-tuple_cube_type!(P1, P2, P3, P4, P5);
-tuple_cube_type!(P1, P2, P3, P4, P5, P6);
+all_tuples!(tuple_cube_type, 0, 12, P);
+all_tuples!(tuple_debug, 0, 12, P);
+all_tuples!(tuple_init, 0, 12, P);
+all_tuples!(tuple_runtime, 0, 12, P);
 
-tuple_init!(P1);
-tuple_init!(P1, P2);
-tuple_init!(P1, P2, P3);
-tuple_init!(P1, P2, P3, P4);
-tuple_init!(P1, P2, P3, P4, P5);
-tuple_init!(P1, P2, P3, P4, P5, P6);
-
-tuple_runtime!(P1);
-tuple_runtime!(P1, P2);
-tuple_runtime!(P1, P2, P3);
-tuple_runtime!(P1, P2, P3, P4);
-tuple_runtime!(P1, P2, P3, P4, P5);
-tuple_runtime!(P1, P2, P3, P4, P5, P6);
+impl<P: CubePrimitive> CubeDebug for P {}
 
 pub trait ExpandElementBaseInit: CubeType {
     fn init_elem(scope: &mut Scope, elem: ExpandElement) -> ExpandElement;
@@ -263,6 +252,12 @@ impl<T: ExpandElementBaseInit> Init for ExpandElementTyped<T> {
     }
 }
 
+impl<T: CubeType> CubeDebug for ExpandElementTyped<T> {
+    fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
+        scope.update_variable_name(*self.expand, name);
+    }
+}
+
 impl<T: CubeType> ExpandElementTyped<T> {
     // Expanded version of vectorization factor.
     pub fn __expand_vectorization_factor_method(self, _scope: &mut Scope) -> u32 {
@@ -375,6 +370,7 @@ impl<T: Init> Init for Vec<T> {
         self.into_iter().map(|e| e.init(scope)).collect()
     }
 }
+impl<T: CubeDebug> CubeDebug for Vec<T> {}
 
 /// Create a constant element of the correct type during expansion.
 pub(crate) fn __expand_new<C: Numeric, Out: Numeric>(

crates/cubecl-core/src/frontend/pipeline.rs

@@ -72,12 +72,13 @@ use std::marker::PhantomData;
 use cubecl_ir::ExpandElement;
 
 use crate::{
-    ir::{Instruction, Item, Operation, PipelineOps, Scope},
+    ir::{Item, PipelineOps, Scope},
     unexpanded,
 };
 
 use super::{
-    CubePrimitive, CubeType, ExpandElementTyped, Init, IntoRuntime, Line, Slice, SliceMut,
+    CubeDebug, CubePrimitive, CubeType, ExpandElementTyped, Init, IntoRuntime, Line, Slice,
+    SliceMut,
 };
 
 /// A mechanism for managing a sequence of `memcpy_async`
@@ -103,6 +104,12 @@ impl<C: CubePrimitive> Init for PipelineExpand<C> {
     }
 }
 
+impl<C: CubePrimitive> CubeDebug for PipelineExpand<C> {
+    fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
+        scope.update_variable_name(*self.elem, name);
+    }
+}
+
 #[derive(Clone)]
 /// Expand type of [Pipeline]
 pub struct PipelineExpand<C: CubePrimitive> {
@@ -204,38 +211,23 @@ impl<C: CubePrimitive> PipelineExpand<C> {
             destination,
         };
 
-        scope.register(Instruction {
-            out: None,
-            operation: Operation::Pipeline(mem_copy),
-        });
+        scope.register(mem_copy);
     }
 
     pub fn __expand_producer_acquire_method(&self, scope: &mut Scope) {
         let pipeline = *self.elem;
-        scope.register(Instruction {
-            out: None,
-            operation: Operation::Pipeline(PipelineOps::ProducerAcquire { pipeline }),
-        });
+        scope.register(PipelineOps::ProducerAcquire { pipeline });
     }
     pub fn __expand_producer_commit_method(&self, scope: &mut Scope) {
         let pipeline = *self.elem;
-        scope.register(Instruction {
-            out: None,
-            operation: Operation::Pipeline(PipelineOps::ProducerCommit { pipeline }),
-        });
+        scope.register(PipelineOps::ProducerCommit { pipeline });
     }
     pub fn __expand_consumer_wait_method(&self, scope: &mut Scope) {
         let pipeline = *self.elem;
-        scope.register(Instruction {
-            out: None,
-            operation: Operation::Pipeline(PipelineOps::ConsumerWait { pipeline }),
-        });
+        scope.register(PipelineOps::ConsumerWait { pipeline });
     }
     pub fn __expand_consumer_release_method(&self, scope: &mut Scope) {
         let pipeline = *self.elem;
-        scope.register(Instruction {
-            out: None,
-            operation: Operation::Pipeline(PipelineOps::ConsumerRelease { pipeline }),
-        });
+        scope.register(PipelineOps::ConsumerRelease { pipeline });
     }
 }