Feat kernel profiler

crates/cubecl-core/src/codegen/mod.rs

@@ -1,6 +1,7 @@
 mod info;
 mod integrator;
 mod metadata;
+mod profiler;
 mod scalars;
 
 mod compiler;
@@ -9,4 +10,5 @@ pub use compiler::*;
 pub use info::*;
 pub use integrator::*;
 pub use metadata::*;
+pub use profiler::*;
 pub use scalars::*;

crates/cubecl-core/src/codegen/profiler.rs

@@ -0,0 +1,173 @@
+use alloc::vec::Vec;
+use cubecl_ir::{
+    AddressSpace, Allocator, Arithmetic, AtomicBinaryOperands, AtomicOp, BinaryOperands,
+    ConstantValue, CountKey, ElemType, IndexOperands, Instruction, Memory, Operation, StorageType,
+    StoreOperands, Trackable, Type, UIntKind, Value,
+};
+use hashbrown::{HashMap, hash_map::Entry};
+
+type BufferId = Value;
+
+const LOCAL_COUNTER_TYPE: Type = Type::Scalar(StorageType::Scalar(ElemType::UInt(UIntKind::U32)));
+
+fn constant_value(amount: usize) -> Value {
+    Value::constant(ConstantValue::UInt(amount as u64), LOCAL_COUNTER_TYPE)
+}
+
+#[derive(Clone, Debug, Default, derive_new::new)]
+pub struct CompilerProfiler {
+    global_counter: Option<BufferId>,
+    locals: HashMap<CountKey, (BufferId, usize)>,
+    current_index: usize,
+}
+
+impl CompilerProfiler {
+    pub fn set_counter(&mut self, counter: BufferId) {
+        self.global_counter = Some(counter)
+    }
+
+    pub fn profile_operation(
+        &mut self,
+        operation: &impl Trackable,
+        out: Option<&Value>,
+        allocator: &Allocator,
+    ) -> Vec<Instruction> {
+        // Terminals carry no result, so this must be checked before touching `out`.
+        if operation.is_terminal() {
+            let mut instructions = Vec::new();
+            for (local_counter, index) in self.locals.values() {
+                instructions.extend(self.flush_instructions(local_counter, index, allocator));
+            }
+            return instructions;
+        }
+
+        let Some(counted) = operation.count(out) else {
+            return Vec::new();
+        };
+
+        let local_counter = self.local_counter(counted.key, allocator);
+        Vec::from(self.increment_local_instructions(
+            &local_counter,
+            counted.amount as usize,
+            allocator,
+        ))
+    }
+
+    pub fn profile(&self, allocator: &Allocator) -> (Vec<Instruction>, Vec<Instruction>) {
+        let mut declares = Vec::new();
+        let mut flushes = Vec::new();
+
+        self.locals.values().for_each(|(local_counter, index)| {
+            declares.extend(self.declare_instructions(local_counter));
+            flushes.extend(self.flush_instructions(local_counter, index, allocator));
+        });
+
+        (declares, flushes)
+    }
+
+    fn increment_local_instructions(
+        &self,
+        local_counter: &BufferId,
+        amount: usize,
+        allocator: &Allocator,
+    ) -> [Instruction; 3] {
+        let loaded = allocator.create_value(LOCAL_COUNTER_TYPE);
+        let added = allocator.create_value(LOCAL_COUNTER_TYPE);
+
+        [
+            Instruction::new(Memory::Load(*local_counter), loaded),
+            Instruction::new(
+                Arithmetic::Add(BinaryOperands {
+                    lhs: loaded,
+                    rhs: constant_value(amount),
+                }),
+                added,
+            ),
+            Instruction::no_out(Memory::Store(StoreOperands {
+                ptr: *local_counter,
+                value: added,
+            })),
+        ]
+    }
+
+    fn declare_instructions(&self, local_counter: &BufferId) -> [Instruction; 2] {
+        [
+            Instruction::new(
+                Operation::DeclareVariable {
+                    value_ty: LOCAL_COUNTER_TYPE,
+                    addr_space: AddressSpace::Local,
+                    alignment: LOCAL_COUNTER_TYPE.align(),
+                },
+                *local_counter,
+            ),
+            Instruction::no_out(Memory::Store(StoreOperands {
+                ptr: *local_counter,
+                value: constant_value(0),
+            })),
+        ]
+    }
+
+    fn flush_instructions(
+        &self,
+        local_counter: &BufferId,
+        index: &usize,
+        allocator: &Allocator,
+    ) -> [Instruction; 3] {
+        let global = self
+            .global_counter
+            .as_ref()
+            .expect("Global counter should be available");
+
+        let loaded = allocator.create_value(LOCAL_COUNTER_TYPE);
+        let slot_ptr =
+            allocator.create_value(Type::pointer(global.value_type(), global.address_space()));
+        let discard = allocator.create_value(LOCAL_COUNTER_TYPE);
+
+        [
+            Instruction::new(Memory::Load(*local_counter), loaded),
+            Instruction::new(
+                Memory::Index(IndexOperands {
+                    list: *global,
+                    index: constant_value(*index),
+                    checked: false,
+                    unroll_factor: 1,
+                }),
+                slot_ptr,
+            ),
+            Instruction::new(
+                AtomicOp::Add(AtomicBinaryOperands {
+                    ptr: slot_ptr,
+                    value: loaded,
+                }),
+                discard,
+            ),
+        ]
+    }
+
+    fn local_counter(&mut self, key: CountKey, allocator: &Allocator) -> BufferId {
+        match self.locals.entry(key) {
+            Entry::Occupied(entry) => entry.get().0,
+            Entry::Vacant(entry) => {
+                let index = self.current_index;
+                self.current_index += 1;
+                let counter =
+                    allocator.create_value(Type::pointer(LOCAL_COUNTER_TYPE, AddressSpace::Local));
+                entry.insert((counter, index));
+                counter
+            }
+        }
+    }
+
+    /// The slot → key decoding table, indexed by dense slot. Its length is the number of used
+    /// counters, i.e. the size the global profiling buffer must be. Empty when nothing was
+    /// tracked.
+    pub fn profile_map(&self) -> Vec<CountKey> {
+        let mut entries: Vec<(usize, CountKey)> = self
+            .locals
+            .iter()
+            .map(|(key, (_, slot))| (*slot, *key))
+            .collect();
+        entries.sort_by_key(|(slot, _)| *slot);
+        entries.into_iter().map(|(_, key)| key).collect()
+    }
+}

crates/cubecl-core/src/compute/builder.rs

@@ -8,7 +8,9 @@ use crate::{
     prelude::KernelDefinition,
 };
 use alloc::collections::BTreeMap;
-use cubecl_ir::{DeviceProperties, Scope, StorageType, TargetProperties, Value};
+use cubecl_ir::{
+    DeviceProperties, Scope, StorageType, TargetProperties, Value, counter_stored_type,
+};
 use cubecl_runtime::config::{
     CubeClRuntimeConfig, RuntimeConfig, compilation::CompilationLogLevel,
 };
@@ -90,7 +92,17 @@ impl KernelBuilder {
     }
 
     /// Build the [kernel definition](KernelDefinition).
-    pub fn build(self, settings: KernelSettings) -> KernelDefinition {
+    pub fn build(mut self, settings: KernelSettings) -> KernelDefinition {
+        if self.profile.enabled {
+            // The profiling counters buffer is a *special trailing binding*: it is bound after
+            // every normal buffer and the `info` struct, and deliberately kept out of
+            // `self.buffers` so it does not contribute to the metadata layout (`num_meta`). It is
+            // only ever accessed via constant slot indices, so it needs no length metadata.
+            let id = self.buffer_id();
+            let buffer = self.scope.global(id, counter_stored_type());
+            self.scope.profile.counters_buffer = Some(buffer);
+        }
+
         let scalars = self
             .scalars
             .into_iter()
@@ -119,8 +131,10 @@ impl KernelBuilder {
             debug == 1
         };
 
+        let profile = CubeClRuntimeConfig::get().profiling.hardware_metrics;
+
         Self {
-            scope: Scope::root(debug),
+            scope: Scope::root(debug, profile),
             buffers: Default::default(),
             scalars: Default::default(),
             tensor_maps: Default::default(),

crates/cubecl-core/src/compute/launcher.rs

@@ -18,7 +18,7 @@ use cubecl_runtime::{
 std::thread_local! {
     static INFO: RefCell<InfoBuilder> = RefCell::new(InfoBuilder::default());
     // Only used for resolving types
-    static SCOPE: RefCell<Scope> = RefCell::new(Scope::root(false));
+    static SCOPE: RefCell<Scope> = RefCell::new(Scope::root(false, false));
 }
 
 /// Prepare a kernel for [launch](KernelLauncher::launch).
@@ -76,7 +76,7 @@ impl<R: Runtime> KernelLauncher<R> {
         let bindings = self.into_bindings();
         let kernel = Box::new(KernelTask::<R::Compiler, K>::new(kernel));
 
-        client.launch(kernel, cube_count, bindings)
+        client.launch(kernel, cube_count, bindings);
     }
 
     /// Launch the kernel without check bounds.
@@ -196,7 +196,7 @@ impl<R: Runtime> KernelLauncher<R> {
             #[cfg(not(feature = "std"))]
             info: InfoBuilder::default(),
             #[cfg(not(feature = "std"))]
-            scope: Scope::root(false),
+            scope: Scope::root(false, false),
         }
     }
 }

crates/cubecl-core/src/post_processing/checked_io.rs

@@ -56,7 +56,8 @@ impl CheckedIoVisitor {
                     if has_length {
                         let list = op.list;
                         let index = op.index;
-                        let scope = Scope::root(false).with_global_state(global_state.clone());
+                        let scope =
+                            Scope::root(false, false).with_global_state(global_state.clone());
 
                         expand_checked_index(
                             &scope,
@@ -88,7 +89,8 @@ impl CheckedIoVisitor {
                     if has_length {
                         let list = op.list;
                         let index = op.index;
-                        let scope = Scope::root(false).with_global_state(global_state.clone());
+                        let scope =
+                            Scope::root(false, false).with_global_state(global_state.clone());
 
                         expand_validate_index(
                             &scope,

crates/cubecl-core/src/post_processing/predicate.rs

@@ -58,7 +58,7 @@ fn run_polyfill<T: CubePrimitive, O: CubePrimitive>(
     out: Value,
     mut polyfill: impl FnMut(&Scope, NativeExpand<T>, u32, u32) -> NativeExpand<O>,
 ) {
-    let scope = Scope::root(false).with_global_state(processing.global_state.clone());
+    let scope = Scope::root(false, false).with_global_state(processing.global_state.clone());
     scope.register_type::<ElemA>(input.storage_type());
     scope.register_size::<SizeA>(input.vector_size());
 

crates/cubecl-core/src/post_processing/saturating.rs

@@ -97,7 +97,7 @@ fn run_polyfill<T: CubePrimitive>(
     out: Value,
     mut polyfill: impl FnMut(&Scope, NativeExpand<T>, NativeExpand<T>) -> NativeExpand<T>,
 ) {
-    let scope = Scope::root(false).with_global_state(processing.global_state.clone());
+    let scope = Scope::root(false, false).with_global_state(processing.global_state.clone());
     scope.register_type::<ElemA>(lhs.storage_type());
     scope.register_size::<SizeA>(lhs.vector_size());
     if let ElemType::Int(kind) = lhs.elem_type() {

crates/cubecl-cpp/src/cuda/processors.rs

@@ -16,8 +16,8 @@ impl Processor for CudaMmaProcessor {
             match instruction.operation {
                 Operation::CoopMma(CoopMma::RowIndex { lane_id, i, matrix }) => {
                     let elems_per_reg = 32 / matrix.storage.elem_type().size_bits();
-                    let scope =
-                        Scope::root(false).with_global_state(processing.global_state.clone());
+                    let scope = Scope::root(false, false)
+                        .with_global_state(processing.global_state.clone());
                     let row_idx: Value = row_index::expand(
                         &scope,
                         lane_id.into(),
@@ -38,8 +38,8 @@ impl Processor for CudaMmaProcessor {
                 }
                 Operation::CoopMma(CoopMma::ColIndex { lane_id, i, matrix }) => {
                     let elems_per_reg = 32 / matrix.storage.elem_type().size_bits();
-                    let scope =
-                        Scope::root(false).with_global_state(processing.global_state.clone());
+                    let scope = Scope::root(false, false)
+                        .with_global_state(processing.global_state.clone());
                     let col_idx: Value = col_index::expand(
                         &scope,
                         lane_id.into(),

crates/cubecl-cpp/src/hip/processors.rs

@@ -15,8 +15,8 @@ impl Processor for HipMmaProcessor {
         for instruction in instructions {
             match instruction.operation {
                 Operation::CoopMma(CoopMma::RowIndex { lane_id, i, matrix }) => {
-                    let scope =
-                        Scope::root(false).with_global_state(processing.global_state.clone());
+                    let scope = Scope::root(false, false)
+                        .with_global_state(processing.global_state.clone());
                     let row_idx: Value =
                         row_index::expand(&scope, lane_id.into(), i.into(), matrix.ident).into();
                     let tmp_processing = scope.process([]);
@@ -30,8 +30,8 @@ impl Processor for HipMmaProcessor {
                     ));
                 }
                 Operation::CoopMma(CoopMma::ColIndex { lane_id, i, matrix }) => {
-                    let scope =
-                        Scope::root(false).with_global_state(processing.global_state.clone());
+                    let scope = Scope::root(false, false)
+                        .with_global_state(processing.global_state.clone());
                     let row_idx: Value =
                         col_index::expand(&scope, lane_id.into(), i.into(), matrix.ident).into();
                     let tmp_processing = scope.process([]);

crates/cubecl-ir/Cargo.toml

@@ -34,6 +34,7 @@ cubecl-common = { path = "../cubecl-common", version = "=0.11.0-pre.1", default-
 cubecl-macros-internal = { path = "../cubecl-macros-internal", version = "=0.11.0-pre.1" }
 
 bumpalo = { workspace = true }
+bytemuck = { workspace = true, features = ["derive"] }
 derive-new = { workspace = true }
 derive_more = { workspace = true, features = ["from", "eq"] }
 enumset = { workspace = true }
@@ -49,3 +50,5 @@ variadics_please = { workspace = true }
 
 hashbrown = { workspace = true }
 portable-atomic = { workspace = true }
+
+paste = { workspace = true }