From c68c5a9daa3849f4788bbe9941be192e062ac7a1 Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Mon, 29 Jun 2026 15:53:16 -0400 Subject: [PATCH 01/14] working base addition profiling --- crates/cubecl-core/src/compute/builder.rs | 14 ++- crates/cubecl-core/src/compute/launcher.rs | 47 +++++++-- .../src/post_processing/checked_io.rs | 6 +- .../src/post_processing/predicate.rs | 2 +- .../src/post_processing/saturating.rs | 2 +- crates/cubecl-cpp/src/cuda/processors.rs | 8 +- crates/cubecl-cpp/src/hip/processors.rs | 8 +- crates/cubecl-ir/src/flop_count.rs | 98 +++++++++++++++++++ crates/cubecl-ir/src/lib.rs | 2 + crates/cubecl-ir/src/scope.rs | 27 ++++- crates/cubecl-opt/src/lib.rs | 4 +- crates/cubecl-runtime/src/client.rs | 4 +- crates/cubecl-runtime/src/config/base.rs | 6 ++ crates/cubecl-runtime/src/config/profiling.rs | 4 + crates/cubecl-runtime/src/kernel.rs | 4 + crates/cubecl-runtime/src/server/base.rs | 7 +- 16 files changed, 212 insertions(+), 31 deletions(-) create mode 100644 crates/cubecl-ir/src/flop_count.rs diff --git a/crates/cubecl-core/src/compute/builder.rs b/crates/cubecl-core/src/compute/builder.rs index 4ac593510d..62d50e4ed9 100644 --- a/crates/cubecl-core/src/compute/builder.rs +++ b/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, FlopCountProcessor, Scope, StorageType, TargetProperties, Value, +}; use cubecl_runtime::config::{ CubeClRuntimeConfig, RuntimeConfig, compilation::CompilationLogLevel, }; @@ -90,7 +92,11 @@ 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 { + self.buffer(FlopCountProcessor::flop_counter_type()); + } + let scalars = self .scalars .into_iter() @@ -119,8 +125,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(), diff --git a/crates/cubecl-core/src/compute/launcher.rs b/crates/cubecl-core/src/compute/launcher.rs index 1eff237aee..7db360c6fd 100644 --- a/crates/cubecl-core/src/compute/launcher.rs +++ b/crates/cubecl-core/src/compute/launcher.rs @@ -6,8 +6,9 @@ use crate::prelude::{BufferArg, TensorArg, TensorMapArg, TensorMapKind}; use crate::{InfoBuilder, KernelSettings, ScalarArgType}; #[cfg(feature = "std")] use core::cell::RefCell; -use cubecl_ir::{AddressType, Scope, StorageType, Type}; -use cubecl_runtime::server::{Binding, CubeCount, TensorMapBinding}; +use cubecl_ir::{AddressType, FlopCountProcessor, Scope, StorageType, Type}; +use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig}; +use cubecl_runtime::server::{Binding, CubeCount, Handle, TensorMapBinding}; use cubecl_runtime::{ client::ComputeClient, kernel::{CubeKernel, KernelTask}, @@ -18,7 +19,7 @@ use cubecl_runtime::{ std::thread_local! { static INFO: RefCell = RefCell::new(InfoBuilder::default()); // Only used for resolving types - static SCOPE: RefCell = RefCell::new(Scope::root(false)); + static SCOPE: RefCell = RefCell::new(Scope::root(false, false)); } /// Prepare a kernel for [launch](KernelLauncher::launch). @@ -65,18 +66,45 @@ impl KernelLauncher { self.with_info(|info| info.scalars.push_raw(bytes, dtype)); } + /// Injects profiling buffers into the bindings before execution. + fn inject_profiling(&mut self, client: &ComputeClient) -> Option { + if !CubeClRuntimeConfig::get().profiling.hardware_metrics { + return None; + } + + let handle = client.create_from_slice(&[0u8; 4]); + let arg = unsafe { BufferArg::from_raw_parts(handle.clone(), 1) }; + self.register_buffer(arg, FlopCountProcessor::flop_counter_type()); + Some(handle) + } + + /// Read back the FLOP counter buffer and report it. + fn report_flop_counter(client: &ComputeClient, handle: Handle) { + let bytes = client.read_one_unchecked(handle); + let count = u32::from_le_bytes(bytes[..4].try_into().expect("Counter should be 4 bytes")); + #[cfg(feature = "std")] + std::eprintln!("[cubecl flop-profile] add count = {count}"); + #[cfg(not(feature = "std"))] + let _ = count; + } + /// Launch the kernel. #[track_caller] pub fn launch( - self, + mut self, cube_count: CubeCount, kernel: K, client: &ComputeClient, ) { + let flop_counter = self.inject_profiling(client); let bindings = self.into_bindings(); let kernel = Box::new(KernelTask::::new(kernel)); - client.launch(kernel, cube_count, bindings) + client.launch(kernel, cube_count, bindings); + + if let Some(handle) = flop_counter { + Self::report_flop_counter(client, handle); + } } /// Launch the kernel without check bounds. @@ -89,17 +117,22 @@ impl KernelLauncher { /// other unpredictable behaviour. #[track_caller] pub unsafe fn launch_unchecked( - self, + mut self, cube_count: CubeCount, kernel: K, client: &ComputeClient, ) { + let flop_counter = self.inject_profiling(client); unsafe { let bindings = self.into_bindings(); let kernel = Box::new(KernelTask::::new(kernel)); client.launch_unchecked(kernel, cube_count, bindings) } + + if let Some(handle) = flop_counter { + Self::report_flop_counter(client, handle); + } } /// We need to create the bindings in the same order they are defined in the compilation step. @@ -196,7 +229,7 @@ impl KernelLauncher { #[cfg(not(feature = "std"))] info: InfoBuilder::default(), #[cfg(not(feature = "std"))] - scope: Scope::root(false), + scope: Scope::root(false, false), } } } diff --git a/crates/cubecl-core/src/post_processing/checked_io.rs b/crates/cubecl-core/src/post_processing/checked_io.rs index 7d8f6ec891..1496923949 100644 --- a/crates/cubecl-core/src/post_processing/checked_io.rs +++ b/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, diff --git a/crates/cubecl-core/src/post_processing/predicate.rs b/crates/cubecl-core/src/post_processing/predicate.rs index 3d2ae420ab..2c5624a678 100644 --- a/crates/cubecl-core/src/post_processing/predicate.rs +++ b/crates/cubecl-core/src/post_processing/predicate.rs @@ -58,7 +58,7 @@ fn run_polyfill( out: Value, mut polyfill: impl FnMut(&Scope, NativeExpand, u32, u32) -> NativeExpand, ) { - 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::(input.storage_type()); scope.register_size::(input.vector_size()); diff --git a/crates/cubecl-core/src/post_processing/saturating.rs b/crates/cubecl-core/src/post_processing/saturating.rs index ea232f6213..e6036bf1dd 100644 --- a/crates/cubecl-core/src/post_processing/saturating.rs +++ b/crates/cubecl-core/src/post_processing/saturating.rs @@ -97,7 +97,7 @@ fn run_polyfill( out: Value, mut polyfill: impl FnMut(&Scope, NativeExpand, NativeExpand) -> NativeExpand, ) { - 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::(lhs.storage_type()); scope.register_size::(lhs.vector_size()); if let ElemType::Int(kind) = lhs.elem_type() { diff --git a/crates/cubecl-cpp/src/cuda/processors.rs b/crates/cubecl-cpp/src/cuda/processors.rs index 5a27a2c33a..46b87c2369 100644 --- a/crates/cubecl-cpp/src/cuda/processors.rs +++ b/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(), diff --git a/crates/cubecl-cpp/src/hip/processors.rs b/crates/cubecl-cpp/src/hip/processors.rs index 2e297117a5..af1528fc69 100644 --- a/crates/cubecl-cpp/src/hip/processors.rs +++ b/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([]); diff --git a/crates/cubecl-ir/src/flop_count.rs b/crates/cubecl-ir/src/flop_count.rs new file mode 100644 index 0000000000..ec58a50123 --- /dev/null +++ b/crates/cubecl-ir/src/flop_count.rs @@ -0,0 +1,98 @@ +//! Experimental FLOP profiling. +//! +//! When profiling is enabled on a [`Scope`], an extra global atomic counter buffer is appended to +//! the kernel and an atomic increment is inserted after each profiled arithmetic operation. This +//! gives a *dynamic* operation count that respects control flow and loops, since the increment +//! lives next to the operation in the same basic block. +//! +//! The processor is applied automatically by [`Scope::process`] for any scope whose +//! [`crate::ProfileInfo`] is enabled, so every backend supports it without backend-specific code. +//! +//! This is a first slice: only [`Arithmetic::Add`] is counted, and the counter increments by one +//! per executed instruction (vectorization width is ignored for now). + +use alloc::vec::Vec; + +use crate::{ + Arithmetic, AtomicBinaryOperands, AtomicOp, ElemType, IndexOperands, Instruction, Memory, + Operation, Processor, Scope, ScopeProcessing, Type, UIntKind, Value, +}; + +/// Inserts an atomic increment of `counter` after each profiled arithmetic operation. +#[derive(Debug)] +pub struct FlopCountProcessor { + /// The global counter buffer (a `&mut [atomic]`). + counter: Value, +} + +impl FlopCountProcessor { + /// Create a processor that increments the given global atomic `counter`. + pub fn new(counter: Value) -> Self { + Self { counter } + } + + /// The element type of the FLOP counter buffer: a single atomic `u32`. + pub fn flop_counter_type() -> Type { + Type::atomic(Type::scalar(ElemType::UInt(UIntKind::U32))) + } + + /// Append `counter[0] += 1` to the processing instruction stream. + fn emit_increment(&self, processing: &mut ScopeProcessing) { + // Build the increment in a throwaway scope that shares the global state (allocator, type + // maps, ...) so the freshly created values don't collide with existing ones. Profiling is + // disabled on this scope so its own `process` call doesn't recurse into FLOP counting. + let scope = Scope::root(false, false).with_global_state(processing.global_state.clone()); + + let u32_ty = Type::scalar(ElemType::UInt(UIntKind::U32)); + + // `&counter[0]` -> pointer to the atomic element. + let elem_ptr = scope.create_value(Type::pointer( + self.counter.value_type(), + self.counter.address_space(), + )); + scope.register(Instruction::new( + Memory::Index(IndexOperands { + list: self.counter, + index: Value::constant(0u32.into(), u32_ty), + unroll_factor: 1, + checked: false, + }), + elem_ptr, + )); + + // `atomicAdd(&counter[0], 1)`. The returned old value is unused. + let old = scope.create_value(u32_ty); + scope.register(Instruction::new( + AtomicOp::Add(AtomicBinaryOperands { + ptr: elem_ptr, + value: Value::constant(1u32.into(), u32_ty), + }), + old, + )); + + let tmp = scope.process([]); + processing.instructions.extend(tmp.instructions); + } +} + +impl Processor for FlopCountProcessor { + fn transform(&self, mut processing: ScopeProcessing) -> ScopeProcessing { + let mut instructions = Vec::new(); + core::mem::swap(&mut processing.instructions, &mut instructions); + + for instruction in instructions { + let is_profiled = matches!( + &instruction.operation, + Operation::Arithmetic(Arithmetic::Add(_)) + ); + + processing.instructions.push(instruction); + + if is_profiled { + self.emit_increment(&mut processing); + } + } + + processing + } +} diff --git a/crates/cubecl-ir/src/lib.rs b/crates/cubecl-ir/src/lib.rs index 8528cf4023..84dbbf5fe9 100644 --- a/crates/cubecl-ir/src/lib.rs +++ b/crates/cubecl-ir/src/lib.rs @@ -16,6 +16,7 @@ mod bitwise; mod branch; mod cmma; mod comparison; +mod flop_count; mod marker; mod memory; mod metadata; @@ -46,6 +47,7 @@ pub use bitwise::*; pub use branch::*; pub use cmma::*; pub use comparison::*; +pub use flop_count::*; pub use marker::*; pub use memory::*; pub use metadata::*; diff --git a/crates/cubecl-ir/src/scope.rs b/crates/cubecl-ir/src/scope.rs index e7d5ba8596..97537d8748 100644 --- a/crates/cubecl-ir/src/scope.rs +++ b/crates/cubecl-ir/src/scope.rs @@ -11,9 +11,9 @@ use hashbrown::{HashMap, HashSet}; use itertools::Itertools; use crate::{ - AddressSpace, AggregateExtractOperands, CubeFnSource, DeviceProperties, FastMath, Function, - OpaqueType, Operation, OperationReflect, Processor, SourceLoc, StorageType, TargetProperties, - TypeHash, arena::DropBump, + AddressSpace, AggregateExtractOperands, CubeFnSource, DeviceProperties, FastMath, + FlopCountProcessor, Function, OpaqueType, Operation, OperationReflect, Processor, SourceLoc, + StorageType, TargetProperties, TypeHash, arena::DropBump, }; use super::{Allocator, Id, Instruction, Type, Value, processing::ScopeProcessing}; @@ -38,6 +38,7 @@ pub struct Scope { pub return_value: Option, pub locals: RefCell>, pub debug: DebugInfo, + pub profile: ProfileInfo, #[cfg_attr(feature = "serde", serde(skip))] pub global_state: GlobalState, @@ -94,6 +95,13 @@ pub struct DebugInfo { pub entry_loc: RefCell>, } +/// Profiling related fields +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq, Eq, TypeHash)] +pub struct ProfileInfo { + pub enabled: bool, +} + /// Modes set and reset during expansion #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, TypeHash)] @@ -136,7 +144,7 @@ impl Scope { /// Create a scope that is at the root of a kernel definition. /// /// A local scope can be created with the [child](Self::child) method. - pub fn root(debug_enabled: bool) -> Self { + pub fn root(debug_enabled: bool, profile_enabled: bool) -> Self { Self { validation_errors: ValidationErrors { errors: Rc::new(RefCell::new(Vec::new())), @@ -152,6 +160,9 @@ impl Scope { source_loc: Default::default(), entry_loc: Default::default(), }, + profile: ProfileInfo { + enabled: profile_enabled, + }, global_state: Default::default(), } } @@ -274,6 +285,7 @@ impl Scope { return_value: None, locals: Default::default(), debug: self.debug.clone(), + profile: self.profile.clone(), global_state: self.global_state.clone(), } } @@ -311,6 +323,13 @@ impl Scope { global_state: self.global_state.clone(), }; + if self.profile.enabled { + let counter = self.global_state.borrow().global_args.last().copied(); + if let Some(counter) = counter { + processing = FlopCountProcessor::new(counter).transform(processing); + } + } + for p in processors { processing = p.transform(processing); } diff --git a/crates/cubecl-opt/src/lib.rs b/crates/cubecl-opt/src/lib.rs index 5161006797..c679ee680c 100644 --- a/crates/cubecl-opt/src/lib.rs +++ b/crates/cubecl-opt/src/lib.rs @@ -183,7 +183,7 @@ impl Default for GlobalState { fn default() -> Self { Self { allocator: Default::default(), - root_scope: Scope::root(false), + root_scope: Scope::root(false, false), buffer_visibility: Default::default(), extra_functions: Default::default(), cube_dim: CubeDim::new_1d(1), @@ -660,7 +660,7 @@ mod test { #[test_log::test] #[ignore = "no good way to assert opt is applied"] fn test_pre() { - let ctx = Scope::root(false); + let ctx = Scope::root(false, false); let u32 = Type::scalar(ElemType::UInt(UIntKind::U32)); let x = ctx.create_value(u32).into(); let cond = ctx.create_value(u32).into(); diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index 8691825978..2987d97849 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -5,8 +5,8 @@ use crate::{ memory_management::{MemoryAllocationMode, MemoryUsage}, runtime::Runtime, server::{ - CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, Handle, IoError, - KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, + CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, Handle, + IoError, KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, MemoryLayoutStrategy, ProfileError, ReduceOperation, ServerCommunication, ServerError, ServerUtilities, }, diff --git a/crates/cubecl-runtime/src/config/base.rs b/crates/cubecl-runtime/src/config/base.rs index 8f8cac3699..c950dece4f 100644 --- a/crates/cubecl-runtime/src/config/base.rs +++ b/crates/cubecl-runtime/src/config/base.rs @@ -133,6 +133,12 @@ impl RuntimeConfig for CubeClRuntimeConfig { } } + if let Ok(val) = std::env::var("CUBECL_PROFILE_FLOPS") { + if val == "1" || val.eq_ignore_ascii_case("true") { + self.profiling.hardware_metrics = true; + } + } + self } } diff --git a/crates/cubecl-runtime/src/config/profiling.rs b/crates/cubecl-runtime/src/config/profiling.rs index 33b1dce7a9..1d95883933 100644 --- a/crates/cubecl-runtime/src/config/profiling.rs +++ b/crates/cubecl-runtime/src/config/profiling.rs @@ -6,6 +6,10 @@ pub struct ProfilingConfig { /// Logger configuration for profiling logs, using profiling-specific log levels. #[serde(default)] pub logger: LoggerConfig, + /// Determines whether hardware metrics are collected during profiling. + /// WARNING: Adds significant overhead to kernels as they are modified to count operations. + #[serde(default)] + pub hardware_metrics: bool, } /// Log levels for profiling in `CubeCL`. diff --git a/crates/cubecl-runtime/src/kernel.rs b/crates/cubecl-runtime/src/kernel.rs index ed0229eec2..25af61f2a7 100644 --- a/crates/cubecl-runtime/src/kernel.rs +++ b/crates/cubecl-runtime/src/kernel.rs @@ -179,6 +179,10 @@ impl CubeTask for KernelTask { let gpu_ir = self.kernel_definition.define(); let entrypoint_name = gpu_ir.options.kernel_name.clone(); let cube_dim = gpu_ir.cube_dim; + + // std::println!("Compiling kernel: {}", entrypoint_name); + // std::println!("IR: {:#?}", gpu_ir); + let lower_level_ir = compiler.compile(gpu_ir, compilation_options, mode, addr_type)?; Ok(CompiledKernel { diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index 4c218296ca..80427f750f 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -96,6 +96,8 @@ pub struct ServerUtilities { pub layout_policy: Server::MemoryLayoutPolicy, /// How to enforce bounds checking on kernels. pub check_mode: BoundsCheckMode, + /// Whether to collect hardware metrics. + pub hardware_metrics: bool, /// A set containing the ids for which the inter-device communication has already been initialized. pub initialized_comms: RwLock>, } @@ -135,6 +137,8 @@ impl ServerUtilities { info: S::Info, allocator: S::MemoryLayoutPolicy, ) -> Self { + let config = CubeClRuntimeConfig::get(); + // Start a tracy client if needed. #[cfg(feature = "profile-tracy")] let client = tracy_client::Client::start(); @@ -159,7 +163,8 @@ impl ServerUtilities { epoch_time: web_time::Instant::now(), info, layout_policy: allocator, - check_mode: CubeClRuntimeConfig::get().compilation.check_mode, + check_mode: config.compilation.check_mode, + hardware_metrics: config.profiling.hardware_metrics, initialized_comms: RwLock::new(HashSet::default()), } } From 6030e85af4bf96e8ba345ab14dbbe4a869c17f61 Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Mon, 29 Jun 2026 16:39:25 -0400 Subject: [PATCH 02/14] support vectorization --- crates/cubecl-core/src/compute/launcher.rs | 18 +-- crates/cubecl-ir/src/flop_count.rs | 14 +-- crates/cubecl-runtime/src/client.rs | 26 +++- crates/cubecl-runtime/src/server/base.rs | 19 ++- crates/cubecl/tests/chose.rs | 136 +++++++++++++++++++++ 5 files changed, 192 insertions(+), 21 deletions(-) create mode 100644 crates/cubecl/tests/chose.rs diff --git a/crates/cubecl-core/src/compute/launcher.rs b/crates/cubecl-core/src/compute/launcher.rs index 7db360c6fd..4fb5327e13 100644 --- a/crates/cubecl-core/src/compute/launcher.rs +++ b/crates/cubecl-core/src/compute/launcher.rs @@ -8,6 +8,7 @@ use crate::{InfoBuilder, KernelSettings, ScalarArgType}; use core::cell::RefCell; use cubecl_ir::{AddressType, FlopCountProcessor, Scope, StorageType, Type}; use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig}; +use cubecl_runtime::id::KernelId; use cubecl_runtime::server::{Binding, CubeCount, Handle, TensorMapBinding}; use cubecl_runtime::{ client::ComputeClient, @@ -78,14 +79,11 @@ impl KernelLauncher { Some(handle) } - /// Read back the FLOP counter buffer and report it. - fn report_flop_counter(client: &ComputeClient, handle: Handle) { + /// Read back the FLOP counter buffer and record it on the client, keyed by kernel id. + fn report_flop_counter(client: &ComputeClient, id: KernelId, handle: Handle) { let bytes = client.read_one_unchecked(handle); let count = u32::from_le_bytes(bytes[..4].try_into().expect("Counter should be 4 bytes")); - #[cfg(feature = "std")] - std::eprintln!("[cubecl flop-profile] add count = {count}"); - #[cfg(not(feature = "std"))] - let _ = count; + client.record_flop_count(id.clone(), count as u64); } /// Launch the kernel. @@ -97,13 +95,15 @@ impl KernelLauncher { client: &ComputeClient, ) { let flop_counter = self.inject_profiling(client); + + let kernel_id = kernel.id(); let bindings = self.into_bindings(); let kernel = Box::new(KernelTask::::new(kernel)); client.launch(kernel, cube_count, bindings); if let Some(handle) = flop_counter { - Self::report_flop_counter(client, handle); + Self::report_flop_counter(client, kernel_id, handle); } } @@ -123,6 +123,8 @@ impl KernelLauncher { client: &ComputeClient, ) { let flop_counter = self.inject_profiling(client); + + let kernel_id = kernel.id(); unsafe { let bindings = self.into_bindings(); let kernel = Box::new(KernelTask::::new(kernel)); @@ -131,7 +133,7 @@ impl KernelLauncher { } if let Some(handle) = flop_counter { - Self::report_flop_counter(client, handle); + Self::report_flop_counter(client, kernel_id, handle); } } diff --git a/crates/cubecl-ir/src/flop_count.rs b/crates/cubecl-ir/src/flop_count.rs index ec58a50123..2de4e8656a 100644 --- a/crates/cubecl-ir/src/flop_count.rs +++ b/crates/cubecl-ir/src/flop_count.rs @@ -37,7 +37,7 @@ impl FlopCountProcessor { } /// Append `counter[0] += 1` to the processing instruction stream. - fn emit_increment(&self, processing: &mut ScopeProcessing) { + fn emit_increment(&self, processing: &mut ScopeProcessing, vector_size: usize) { // Build the increment in a throwaway scope that shares the global state (allocator, type // maps, ...) so the freshly created values don't collide with existing ones. Profiling is // disabled on this scope so its own `process` call doesn't recurse into FLOP counting. @@ -65,7 +65,7 @@ impl FlopCountProcessor { scope.register(Instruction::new( AtomicOp::Add(AtomicBinaryOperands { ptr: elem_ptr, - value: Value::constant(1u32.into(), u32_ty), + value: Value::constant(vector_size.into(), u32_ty), }), old, )); @@ -81,15 +81,15 @@ impl Processor for FlopCountProcessor { core::mem::swap(&mut processing.instructions, &mut instructions); for instruction in instructions { - let is_profiled = matches!( - &instruction.operation, - Operation::Arithmetic(Arithmetic::Add(_)) - ); + let (is_profiled, vector_size) = match &instruction.operation { + Operation::Arithmetic(Arithmetic::Add(ops)) => (true, ops.lhs.vector_size()), + _ => (false, 1), + }; processing.instructions.push(instruction); if is_profiled { - self.emit_increment(&mut processing); + self.emit_increment(&mut processing, vector_size); } } diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index 2987d97849..babbb5363e 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -1,12 +1,13 @@ use crate::{ config::{TypeNameFormatLevel, type_name_format}, + id::KernelId, kernel::KernelMetadata, logging::ProfileLevel, memory_management::{MemoryAllocationMode, MemoryUsage}, runtime::Runtime, server::{ - CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, Handle, - IoError, KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, + CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, FlopRecord, + Handle, IoError, KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, MemoryLayoutStrategy, ProfileError, ReduceOperation, ServerCommunication, ServerError, ServerUtilities, }, @@ -23,6 +24,7 @@ use cubecl_common::{ device_handle::{CallResultExt, DeviceHandle}, future::DynFut, profile::ProfileDuration, + stub::RwLock, }; use cubecl_ir::{DeviceProperties, ElemType, VectorSize, features::Features}; use cubecl_zspace::Shape; @@ -30,6 +32,7 @@ use cubecl_zspace::Shape; #[allow(unused)] use cubecl_common::profile::TimingMethod; use cubecl_common::stream_id::StreamId; +use hashbrown::HashMap; /// The `ComputeClient` is the entry point to require tasks from the `ComputeServer`. /// It should be obtained for a specific device via the Compute struct. @@ -730,6 +733,25 @@ impl ComputeClient { handle } + /// Record a FLOP count for the given kernel, stored per-device in + /// [`ServerUtilities::metrics`]. Used by the launcher when `hardware_metrics` profiling is on. + pub fn record_flop_count(&self, id: KernelId, count: u64) { + let mut metrics = self.utilities.metrics.write().unwrap(); + let record = metrics.entry(id).or_default(); + record.last = count; + record.samples += 1; + } + + /// Returns the recorded FLOP metrics for the given kernel, if any. + pub fn flop_count(&self, id: &KernelId) -> Option { + self.utilities.metrics.read().unwrap().get(id).copied() + } + + /// Returns a reference to the [`RwLock`] containing the recorded FLOP metrics for all kernels. + pub fn metrics(&self) -> &RwLock> { + &self.utilities.metrics + } + #[track_caller] #[cfg_attr(feature = "tracing", tracing::instrument(level="trace", skip(self, kernel, bindings), diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index 80427f750f..9948b34509 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -3,6 +3,7 @@ use crate::{ client::ComputeClient, compiler::CompilationError, config::{CubeClRuntimeConfig, RuntimeConfig, compilation::BoundsCheckMode}, + id::KernelId, kernel::KernelMetadata, logging::ServerLogger, memory_management::{ManagedMemoryHandle, MemoryAllocationMode, MemoryUsage}, @@ -33,7 +34,7 @@ use cubecl_common::{ }; use cubecl_ir::{DeviceProperties, ElemType, StorageType}; use cubecl_zspace::{Shape, Strides, metadata::Metadata}; -use hashbrown::HashSet; +use hashbrown::{HashMap, HashSet}; use itertools::Itertools; use thiserror::Error; @@ -76,6 +77,16 @@ impl core::fmt::Debug for ProfileError { } } +/// Profiling metrics recorded for a single kernel, keyed by [`KernelId`] in +/// [`ServerUtilities::metrics`]. Populated when `hardware_metrics` profiling is enabled. +#[derive(Debug, Default, Clone, Copy)] +pub struct FlopRecord { + /// Most recently recorded FLOP count for the kernel. + pub last: u64, + /// Number of launches recorded for the kernel. + pub samples: u64, +} + /// Contains many different types that are useful for server implementations and compute clients. pub struct ServerUtilities { /// The time when `profile-tracy` is activated. @@ -96,10 +107,10 @@ pub struct ServerUtilities { pub layout_policy: Server::MemoryLayoutPolicy, /// How to enforce bounds checking on kernels. pub check_mode: BoundsCheckMode, - /// Whether to collect hardware metrics. - pub hardware_metrics: bool, /// A set containing the ids for which the inter-device communication has already been initialized. pub initialized_comms: RwLock>, + /// FLOP-profiling metrics recorded per kernel. Populated when `hardware_metrics` is enabled. + pub metrics: RwLock>, } /// Defines how the memory layout is determined. @@ -164,8 +175,8 @@ impl ServerUtilities { info, layout_policy: allocator, check_mode: config.compilation.check_mode, - hardware_metrics: config.profiling.hardware_metrics, initialized_comms: RwLock::new(HashSet::default()), + metrics: RwLock::new(HashMap::default()), } } } diff --git a/crates/cubecl/tests/chose.rs b/crates/cubecl/tests/chose.rs new file mode 100644 index 0000000000..7bf358f710 --- /dev/null +++ b/crates/cubecl/tests/chose.rs @@ -0,0 +1,136 @@ +use cubecl::{CubeDim, CubeElement, Runtime, TestRuntime, cube, prelude::*}; + +#[cube(launch)] +fn test_chose_cube(lhs: f32, rhs: f32, out: &mut [f32]) { + if UNIT_POS == 0 { + out[0] = lhs + rhs + lhs; + out[0] = lhs + rhs; + out[0] = lhs + rhs; + } +} + +#[test] +fn test_chose_1() { + let client = TestRuntime::client(&Default::default()); + + let lhs = 1.0; + let rhs = 2.0; + let output = client.empty(core::mem::size_of::()); + + test_chose_cube::launch( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(1), + lhs, + rhs, + unsafe { BufferArg::from_raw_parts(output.clone(), 1) }, + ); + + let bytes = client.read_one(output).unwrap(); + let result = f32::from_bytes(&bytes); + + println!("{:#?}", client.metrics().read().unwrap()); + + assert_eq!(result[0], lhs + rhs); +} + +#[cube(launch)] +fn test_chose_cube_tensor(lhs: &Tensor, rhs: &Tensor, out: &mut Tensor) { + if ABSOLUTE_POS < out.len() { + out[ABSOLUTE_POS] = lhs[ABSOLUTE_POS] + rhs[ABSOLUTE_POS]; + } +} + +#[test] +fn test_chose_tensor() { + let client = TestRuntime::client(&Default::default()); + + let lhs = client.create_from_slice(f32::as_bytes(&[1.0, 2.0, 3.0])); + let rhs = client.create_from_slice(f32::as_bytes(&[10.0, 20.0, 30.0])); + let output = client.empty(3 * core::mem::size_of::()); + + test_chose_cube_tensor::launch( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(3), + // from_raw_parts(handle, strides, shape) — contiguous 1D tensor of length 3 + unsafe { TensorArg::from_raw_parts(lhs, [1].into(), [3].into()) }, + unsafe { TensorArg::from_raw_parts(rhs, [1].into(), [3].into()) }, + unsafe { TensorArg::from_raw_parts(output.clone(), [1].into(), [3].into()) }, + ); + + let bytes = client.read_one(output).unwrap(); + let result = f32::from_bytes(&bytes); + + assert_eq!(result, &[11.0, 22.0, 33.0]); +} + +#[cube(launch)] +fn test_chose_cube_shape(input: &Tensor, out: &mut Tensor) { + if UNIT_POS == 0 { + out[0] = input.shape(0) as u32; + out[1] = input.shape(1) as u32; + } +} + +#[test] +fn test_chose_shape() { + let client = TestRuntime::client(&Default::default()); + let input = client.create_from_slice(f32::as_bytes(&[0.0; 12])); + let output = client.empty(2 * core::mem::size_of::()); + + test_chose_cube_shape::launch( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(1), + unsafe { TensorArg::from_raw_parts(input, [4, 1].into(), [3, 4].into()) }, + unsafe { TensorArg::from_raw_parts(output.clone(), [2, 1].into(), [1, 2].into()) }, + ); + + let bytes = client.read_one(output).unwrap(); + let result = u32::from_bytes(&bytes); + // input shape is [3, 4] + assert_eq!(result, &[3, 4]); +} + +#[cube(launch)] +fn test_chose_vectorized( + input: &Tensor>, + scalar: f32, + out: &mut Tensor>, +) { + if ABSOLUTE_POS < out.len() { + out[ABSOLUTE_POS] = input[ABSOLUTE_POS] + Vector::cast_from(scalar); + } +} + +#[test] +fn test_chose_cube_vectorized() { + let client = TestRuntime::client(&Default::default()); + + let vectorization = 2; + let input = client.create_from_slice(f32::as_bytes(&[ + 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., + ])); + let output = client.empty(12 * core::mem::size_of::()); + + test_chose_vectorized::launch::( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(6), + vectorization, + unsafe { TensorArg::from_raw_parts(input, [4, 1].into(), [3, 4].into()) }, + 10.0, + unsafe { TensorArg::from_raw_parts(output.clone(), [4, 1].into(), [3, 4].into()) }, + ); + + let bytes = client.read_one(output).unwrap(); + let result = f32::from_bytes(&bytes); + + println!("{:#?}", client.metrics().read().unwrap()); + + assert_eq!( + result, + &[11., 12., 13., 14., 15., 16., 17., 18., 19., 20., 21., 22.] + ); +} From 7795e3ad968f0685d29bf50acc7b88b979680501 Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Tue, 30 Jun 2026 14:02:27 -0400 Subject: [PATCH 03/14] working? --- crates/cubecl-core/src/compute/builder.rs | 6 +- crates/cubecl-core/src/compute/launcher.rs | 29 +-- crates/cubecl-ir/Cargo.toml | 3 + crates/cubecl-ir/src/flop_count.rs | 230 +++++++++++++++++---- crates/cubecl-ir/src/scope.rs | 16 +- crates/cubecl-runtime/src/client.rs | 24 ++- crates/cubecl-runtime/src/server/base.rs | 21 +- crates/cubecl/tests/chose.rs | 61 ++++-- 8 files changed, 292 insertions(+), 98 deletions(-) diff --git a/crates/cubecl-core/src/compute/builder.rs b/crates/cubecl-core/src/compute/builder.rs index 62d50e4ed9..74a2667198 100644 --- a/crates/cubecl-core/src/compute/builder.rs +++ b/crates/cubecl-core/src/compute/builder.rs @@ -8,9 +8,7 @@ use crate::{ prelude::KernelDefinition, }; use alloc::collections::BTreeMap; -use cubecl_ir::{ - DeviceProperties, FlopCountProcessor, Scope, StorageType, TargetProperties, Value, -}; +use cubecl_ir::{DeviceProperties, OpsCounts, Scope, StorageType, TargetProperties, Value}; use cubecl_runtime::config::{ CubeClRuntimeConfig, RuntimeConfig, compilation::CompilationLogLevel, }; @@ -94,7 +92,7 @@ impl KernelBuilder { /// Build the [kernel definition](KernelDefinition). pub fn build(mut self, settings: KernelSettings) -> KernelDefinition { if self.profile.enabled { - self.buffer(FlopCountProcessor::flop_counter_type()); + self.buffer(OpsCounts::stored_type()); } let scalars = self diff --git a/crates/cubecl-core/src/compute/launcher.rs b/crates/cubecl-core/src/compute/launcher.rs index 4fb5327e13..d82d800531 100644 --- a/crates/cubecl-core/src/compute/launcher.rs +++ b/crates/cubecl-core/src/compute/launcher.rs @@ -1,4 +1,5 @@ use alloc::{boxed::Box, vec::Vec}; +use bytemuck::Zeroable; use core::marker::PhantomData; use crate::Runtime; @@ -6,7 +7,7 @@ use crate::prelude::{BufferArg, TensorArg, TensorMapArg, TensorMapKind}; use crate::{InfoBuilder, KernelSettings, ScalarArgType}; #[cfg(feature = "std")] use core::cell::RefCell; -use cubecl_ir::{AddressType, FlopCountProcessor, Scope, StorageType, Type}; +use cubecl_ir::{AddressType, OpsCounts, Scope, StorageType, Type}; use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig}; use cubecl_runtime::id::KernelId; use cubecl_runtime::server::{Binding, CubeCount, Handle, TensorMapBinding}; @@ -67,23 +68,25 @@ impl KernelLauncher { self.with_info(|info| info.scalars.push_raw(bytes, dtype)); } - /// Injects profiling buffers into the bindings before execution. + /// Injects the profiling counters buffer into the bindings before execution. fn inject_profiling(&mut self, client: &ComputeClient) -> Option { if !CubeClRuntimeConfig::get().profiling.hardware_metrics { return None; } - let handle = client.create_from_slice(&[0u8; 4]); - let arg = unsafe { BufferArg::from_raw_parts(handle.clone(), 1) }; - self.register_buffer(arg, FlopCountProcessor::flop_counter_type()); + let handle = client.create_from_slice(bytemuck::bytes_of(&OpsCounts::zeroed())); + let arg = unsafe { BufferArg::from_raw_parts(handle.clone(), OpsCounts::LEN) }; + self.register_buffer(arg, OpsCounts::stored_type()); Some(handle) } - /// Read back the FLOP counter buffer and record it on the client, keyed by kernel id. - fn report_flop_counter(client: &ComputeClient, id: KernelId, handle: Handle) { + /// Read back the per-op FLOP counters and record them on the client, keyed by kernel id. The + /// dense slot buffer is turned into a sparse, name-keyed map here so storage stays legible. + fn report_profiling(client: &ComputeClient, id: KernelId, handle: Handle) { let bytes = client.read_one_unchecked(handle); - let count = u32::from_le_bytes(bytes[..4].try_into().expect("Counter should be 4 bytes")); - client.record_flop_count(id.clone(), count as u64); + let ops_counts: OpsCounts = bytemuck::pod_read_unaligned(&bytes); + + client.record_flop_count(id, ops_counts); } /// Launch the kernel. @@ -103,7 +106,7 @@ impl KernelLauncher { client.launch(kernel, cube_count, bindings); if let Some(handle) = flop_counter { - Self::report_flop_counter(client, kernel_id, handle); + Self::report_profiling(client, kernel_id, handle); } } @@ -122,7 +125,7 @@ impl KernelLauncher { kernel: K, client: &ComputeClient, ) { - let flop_counter = self.inject_profiling(client); + let profile = self.inject_profiling(client); let kernel_id = kernel.id(); unsafe { @@ -132,8 +135,8 @@ impl KernelLauncher { client.launch_unchecked(kernel, cube_count, bindings) } - if let Some(handle) = flop_counter { - Self::report_flop_counter(client, kernel_id, handle); + if let Some(handle) = profile { + Self::report_profiling(client, kernel_id, handle); } } diff --git a/crates/cubecl-ir/Cargo.toml b/crates/cubecl-ir/Cargo.toml index 82a9373a98..d629d009c1 100644 --- a/crates/cubecl-ir/Cargo.toml +++ b/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 } diff --git a/crates/cubecl-ir/src/flop_count.rs b/crates/cubecl-ir/src/flop_count.rs index 2de4e8656a..4aa14c7761 100644 --- a/crates/cubecl-ir/src/flop_count.rs +++ b/crates/cubecl-ir/src/flop_count.rs @@ -1,43 +1,188 @@ -//! Experimental FLOP profiling. -//! -//! When profiling is enabled on a [`Scope`], an extra global atomic counter buffer is appended to -//! the kernel and an atomic increment is inserted after each profiled arithmetic operation. This -//! gives a *dynamic* operation count that respects control flow and loops, since the increment -//! lives next to the operation in the same basic block. -//! -//! The processor is applied automatically by [`Scope::process`] for any scope whose -//! [`crate::ProfileInfo`] is enabled, so every backend supports it without backend-specific code. -//! -//! This is a first slice: only [`Arithmetic::Add`] is counted, and the counter increments by one -//! per executed instruction (vectorization width is ignored for now). - +use alloc::format; +use alloc::string::String; +use alloc::string::ToString; use alloc::vec::Vec; +use paste::paste; use crate::{ - Arithmetic, AtomicBinaryOperands, AtomicOp, ElemType, IndexOperands, Instruction, Memory, - Operation, Processor, Scope, ScopeProcessing, Type, UIntKind, Value, + ArithmeticOpCode, AtomicBinaryOperands, AtomicOp, ElemType, FloatKind, IndexOperands, + Instruction, IntKind, Memory, Operation, OperationReflect, Processor, Scope, ScopeProcessing, + Type, UIntKind, Value, }; -/// Inserts an atomic increment of `counter` after each profiled arithmetic operation. -#[derive(Debug)] -pub struct FlopCountProcessor { - /// The global counter buffer (a `&mut [atomic]`). - counter: Value, +macro_rules! arith_ops_count { + ($($op:ident),*) => { + paste! { + #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] + #[repr(C)] + pub struct ArithOpsCount { + $(pub [< $op:lower >]: u32),* + } + + impl ArithOpsCount { + pub const LEN: usize = size_of::() / size_of::(); + + pub fn offset(op_code: &ArithmeticOpCode) -> Option { + enum OpIndex { + $($op),* + } + + Some(match op_code { + $(ArithmeticOpCode::$op => OpIndex::$op as usize),* + }) + } + + /// Whether every counter in this block is zero. + pub fn is_empty(&self) -> bool { + true $(&& self.[< $op:lower >] == 0)* + } + } + + impl core::fmt::Display for ArithOpsCount { + /// Prints each non-zero op as `Name: count`, comma separated. Zero counts are + /// omitted, so an all-zero block prints nothing. + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut first = true; + $( + if self.[< $op:lower >] != 0 { + if !first { + write!(f, ", ")?; + } + write!(f, "{}: {}", stringify!($op), self.[< $op:lower >])?; + first = false; + } + )* + Ok(()) + } + } + } + } } -impl FlopCountProcessor { - /// Create a processor that increments the given global atomic `counter`. - pub fn new(counter: Value) -> Self { - Self { counter } +macro_rules! ops_counts { + ( + ariths : [$($arith:ident),*], + elems: [$($elems:ty),*] + ) => { + paste! { + arith_ops_count!($($arith),*); + + #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] + #[repr(C)] + pub struct OpsCount { + pub arith: ArithOpsCount, + } + + impl OpsCount { + pub const LEN: usize = size_of::() / size_of::(); + + pub fn offset(op: &Operation) -> Option { + match op { + Operation::Arithmetic(arith) => { + let base = core::mem::offset_of!(OpsCount, arith) / size_of::(); + ArithOpsCount::offset(&arith.op_code()).map(|o| base + o) + } + _ => None, + } + } + + /// Whether every counter in this block is zero. + pub fn is_empty(&self) -> bool { + self.arith.is_empty() + } + } + + impl core::fmt::Display for OpsCount { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "{}", self.arith) + } + } + + #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] + #[repr(C)] + pub struct OpsCounts { + $(pub [< $elems:lower >]: OpsCount),* + } + + impl core::fmt::Display for OpsCounts { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut any = false; + $( + if !self.[< $elems:lower >].is_empty() { + writeln!(f, " {}: {}", stringify!($elems), self.[< $elems:lower >])?; + any = true; + } + )* + if !any { + writeln!(f, " (none)")?; + } + Ok(()) + } + } + } } +} - /// The element type of the FLOP counter buffer: a single atomic `u32`. - pub fn flop_counter_type() -> Type { +ops_counts!( + ariths: [Add, SaturatingAdd, Fma, Sub, SaturatingSub, Mul, Div, Abs, Exp, Log, Log1p, Expm1, Cos, Sin, Tan, Tanh, Sinh, Cosh, ArcCos, ArcSin, ArcTan, ArcSinh, ArcCosh, ArcTanh, Degrees, Radians, ArcTan2, Powf, Powi, Hypot, Rhypot, Sqrt, InverseSqrt, Round, Floor, Ceil, Trunc, Erf, Recip, Clamp, Neg, Max, Min, Rem, ModFloor, Magnitude, Normalize, Dot, MulHi, VectorSum], + elems: [Bool, E2M1, E2M3, E3M2, E4M3, E5M2, UE8M0, F16, BF16, Flex32, F32, TF32, F64, Int8, Int16, Int32, Int64, Uint8, Uint16, Uint32, Uint64] +); + +impl OpsCounts { + pub const LEN: usize = size_of::() / size_of::(); + + pub fn stored_type() -> Type { Type::atomic(Type::scalar(ElemType::UInt(UIntKind::U32))) } - /// Append `counter[0] += 1` to the processing instruction stream. - fn emit_increment(&self, processing: &mut ScopeProcessing, vector_size: usize) { + pub fn elem_index(elem: &ElemType) -> usize { + match elem { + ElemType::Bool => 0, + ElemType::Float(FloatKind::E2M1) => 1, + ElemType::Float(FloatKind::E2M3) => 2, + ElemType::Float(FloatKind::E3M2) => 3, + ElemType::Float(FloatKind::E4M3) => 4, + ElemType::Float(FloatKind::E5M2) => 5, + ElemType::Float(FloatKind::UE8M0) => 6, + ElemType::Float(FloatKind::F16) => 7, + ElemType::Float(FloatKind::BF16) => 8, + ElemType::Float(FloatKind::Flex32) => 9, + ElemType::Float(FloatKind::F32) => 10, + ElemType::Float(FloatKind::TF32) => 11, + ElemType::Float(FloatKind::F64) => 12, + ElemType::Int(IntKind::I8) => 13, + ElemType::Int(IntKind::I16) => 14, + ElemType::Int(IntKind::I32) => 15, + ElemType::Int(IntKind::I64) => 16, + ElemType::UInt(UIntKind::U8) => 17, + ElemType::UInt(UIntKind::U16) => 18, + ElemType::UInt(UIntKind::U32) => 19, + ElemType::UInt(UIntKind::U64) => 20, + } + } + + pub fn offset(elem: &ElemType, op: &Operation) -> Option { + OpsCount::offset(op).map(|within| Self::elem_index(elem) * OpsCount::LEN + within) + } + + pub fn tracked_name(elem: &ElemType, op: &Operation) -> String { + format!("{:?}_{:?}", elem.to_string(), op.to_string()) + } +} + +#[derive(Debug)] +pub struct OpsCountsProcessor { + ops_counts: Value, +} + +impl OpsCountsProcessor { + /// Create a processor that increments the given global atomic `counter` array. + pub fn new(ops_counts: Value) -> Self { + Self { ops_counts } + } + + /// Append `counter[slot] += amount` to the processing instruction stream. + fn emit_increment(&self, processing: &mut ScopeProcessing, slot: u32, amount: u32) { // Build the increment in a throwaway scope that shares the global state (allocator, type // maps, ...) so the freshly created values don't collide with existing ones. Profiling is // disabled on this scope so its own `process` call doesn't recurse into FLOP counting. @@ -45,27 +190,27 @@ impl FlopCountProcessor { let u32_ty = Type::scalar(ElemType::UInt(UIntKind::U32)); - // `&counter[0]` -> pointer to the atomic element. + // `&counter[slot]` -> pointer to the atomic element for this op. let elem_ptr = scope.create_value(Type::pointer( - self.counter.value_type(), - self.counter.address_space(), + self.ops_counts.value_type(), + self.ops_counts.address_space(), )); scope.register(Instruction::new( Memory::Index(IndexOperands { - list: self.counter, - index: Value::constant(0u32.into(), u32_ty), + list: self.ops_counts, + index: Value::constant(slot.into(), u32_ty), unroll_factor: 1, checked: false, }), elem_ptr, )); - // `atomicAdd(&counter[0], 1)`. The returned old value is unused. + // `atomicAdd(&counter[slot], amount)`. The returned old value is unused. let old = scope.create_value(u32_ty); scope.register(Instruction::new( AtomicOp::Add(AtomicBinaryOperands { ptr: elem_ptr, - value: Value::constant(vector_size.into(), u32_ty), + value: Value::constant(amount.into(), u32_ty), }), old, )); @@ -75,21 +220,24 @@ impl FlopCountProcessor { } } -impl Processor for FlopCountProcessor { +impl Processor for OpsCountsProcessor { fn transform(&self, mut processing: ScopeProcessing) -> ScopeProcessing { let mut instructions = Vec::new(); core::mem::swap(&mut processing.instructions, &mut instructions); for instruction in instructions { - let (is_profiled, vector_size) = match &instruction.operation { - Operation::Arithmetic(Arithmetic::Add(ops)) => (true, ops.lhs.vector_size()), - _ => (false, 1), + let metric = match &instruction.operation { + Operation::Arithmetic(_) => { + OpsCounts::offset(&instruction.out().elem_type(), &instruction.operation) + .map(|slot| (slot as u32, instruction.out().vector_size() as u32)) + } + _ => None, }; processing.instructions.push(instruction); - if is_profiled { - self.emit_increment(&mut processing, vector_size); + if let Some((slot, amount)) = metric { + self.emit_increment(&mut processing, slot, amount); } } diff --git a/crates/cubecl-ir/src/scope.rs b/crates/cubecl-ir/src/scope.rs index 97537d8748..3ef7cb4745 100644 --- a/crates/cubecl-ir/src/scope.rs +++ b/crates/cubecl-ir/src/scope.rs @@ -11,9 +11,9 @@ use hashbrown::{HashMap, HashSet}; use itertools::Itertools; use crate::{ - AddressSpace, AggregateExtractOperands, CubeFnSource, DeviceProperties, FastMath, - FlopCountProcessor, Function, OpaqueType, Operation, OperationReflect, Processor, SourceLoc, - StorageType, TargetProperties, TypeHash, arena::DropBump, + AddressSpace, AggregateExtractOperands, CubeFnSource, DeviceProperties, FastMath, Function, + OpaqueType, Operation, OperationReflect, OpsCountsProcessor, Processor, SourceLoc, StorageType, + TargetProperties, TypeHash, arena::DropBump, }; use super::{Allocator, Id, Instruction, Type, Value, processing::ScopeProcessing}; @@ -323,17 +323,17 @@ impl Scope { global_state: self.global_state.clone(), }; + for p in processors { + processing = p.transform(processing); + } + if self.profile.enabled { let counter = self.global_state.borrow().global_args.last().copied(); if let Some(counter) = counter { - processing = FlopCountProcessor::new(counter).transform(processing); + processing = OpsCountsProcessor::new(counter).transform(processing); } } - for p in processors { - processing = p.transform(processing); - } - processing } diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index babbb5363e..e21b817d33 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -26,7 +26,7 @@ use cubecl_common::{ profile::ProfileDuration, stub::RwLock, }; -use cubecl_ir::{DeviceProperties, ElemType, VectorSize, features::Features}; +use cubecl_ir::{DeviceProperties, ElemType, OpsCounts, VectorSize, features::Features}; use cubecl_zspace::Shape; #[allow(unused)] @@ -733,18 +733,20 @@ impl ComputeClient { handle } - /// Record a FLOP count for the given kernel, stored per-device in + /// Record per-operation FLOP counts for the given kernel, stored per-device in /// [`ServerUtilities::metrics`]. Used by the launcher when `hardware_metrics` profiling is on. - pub fn record_flop_count(&self, id: KernelId, count: u64) { + pub fn record_flop_count(&self, id: KernelId, counts: OpsCounts) { let mut metrics = self.utilities.metrics.write().unwrap(); - let record = metrics.entry(id).or_default(); - record.last = count; - record.samples += 1; - } - - /// Returns the recorded FLOP metrics for the given kernel, if any. - pub fn flop_count(&self, id: &KernelId) -> Option { - self.utilities.metrics.read().unwrap().get(id).copied() + metrics + .entry(id) + .and_modify(|record| { + record.last = counts; + record.samples += 1; + }) + .or_insert(FlopRecord { + last: counts, + samples: 1, + }); } /// Returns a reference to the [`RwLock`] containing the recorded FLOP metrics for all kernels. diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index 9948b34509..a7be990f26 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -32,7 +32,7 @@ use cubecl_common::{ stream_id::StreamId, stub::RwLock, }; -use cubecl_ir::{DeviceProperties, ElemType, StorageType}; +use cubecl_ir::{DeviceProperties, ElemType, OpsCounts, StorageType}; use cubecl_zspace::{Shape, Strides, metadata::Metadata}; use hashbrown::{HashMap, HashSet}; use itertools::Itertools; @@ -79,12 +79,23 @@ impl core::fmt::Debug for ProfileError { /// Profiling metrics recorded for a single kernel, keyed by [`KernelId`] in /// [`ServerUtilities::metrics`]. Populated when `hardware_metrics` profiling is enabled. -#[derive(Debug, Default, Clone, Copy)] +#[derive(Debug, Clone)] pub struct FlopRecord { - /// Most recently recorded FLOP count for the kernel. - pub last: u64, + /// Most recent per-operation counts, keyed by op name (e.g. `"Add"`). Only operations that + /// actually executed appear. Sorted for deterministic, serialization-friendly output. + pub last: OpsCounts, /// Number of launches recorded for the kernel. - pub samples: u64, + pub samples: u32, +} + +impl std::fmt::Display for FlopRecord { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + writeln!(f, "FlopRecord:")?; + writeln!(f, " samples: {}", self.samples)?; + writeln!(f, " last:")?; + write!(f, "{}", self.last)?; + Ok(()) + } } /// Contains many different types that are useful for server implementations and compute clients. diff --git a/crates/cubecl/tests/chose.rs b/crates/cubecl/tests/chose.rs index 7bf358f710..0b651679e0 100644 --- a/crates/cubecl/tests/chose.rs +++ b/crates/cubecl/tests/chose.rs @@ -29,7 +29,17 @@ fn test_chose_1() { let bytes = client.read_one(output).unwrap(); let result = f32::from_bytes(&bytes); - println!("{:#?}", client.metrics().read().unwrap()); + println!( + "{:#}", + client + .metrics() + .read() + .unwrap() + .iter() + .take(1) + .collect::>()[0] + .1 + ); assert_eq!(result[0], lhs + rhs); } @@ -93,7 +103,7 @@ fn test_chose_shape() { assert_eq!(result, &[3, 4]); } -#[cube(launch)] +#[cube(launch_unchecked)] fn test_chose_vectorized( input: &Tensor>, scalar: f32, @@ -101,6 +111,10 @@ fn test_chose_vectorized( ) { if ABSOLUTE_POS < out.len() { out[ABSOLUTE_POS] = input[ABSOLUTE_POS] + Vector::cast_from(scalar); + out[ABSOLUTE_POS] = input[ABSOLUTE_POS] - Vector::cast_from(scalar); + out[ABSOLUTE_POS] = input[ABSOLUTE_POS] * Vector::cast_from(scalar); + out[ABSOLUTE_POS] = input[ABSOLUTE_POS] / Vector::cast_from(scalar); + out[ABSOLUTE_POS] = input[ABSOLUTE_POS].sin(); } } @@ -109,28 +123,43 @@ fn test_chose_cube_vectorized() { let client = TestRuntime::client(&Default::default()); let vectorization = 2; - let input = client.create_from_slice(f32::as_bytes(&[ - 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., - ])); + let input_values = [1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12.]; + let input = client.create_from_slice(f32::as_bytes(&input_values)); + let scale = 10.0; let output = client.empty(12 * core::mem::size_of::()); - test_chose_vectorized::launch::( - &client, - CubeCount::Static(1, 1, 1), - CubeDim::new_1d(6), - vectorization, - unsafe { TensorArg::from_raw_parts(input, [4, 1].into(), [3, 4].into()) }, - 10.0, - unsafe { TensorArg::from_raw_parts(output.clone(), [4, 1].into(), [3, 4].into()) }, - ); + unsafe { + test_chose_vectorized::launch_unchecked::( + &client, + CubeCount::Static(2, 1, 1), + CubeDim::new_1d((input_values.len() / vectorization) as u32), + vectorization, + unsafe { TensorArg::from_raw_parts(input, [4, 1].into(), [3, 4].into()) }, + scale, + unsafe { TensorArg::from_raw_parts(output.clone(), [4, 1].into(), [3, 4].into()) }, + ); + } let bytes = client.read_one(output).unwrap(); let result = f32::from_bytes(&bytes); - println!("{:#?}", client.metrics().read().unwrap()); + println!( + "{:#}", + client + .metrics() + .read() + .unwrap() + .iter() + .take(1) + .collect::>()[0] + .1 + ); assert_eq!( result, - &[11., 12., 13., 14., 15., 16., 17., 18., 19., 20., 21., 22.] + input_values + .into_iter() + .map(|v| v * scale) + .collect::>() ); } From 3acde7817a3706cc447d2ac554cd7b902b1297e2 Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Tue, 30 Jun 2026 16:11:04 -0400 Subject: [PATCH 04/14] better display and support memory ops --- crates/cubecl-ir/src/flop_count.rs | 184 +++++++++++++++++++++++------ 1 file changed, 149 insertions(+), 35 deletions(-) diff --git a/crates/cubecl-ir/src/flop_count.rs b/crates/cubecl-ir/src/flop_count.rs index 4aa14c7761..249f9259af 100644 --- a/crates/cubecl-ir/src/flop_count.rs +++ b/crates/cubecl-ir/src/flop_count.rs @@ -4,43 +4,64 @@ use alloc::string::ToString; use alloc::vec::Vec; use paste::paste; +use crate::AddressSpace; use crate::{ ArithmeticOpCode, AtomicBinaryOperands, AtomicOp, ElemType, FloatKind, IndexOperands, - Instruction, IntKind, Memory, Operation, OperationReflect, Processor, Scope, ScopeProcessing, - Type, UIntKind, Value, + Instruction, IntKind, Memory, MemoryOpCode, Operation, OperationReflect, Processor, Scope, + ScopeProcessing, Type, UIntKind, Value, }; -macro_rules! arith_ops_count { - ($($op:ident),*) => { +fn display_u32(value: u32) -> String { + let s = value.to_string(); + s.as_bytes() + .rchunks(3) + .rev() + .map(|chunk| std::str::from_utf8(chunk).unwrap()) + .collect::>() + .join(" ") +} + +macro_rules! ops_count { + ({$($op:ident),*}, $name:ident, $unit:literal) => { paste! { #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] #[repr(C)] - pub struct ArithOpsCount { + pub struct [< $name OpsCount >] { $(pub [< $op:lower >]: u32),* } - impl ArithOpsCount { + impl [< $name OpsCount >] { pub const LEN: usize = size_of::() / size_of::(); - pub fn offset(op_code: &ArithmeticOpCode) -> Option { + pub fn offset(op_code: &[< $name OpCode >]) -> Option { enum OpIndex { $($op),* } Some(match op_code { - $(ArithmeticOpCode::$op => OpIndex::$op as usize),* + $([< $name OpCode >]::$op => OpIndex::$op as usize),* }) } - /// Whether every counter in this block is zero. pub fn is_empty(&self) -> bool { true $(&& self.[< $op:lower >] == 0)* } + + pub const KIND: &'static str = stringify!($name); + pub const UNIT: &'static str = $unit; + + pub fn entries(&self) -> Vec<(&'static str, u32)> { + let mut entries = Vec::new(); + $( + if self.[< $op:lower >] != 0 { + entries.push((stringify!($op), self.[< $op:lower >])); + } + )* + entries + } } - impl core::fmt::Display for ArithOpsCount { - /// Prints each non-zero op as `Name: count`, comma separated. Zero counts are - /// omitted, so an all-zero block prints nothing. + impl core::fmt::Display for [< $name OpsCount >] { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { let mut first = true; $( @@ -48,7 +69,7 @@ macro_rules! arith_ops_count { if !first { write!(f, ", ")?; } - write!(f, "{}: {}", stringify!($op), self.[< $op:lower >])?; + write!(f, "{}: {}{}", stringify!($op), display_u32(self.[< $op:lower >]), $unit)?; first = false; } )* @@ -62,15 +83,18 @@ macro_rules! arith_ops_count { macro_rules! ops_counts { ( ariths : [$($arith:ident),*], + memories : [$($memory:ident),*], elems: [$($elems:ty),*] ) => { paste! { - arith_ops_count!($($arith),*); + ops_count!({$($arith),*}, Arithmetic, "ops"); + ops_count!({$($memory),*}, Memory, "bytes"); #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] #[repr(C)] pub struct OpsCount { - pub arith: ArithOpsCount, + pub arith: ArithmeticOpsCount, + pub memory: MemoryOpsCount, } impl OpsCount { @@ -80,21 +104,37 @@ macro_rules! ops_counts { match op { Operation::Arithmetic(arith) => { let base = core::mem::offset_of!(OpsCount, arith) / size_of::(); - ArithOpsCount::offset(&arith.op_code()).map(|o| base + o) + ArithmeticOpsCount::offset(&arith.op_code()).map(|o| base + o) + } + Operation::Memory(memory) => { + let base = core::mem::offset_of!(OpsCount, memory) / size_of::(); + MemoryOpsCount::offset(&memory.op_code()).map(|o| base + o) } _ => None, } } - /// Whether every counter in this block is zero. pub fn is_empty(&self) -> bool { - self.arith.is_empty() + self.arith.is_empty() && self.memory.is_empty() + } + + pub fn entries(&self) -> Vec<(&'static str, &'static str, &'static str, u32)> { + let mut entries = Vec::new(); + for (op, count) in self.arith.entries() { + entries.push((ArithmeticOpsCount::KIND, ArithmeticOpsCount::UNIT, op, count)); + } + for (op, count) in self.memory.entries() { + entries.push((MemoryOpsCount::KIND, MemoryOpsCount::UNIT, op, count)); + } + entries } } impl core::fmt::Display for OpsCount { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - write!(f, "{}", self.arith) + write!(f, "{}", self.arith)?; + write!(f, ", {}", self.memory)?; + Ok(()) } } @@ -106,16 +146,69 @@ macro_rules! ops_counts { impl core::fmt::Display for OpsCounts { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - let mut any = false; + let mut entries = Vec::new(); $( - if !self.[< $elems:lower >].is_empty() { - writeln!(f, " {}: {}", stringify!($elems), self.[< $elems:lower >])?; - any = true; + for (kind, unit, op, count) in self.[< $elems:lower >].entries() { + entries.push(( + stringify!($elems), + kind, + op, + display_u32(count), + unit + )); } )* - if !any { - writeln!(f, " (none)")?; + + if entries.is_empty() { + return writeln!(f, " (none)"); + } + + let max_op_len = entries.iter().map(|e| e.2.len()).max().unwrap_or(0); + let max_count_len = entries.iter().map(|e| e.3.len()).max().unwrap_or(0); + + let mut current_elem = ""; + let mut current_kind = ""; + let mut is_last_kind = false; + + for i in 0..entries.len() { + let (elem, kind, op, ref count_str, unit) = entries[i]; + + if elem != current_elem { + writeln!(f, "{}", elem)?; + current_elem = elem; + current_kind = ""; + } + + if kind != current_kind { + is_last_kind = !entries[i + 1..] + .iter() + .any(|e| e.0 == elem && e.1 != kind); + + let kind_prefix = if is_last_kind { "└── " } else { "├── " }; + writeln!(f, "{}{}", kind_prefix, kind)?; + current_kind = kind; + } + + let is_last_op = i + 1 == entries.len() + || entries[i + 1].0 != elem + || entries[i + 1].1 != kind; + + let kind_indent = if is_last_kind { " " } else { "│ " }; + let op_prefix = if is_last_op { "└── " } else { "├── " }; + + writeln!( + f, + "{}{}{:count_width$} {}", + kind_indent, + op_prefix, + format!("{}:", op), + count_str, + unit, + width = max_op_len + 1, + count_width = max_count_len + )?; } + Ok(()) } } @@ -125,6 +218,7 @@ macro_rules! ops_counts { ops_counts!( ariths: [Add, SaturatingAdd, Fma, Sub, SaturatingSub, Mul, Div, Abs, Exp, Log, Log1p, Expm1, Cos, Sin, Tan, Tanh, Sinh, Cosh, ArcCos, ArcSin, ArcTan, ArcSinh, ArcCosh, ArcTanh, Degrees, Radians, ArcTan2, Powf, Powi, Hypot, Rhypot, Sqrt, InverseSqrt, Round, Floor, Ceil, Trunc, Erf, Recip, Clamp, Neg, Max, Min, Rem, ModFloor, Magnitude, Normalize, Dot, MulHi, VectorSum], + memories: [Index, Load, Store, CopyMemory], elems: [Bool, E2M1, E2M3, E3M2, E4M3, E5M2, UE8M0, F16, BF16, Flex32, F32, TF32, F64, Int8, Int16, Int32, Int64, Uint8, Uint16, Uint32, Uint64] ); @@ -176,21 +270,15 @@ pub struct OpsCountsProcessor { } impl OpsCountsProcessor { - /// Create a processor that increments the given global atomic `counter` array. pub fn new(ops_counts: Value) -> Self { Self { ops_counts } } - /// Append `counter[slot] += amount` to the processing instruction stream. - fn emit_increment(&self, processing: &mut ScopeProcessing, slot: u32, amount: u32) { - // Build the increment in a throwaway scope that shares the global state (allocator, type - // maps, ...) so the freshly created values don't collide with existing ones. Profiling is - // disabled on this scope so its own `process` call doesn't recurse into FLOP counting. + fn emit_increment(&self, processing: &mut ScopeProcessing, index: u32, amount: u32) { let scope = Scope::root(false, false).with_global_state(processing.global_state.clone()); let u32_ty = Type::scalar(ElemType::UInt(UIntKind::U32)); - // `&counter[slot]` -> pointer to the atomic element for this op. let elem_ptr = scope.create_value(Type::pointer( self.ops_counts.value_type(), self.ops_counts.address_space(), @@ -198,14 +286,13 @@ impl OpsCountsProcessor { scope.register(Instruction::new( Memory::Index(IndexOperands { list: self.ops_counts, - index: Value::constant(slot.into(), u32_ty), + index: Value::constant(index.into(), u32_ty), unroll_factor: 1, checked: false, }), elem_ptr, )); - // `atomicAdd(&counter[slot], amount)`. The returned old value is unused. let old = scope.create_value(u32_ty); scope.register(Instruction::new( AtomicOp::Add(AtomicBinaryOperands { @@ -229,7 +316,34 @@ impl Processor for OpsCountsProcessor { let metric = match &instruction.operation { Operation::Arithmetic(_) => { OpsCounts::offset(&instruction.out().elem_type(), &instruction.operation) - .map(|slot| (slot as u32, instruction.out().vector_size() as u32)) + .map(|index| (index as u32, instruction.out().vector_size() as u32)) + } + Operation::Memory(memory) => { + let (value, target_address_spaces) = match memory { + Memory::Index(_) => ( + instruction.out(), + std::vec![instruction.out().address_space()], + ), + Memory::Load(variable) => (*variable, std::vec![variable.address_space()]), + Memory::Store(op) => (op.value, std::vec![op.ptr.address_space()]), + Memory::CopyMemory(op) => ( + op.source, + std::vec![op.source.address_space(), op.target.address_space()], + ), + }; + + if target_address_spaces.iter().any(|addr_space| { + !matches!(addr_space, AddressSpace::Local | AddressSpace::Shared) + }) { + OpsCounts::offset(&value.elem_type(), &instruction.operation).map(|index| { + ( + index as u32, + (value.vector_size() * value.elem_type().size()) as u32, + ) + }) + } else { + None + } } _ => None, }; From 01caea366b484944e133c7b8a6c64f892fba1082 Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Thu, 2 Jul 2026 12:26:04 -0400 Subject: [PATCH 05/14] thread local counters Co-authored-by: SamuelBelanger --- .../src/{flop_count.rs => counters.rs} | 152 ++++++------------ crates/cubecl-ir/src/lib.rs | 6 +- crates/cubecl-ir/src/profiler.rs | 72 +++++++++ crates/cubecl-ir/src/scope.rs | 11 +- .../cubecl-wgpu/src/compiler/wgsl/compiler.rs | 41 ++++- .../src/compiler/wgsl/instructions.rs | 2 + 6 files changed, 170 insertions(+), 114 deletions(-) rename crates/cubecl-ir/src/{flop_count.rs => counters.rs} (72%) create mode 100644 crates/cubecl-ir/src/profiler.rs diff --git a/crates/cubecl-ir/src/flop_count.rs b/crates/cubecl-ir/src/counters.rs similarity index 72% rename from crates/cubecl-ir/src/flop_count.rs rename to crates/cubecl-ir/src/counters.rs index 249f9259af..d121a72461 100644 --- a/crates/cubecl-ir/src/flop_count.rs +++ b/crates/cubecl-ir/src/counters.rs @@ -6,9 +6,8 @@ use paste::paste; use crate::AddressSpace; use crate::{ - ArithmeticOpCode, AtomicBinaryOperands, AtomicOp, ElemType, FloatKind, IndexOperands, - Instruction, IntKind, Memory, MemoryOpCode, Operation, OperationReflect, Processor, Scope, - ScopeProcessing, Type, UIntKind, Value, + ArithmeticOpCode, ElemType, FloatKind, IntKind, Memory, MemoryOpCode, Operation, + OperationReflect, Type, UIntKind, Value, }; fn display_u32(value: u32) -> String { @@ -21,6 +20,54 @@ fn display_u32(value: u32) -> String { .join(" ") } +#[derive(derive_new::new)] +pub struct CountInfo { + pub slot: u32, + pub amount: u32, +} + +pub struct OpsCounter; + +impl OpsCounter { + pub fn count(operation: &Operation, out: Option<&Value>) -> Option { + match operation { + Operation::Arithmetic(_) => { + let out = out?; + OpsCounts::offset(&out.elem_type(), operation) + .map(|index| CountInfo::new(index as u32, out.vector_size() as u32)) + } + Operation::Memory(memory) => { + let (value, target_address_spaces) = match memory { + Memory::Index(_) => { + let out = out?; + (out, std::vec![out.address_space()]) + } + Memory::Load(variable) => (variable, std::vec![variable.address_space()]), + Memory::Store(op) => (&op.value, std::vec![op.ptr.address_space()]), + Memory::CopyMemory(op) => ( + &op.source, + std::vec![op.source.address_space(), op.target.address_space()], + ), + }; + + if target_address_spaces.iter().any(|addr_space| { + !matches!(addr_space, AddressSpace::Local | AddressSpace::Shared) + }) { + OpsCounts::offset(&value.elem_type(), operation).map(|index| { + CountInfo::new( + index as u32, + (value.vector_size() * value.elem_type().size()) as u32, + ) + }) + } else { + None + } + } + _ => None, + } + } +} + macro_rules! ops_count { ({$($op:ident),*}, $name:ident, $unit:literal) => { paste! { @@ -258,103 +305,4 @@ impl OpsCounts { pub fn offset(elem: &ElemType, op: &Operation) -> Option { OpsCount::offset(op).map(|within| Self::elem_index(elem) * OpsCount::LEN + within) } - - pub fn tracked_name(elem: &ElemType, op: &Operation) -> String { - format!("{:?}_{:?}", elem.to_string(), op.to_string()) - } -} - -#[derive(Debug)] -pub struct OpsCountsProcessor { - ops_counts: Value, -} - -impl OpsCountsProcessor { - pub fn new(ops_counts: Value) -> Self { - Self { ops_counts } - } - - fn emit_increment(&self, processing: &mut ScopeProcessing, index: u32, amount: u32) { - let scope = Scope::root(false, false).with_global_state(processing.global_state.clone()); - - let u32_ty = Type::scalar(ElemType::UInt(UIntKind::U32)); - - let elem_ptr = scope.create_value(Type::pointer( - self.ops_counts.value_type(), - self.ops_counts.address_space(), - )); - scope.register(Instruction::new( - Memory::Index(IndexOperands { - list: self.ops_counts, - index: Value::constant(index.into(), u32_ty), - unroll_factor: 1, - checked: false, - }), - elem_ptr, - )); - - let old = scope.create_value(u32_ty); - scope.register(Instruction::new( - AtomicOp::Add(AtomicBinaryOperands { - ptr: elem_ptr, - value: Value::constant(amount.into(), u32_ty), - }), - old, - )); - - let tmp = scope.process([]); - processing.instructions.extend(tmp.instructions); - } -} - -impl Processor for OpsCountsProcessor { - fn transform(&self, mut processing: ScopeProcessing) -> ScopeProcessing { - let mut instructions = Vec::new(); - core::mem::swap(&mut processing.instructions, &mut instructions); - - for instruction in instructions { - let metric = match &instruction.operation { - Operation::Arithmetic(_) => { - OpsCounts::offset(&instruction.out().elem_type(), &instruction.operation) - .map(|index| (index as u32, instruction.out().vector_size() as u32)) - } - Operation::Memory(memory) => { - let (value, target_address_spaces) = match memory { - Memory::Index(_) => ( - instruction.out(), - std::vec![instruction.out().address_space()], - ), - Memory::Load(variable) => (*variable, std::vec![variable.address_space()]), - Memory::Store(op) => (op.value, std::vec![op.ptr.address_space()]), - Memory::CopyMemory(op) => ( - op.source, - std::vec![op.source.address_space(), op.target.address_space()], - ), - }; - - if target_address_spaces.iter().any(|addr_space| { - !matches!(addr_space, AddressSpace::Local | AddressSpace::Shared) - }) { - OpsCounts::offset(&value.elem_type(), &instruction.operation).map(|index| { - ( - index as u32, - (value.vector_size() * value.elem_type().size()) as u32, - ) - }) - } else { - None - } - } - _ => None, - }; - - processing.instructions.push(instruction); - - if let Some((slot, amount)) = metric { - self.emit_increment(&mut processing, slot, amount); - } - } - - processing - } } diff --git a/crates/cubecl-ir/src/lib.rs b/crates/cubecl-ir/src/lib.rs index 84dbbf5fe9..0748b0cc01 100644 --- a/crates/cubecl-ir/src/lib.rs +++ b/crates/cubecl-ir/src/lib.rs @@ -16,7 +16,7 @@ mod bitwise; mod branch; mod cmma; mod comparison; -mod flop_count; +mod counters; mod marker; mod memory; mod metadata; @@ -25,6 +25,7 @@ mod operation; mod operator; mod plane; mod processing; +pub mod profiler; mod properties; mod reflect; mod runtime_properties; @@ -47,7 +48,7 @@ pub use bitwise::*; pub use branch::*; pub use cmma::*; pub use comparison::*; -pub use flop_count::*; +pub use counters::*; pub use marker::*; pub use memory::*; pub use metadata::*; @@ -56,6 +57,7 @@ pub use operation::*; pub use operator::*; pub use plane::*; pub use processing::*; +pub use profiler::*; pub use properties::*; pub use reflect::*; pub use runtime_properties::*; diff --git a/crates/cubecl-ir/src/profiler.rs b/crates/cubecl-ir/src/profiler.rs new file mode 100644 index 0000000000..8ba16bd522 --- /dev/null +++ b/crates/cubecl-ir/src/profiler.rs @@ -0,0 +1,72 @@ +use crate::{Branch, CountInfo, Operation, OpsCounter, Value}; +use alloc::vec::Vec; +use hashbrown::HashSet; + +/// A reusable profiler component that tracks used flop slots during backend compilation. +#[derive(Default, Clone, Debug)] +pub struct CompilerProfiler { + flops_used: HashSet, + compiler: C, +} + +pub trait ProfileCompiler { + type I; + + fn increment_instruction(&self, slot: u32, amount: u32) -> Self::I; + fn prefix_instruction(&self, slot: u32) -> Self::I; + fn suffix_instruction(&self, slot: u32) -> Self::I; +} + +impl CompilerProfiler { + pub fn new(compiler: C) -> Self { + Self { + flops_used: HashSet::default(), + compiler, + } + } + + /// Mutable access to the backend compiler, e.g. to configure it before emitting prefix/suffix. + pub fn compiler_mut(&mut self) -> &mut C { + &mut self.compiler + } + + /// Intercepts an operation, tracking any used counters (flops, etc). + /// The backend provides a formatting closure to emit its specific increment instruction. + pub fn profile_operation( + &mut self, + operation: &Operation, + out: Option<&Value>, + instructions: &mut Vec<::I>, + ) { + if matches!(operation, Operation::Branch(Branch::Return)) {} + + if let Some(CountInfo { slot, amount }) = OpsCounter::count(operation, out) { + self.flops_used.insert(slot); + instructions.push(self.compiler.increment_instruction(slot, amount)); + } + } + + /// Iterates over all used profiling slots and generates initialization (prefix) + /// and aggregation (suffix) instructions via the provided formatting closures. + pub fn profile(&self, instructions: &mut Vec<::I>) { + let mut prefix = Vec::new(); + for slot in self.used_slots() { + prefix.push(self.compiler.prefix_instruction(*slot)); + } + + let mut old_instructions = core::mem::take(instructions); + + instructions.extend(prefix); + instructions.append(&mut old_instructions); + + for slot in self.used_slots() { + instructions.push(self.compiler.suffix_instruction(*slot)); + } + } + + /// Returns the flop slots that were used during compilation. + /// The backend can use this to emit initialization and flush instructions. + pub fn used_slots(&self) -> &HashSet { + &self.flops_used + } +} diff --git a/crates/cubecl-ir/src/scope.rs b/crates/cubecl-ir/src/scope.rs index 3ef7cb4745..450ecc9209 100644 --- a/crates/cubecl-ir/src/scope.rs +++ b/crates/cubecl-ir/src/scope.rs @@ -12,8 +12,8 @@ use itertools::Itertools; use crate::{ AddressSpace, AggregateExtractOperands, CubeFnSource, DeviceProperties, FastMath, Function, - OpaqueType, Operation, OperationReflect, OpsCountsProcessor, Processor, SourceLoc, StorageType, - TargetProperties, TypeHash, arena::DropBump, + OpaqueType, Operation, OperationReflect, Processor, SourceLoc, StorageType, TargetProperties, + TypeHash, arena::DropBump, }; use super::{Allocator, Id, Instruction, Type, Value, processing::ScopeProcessing}; @@ -327,13 +327,6 @@ impl Scope { processing = p.transform(processing); } - if self.profile.enabled { - let counter = self.global_state.borrow().global_args.last().copied(); - if let Some(counter) = counter { - processing = OpsCountsProcessor::new(counter).transform(processing); - } - } - processing } diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs index ade0f21cb6..fcc376b18a 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs @@ -19,6 +19,7 @@ use cubecl_core::{ }; use cubecl_core::{post_processing::disaggregate::DisaggregateVisitor, prelude::*}; use cubecl_ir::AddressSpace; +use cubecl_ir::ProfileCompiler; use cubecl_runtime::compiler::CompilationError; use cubecl_runtime::kernel; use hashbrown::HashMap; @@ -50,6 +51,33 @@ pub struct WgslCompiler { strategy: ExecutionMode, subgroup_instructions_used: bool, f16_used: bool, + profiler: cubecl_core::ir::CompilerProfiler, +} + +#[derive(Clone, Default)] +pub struct WgslCompilerProfiler { + /// WGSL name of the global atomic counter buffer (e.g. `val_12`), set before emitting suffixes. + counter: Option, +} + +impl ProfileCompiler for WgslCompilerProfiler { + type I = wgsl::Instruction; + + fn increment_instruction(&self, slot: u32, amount: u32) -> Self::I { + wgsl::Instruction::Custom(format!("flop_{slot} += {amount}u;")) + } + + fn prefix_instruction(&self, slot: u32) -> Self::I { + wgsl::Instruction::Custom(format!("var flop_{slot} = 0u;")) + } + + fn suffix_instruction(&self, slot: u32) -> Self::I { + let counter = self + .counter + .as_deref() + .expect("The counter buffer should be set when profiling is enabled"); + wgsl::Instruction::Custom(format!("atomicAdd(&{counter}[{slot}], flop_{slot});")) + } } impl core::fmt::Debug for WgslCompiler { @@ -130,7 +158,14 @@ impl WgslCompiler { .resize(value.num_global_buffers(), Visibility::Read); let address_type = self.compile_storage_type(address_type); - let instructions = self.compile_scope(&value.body); + let mut instructions = self.compile_scope(&value.body); + if value.body.profile.enabled { + // The counter buffer is the last one, appended by the kernel builder when profiling. + if let Some(counter) = value.buffers.last() { + self.profiler.compiler_mut().counter = Some(format!("val_{}", counter.value.id())); + } + self.profiler.profile(&mut instructions); + } let extensions = register_extensions(&instructions); let body = wgsl::Body { instructions, @@ -311,6 +346,10 @@ impl WgslCompiler { out: Option, scope: &cube::Scope, ) { + if scope.profile.enabled { + self.profiler + .profile_operation(&operation, out.as_ref(), instructions); + } match operation { cube::Operation::Copy(value) => instructions.push(wgsl::Instruction::Assign { input: self.compile_value(value), diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs b/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs index adc80ae515..de9885fc47 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs @@ -12,6 +12,7 @@ use std::fmt::Display; #[derive(Debug, Clone)] #[allow(dead_code)] // Some variants might not be used with different flags pub enum Instruction { + Custom(String), DeclareVariable { val: Value, value_ty: Item, @@ -482,6 +483,7 @@ pub enum Instruction { impl Display for Instruction { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + Instruction::Custom(s) => writeln!(f, "{s}"), Instruction::DeclareVariable { val, value_ty } => { writeln!(f, "var {val}_store: {value_ty};")?; writeln!(f, "let {val} = &{val}_store;") From b173451e08f2a6a038b7a6a2061103748f0e8993 Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Thu, 2 Jul 2026 15:56:10 -0400 Subject: [PATCH 06/14] working wgpu clean Co-authored-by: SamuelBelanger --- crates/cubecl-core/src/codegen/mod.rs | 2 + crates/cubecl-core/src/codegen/profiler.rs | 93 +++++++++++++++++++ crates/cubecl-core/src/compute/builder.rs | 3 +- crates/cubecl-core/src/compute/launcher.rs | 4 +- crates/cubecl-ir/src/lib.rs | 2 - crates/cubecl-ir/src/profiler.rs | 72 -------------- crates/cubecl-ir/src/scope.rs | 2 + .../cubecl-wgpu/src/compiler/wgsl/compiler.rs | 63 ++++++------- crates/cubecl-wgpu/src/compiler/wgsl/mod.rs | 2 + .../cubecl-wgpu/src/compiler/wgsl/profile.rs | 23 +++++ crates/cubecl/tests/chose.rs | 5 +- 11 files changed, 156 insertions(+), 115 deletions(-) create mode 100644 crates/cubecl-core/src/codegen/profiler.rs delete mode 100644 crates/cubecl-ir/src/profiler.rs create mode 100644 crates/cubecl-wgpu/src/compiler/wgsl/profile.rs diff --git a/crates/cubecl-core/src/codegen/mod.rs b/crates/cubecl-core/src/codegen/mod.rs index 13cc7ebcd4..f5d58bbee8 100644 --- a/crates/cubecl-core/src/codegen/mod.rs +++ b/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::*; diff --git a/crates/cubecl-core/src/codegen/profiler.rs b/crates/cubecl-core/src/codegen/profiler.rs new file mode 100644 index 0000000000..90fd2a4d5e --- /dev/null +++ b/crates/cubecl-core/src/codegen/profiler.rs @@ -0,0 +1,93 @@ +use alloc::{fmt::Display, vec::Vec}; +use cubecl_ir::{Branch, CountInfo, Operation, OpsCounter, Value}; +use hashbrown::HashSet; + +/// A reusable profiler component that tracks used flop slots during backend compilation. +#[derive(Clone, Debug)] +pub struct CompilerProfiler { + flops_used: HashSet, + compiler: C, + counter: Option, +} + +impl Default for CompilerProfiler { + fn default() -> Self { + Self { + flops_used: HashSet::default(), + compiler: C::default(), + counter: Option::default(), + } + } +} + +pub trait CompilerProfilerInstructions: Default { + type Instruction; + type CounterVariable: Display; + + /// Returns an instruction that increments the value of the given slot by the given amount. + fn increment_instruction(&self, slot: u32, amount: u32) -> Self::Instruction; + /// Returns an instruction that prefixes the given slot, e.g. to initialize it to zero. + fn declare_instruction(&self, slot: u32) -> Self::Instruction; + /// Returns an instruction that suffixes the given slot, e.g. to aggregate its value. + fn flush_instruction(&self, slot: u32, counter: &Self::CounterVariable) -> Self::Instruction; +} + +impl CompilerProfiler { + pub fn new(compiler: C) -> Self { + Self { + flops_used: HashSet::default(), + compiler, + counter: None, + } + } + + pub fn set_counter(&mut self, counter: C::CounterVariable) { + self.counter = Some(counter) + } + + /// Intercepts an operation, tracking any used counters (flops, etc). + /// The backend provides a formatting closure to emit its specific increment instruction. + pub fn profile_operation( + &mut self, + operation: &Operation, + out: Option<&Value>, + instructions: &mut Vec<::Instruction>, + ) { + if let Some(counter) = &self.counter { + if matches!(operation, Operation::Branch(Branch::Return)) { + for slot in &self.flops_used { + instructions.push(self.compiler.flush_instruction(*slot, counter)); + } + return; + } + + if let Some(CountInfo { slot, amount }) = OpsCounter::count(operation, out) { + self.flops_used.insert(slot); + instructions.push(self.compiler.increment_instruction(slot, amount)); + } + } + } + + /// Iterates over all used profiling slots and generates initialization (prefix) + /// and aggregation (suffix) instructions via the provided formatting closures. + pub fn profile( + &self, + instructions: &mut Vec<::Instruction>, + ) { + if let Some(counter) = &self.counter { + let mut declares = Vec::new(); + for slot in &self.flops_used { + declares.push(self.compiler.declare_instruction(*slot)); + } + + let mut old_instructions = core::mem::take(instructions); + + instructions.extend(declares); + instructions.append(&mut old_instructions); + + for slot in &self.flops_used { + instructions.push(self.compiler.flush_instruction(*slot, counter)); + } + } + } +} diff --git a/crates/cubecl-core/src/compute/builder.rs b/crates/cubecl-core/src/compute/builder.rs index 74a2667198..66d71343f0 100644 --- a/crates/cubecl-core/src/compute/builder.rs +++ b/crates/cubecl-core/src/compute/builder.rs @@ -92,7 +92,8 @@ impl KernelBuilder { /// Build the [kernel definition](KernelDefinition). pub fn build(mut self, settings: KernelSettings) -> KernelDefinition { if self.profile.enabled { - self.buffer(OpsCounts::stored_type()); + let buffer = self.buffer(OpsCounts::stored_type()); + self.scope.profile.counters_buffer = Some(buffer); } let scalars = self diff --git a/crates/cubecl-core/src/compute/launcher.rs b/crates/cubecl-core/src/compute/launcher.rs index d82d800531..8fa3d471d4 100644 --- a/crates/cubecl-core/src/compute/launcher.rs +++ b/crates/cubecl-core/src/compute/launcher.rs @@ -97,7 +97,7 @@ impl KernelLauncher { kernel: K, client: &ComputeClient, ) { - let flop_counter = self.inject_profiling(client); + let profile = self.inject_profiling(client); let kernel_id = kernel.id(); let bindings = self.into_bindings(); @@ -105,7 +105,7 @@ impl KernelLauncher { client.launch(kernel, cube_count, bindings); - if let Some(handle) = flop_counter { + if let Some(handle) = profile { Self::report_profiling(client, kernel_id, handle); } } diff --git a/crates/cubecl-ir/src/lib.rs b/crates/cubecl-ir/src/lib.rs index 0748b0cc01..de0b2fbad5 100644 --- a/crates/cubecl-ir/src/lib.rs +++ b/crates/cubecl-ir/src/lib.rs @@ -25,7 +25,6 @@ mod operation; mod operator; mod plane; mod processing; -pub mod profiler; mod properties; mod reflect; mod runtime_properties; @@ -57,7 +56,6 @@ pub use operation::*; pub use operator::*; pub use plane::*; pub use processing::*; -pub use profiler::*; pub use properties::*; pub use reflect::*; pub use runtime_properties::*; diff --git a/crates/cubecl-ir/src/profiler.rs b/crates/cubecl-ir/src/profiler.rs deleted file mode 100644 index 8ba16bd522..0000000000 --- a/crates/cubecl-ir/src/profiler.rs +++ /dev/null @@ -1,72 +0,0 @@ -use crate::{Branch, CountInfo, Operation, OpsCounter, Value}; -use alloc::vec::Vec; -use hashbrown::HashSet; - -/// A reusable profiler component that tracks used flop slots during backend compilation. -#[derive(Default, Clone, Debug)] -pub struct CompilerProfiler { - flops_used: HashSet, - compiler: C, -} - -pub trait ProfileCompiler { - type I; - - fn increment_instruction(&self, slot: u32, amount: u32) -> Self::I; - fn prefix_instruction(&self, slot: u32) -> Self::I; - fn suffix_instruction(&self, slot: u32) -> Self::I; -} - -impl CompilerProfiler { - pub fn new(compiler: C) -> Self { - Self { - flops_used: HashSet::default(), - compiler, - } - } - - /// Mutable access to the backend compiler, e.g. to configure it before emitting prefix/suffix. - pub fn compiler_mut(&mut self) -> &mut C { - &mut self.compiler - } - - /// Intercepts an operation, tracking any used counters (flops, etc). - /// The backend provides a formatting closure to emit its specific increment instruction. - pub fn profile_operation( - &mut self, - operation: &Operation, - out: Option<&Value>, - instructions: &mut Vec<::I>, - ) { - if matches!(operation, Operation::Branch(Branch::Return)) {} - - if let Some(CountInfo { slot, amount }) = OpsCounter::count(operation, out) { - self.flops_used.insert(slot); - instructions.push(self.compiler.increment_instruction(slot, amount)); - } - } - - /// Iterates over all used profiling slots and generates initialization (prefix) - /// and aggregation (suffix) instructions via the provided formatting closures. - pub fn profile(&self, instructions: &mut Vec<::I>) { - let mut prefix = Vec::new(); - for slot in self.used_slots() { - prefix.push(self.compiler.prefix_instruction(*slot)); - } - - let mut old_instructions = core::mem::take(instructions); - - instructions.extend(prefix); - instructions.append(&mut old_instructions); - - for slot in self.used_slots() { - instructions.push(self.compiler.suffix_instruction(*slot)); - } - } - - /// Returns the flop slots that were used during compilation. - /// The backend can use this to emit initialization and flush instructions. - pub fn used_slots(&self) -> &HashSet { - &self.flops_used - } -} diff --git a/crates/cubecl-ir/src/scope.rs b/crates/cubecl-ir/src/scope.rs index 450ecc9209..7988592b24 100644 --- a/crates/cubecl-ir/src/scope.rs +++ b/crates/cubecl-ir/src/scope.rs @@ -100,6 +100,7 @@ pub struct DebugInfo { #[derive(Debug, Clone, PartialEq, Eq, TypeHash)] pub struct ProfileInfo { pub enabled: bool, + pub counters_buffer: Option, } /// Modes set and reset during expansion @@ -162,6 +163,7 @@ impl Scope { }, profile: ProfileInfo { enabled: profile_enabled, + counters_buffer: None, }, global_state: Default::default(), } diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs index fcc376b18a..d251a757b3 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs @@ -1,9 +1,11 @@ use super::Item; use super::Subgroup; use super::shader::ComputeShader; +use crate::compiler::wgsl::WgslCompilerProfiler; use crate::compiler::wgsl::{self, SharedValue}; use cubecl_common::backtrace::BackTrace; +use cubecl_core::CompilerProfiler; use cubecl_core::ir::{Processor, UIntKind}; use cubecl_core::{ Info, @@ -19,7 +21,6 @@ use cubecl_core::{ }; use cubecl_core::{post_processing::disaggregate::DisaggregateVisitor, prelude::*}; use cubecl_ir::AddressSpace; -use cubecl_ir::ProfileCompiler; use cubecl_runtime::compiler::CompilationError; use cubecl_runtime::kernel; use hashbrown::HashMap; @@ -51,33 +52,7 @@ pub struct WgslCompiler { strategy: ExecutionMode, subgroup_instructions_used: bool, f16_used: bool, - profiler: cubecl_core::ir::CompilerProfiler, -} - -#[derive(Clone, Default)] -pub struct WgslCompilerProfiler { - /// WGSL name of the global atomic counter buffer (e.g. `val_12`), set before emitting suffixes. - counter: Option, -} - -impl ProfileCompiler for WgslCompilerProfiler { - type I = wgsl::Instruction; - - fn increment_instruction(&self, slot: u32, amount: u32) -> Self::I { - wgsl::Instruction::Custom(format!("flop_{slot} += {amount}u;")) - } - - fn prefix_instruction(&self, slot: u32) -> Self::I { - wgsl::Instruction::Custom(format!("var flop_{slot} = 0u;")) - } - - fn suffix_instruction(&self, slot: u32) -> Self::I { - let counter = self - .counter - .as_deref() - .expect("The counter buffer should be set when profiling is enabled"); - wgsl::Instruction::Custom(format!("atomicAdd(&{counter}[{slot}], flop_{slot});")) - } + profiler: CompilerProfiler, } impl core::fmt::Debug for WgslCompiler { @@ -157,16 +132,13 @@ impl WgslCompiler { self.buffer_vis .resize(value.num_global_buffers(), Visibility::Read); - let address_type = self.compile_storage_type(address_type); + self.setup_profiler(&value.body); let mut instructions = self.compile_scope(&value.body); - if value.body.profile.enabled { - // The counter buffer is the last one, appended by the kernel builder when profiling. - if let Some(counter) = value.buffers.last() { - self.profiler.compiler_mut().counter = Some(format!("val_{}", counter.value.id())); - } - self.profiler.profile(&mut instructions); - } + self.profile(&value.body, &mut instructions); + + let address_type = self.compile_storage_type(address_type); let extensions = register_extensions(&instructions); + let body = wgsl::Body { instructions, id: self.id, @@ -325,6 +297,25 @@ impl WgslCompiler { self.compile_value(val) } + fn setup_profiler(&mut self, scope: &cube::Scope) { + if scope.profile.enabled { + let counter = self.compile_value( + scope + .profile + .counters_buffer + .expect("Profiling counters buffer should be initialized"), + ); + + self.profiler.set_counter(counter); + } + } + + fn profile(&self, scope: &cube::Scope, mut instructions: &mut Vec) { + if scope.profile.enabled { + self.profiler.profile(&mut instructions); + } + } + fn compile_scope(&mut self, scope: &cube::Scope) -> Vec { let mut instructions = Vec::new(); diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs b/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs index d549832b14..ec10c6e5cd 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs @@ -3,6 +3,7 @@ mod body; mod compiler; mod extension; mod instructions; +mod profile; pub(crate) mod shader; mod subgroup; @@ -11,5 +12,6 @@ pub(crate) use body::*; pub use compiler::*; pub(crate) use extension::*; pub(crate) use instructions::*; +pub(crate) use profile::*; pub(crate) use shader::*; pub(crate) use subgroup::*; diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/profile.rs b/crates/cubecl-wgpu/src/compiler/wgsl/profile.rs new file mode 100644 index 0000000000..5f4676ab1d --- /dev/null +++ b/crates/cubecl-wgpu/src/compiler/wgsl/profile.rs @@ -0,0 +1,23 @@ +use cubecl_core::CompilerProfilerInstructions; + +use crate::compiler::wgsl; + +#[derive(Clone, Default)] +pub struct WgslCompilerProfiler; + +impl CompilerProfilerInstructions for WgslCompilerProfiler { + type Instruction = wgsl::Instruction; + type CounterVariable = wgsl::Value; + + fn increment_instruction(&self, slot: u32, amount: u32) -> Self::Instruction { + wgsl::Instruction::Custom(format!("flop_{slot} += {amount}u;")) + } + + fn declare_instruction(&self, slot: u32) -> Self::Instruction { + wgsl::Instruction::Custom(format!("var flop_{slot} = 0u;")) + } + + fn flush_instruction(&self, slot: u32, counter: &Self::CounterVariable) -> Self::Instruction { + wgsl::Instruction::Custom(format!("atomicAdd(&{counter}[{slot}], flop_{slot});")) + } +} diff --git a/crates/cubecl/tests/chose.rs b/crates/cubecl/tests/chose.rs index 0b651679e0..bbbfb10fb2 100644 --- a/crates/cubecl/tests/chose.rs +++ b/crates/cubecl/tests/chose.rs @@ -5,6 +5,7 @@ fn test_chose_cube(lhs: f32, rhs: f32, out: &mut [f32]) { if UNIT_POS == 0 { out[0] = lhs + rhs + lhs; out[0] = lhs + rhs; + terminate!(); out[0] = lhs + rhs; } } @@ -134,9 +135,9 @@ fn test_chose_cube_vectorized() { CubeCount::Static(2, 1, 1), CubeDim::new_1d((input_values.len() / vectorization) as u32), vectorization, - unsafe { TensorArg::from_raw_parts(input, [4, 1].into(), [3, 4].into()) }, + TensorArg::from_raw_parts(input, [4, 1].into(), [3, 4].into()), scale, - unsafe { TensorArg::from_raw_parts(output.clone(), [4, 1].into(), [3, 4].into()) }, + TensorArg::from_raw_parts(output.clone(), [4, 1].into(), [3, 4].into()), ); } From 03b02293d81669c1c6b5e57dc6f1b8fbdb1c813d Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Thu, 2 Jul 2026 17:11:46 -0400 Subject: [PATCH 07/14] EnumCounts derive Co-authored-by: SamuelBelanger --- crates/cubecl-ir/Cargo.toml | 2 + crates/cubecl-ir/src/lib.rs | 4 +- crates/cubecl-ir/src/type.rs | 46 +++++- .../cubecl-macros-internal/src/enum_counts.rs | 138 ++++++++++++++++++ crates/cubecl-macros-internal/src/lib.rs | 15 ++ 5 files changed, 200 insertions(+), 5 deletions(-) create mode 100644 crates/cubecl-macros-internal/src/enum_counts.rs diff --git a/crates/cubecl-ir/Cargo.toml b/crates/cubecl-ir/Cargo.toml index d629d009c1..f8861ac16e 100644 --- a/crates/cubecl-ir/Cargo.toml +++ b/crates/cubecl-ir/Cargo.toml @@ -52,3 +52,5 @@ hashbrown = { workspace = true } portable-atomic = { workspace = true } paste = { workspace = true } + +strum = { version = "0.28.0", features = ["derive"] } diff --git a/crates/cubecl-ir/src/lib.rs b/crates/cubecl-ir/src/lib.rs index de0b2fbad5..f50cd3df94 100644 --- a/crates/cubecl-ir/src/lib.rs +++ b/crates/cubecl-ir/src/lib.rs @@ -66,5 +66,7 @@ pub use tma::*; pub use r#type::*; pub use variable::*; -pub(crate) use cubecl_macros_internal::{OperationArgs, OperationCode, OperationReflect, TypeHash}; +pub(crate) use cubecl_macros_internal::{ + EnumCounts, OperationArgs, OperationCode, OperationReflect, TypeHash, +}; pub use type_hash::TypeHash; diff --git a/crates/cubecl-ir/src/type.rs b/crates/cubecl-ir/src/type.rs index cf590781b0..6bd3844825 100644 --- a/crates/cubecl-ir/src/type.rs +++ b/crates/cubecl-ir/src/type.rs @@ -1,5 +1,5 @@ use super::{ConstantValue, Value, ValueKind}; -use crate::{BarrierLevel, ClampMode, Id, MatrixType, TypeHash}; +use crate::{BarrierLevel, ClampMode, EnumCounts, Id, MatrixType, TypeHash}; use core::fmt::Display; use cubecl_common::{ e2m1, e2m1x2, e2m3, e3m2, e4m3, e5m2, flex32, @@ -12,7 +12,7 @@ use half::{bf16, f16}; pub use internment::Intern; #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, EnumCounts)] #[allow(missing_docs)] pub enum FloatKind { /// FP4, 2 bit exponent, 1 bit mantissa @@ -38,7 +38,7 @@ pub enum FloatKind { } #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, EnumCounts)] #[allow(missing_docs)] pub enum IntKind { I8, @@ -48,7 +48,7 @@ pub enum IntKind { } #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, EnumCounts)] #[allow(missing_docs)] pub enum UIntKind { U8, @@ -1162,3 +1162,41 @@ impl_into_value!( usize => UIntKind::U32, isize => IntKind::I32, ); + +#[cfg(test)] +mod enum_counts_tests { + use super::*; + + #[test] + fn count_and_index_are_consistent() { + assert_eq!(FloatKind::COUNT, 12); + assert_eq!(FloatKind::VARIANTS.len(), FloatKind::COUNT); + for (i, kind) in FloatKind::VARIANTS.into_iter().enumerate() { + assert_eq!(kind.index(), i); + } + } + + #[test] + fn named_fields_match_indexed_access() { + let mut counts = FloatKindCounts::::default(); + counts.f16 = 7; + counts[FloatKind::F32] = 9; + + assert_eq!(counts[FloatKind::F16], 7); + assert_eq!(counts.f32, 9); + // as_slice is in declaration order, so index() addresses the same entry. + assert_eq!(counts.as_slice()[FloatKind::F16.index()], 7); + assert_eq!(counts.iter().filter(|(_, v)| **v != 0).count(), 2); + } + + #[test] + fn is_pod() { + // Round-trips through bytes: relies on the generated `#[repr(C)]` + `Pod` impl. + let mut counts = UIntKindCounts::::default(); + counts.u32 = 42; + let bytes = bytemuck::bytes_of(&counts); + assert_eq!(bytes.len(), UIntKind::COUNT * size_of::()); + let back: UIntKindCounts = *bytemuck::from_bytes(bytes); + assert_eq!(back.u32, 42); + } +} diff --git a/crates/cubecl-macros-internal/src/enum_counts.rs b/crates/cubecl-macros-internal/src/enum_counts.rs new file mode 100644 index 0000000000..f4dcc30d6e --- /dev/null +++ b/crates/cubecl-macros-internal/src/enum_counts.rs @@ -0,0 +1,138 @@ +use proc_macro2::TokenStream; +use quote::{format_ident, quote}; +use syn::{Data, DeriveInput, Fields, Ident}; + +/// Generates a flat, named-field companion for a fieldless enum, plus variant reflection on the +/// enum itself. For an enum `Kind { A, B }` it emits: +/// +/// - `impl Kind { const COUNT; const VARIANTS; const fn index(self) }` +/// - `struct KindCounts { pub a: T, pub b: T }` (`#[repr(C)]`, `Pod`/`Zeroable` when `T` is), +/// with `get`/`get_mut`, `iter`, `as_slice`, and `Index`/`IndexMut`. +/// +/// Field order matches declaration order, so `KindCounts::as_slice()[k.index()]` is the entry for +/// `k`. Adding a variant automatically extends both the reflection and the companion struct. +pub fn enum_counts_impl(input: DeriveInput) -> syn::Result { + let enum_ident = &input.ident; + let vis = &input.vis; + + let Data::Enum(data) = &input.data else { + return Err(syn::Error::new_spanned( + enum_ident, + "EnumCounts can only be derived on enums", + )); + }; + + if !input.generics.params.is_empty() { + return Err(syn::Error::new_spanned( + &input.generics, + "EnumCounts does not support generic enums", + )); + } + + let mut variants = Vec::new(); + let mut fields = Vec::new(); + for variant in &data.variants { + if !matches!(variant.fields, Fields::Unit) { + return Err(syn::Error::new_spanned( + &variant.ident, + "EnumCounts requires fieldless (unit) variants", + )); + } + let ident = &variant.ident; + fields.push(Ident::new(&ident.to_string().to_lowercase(), ident.span())); + variants.push(ident.clone()); + } + + let count = variants.len(); + let counts_ident = format_ident!("{}Counts", enum_ident); + + let index_arms = variants + .iter() + .enumerate() + .map(|(i, v)| quote! { #enum_ident::#v => #i }); + let get_arms = variants + .iter() + .zip(&fields) + .map(|(v, f)| quote! { #enum_ident::#v => &self.#f }); + let get_mut_arms = variants + .iter() + .zip(&fields) + .map(|(v, f)| quote! { #enum_ident::#v => &mut self.#f }); + let entries = variants + .iter() + .zip(&fields) + .map(|(v, f)| quote! { (#enum_ident::#v, &self.#f) }); + let variants_list = variants.iter().map(|v| quote! { #enum_ident::#v }); + let struct_fields = fields.iter().map(|f| quote! { pub #f: T }); + + Ok(quote! { + impl #enum_ident { + /// Number of variants. + pub const COUNT: usize = #count; + /// All variants, in declaration order. + pub const VARIANTS: [#enum_ident; #count] = [ #(#variants_list),* ]; + + /// Stable index of this variant (declaration order, `0..COUNT`). + pub const fn index(self) -> usize { + match self { #(#index_arms),* } + } + } + + #[doc = concat!("Per-variant values keyed by [`", stringify!(#enum_ident), "`].")] + #[repr(C)] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] + #vis struct #counts_ident { + #(#struct_fields),* + } + + // Safe: `#[repr(C)]` with `COUNT` fields all of type `T` (no padding), so the layout is + // exactly `[T; COUNT]`. + unsafe impl ::bytemuck::Zeroable for #counts_ident {} + unsafe impl ::bytemuck::Pod for #counts_ident {} + + impl #counts_ident { + /// Number of entries (equals `COUNT` of the key enum). + pub const LEN: usize = #count; + + /// The entry for `key`. + pub fn get(&self, key: #enum_ident) -> &T { + match key { #(#get_arms),* } + } + + /// The mutable entry for `key`. + pub fn get_mut(&mut self, key: #enum_ident) -> &mut T { + match key { #(#get_mut_arms),* } + } + + /// Iterate `(variant, &value)` in declaration order. + pub fn iter(&self) -> impl Iterator { + [ #(#entries),* ].into_iter() + } + + /// Contiguous view in declaration order. + pub fn as_slice(&self) -> &[T] { + unsafe { ::core::slice::from_raw_parts((self as *const Self).cast::(), Self::LEN) } + } + + /// Mutable contiguous view in declaration order. + pub fn as_mut_slice(&mut self) -> &mut [T] { + unsafe { + ::core::slice::from_raw_parts_mut((self as *mut Self).cast::(), Self::LEN) + } + } + } + + impl ::core::ops::Index<#enum_ident> for #counts_ident { + type Output = T; + fn index(&self, key: #enum_ident) -> &T { + self.get(key) + } + } + + impl ::core::ops::IndexMut<#enum_ident> for #counts_ident { + fn index_mut(&mut self, key: #enum_ident) -> &mut T { + self.get_mut(key) + } + } + }) +} diff --git a/crates/cubecl-macros-internal/src/lib.rs b/crates/cubecl-macros-internal/src/lib.rs index 24774d9022..6a9c9b2ead 100644 --- a/crates/cubecl-macros-internal/src/lib.rs +++ b/crates/cubecl-macros-internal/src/lib.rs @@ -2,9 +2,11 @@ use generate::{ op_args::generate_op_args, operation::{generate_opcode, generate_operation}, }; +use enum_counts::enum_counts_impl; use proc_macro::TokenStream; use type_hash::type_hash_impl; +mod enum_counts; mod generate; mod parse; mod type_hash; @@ -72,3 +74,16 @@ pub fn derive_type_hash(input: proc_macro::TokenStream) -> proc_macro::TokenStre let input = syn::parse(input).unwrap(); type_hash_impl(input).into() } + +/// Generates variant reflection (`COUNT`, `VARIANTS`, `index`) for a fieldless enum, plus a flat +/// `#[repr(C)]` named-field companion `{Enum}Counts` (one field per variant) with `get`/`get_mut`, +/// `iter`, `as_slice`, `Index`/`IndexMut`, and `Pod`/`Zeroable` when `T` is. Adding a variant extends +/// both automatically. +#[proc_macro_derive(EnumCounts)] +pub fn derive_enum_counts(input: TokenStream) -> TokenStream { + let input = syn::parse(input).unwrap(); + match enum_counts_impl(input) { + Ok(tokens) => tokens.into(), + Err(e) => e.into_compile_error().into(), + } +} From c86965518fcb8a9e4af4ba984a9bc367c959d42a Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Fri, 3 Jul 2026 15:06:20 -0400 Subject: [PATCH 08/14] cleanup trackable Co-authored-by: SamuelBelanger --- crates/cubecl-core/src/codegen/profiler.rs | 207 ++++++++---- crates/cubecl-ir/Cargo.toml | 2 - crates/cubecl-ir/src/count.rs | 208 ++++++++++++ crates/cubecl-ir/src/counters.rs | 308 ------------------ crates/cubecl-ir/src/lib.rs | 4 +- .../src/generate/operation.rs | 31 ++ crates/cubecl-runtime/src/client.rs | 2 +- crates/cubecl-runtime/src/server/base.rs | 4 +- .../cubecl-wgpu/src/compiler/wgsl/compiler.rs | 72 +++- .../src/compiler/wgsl/instructions.rs | 2 - crates/cubecl-wgpu/src/compiler/wgsl/mod.rs | 2 - .../cubecl-wgpu/src/compiler/wgsl/profile.rs | 23 -- crates/cubecl/tests/chose.rs | 2 +- 13 files changed, 436 insertions(+), 431 deletions(-) create mode 100644 crates/cubecl-ir/src/count.rs delete mode 100644 crates/cubecl-ir/src/counters.rs delete mode 100644 crates/cubecl-wgpu/src/compiler/wgsl/profile.rs diff --git a/crates/cubecl-core/src/codegen/profiler.rs b/crates/cubecl-core/src/codegen/profiler.rs index 90fd2a4d5e..c55692e347 100644 --- a/crates/cubecl-core/src/codegen/profiler.rs +++ b/crates/cubecl-core/src/codegen/profiler.rs @@ -1,92 +1,159 @@ -use alloc::{fmt::Display, vec::Vec}; -use cubecl_ir::{Branch, CountInfo, Operation, OpsCounter, Value}; -use hashbrown::HashSet; - -/// A reusable profiler component that tracks used flop slots during backend compilation. -#[derive(Clone, Debug)] -pub struct CompilerProfiler { - flops_used: HashSet, - compiler: C, - counter: Option, -} +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}; -impl Default for CompilerProfiler { - fn default() -> Self { - Self { - flops_used: HashSet::default(), - compiler: C::default(), - counter: Option::default(), - } - } -} +type BufferId = Value; -pub trait CompilerProfilerInstructions: Default { - type Instruction; - type CounterVariable: Display; +const LOCAL_COUNTER_TYPE: Type = Type::Scalar(StorageType::Scalar(ElemType::UInt(UIntKind::U32))); - /// Returns an instruction that increments the value of the given slot by the given amount. - fn increment_instruction(&self, slot: u32, amount: u32) -> Self::Instruction; - /// Returns an instruction that prefixes the given slot, e.g. to initialize it to zero. - fn declare_instruction(&self, slot: u32) -> Self::Instruction; - /// Returns an instruction that suffixes the given slot, e.g. to aggregate its value. - fn flush_instruction(&self, slot: u32, counter: &Self::CounterVariable) -> Self::Instruction; +fn constant_value(amount: usize) -> Value { + Value::constant(ConstantValue::UInt(amount as u64), LOCAL_COUNTER_TYPE) } -impl CompilerProfiler { - pub fn new(compiler: C) -> Self { - Self { - flops_used: HashSet::default(), - compiler, - counter: None, - } - } +#[derive(Clone, Debug, Default, derive_new::new)] +pub struct CompilerProfiler { + global_counter: Option, + locals: HashMap, + current_index: usize, +} - pub fn set_counter(&mut self, counter: C::CounterVariable) { - self.counter = Some(counter) +impl CompilerProfiler { + pub fn set_counter(&mut self, counter: BufferId) { + self.global_counter = Some(counter) } - /// Intercepts an operation, tracking any used counters (flops, etc). - /// The backend provides a formatting closure to emit its specific increment instruction. pub fn profile_operation( &mut self, - operation: &Operation, + operation: &impl Trackable, out: Option<&Value>, - instructions: &mut Vec<::Instruction>, - ) { - if let Some(counter) = &self.counter { - if matches!(operation, Operation::Branch(Branch::Return)) { - for slot in &self.flops_used { - instructions.push(self.compiler.flush_instruction(*slot, counter)); - } - return; - } - - if let Some(CountInfo { slot, amount }) = OpsCounter::count(operation, out) { - self.flops_used.insert(slot); - instructions.push(self.compiler.increment_instruction(slot, amount)); + allocator: &Allocator, + ) -> Vec { + // 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, + )) } - /// Iterates over all used profiling slots and generates initialization (prefix) - /// and aggregation (suffix) instructions via the provided formatting closures. - pub fn profile( + pub fn profile(&self, allocator: &Allocator) -> (Vec, Vec) { + 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, - instructions: &mut Vec<::Instruction>, - ) { - if let Some(counter) = &self.counter { - let mut declares = Vec::new(); - for slot in &self.flops_used { - declares.push(self.compiler.declare_instruction(*slot)); - } + 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); - let mut old_instructions = core::mem::take(instructions); + [ + 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, + })), + ] + } - instructions.extend(declares); - instructions.append(&mut old_instructions); + 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, + ), + ] + } - for slot in &self.flops_used { - instructions.push(self.compiler.flush_instruction(*slot, counter)); + 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 } } } diff --git a/crates/cubecl-ir/Cargo.toml b/crates/cubecl-ir/Cargo.toml index f8861ac16e..d629d009c1 100644 --- a/crates/cubecl-ir/Cargo.toml +++ b/crates/cubecl-ir/Cargo.toml @@ -52,5 +52,3 @@ hashbrown = { workspace = true } portable-atomic = { workspace = true } paste = { workspace = true } - -strum = { version = "0.28.0", features = ["derive"] } diff --git a/crates/cubecl-ir/src/count.rs b/crates/cubecl-ir/src/count.rs new file mode 100644 index 0000000000..24d649e208 --- /dev/null +++ b/crates/cubecl-ir/src/count.rs @@ -0,0 +1,208 @@ +use hashbrown::HashMap; +use internment::Intern; + +use crate::{ + AddressSpace, Arithmetic, Branch, ElemType, Memory, Operation, OperationReflect, Type, + UIntKind, Value, +}; + +use alloc::{format, string::String, vec::Vec}; + +/// The unit a counted operation is measured in. +/// +/// Part of [`CountKey`], so op counts and byte counts never share a slot even if two ops happened +/// to share a name. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum CountUnit { + /// A number of scalar operations. + Ops, + /// A number of bytes moved. + Bytes, +} + +impl CountUnit { + /// Short suffix used when displaying counts (e.g. `"ops"`, `"bytes"`). + pub fn suffix(self) -> &'static str { + match self { + CountUnit::Ops => "ops", + CountUnit::Bytes => "bytes", + } + } +} + +/// The identity a count accumulates under: the operation name, the element type it acts on, and +/// the unit it is measured in. +/// +/// `name` is interned, so equality/hashing are pointer-based — a collision-free, hash-like +/// identity — while the readable name is kept for display. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct CountKey { + /// Interned operation name, e.g. `"Add"` or `"Load"`. + pub name: Intern, + /// The element type the operation acts on. + pub elem: ElemType, + /// What the operation is measured in. + pub unit: CountUnit, +} + +/// One operation's contribution to the counts: what to accumulate, and how much. +pub struct Counted { + /// The identity to accumulate under. + pub key: CountKey, + /// The amount to add: an op count for [`CountUnit::Ops`], a byte count for + /// [`CountUnit::Bytes`]. + pub amount: u32, +} + +/// How an operation participates in profiling. +/// +/// The counting layer only ever sees this trait — never the concrete operation enums or opcodes. +/// The current IR implements it on [`Operation`] and its sub-enums; a future op representation +/// (e.g. pliron) implements it per op type without touching anything downstream. +/// +/// [`count`](Trackable::count) is the single place that inspects the operation's result (`out`), +/// because in this IR the element type and vectorization live on the result value, not on the op +/// node. Structural queries like [`is_terminal`](Trackable::is_terminal) need no result and stay +/// separate, so callers can test them without a result in hand. +pub trait Trackable { + /// This operation's counting contribution given its result value, or `None` if it is not + /// counted. `out` is `None` for operations that produce no value. + fn count(&self, out: Option<&Value>) -> Option; + + /// Whether this operation ends a block and should trigger a flush of the accumulated counts. + fn is_terminal(&self) -> bool { + false + } +} + +impl Trackable for Operation { + fn count(&self, out: Option<&Value>) -> Option { + match self { + Operation::Arithmetic(op) => op.count(out), + Operation::Memory(op) => op.count(out), + _ => None, + } + } + + fn is_terminal(&self) -> bool { + match self { + Operation::Branch(op) => op.is_terminal(), + _ => false, + } + } +} + +impl Trackable for Arithmetic { + fn count(&self, out: Option<&Value>) -> Option { + // Arithmetic is billed against its result's type and vectorization. + let out = out?; + Some(Counted { + key: CountKey { + name: Intern::from(self.op_code().name()), + elem: out.elem_type(), + unit: CountUnit::Ops, + }, + amount: out.vector_size() as u32, + }) + } +} + +impl Trackable for Memory { + fn count(&self, out: Option<&Value>) -> Option { + // The value whose element type/width we bill, plus the address space(s) the op touches. + let (value, spaces): (&Value, [Option; 2]) = match self { + Memory::Index(_) => { + let out = out?; + (out, [Some(out.address_space()), None]) + } + Memory::Load(variable) => (variable, [Some(variable.address_space()), None]), + Memory::Store(op) => (&op.value, [Some(op.ptr.address_space()), None]), + Memory::CopyMemory(op) => ( + &op.source, + [ + Some(op.source.address_space()), + Some(op.target.address_space()), + ], + ), + }; + + // Local and shared traffic is free; only other address spaces count as moved bytes. + let billable = spaces + .iter() + .flatten() + .any(|space| !matches!(space, AddressSpace::Local | AddressSpace::Shared)); + if !billable { + return None; + } + + let elem = value.elem_type(); + Some(Counted { + key: CountKey { + name: Intern::from(self.op_code().name()), + elem, + unit: CountUnit::Bytes, + }, + amount: (value.vector_size() * elem.size()) as u32, + }) + } +} + +impl Trackable for Branch { + fn count(&self, _out: Option<&Value>) -> Option { + None + } + + fn is_terminal(&self) -> bool { + matches!(self, Branch::Return) + } +} + +pub fn counter_stored_type() -> Type { + Type::atomic(Type::scalar(ElemType::UInt(UIntKind::U32))) +} + +fn display_u32(value: u32) -> String { + format!("{value}") + .as_bytes() + .rchunks(3) + .rev() + .map(|c| core::str::from_utf8(c).unwrap()) + .collect::>() + .join(" ") +} + +pub fn fmt_counts( + counts: &HashMap, + f: &mut core::fmt::Formatter<'_>, +) -> core::fmt::Result { + let mut entries: Vec<(ElemType, CountUnit, &str, u32)> = counts + .iter() + .filter(|(_, c)| **c != 0) + .map(|(k, &c)| (k.elem, k.unit, &*k.name, c)) + .collect(); + entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)).then(a.2.cmp(b.2))); + + if entries.is_empty() { + return writeln!(f, " (none)"); + } + let mut cur_elem: Option = None; + let mut cur_unit: Option = None; + for (elem, unit, name, count) in entries { + if cur_elem != Some(elem) { + writeln!(f, "{elem}")?; + cur_elem = Some(elem); + cur_unit = None; + } + if cur_unit != Some(unit) { + writeln!(f, "└── {}", unit.suffix())?; + cur_unit = Some(unit); + } + writeln!( + f, + " └── {name}: {} {}", + display_u32(count), + unit.suffix() + )?; + } + Ok(()) +} diff --git a/crates/cubecl-ir/src/counters.rs b/crates/cubecl-ir/src/counters.rs deleted file mode 100644 index d121a72461..0000000000 --- a/crates/cubecl-ir/src/counters.rs +++ /dev/null @@ -1,308 +0,0 @@ -use alloc::format; -use alloc::string::String; -use alloc::string::ToString; -use alloc::vec::Vec; -use paste::paste; - -use crate::AddressSpace; -use crate::{ - ArithmeticOpCode, ElemType, FloatKind, IntKind, Memory, MemoryOpCode, Operation, - OperationReflect, Type, UIntKind, Value, -}; - -fn display_u32(value: u32) -> String { - let s = value.to_string(); - s.as_bytes() - .rchunks(3) - .rev() - .map(|chunk| std::str::from_utf8(chunk).unwrap()) - .collect::>() - .join(" ") -} - -#[derive(derive_new::new)] -pub struct CountInfo { - pub slot: u32, - pub amount: u32, -} - -pub struct OpsCounter; - -impl OpsCounter { - pub fn count(operation: &Operation, out: Option<&Value>) -> Option { - match operation { - Operation::Arithmetic(_) => { - let out = out?; - OpsCounts::offset(&out.elem_type(), operation) - .map(|index| CountInfo::new(index as u32, out.vector_size() as u32)) - } - Operation::Memory(memory) => { - let (value, target_address_spaces) = match memory { - Memory::Index(_) => { - let out = out?; - (out, std::vec![out.address_space()]) - } - Memory::Load(variable) => (variable, std::vec![variable.address_space()]), - Memory::Store(op) => (&op.value, std::vec![op.ptr.address_space()]), - Memory::CopyMemory(op) => ( - &op.source, - std::vec![op.source.address_space(), op.target.address_space()], - ), - }; - - if target_address_spaces.iter().any(|addr_space| { - !matches!(addr_space, AddressSpace::Local | AddressSpace::Shared) - }) { - OpsCounts::offset(&value.elem_type(), operation).map(|index| { - CountInfo::new( - index as u32, - (value.vector_size() * value.elem_type().size()) as u32, - ) - }) - } else { - None - } - } - _ => None, - } - } -} - -macro_rules! ops_count { - ({$($op:ident),*}, $name:ident, $unit:literal) => { - paste! { - #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] - #[repr(C)] - pub struct [< $name OpsCount >] { - $(pub [< $op:lower >]: u32),* - } - - impl [< $name OpsCount >] { - pub const LEN: usize = size_of::() / size_of::(); - - pub fn offset(op_code: &[< $name OpCode >]) -> Option { - enum OpIndex { - $($op),* - } - - Some(match op_code { - $([< $name OpCode >]::$op => OpIndex::$op as usize),* - }) - } - - pub fn is_empty(&self) -> bool { - true $(&& self.[< $op:lower >] == 0)* - } - - pub const KIND: &'static str = stringify!($name); - pub const UNIT: &'static str = $unit; - - pub fn entries(&self) -> Vec<(&'static str, u32)> { - let mut entries = Vec::new(); - $( - if self.[< $op:lower >] != 0 { - entries.push((stringify!($op), self.[< $op:lower >])); - } - )* - entries - } - } - - impl core::fmt::Display for [< $name OpsCount >] { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - let mut first = true; - $( - if self.[< $op:lower >] != 0 { - if !first { - write!(f, ", ")?; - } - write!(f, "{}: {}{}", stringify!($op), display_u32(self.[< $op:lower >]), $unit)?; - first = false; - } - )* - Ok(()) - } - } - } - } -} - -macro_rules! ops_counts { - ( - ariths : [$($arith:ident),*], - memories : [$($memory:ident),*], - elems: [$($elems:ty),*] - ) => { - paste! { - ops_count!({$($arith),*}, Arithmetic, "ops"); - ops_count!({$($memory),*}, Memory, "bytes"); - - #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] - #[repr(C)] - pub struct OpsCount { - pub arith: ArithmeticOpsCount, - pub memory: MemoryOpsCount, - } - - impl OpsCount { - pub const LEN: usize = size_of::() / size_of::(); - - pub fn offset(op: &Operation) -> Option { - match op { - Operation::Arithmetic(arith) => { - let base = core::mem::offset_of!(OpsCount, arith) / size_of::(); - ArithmeticOpsCount::offset(&arith.op_code()).map(|o| base + o) - } - Operation::Memory(memory) => { - let base = core::mem::offset_of!(OpsCount, memory) / size_of::(); - MemoryOpsCount::offset(&memory.op_code()).map(|o| base + o) - } - _ => None, - } - } - - pub fn is_empty(&self) -> bool { - self.arith.is_empty() && self.memory.is_empty() - } - - pub fn entries(&self) -> Vec<(&'static str, &'static str, &'static str, u32)> { - let mut entries = Vec::new(); - for (op, count) in self.arith.entries() { - entries.push((ArithmeticOpsCount::KIND, ArithmeticOpsCount::UNIT, op, count)); - } - for (op, count) in self.memory.entries() { - entries.push((MemoryOpsCount::KIND, MemoryOpsCount::UNIT, op, count)); - } - entries - } - } - - impl core::fmt::Display for OpsCount { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - write!(f, "{}", self.arith)?; - write!(f, ", {}", self.memory)?; - Ok(()) - } - } - - #[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Zeroable, bytemuck::Pod)] - #[repr(C)] - pub struct OpsCounts { - $(pub [< $elems:lower >]: OpsCount),* - } - - impl core::fmt::Display for OpsCounts { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - let mut entries = Vec::new(); - $( - for (kind, unit, op, count) in self.[< $elems:lower >].entries() { - entries.push(( - stringify!($elems), - kind, - op, - display_u32(count), - unit - )); - } - )* - - if entries.is_empty() { - return writeln!(f, " (none)"); - } - - let max_op_len = entries.iter().map(|e| e.2.len()).max().unwrap_or(0); - let max_count_len = entries.iter().map(|e| e.3.len()).max().unwrap_or(0); - - let mut current_elem = ""; - let mut current_kind = ""; - let mut is_last_kind = false; - - for i in 0..entries.len() { - let (elem, kind, op, ref count_str, unit) = entries[i]; - - if elem != current_elem { - writeln!(f, "{}", elem)?; - current_elem = elem; - current_kind = ""; - } - - if kind != current_kind { - is_last_kind = !entries[i + 1..] - .iter() - .any(|e| e.0 == elem && e.1 != kind); - - let kind_prefix = if is_last_kind { "└── " } else { "├── " }; - writeln!(f, "{}{}", kind_prefix, kind)?; - current_kind = kind; - } - - let is_last_op = i + 1 == entries.len() - || entries[i + 1].0 != elem - || entries[i + 1].1 != kind; - - let kind_indent = if is_last_kind { " " } else { "│ " }; - let op_prefix = if is_last_op { "└── " } else { "├── " }; - - writeln!( - f, - "{}{}{:count_width$} {}", - kind_indent, - op_prefix, - format!("{}:", op), - count_str, - unit, - width = max_op_len + 1, - count_width = max_count_len - )?; - } - - Ok(()) - } - } - } - } -} - -ops_counts!( - ariths: [Add, SaturatingAdd, Fma, Sub, SaturatingSub, Mul, Div, Abs, Exp, Log, Log1p, Expm1, Cos, Sin, Tan, Tanh, Sinh, Cosh, ArcCos, ArcSin, ArcTan, ArcSinh, ArcCosh, ArcTanh, Degrees, Radians, ArcTan2, Powf, Powi, Hypot, Rhypot, Sqrt, InverseSqrt, Round, Floor, Ceil, Trunc, Erf, Recip, Clamp, Neg, Max, Min, Rem, ModFloor, Magnitude, Normalize, Dot, MulHi, VectorSum], - memories: [Index, Load, Store, CopyMemory], - elems: [Bool, E2M1, E2M3, E3M2, E4M3, E5M2, UE8M0, F16, BF16, Flex32, F32, TF32, F64, Int8, Int16, Int32, Int64, Uint8, Uint16, Uint32, Uint64] -); - -impl OpsCounts { - pub const LEN: usize = size_of::() / size_of::(); - - pub fn stored_type() -> Type { - Type::atomic(Type::scalar(ElemType::UInt(UIntKind::U32))) - } - - pub fn elem_index(elem: &ElemType) -> usize { - match elem { - ElemType::Bool => 0, - ElemType::Float(FloatKind::E2M1) => 1, - ElemType::Float(FloatKind::E2M3) => 2, - ElemType::Float(FloatKind::E3M2) => 3, - ElemType::Float(FloatKind::E4M3) => 4, - ElemType::Float(FloatKind::E5M2) => 5, - ElemType::Float(FloatKind::UE8M0) => 6, - ElemType::Float(FloatKind::F16) => 7, - ElemType::Float(FloatKind::BF16) => 8, - ElemType::Float(FloatKind::Flex32) => 9, - ElemType::Float(FloatKind::F32) => 10, - ElemType::Float(FloatKind::TF32) => 11, - ElemType::Float(FloatKind::F64) => 12, - ElemType::Int(IntKind::I8) => 13, - ElemType::Int(IntKind::I16) => 14, - ElemType::Int(IntKind::I32) => 15, - ElemType::Int(IntKind::I64) => 16, - ElemType::UInt(UIntKind::U8) => 17, - ElemType::UInt(UIntKind::U16) => 18, - ElemType::UInt(UIntKind::U32) => 19, - ElemType::UInt(UIntKind::U64) => 20, - } - } - - pub fn offset(elem: &ElemType, op: &Operation) -> Option { - OpsCount::offset(op).map(|within| Self::elem_index(elem) * OpsCount::LEN + within) - } -} diff --git a/crates/cubecl-ir/src/lib.rs b/crates/cubecl-ir/src/lib.rs index f50cd3df94..3d0341276d 100644 --- a/crates/cubecl-ir/src/lib.rs +++ b/crates/cubecl-ir/src/lib.rs @@ -16,7 +16,7 @@ mod bitwise; mod branch; mod cmma; mod comparison; -mod counters; +mod count; mod marker; mod memory; mod metadata; @@ -47,7 +47,7 @@ pub use bitwise::*; pub use branch::*; pub use cmma::*; pub use comparison::*; -pub use counters::*; +pub use count::*; pub use marker::*; pub use memory::*; pub use metadata::*; diff --git a/crates/cubecl-macros-internal/src/generate/operation.rs b/crates/cubecl-macros-internal/src/generate/operation.rs index 84813e74cd..cb57a5b16d 100644 --- a/crates/cubecl-macros-internal/src/generate/operation.rs +++ b/crates/cubecl-macros-internal/src/generate/operation.rs @@ -347,16 +347,45 @@ impl Operation { } } } + + /// Generates a `name()` method on the opcode enum returning a stable, human-readable + /// identity for each variant (e.g. `"Add"`), delegating through nested opcodes. + fn generate_opcode_name(&self) -> TokenStream { + let name = &self.opcode_name; + let variants = self.variants(); + let match_variants = variants.iter().map(|variant| { + let ident = &variant.ident; + if variant.nested.is_present() { + quote![Self::#ident(child) => child.name()] + } else { + quote![Self::#ident => stringify!(#ident)] + } + }); + quote! { + impl #name { + /// Stable, human-readable name of this opcode variant (e.g. `"Add"`). + /// + /// Nested opcodes delegate to their child's name. + pub fn name(&self) -> &'static str { + match self { + #(#match_variants),* + } + } + } + } + } } pub fn generate_operation(input: DeriveInput) -> syn::Result { let operation = Operation::from_derive_input(&input)?; let opcode = operation.generate_opcode(); + let opcode_name = operation.generate_opcode_name(); let operation_impl = operation.generate_operation_impl(); Ok(quote! { #opcode + #opcode_name #operation_impl }) } @@ -377,10 +406,12 @@ pub fn generate_opcode(input: DeriveInput) -> syn::Result { let generics = &operation.generics; let opcode_name = &operation.opcode_name; let opcode = operation.generate_opcode(); + let opcode_name_impl = operation.generate_opcode_name(); let match_opcode = operation.generate_opcode_impl(); Ok(quote! { #opcode + #opcode_name_impl impl #generics #name { fn __match_opcode(&self) -> #opcode_name { diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index e21b817d33..955139ec9d 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -26,7 +26,7 @@ use cubecl_common::{ profile::ProfileDuration, stub::RwLock, }; -use cubecl_ir::{DeviceProperties, ElemType, OpsCounts, VectorSize, features::Features}; +use cubecl_ir::{CountUnit, DeviceProperties, ElemType, VectorSize, features::Features}; use cubecl_zspace::Shape; #[allow(unused)] diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index a7be990f26..06a53b1bd9 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -32,7 +32,7 @@ use cubecl_common::{ stream_id::StreamId, stub::RwLock, }; -use cubecl_ir::{DeviceProperties, ElemType, OpsCounts, StorageType}; +use cubecl_ir::{CountKey, DeviceProperties, ElemType, OpsCounts, StorageType}; use cubecl_zspace::{Shape, Strides, metadata::Metadata}; use hashbrown::{HashMap, HashSet}; use itertools::Itertools; @@ -83,7 +83,7 @@ impl core::fmt::Debug for ProfileError { pub struct FlopRecord { /// Most recent per-operation counts, keyed by op name (e.g. `"Add"`). Only operations that /// actually executed appear. Sorted for deterministic, serialization-friendly output. - pub last: OpsCounts, + pub last: HashMap, /// Number of launches recorded for the kernel. pub samples: u32, } diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs index d251a757b3..54e6e77608 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs @@ -1,7 +1,6 @@ use super::Item; use super::Subgroup; use super::shader::ComputeShader; -use crate::compiler::wgsl::WgslCompilerProfiler; use crate::compiler::wgsl::{self, SharedValue}; use cubecl_common::backtrace::BackTrace; @@ -52,7 +51,7 @@ pub struct WgslCompiler { strategy: ExecutionMode, subgroup_instructions_used: bool, f16_used: bool, - profiler: CompilerProfiler, + profiler: CompilerProfiler, } impl core::fmt::Debug for WgslCompiler { @@ -298,21 +297,42 @@ impl WgslCompiler { } fn setup_profiler(&mut self, scope: &cube::Scope) { - if scope.profile.enabled { - let counter = self.compile_value( - scope - .profile - .counters_buffer - .expect("Profiling counters buffer should be initialized"), - ); - - self.profiler.set_counter(counter); + if !scope.profile.enabled { + return; } + + let counter = scope + .profile + .counters_buffer + .expect("Profiling counters buffer should be initialized"); + + self.profiler.set_counter(counter); } - fn profile(&self, scope: &cube::Scope, mut instructions: &mut Vec) { - if scope.profile.enabled { - self.profiler.profile(&mut instructions); + fn profile(&mut self, scope: &cube::Scope, instructions: &mut Vec) { + if !scope.profile.enabled { + return; + } + + let (declare_instructions, flush_instructions) = + self.profiler.profile(&scope.state().allocator); + + let mut declare_wgsl_instructions = Vec::new(); + + for instruction in declare_instructions { + self.compile_operation( + &mut declare_wgsl_instructions, + instruction.operation, + instruction.out, + scope, + ); + } + + let mut old_instructions = core::mem::replace(instructions, declare_wgsl_instructions); + instructions.append(&mut old_instructions); + + for instruction in flush_instructions { + self.compile_operation(instructions, instruction.operation, instruction.out, scope); } } @@ -325,12 +345,12 @@ impl WgslCompiler { processing .instructions .into_iter() - .for_each(|op| self.compile_operation(&mut instructions, op.operation, op.out, scope)); + .for_each(|op| self.process_operation(&mut instructions, op.operation, op.out, scope)); instructions } - fn compile_operation( + fn process_operation( &mut self, instructions: &mut Vec, operation: cube::Operation, @@ -338,9 +358,25 @@ impl WgslCompiler { scope: &cube::Scope, ) { if scope.profile.enabled { - self.profiler - .profile_operation(&operation, out.as_ref(), instructions); + let new_instructions = + self.profiler + .profile_operation(&operation, out.as_ref(), &scope.state().allocator); + + for instruction in new_instructions { + self.compile_operation(instructions, instruction.operation, instruction.out, scope); + } } + + self.compile_operation(instructions, operation, out, scope); + } + + fn compile_operation( + &mut self, + instructions: &mut Vec, + operation: cube::Operation, + out: Option, + scope: &cube::Scope, + ) { match operation { cube::Operation::Copy(value) => instructions.push(wgsl::Instruction::Assign { input: self.compile_value(value), diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs b/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs index de9885fc47..adc80ae515 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs @@ -12,7 +12,6 @@ use std::fmt::Display; #[derive(Debug, Clone)] #[allow(dead_code)] // Some variants might not be used with different flags pub enum Instruction { - Custom(String), DeclareVariable { val: Value, value_ty: Item, @@ -483,7 +482,6 @@ pub enum Instruction { impl Display for Instruction { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Instruction::Custom(s) => writeln!(f, "{s}"), Instruction::DeclareVariable { val, value_ty } => { writeln!(f, "var {val}_store: {value_ty};")?; writeln!(f, "let {val} = &{val}_store;") diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs b/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs index ec10c6e5cd..d549832b14 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/mod.rs @@ -3,7 +3,6 @@ mod body; mod compiler; mod extension; mod instructions; -mod profile; pub(crate) mod shader; mod subgroup; @@ -12,6 +11,5 @@ pub(crate) use body::*; pub use compiler::*; pub(crate) use extension::*; pub(crate) use instructions::*; -pub(crate) use profile::*; pub(crate) use shader::*; pub(crate) use subgroup::*; diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/profile.rs b/crates/cubecl-wgpu/src/compiler/wgsl/profile.rs deleted file mode 100644 index 5f4676ab1d..0000000000 --- a/crates/cubecl-wgpu/src/compiler/wgsl/profile.rs +++ /dev/null @@ -1,23 +0,0 @@ -use cubecl_core::CompilerProfilerInstructions; - -use crate::compiler::wgsl; - -#[derive(Clone, Default)] -pub struct WgslCompilerProfiler; - -impl CompilerProfilerInstructions for WgslCompilerProfiler { - type Instruction = wgsl::Instruction; - type CounterVariable = wgsl::Value; - - fn increment_instruction(&self, slot: u32, amount: u32) -> Self::Instruction { - wgsl::Instruction::Custom(format!("flop_{slot} += {amount}u;")) - } - - fn declare_instruction(&self, slot: u32) -> Self::Instruction { - wgsl::Instruction::Custom(format!("var flop_{slot} = 0u;")) - } - - fn flush_instruction(&self, slot: u32, counter: &Self::CounterVariable) -> Self::Instruction { - wgsl::Instruction::Custom(format!("atomicAdd(&{counter}[{slot}], flop_{slot});")) - } -} diff --git a/crates/cubecl/tests/chose.rs b/crates/cubecl/tests/chose.rs index bbbfb10fb2..837dcabe69 100644 --- a/crates/cubecl/tests/chose.rs +++ b/crates/cubecl/tests/chose.rs @@ -5,7 +5,7 @@ fn test_chose_cube(lhs: f32, rhs: f32, out: &mut [f32]) { if UNIT_POS == 0 { out[0] = lhs + rhs + lhs; out[0] = lhs + rhs; - terminate!(); + // terminate!(); out[0] = lhs + rhs; } } From 3825c01d3a1f97222a081aeb85c63e80c47efc1f Mon Sep 17 00:00:00 2001 From: Thierry Cantin-Demers Date: Fri, 3 Jul 2026 17:24:41 -0400 Subject: [PATCH 09/14] global buffer post alloc and remove opscounts --- crates/cubecl-core/src/codegen/profiler.rs | 13 ++ crates/cubecl-core/src/compute/builder.rs | 11 +- crates/cubecl-core/src/compute/launcher.rs | 46 +---- crates/cubecl-ir/src/count.rs | 179 ++++++++++++++++-- crates/cubecl-runtime/src/client.rs | 18 +- crates/cubecl-runtime/src/server/base.rs | 5 +- crates/cubecl-wgpu/src/backend/wgsl.rs | 4 + .../cubecl-wgpu/src/compiler/wgsl/compiler.rs | 30 +++ .../cubecl-wgpu/src/compiler/wgsl/shader.rs | 26 ++- crates/cubecl-wgpu/src/compute/mem_manager.rs | 19 ++ crates/cubecl-wgpu/src/compute/schedule.rs | 6 + crates/cubecl-wgpu/src/compute/server.rs | 101 +++++++++- 12 files changed, 367 insertions(+), 91 deletions(-) diff --git a/crates/cubecl-core/src/codegen/profiler.rs b/crates/cubecl-core/src/codegen/profiler.rs index c55692e347..8b4528de4b 100644 --- a/crates/cubecl-core/src/codegen/profiler.rs +++ b/crates/cubecl-core/src/codegen/profiler.rs @@ -157,4 +157,17 @@ impl CompilerProfiler { } } } + + /// 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 { + let mut entries: Vec<(usize, CountKey)> = self + .locals + .iter() + .map(|(key, (_, slot))| (*slot, key.clone())) + .collect(); + entries.sort_by_key(|(slot, _)| *slot); + entries.into_iter().map(|(_, key)| key).collect() + } } diff --git a/crates/cubecl-core/src/compute/builder.rs b/crates/cubecl-core/src/compute/builder.rs index 66d71343f0..a6c1f9504c 100644 --- a/crates/cubecl-core/src/compute/builder.rs +++ b/crates/cubecl-core/src/compute/builder.rs @@ -8,7 +8,9 @@ use crate::{ prelude::KernelDefinition, }; use alloc::collections::BTreeMap; -use cubecl_ir::{DeviceProperties, OpsCounts, Scope, StorageType, TargetProperties, Value}; +use cubecl_ir::{ + DeviceProperties, Scope, StorageType, TargetProperties, Value, counter_stored_type, +}; use cubecl_runtime::config::{ CubeClRuntimeConfig, RuntimeConfig, compilation::CompilationLogLevel, }; @@ -92,7 +94,12 @@ impl KernelBuilder { /// Build the [kernel definition](KernelDefinition). pub fn build(mut self, settings: KernelSettings) -> KernelDefinition { if self.profile.enabled { - let buffer = self.buffer(OpsCounts::stored_type()); + // 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); } diff --git a/crates/cubecl-core/src/compute/launcher.rs b/crates/cubecl-core/src/compute/launcher.rs index 8fa3d471d4..013392960e 100644 --- a/crates/cubecl-core/src/compute/launcher.rs +++ b/crates/cubecl-core/src/compute/launcher.rs @@ -1,5 +1,4 @@ use alloc::{boxed::Box, vec::Vec}; -use bytemuck::Zeroable; use core::marker::PhantomData; use crate::Runtime; @@ -7,10 +6,8 @@ use crate::prelude::{BufferArg, TensorArg, TensorMapArg, TensorMapKind}; use crate::{InfoBuilder, KernelSettings, ScalarArgType}; #[cfg(feature = "std")] use core::cell::RefCell; -use cubecl_ir::{AddressType, OpsCounts, Scope, StorageType, Type}; -use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig}; -use cubecl_runtime::id::KernelId; -use cubecl_runtime::server::{Binding, CubeCount, Handle, TensorMapBinding}; +use cubecl_ir::{AddressType, Scope, StorageType, Type}; +use cubecl_runtime::server::{Binding, CubeCount, TensorMapBinding}; use cubecl_runtime::{ client::ComputeClient, kernel::{CubeKernel, KernelTask}, @@ -68,46 +65,18 @@ impl KernelLauncher { self.with_info(|info| info.scalars.push_raw(bytes, dtype)); } - /// Injects the profiling counters buffer into the bindings before execution. - fn inject_profiling(&mut self, client: &ComputeClient) -> Option { - if !CubeClRuntimeConfig::get().profiling.hardware_metrics { - return None; - } - - let handle = client.create_from_slice(bytemuck::bytes_of(&OpsCounts::zeroed())); - let arg = unsafe { BufferArg::from_raw_parts(handle.clone(), OpsCounts::LEN) }; - self.register_buffer(arg, OpsCounts::stored_type()); - Some(handle) - } - - /// Read back the per-op FLOP counters and record them on the client, keyed by kernel id. The - /// dense slot buffer is turned into a sparse, name-keyed map here so storage stays legible. - fn report_profiling(client: &ComputeClient, id: KernelId, handle: Handle) { - let bytes = client.read_one_unchecked(handle); - let ops_counts: OpsCounts = bytemuck::pod_read_unaligned(&bytes); - - client.record_flop_count(id, ops_counts); - } - /// Launch the kernel. #[track_caller] pub fn launch( - mut self, + self, cube_count: CubeCount, kernel: K, client: &ComputeClient, ) { - let profile = self.inject_profiling(client); - - let kernel_id = kernel.id(); let bindings = self.into_bindings(); let kernel = Box::new(KernelTask::::new(kernel)); client.launch(kernel, cube_count, bindings); - - if let Some(handle) = profile { - Self::report_profiling(client, kernel_id, handle); - } } /// Launch the kernel without check bounds. @@ -120,24 +89,17 @@ impl KernelLauncher { /// other unpredictable behaviour. #[track_caller] pub unsafe fn launch_unchecked( - mut self, + self, cube_count: CubeCount, kernel: K, client: &ComputeClient, ) { - let profile = self.inject_profiling(client); - - let kernel_id = kernel.id(); unsafe { let bindings = self.into_bindings(); let kernel = Box::new(KernelTask::::new(kernel)); client.launch_unchecked(kernel, cube_count, bindings) } - - if let Some(handle) = profile { - Self::report_profiling(client, kernel_id, handle); - } } /// We need to create the bindings in the same order they are defined in the compilation step. diff --git a/crates/cubecl-ir/src/count.rs b/crates/cubecl-ir/src/count.rs index 24d649e208..b946ead09f 100644 --- a/crates/cubecl-ir/src/count.rs +++ b/crates/cubecl-ir/src/count.rs @@ -30,6 +30,20 @@ impl CountUnit { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct OpPath(Intern<[Intern]>); + +impl OpPath { + pub fn new(segments: &[&str]) -> Self { + let segments: Vec> = segments.iter().copied().map(Intern::from).collect(); + Self(Intern::from(segments.as_slice())) + } + /// The segments, outermost first. + pub fn segments(&self) -> &[Intern] { + &self.0 + } +} + /// The identity a count accumulates under: the operation name, the element type it acts on, and /// the unit it is measured in. /// @@ -38,7 +52,7 @@ impl CountUnit { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct CountKey { /// Interned operation name, e.g. `"Add"` or `"Load"`. - pub name: Intern, + pub name: OpPath, /// The element type the operation acts on. pub elem: ElemType, /// What the operation is measured in. @@ -98,7 +112,7 @@ impl Trackable for Arithmetic { let out = out?; Some(Counted { key: CountKey { - name: Intern::from(self.op_code().name()), + name: OpPath::new(&[family::(), self.op_code().name()]), elem: out.elem_type(), unit: CountUnit::Ops, }, @@ -138,7 +152,7 @@ impl Trackable for Memory { let elem = value.elem_type(); Some(Counted { key: CountKey { - name: Intern::from(self.op_code().name()), + name: OpPath::new(&[family::(), self.op_code().name()]), elem, unit: CountUnit::Bytes, }, @@ -171,38 +185,163 @@ fn display_u32(value: u32) -> String { .join(" ") } +fn family() -> &'static str { + let name = core::any::type_name::(); + name.rsplit("::").next().unwrap_or(name) +} + pub fn fmt_counts( counts: &HashMap, f: &mut core::fmt::Formatter<'_>, ) -> core::fmt::Result { - let mut entries: Vec<(ElemType, CountUnit, &str, u32)> = counts + let mut entries: Vec<(ElemType, CountUnit, &[Intern], u32)> = counts .iter() .filter(|(_, c)| **c != 0) - .map(|(k, &c)| (k.elem, k.unit, &*k.name, c)) + .map(|(k, &c)| (k.elem, k.unit, k.name.segments(), c)) .collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)).then(a.2.cmp(b.2))); + + // Sort by Elem, then strictly by Path (names), then Unit + // Grouping by path ensures our tree algorithm correctly renders branches. + entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2)).then(a.1.cmp(&b.1))); if entries.is_empty() { return writeln!(f, " (none)"); } + + // 1. Calculate max lengths for vertical column alignment + let mut max_op_len = 0; + let mut max_count_len = 0; + let mut formatted_counts = Vec::with_capacity(entries.len()); + + for entry in &entries { + let names = entry.2; + let count = entry.3; + + // The operation name is the leaf node (last item in the path) + let op_len = names.last().map(|s| s.len()).unwrap_or(6); // 6 is "(root)".len() + if op_len > max_op_len { + max_op_len = op_len; + } + + let count_str = display_u32(count); + if count_str.len() > max_count_len { + max_count_len = count_str.len(); + } + formatted_counts.push(count_str); + } + let mut cur_elem: Option = None; - let mut cur_unit: Option = None; - for (elem, unit, name, count) in entries { - if cur_elem != Some(elem) { + let mut is_last_at_depth: Vec = Vec::new(); // Tracks indentation levels + + // 2. Render the tree + for (entry_index, entry) in entries.iter().enumerate() { + let elem = &entry.0; + let unit = &entry.1; + let names = entry.2; + + // Print root elements without indentation + if cur_elem.as_ref() != Some(elem) { writeln!(f, "{elem}")?; - cur_elem = Some(elem); - cur_unit = None; + cur_elem = Some(elem.clone()); + is_last_at_depth.clear(); } - if cur_unit != Some(unit) { - writeln!(f, "└── {}", unit.suffix())?; - cur_unit = Some(unit); + + let mut common_len = 0; + if entry_index > 0 && entries[entry_index - 1].0 == *elem { + let prev_names = entries[entry_index - 1].2; + for (p, n) in prev_names.iter().zip(names.iter()) { + if p == n { + common_len += 1; + } else { + break; + } + } + } + + // If the path is an exact duplicate of the previous entry (e.g. same path, different unit) + // Step back one depth level to ensure the leaf node prints again. + if common_len == names.len() && !names.is_empty() { + common_len -= 1; + } + + let count_str = &formatted_counts[entry_index]; + + if names.is_empty() { + writeln!( + f, + "└── {:count_width$} {}", + "(root):", + count_str, + unit.suffix(), + width = max_op_len + 1, + count_width = max_count_len + )?; + continue; + } + + for depth in common_len..names.len() { + // Print background indentation structure mapped from parent relationships + for p in 0..depth { + if *is_last_at_depth.get(p).unwrap_or(&true) { + write!(f, " ")?; + } else { + write!(f, "│ ")?; + } + } + + let parent_path = &names[0..depth]; + let mut is_last = true; + + // Look ahead to evaluate if the current node has remaining siblings + for next_entry in &entries[entry_index + 1..] { + if next_entry.0 != *elem { + break; + } + let next_names = next_entry.2; + + if next_names.starts_with(parent_path) { + if next_names.len() > depth && next_names[depth] != names[depth] { + is_last = false; // Found a sibling + break; + } + if depth == names.len() - 1 && next_names == names { + is_last = false; // Identical leaf node (usually from duplicate/multiple units) + break; + } + } else { + break; // Out of the shared parent branch + } + } + + // Record this depth's 'last' status for its children to use + if depth >= is_last_at_depth.len() { + is_last_at_depth.push(is_last); + } else { + is_last_at_depth[depth] = is_last; + } + is_last_at_depth.truncate(depth + 1); + + let prefix = if is_last { "└── " } else { "├── " }; + + if depth == names.len() - 1 { + // Leaf Node: Append the colon and format with column alignment + let op_with_colon = format!("{}:", names[depth]); + writeln!( + f, + "{}{:count_width$} {}", + prefix, + op_with_colon, + count_str, + unit.suffix(), + width = max_op_len + 1, + count_width = max_count_len + )?; + } else { + // Intermediate Parent Node + writeln!(f, "{}{}", prefix, names[depth])?; + } } - writeln!( - f, - " └── {name}: {} {}", - display_u32(count), - unit.suffix() - )?; } + Ok(()) } diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index 955139ec9d..5b24a1d483 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -26,7 +26,7 @@ use cubecl_common::{ profile::ProfileDuration, stub::RwLock, }; -use cubecl_ir::{CountUnit, DeviceProperties, ElemType, VectorSize, features::Features}; +use cubecl_ir::{DeviceProperties, ElemType, VectorSize, features::Features}; use cubecl_zspace::Shape; #[allow(unused)] @@ -733,22 +733,6 @@ impl ComputeClient { handle } - /// Record per-operation FLOP counts for the given kernel, stored per-device in - /// [`ServerUtilities::metrics`]. Used by the launcher when `hardware_metrics` profiling is on. - pub fn record_flop_count(&self, id: KernelId, counts: OpsCounts) { - let mut metrics = self.utilities.metrics.write().unwrap(); - metrics - .entry(id) - .and_modify(|record| { - record.last = counts; - record.samples += 1; - }) - .or_insert(FlopRecord { - last: counts, - samples: 1, - }); - } - /// Returns a reference to the [`RwLock`] containing the recorded FLOP metrics for all kernels. pub fn metrics(&self) -> &RwLock> { &self.utilities.metrics diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index 06a53b1bd9..71c8e14afa 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -32,7 +32,7 @@ use cubecl_common::{ stream_id::StreamId, stub::RwLock, }; -use cubecl_ir::{CountKey, DeviceProperties, ElemType, OpsCounts, StorageType}; +use cubecl_ir::{CountKey, DeviceProperties, ElemType, StorageType, fmt_counts}; use cubecl_zspace::{Shape, Strides, metadata::Metadata}; use hashbrown::{HashMap, HashSet}; use itertools::Itertools; @@ -93,8 +93,7 @@ impl std::fmt::Display for FlopRecord { writeln!(f, "FlopRecord:")?; writeln!(f, " samples: {}", self.samples)?; writeln!(f, " last:")?; - write!(f, "{}", self.last)?; - Ok(()) + fmt_counts(&self.last, f) } } diff --git a/crates/cubecl-wgpu/src/backend/wgsl.rs b/crates/cubecl-wgpu/src/backend/wgsl.rs index 54e50d665e..b4e0b7cc15 100644 --- a/crates/cubecl-wgpu/src/backend/wgsl.rs +++ b/crates/cubecl-wgpu/src/backend/wgsl.rs @@ -28,6 +28,10 @@ pub fn bindings( if !args.info.data.is_empty() { bindings.push(Visibility::Read); } + // The profiling counters buffer is bound after `info` and written via atomics. + if repr.profile_counter.is_some() { + bindings.push(Visibility::ReadWrite); + } (bindings, 0) } diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs index 54e6e77608..25a0b5a2f4 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs @@ -131,6 +131,18 @@ impl WgslCompiler { self.buffer_vis .resize(value.num_global_buffers(), Visibility::Read); + // The counters buffer is a global (accessed via atomics) but lives outside `value.buffers`, + // so extend `buffer_vis` to cover its id with read-write visibility. + if let Some(counter) = value.body.profile.counters_buffer + && let cube::AddressSpace::Global(id) = counter.address_space() + { + let id = id as usize; + if id >= self.buffer_vis.len() { + self.buffer_vis.resize(id + 1, Visibility::Read); + } + self.buffer_vis[id] = Visibility::ReadWrite; + } + self.setup_profiler(&value.body); let mut instructions = self.compile_scope(&value.body); self.profile(&value.body, &mut instructions); @@ -144,6 +156,22 @@ impl WgslCompiler { address_type, }; + // The counters buffer is a special trailing binding (not in `value.buffers`), with the + // decode map alongside it so the server can turn the read-back slots into named counts. + let profile_counter = value + .body + .profile + .counters_buffer + .map(|counter| wgsl::KernelArg { + id: match counter.address_space() { + cube::AddressSpace::Global(id) => id, + _ => 0, + }, + visibility: Visibility::ReadWrite, + value: self.compile_value(counter), + }); + let profile_map = self.profiler.profile_map(); + Ok(wgsl::ComputeShader { address_type, buffers: value @@ -185,6 +213,8 @@ impl WgslCompiler { workgroup_size_no_axis: self.workgroup_size_no_axis, subgroup_instructions_used: self.subgroup_instructions_used, f16_used: self.f16_used, + profile_counter, + profile_map, kernel_name: value.options.kernel_name, }) } diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/shader.rs b/crates/cubecl-wgpu/src/compiler/wgsl/shader.rs index 83d2aaf7a9..f810e61df5 100644 --- a/crates/cubecl-wgpu/src/compiler/wgsl/shader.rs +++ b/crates/cubecl-wgpu/src/compiler/wgsl/shader.rs @@ -1,7 +1,7 @@ use crate::compiler::wgsl::Value; use super::{Body, Elem, Extension, Item}; -use cubecl_core::{CubeDim, Info, ir::Id, prelude::Visibility}; +use cubecl_core::{CubeDim, Info, ir::CountKey, ir::Id, prelude::Visibility}; use std::fmt::Display; #[derive(Debug, PartialEq, Eq, Clone)] @@ -53,6 +53,12 @@ pub struct ComputeShader { pub kernel_name: String, pub subgroup_instructions_used: bool, pub f16_used: bool, + /// Profiling counters buffer, bound after every normal buffer and `info` when + /// `hardware_metrics` profiling is on. Kept out of `buffers` so it carries no metadata. + pub profile_counter: Option, + /// Decodes the counter buffer slots back to `(op, elem, unit)` after readback. Its length is + /// the counter buffer size. Empty when profiling is off. + pub profile_map: Vec, } impl ComputeShader { @@ -105,6 +111,20 @@ var<{location}, {visibility}> info: info_st; )?; } + // The profiling counters buffer is bound after every normal buffer and `info`. + if let Some(counter) = &self.profile_counter { + let binding = self.buffers.len() + self.info.has_info() as usize; + let ty = counter.value.item().unwrap_ptr(); + write!( + f, + "@group(0) +@binding({binding}) +var {}_store: {ty}; +\n", + counter.value, + )?; + } + for value in self.shared_values.iter() { let location = "workgroup"; write!( @@ -189,6 +209,10 @@ fn {}( writeln!(f, "let {value} = &{value}_store;")?; } + if let Some(counter) = &self.profile_counter { + writeln!(f, "let {0} = &{0}_store;", counter.value)?; + } + for SharedValue { value, .. } in self.shared_values.iter() { writeln!(f, "let {value} = &{value}_store;")?; } diff --git a/crates/cubecl-wgpu/src/compute/mem_manager.rs b/crates/cubecl-wgpu/src/compute/mem_manager.rs index 797523fcb2..438c2b6a69 100644 --- a/crates/cubecl-wgpu/src/compute/mem_manager.rs +++ b/crates/cubecl-wgpu/src/compute/mem_manager.rs @@ -113,6 +113,25 @@ impl WgpuMemManager { Ok((resource, binding)) } + /// Reserve a storage buffer (usable as a compute binding) and return both a resource and a + /// binding, so it can be bound, zeroed, and later read back. + pub(crate) fn reserve_storage( + &mut self, + size: u64, + ) -> Result { + let handle = self.memory_pool.reserve(size)?; + Ok(MemoryHandle::binding(handle)) + } + + /// Resolve a storage-pool binding into a resource (may be called several times for the same + /// binding, e.g. to bind and to read back). + pub(crate) fn storage_resource( + &mut self, + binding: ManagedMemoryBinding, + ) -> Result { + self.memory_pool.get_resource(binding, None, None) + } + pub(crate) fn get_resource(&mut self, binding: Binding) -> Result { self.memory_pool .get_resource(binding.memory, binding.offset_start, binding.offset_end) diff --git a/crates/cubecl-wgpu/src/compute/schedule.rs b/crates/cubecl-wgpu/src/compute/schedule.rs index b0e71c4e6d..dabebf94d2 100644 --- a/crates/cubecl-wgpu/src/compute/schedule.rs +++ b/crates/cubecl-wgpu/src/compute/schedule.rs @@ -59,6 +59,8 @@ pub struct BindingsResource { /// Which compiler was used. This determines the passing strategy of params. /// WGSL and metal use bindings, Vulkan uses buffer addresses sent via a uniform buffer. pub compiler_info: CompilerInfo, + /// The profiling counters buffer, bound after `info` when profiling is on. + pub profile_counter: Option, } /// Represents a WGPU backend for scheduling tasks on streams. @@ -175,6 +177,10 @@ impl BindingsResource { if let Some(info) = info { self.resources.push(info); } + // The counters buffer is bound last, after every buffer and `info`. + if let Some(counter) = self.profile_counter { + self.resources.push(counter); + } (self.resources, vec![], None) } } diff --git a/crates/cubecl-wgpu/src/compute/server.rs b/crates/cubecl-wgpu/src/compute/server.rs index 75ab4be0c5..a2dd43eea7 100644 --- a/crates/cubecl-wgpu/src/compute/server.rs +++ b/crates/cubecl-wgpu/src/compute/server.rs @@ -17,16 +17,16 @@ use cubecl_core::{ future::DynFut, prelude::*, server::{ - CopyDescriptor, IoError, KernelArguments, LaunchError, ProfileError, ProfilingToken, - ServerCommunication, ServerError, ServerUtilities, + CopyDescriptor, FlopRecord, IoError, KernelArguments, LaunchError, ProfileError, + ProfilingToken, ServerCommunication, ServerError, ServerUtilities, }, zspace::{Strides, strides}, }; #[cfg(feature = "spirv")] use cubecl_core::{cache::CacheOption, compilation_cache::CompilationCache, hash::StableHash}; -use cubecl_ir::MemoryDeviceProperties; +use cubecl_ir::{CountKey, MemoryDeviceProperties}; use cubecl_runtime::allocator::ContiguousMemoryLayoutPolicy; -use cubecl_runtime::memory_management::{ManagedMemoryHandle, MemoryUsage}; +use cubecl_runtime::memory_management::{ManagedMemoryBinding, ManagedMemoryHandle, MemoryUsage}; use cubecl_runtime::{ compiler::CubeTask, config::{CubeClRuntimeConfig, RuntimeConfig}, @@ -62,6 +62,9 @@ pub struct WgpuServer { // A buffer that can be used to store stream id without extra allocations. streams_pool: Vec, pipelines: HashMap, CompilerInfo)>, + /// Slot → key decode tables for profiled kernels, cached at compile time and used to turn the + /// read-back counters buffer into named counts. + profile_maps: HashMap>>, scheduler: SchedulerMultiStream, #[cfg(feature = "spirv")] pub(crate) spirv_cache: @@ -112,6 +115,7 @@ impl WgpuServer { streams_pool: Vec::new(), device, pipelines: HashMap::new(), + profile_maps: HashMap::new(), scheduler: SchedulerMultiStream::new( utilities.logger.clone(), backend_scheduler, @@ -159,6 +163,7 @@ impl WgpuServer { resources, info: bindings.info, compiler_info, + profile_counter: None, }) } @@ -200,6 +205,14 @@ impl WgpuServer { let (compiler_info, auto_repr) = compiler.normalize_repr(compiled.repr); let repr = auto_repr.as_ref().map(|r| r.as_ref()); + // Cache the profiling decode table produced during compilation, keyed by kernel id. + if let Some(crate::AutoRepresentation::Wgsl(shader)) = auto_repr.as_ref() + && !shader.profile_map.is_empty() + { + self.profile_maps + .insert(kernel_id.clone(), Arc::new(shader.profile_map.clone())); + } + // /!\ Do not delete the following commented code. // This is useful while working on the metal compiler. // Also the errors are printed nicely which is not the case when this is the runtime @@ -253,6 +266,44 @@ impl WgpuServer { Ok((pipeline, compiler_info)) } + + /// Read back the counters buffer after the profiled dispatch, decode its slots through `map`, + /// and record the per-op counts on the shared metrics, keyed by `kernel_id`. + fn record_profile( + &mut self, + kernel_id: KernelId, + binding: ManagedMemoryBinding, + map: &[CountKey], + stream_id: StreamId, + ) { + // Flush the zeroing write and the dispatch before reading. + self.scheduler.execute_streams(vec![stream_id]); + let stream = self.scheduler.stream(&stream_id); + let Ok(resource) = stream.mem_manage.storage_resource(binding) else { + return; + }; + let read = + stream.read_resources(vec![(resource, Shape::new([map.len()]), size_of::())]); + let Ok(mut bytes) = cubecl_common::future::block_on(read) else { + return; + }; + let bytes = bytes.remove(0); + let data: &[u32] = bytemuck::cast_slice(&bytes); + let counts: HashMap = + map.iter().cloned().zip(data.iter().copied()).collect(); + + let mut metrics = self.utilities.metrics.write().unwrap(); + metrics + .entry(kernel_id) + .and_modify(|record| { + record.last = counts.clone(); + record.samples += 1; + }) + .or_insert(FlopRecord { + last: counts, + samples: 1, + }); + } } impl ComputeServer for WgpuServer { @@ -372,6 +423,9 @@ impl ComputeServer for WgpuServer { mode: ExecutionMode, stream_id: StreamId, ) { + let mut kernel_id = kernel.id(); + kernel_id.mode(mode); + let (pipeline, compiler_info) = match self.pipeline(kernel, &args, mode) { Ok(val) => val, Err(err) => { @@ -387,7 +441,7 @@ impl ComputeServer for WgpuServer { .iter() .for_each(|b| self.streams_pool.push(b.stream)); - let resources = match self.prepare_bindings(args, compiler_info) { + let mut resources = match self.prepare_bindings(args, compiler_info) { Ok(val) => val, Err(err) => { // We make the stream that would execute the kernel in error. @@ -396,13 +450,48 @@ impl ComputeServer for WgpuServer { return; } }; + + // If this kernel is profiled, allocate a zeroed counters buffer and bind it last. + let profile = self.profile_maps.get(&kernel_id).cloned(); + let counter_binding = profile.as_ref().map(|map| { + let size = (map.len() * size_of::()) as u64; + let stream = self.scheduler.stream(&stream_id); + let binding = stream + .mem_manage + .reserve_storage(size) + .expect("The counters buffer should allocate"); + // One resource to zero it, another to bind it (`WgpuResource` is not `Clone`). + let zero_resource = stream + .mem_manage + .storage_resource(binding.clone()) + .expect("The counters resource should resolve"); + let bind_resource = stream + .mem_manage + .storage_resource(binding.clone()) + .expect("The counters resource should resolve"); + // Zero it before the kernel accumulates into it. + self.scheduler.register( + stream_id, + ScheduleTask::Write { + data: Bytes::from_bytes_vec(vec![0u8; size as usize]), + buffer: zero_resource, + }, + &[], + ); + resources.profile_counter = Some(bind_resource); + binding + }); + let task = ScheduleTask::Execute { pipeline, count, resources, }; - self.scheduler.register(stream_id, task, &self.streams_pool); + + if let (Some(binding), Some(map)) = (counter_binding, profile) { + self.record_profile(kernel_id, binding, &map, stream_id); + } } fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError> { From d251e94871ec9aafd9ac2441c22e7645118f1b39 Mon Sep 17 00:00:00 2001 From: SamuelBelanger Date: Mon, 6 Jul 2026 09:02:38 -0400 Subject: [PATCH 10/14] refactor tests naming --- crates/cubecl-runtime/src/kernel.rs | 3 --- crates/cubecl/tests/{chose.rs => profiler.rs} | 26 +++++++++---------- 2 files changed, 13 insertions(+), 16 deletions(-) rename crates/cubecl/tests/{chose.rs => profiler.rs} (88%) diff --git a/crates/cubecl-runtime/src/kernel.rs b/crates/cubecl-runtime/src/kernel.rs index 25af61f2a7..aafd79153a 100644 --- a/crates/cubecl-runtime/src/kernel.rs +++ b/crates/cubecl-runtime/src/kernel.rs @@ -180,9 +180,6 @@ impl CubeTask for KernelTask { let entrypoint_name = gpu_ir.options.kernel_name.clone(); let cube_dim = gpu_ir.cube_dim; - // std::println!("Compiling kernel: {}", entrypoint_name); - // std::println!("IR: {:#?}", gpu_ir); - let lower_level_ir = compiler.compile(gpu_ir, compilation_options, mode, addr_type)?; Ok(CompiledKernel { diff --git a/crates/cubecl/tests/chose.rs b/crates/cubecl/tests/profiler.rs similarity index 88% rename from crates/cubecl/tests/chose.rs rename to crates/cubecl/tests/profiler.rs index 837dcabe69..8e6a19b2e7 100644 --- a/crates/cubecl/tests/chose.rs +++ b/crates/cubecl/tests/profiler.rs @@ -1,24 +1,24 @@ use cubecl::{CubeDim, CubeElement, Runtime, TestRuntime, cube, prelude::*}; #[cube(launch)] -fn test_chose_cube(lhs: f32, rhs: f32, out: &mut [f32]) { +fn profile_basic_kernel(lhs: f32, rhs: f32, out: &mut [f32]) { if UNIT_POS == 0 { out[0] = lhs + rhs + lhs; out[0] = lhs + rhs; - // terminate!(); + terminate!(); out[0] = lhs + rhs; } } #[test] -fn test_chose_1() { +fn test_profile_basic() { let client = TestRuntime::client(&Default::default()); let lhs = 1.0; let rhs = 2.0; let output = client.empty(core::mem::size_of::()); - test_chose_cube::launch( + profile_basic_kernel::launch( &client, CubeCount::Static(1, 1, 1), CubeDim::new_1d(1), @@ -46,21 +46,21 @@ fn test_chose_1() { } #[cube(launch)] -fn test_chose_cube_tensor(lhs: &Tensor, rhs: &Tensor, out: &mut Tensor) { +fn profile_tensor_kernel(lhs: &Tensor, rhs: &Tensor, out: &mut Tensor) { if ABSOLUTE_POS < out.len() { out[ABSOLUTE_POS] = lhs[ABSOLUTE_POS] + rhs[ABSOLUTE_POS]; } } #[test] -fn test_chose_tensor() { +fn test_profile_tensor() { let client = TestRuntime::client(&Default::default()); let lhs = client.create_from_slice(f32::as_bytes(&[1.0, 2.0, 3.0])); let rhs = client.create_from_slice(f32::as_bytes(&[10.0, 20.0, 30.0])); let output = client.empty(3 * core::mem::size_of::()); - test_chose_cube_tensor::launch( + profile_tensor_kernel::launch( &client, CubeCount::Static(1, 1, 1), CubeDim::new_1d(3), @@ -77,7 +77,7 @@ fn test_chose_tensor() { } #[cube(launch)] -fn test_chose_cube_shape(input: &Tensor, out: &mut Tensor) { +fn profile_shape_kernel(input: &Tensor, out: &mut Tensor) { if UNIT_POS == 0 { out[0] = input.shape(0) as u32; out[1] = input.shape(1) as u32; @@ -85,12 +85,12 @@ fn test_chose_cube_shape(input: &Tensor, out: &mut Tensor) { } #[test] -fn test_chose_shape() { +fn test_profile_shape() { let client = TestRuntime::client(&Default::default()); let input = client.create_from_slice(f32::as_bytes(&[0.0; 12])); let output = client.empty(2 * core::mem::size_of::()); - test_chose_cube_shape::launch( + profile_shape_kernel::launch( &client, CubeCount::Static(1, 1, 1), CubeDim::new_1d(1), @@ -105,7 +105,7 @@ fn test_chose_shape() { } #[cube(launch_unchecked)] -fn test_chose_vectorized( +fn profile_vectorized_kernel( input: &Tensor>, scalar: f32, out: &mut Tensor>, @@ -120,7 +120,7 @@ fn test_chose_vectorized( } #[test] -fn test_chose_cube_vectorized() { +fn test_profile_vectorized() { let client = TestRuntime::client(&Default::default()); let vectorization = 2; @@ -130,7 +130,7 @@ fn test_chose_cube_vectorized() { let output = client.empty(12 * core::mem::size_of::()); unsafe { - test_chose_vectorized::launch_unchecked::( + profile_vectorized_kernel::launch_unchecked::( &client, CubeCount::Static(2, 1, 1), CubeDim::new_1d((input_values.len() / vectorization) as u32), From ae02724ccb81468c81f92ba757707aadf6749faf Mon Sep 17 00:00:00 2001 From: SamuelBelanger Date: Mon, 6 Jul 2026 09:12:01 -0400 Subject: [PATCH 11/14] fix std fmt warning --- crates/cubecl-runtime/src/server/base.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index 71c8e14afa..c3a932b07c 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -14,6 +14,7 @@ use crate::{ }; use ahash::AHasher; use alloc::boxed::Box; +use alloc::fmt; #[cfg(feature = "profile-tracy")] use alloc::format; use alloc::string::String; @@ -88,7 +89,7 @@ pub struct FlopRecord { pub samples: u32, } -impl std::fmt::Display for FlopRecord { +impl fmt::Display for FlopRecord { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { writeln!(f, "FlopRecord:")?; writeln!(f, " samples: {}", self.samples)?; From 840f4ce2a16c0140a5e3bf810d85ca152dd69812 Mon Sep 17 00:00:00 2001 From: SamuelBelanger Date: Mon, 6 Jul 2026 13:38:38 -0400 Subject: [PATCH 12/14] fix cargo fmt --- crates/cubecl-macros-internal/src/lib.rs | 2 +- crates/cubecl-wgpu/src/compute/mem_manager.rs | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/crates/cubecl-macros-internal/src/lib.rs b/crates/cubecl-macros-internal/src/lib.rs index 6a9c9b2ead..ecbb6244af 100644 --- a/crates/cubecl-macros-internal/src/lib.rs +++ b/crates/cubecl-macros-internal/src/lib.rs @@ -1,8 +1,8 @@ +use enum_counts::enum_counts_impl; use generate::{ op_args::generate_op_args, operation::{generate_opcode, generate_operation}, }; -use enum_counts::enum_counts_impl; use proc_macro::TokenStream; use type_hash::type_hash_impl; diff --git a/crates/cubecl-wgpu/src/compute/mem_manager.rs b/crates/cubecl-wgpu/src/compute/mem_manager.rs index 438c2b6a69..e6dfc4f8db 100644 --- a/crates/cubecl-wgpu/src/compute/mem_manager.rs +++ b/crates/cubecl-wgpu/src/compute/mem_manager.rs @@ -115,10 +115,7 @@ impl WgpuMemManager { /// Reserve a storage buffer (usable as a compute binding) and return both a resource and a /// binding, so it can be bound, zeroed, and later read back. - pub(crate) fn reserve_storage( - &mut self, - size: u64, - ) -> Result { + pub(crate) fn reserve_storage(&mut self, size: u64) -> Result { let handle = self.memory_pool.reserve(size)?; Ok(MemoryHandle::binding(handle)) } From bed7d1a84eb7a04021e601ae6d4a9799516ccd60 Mon Sep 17 00:00:00 2001 From: SamuelBelanger Date: Mon, 6 Jul 2026 13:54:42 -0400 Subject: [PATCH 13/14] fix lint --- crates/cubecl-ir/src/count.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/cubecl-ir/src/count.rs b/crates/cubecl-ir/src/count.rs index b946ead09f..611e73c604 100644 --- a/crates/cubecl-ir/src/count.rs +++ b/crates/cubecl-ir/src/count.rs @@ -202,7 +202,7 @@ pub fn fmt_counts( // Sort by Elem, then strictly by Path (names), then Unit // Grouping by path ensures our tree algorithm correctly renders branches. - entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2)).then(a.1.cmp(&b.1))); + entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(b.2)).then(a.1.cmp(&b.1))); if entries.is_empty() { return writeln!(f, " (none)"); @@ -242,7 +242,7 @@ pub fn fmt_counts( // Print root elements without indentation if cur_elem.as_ref() != Some(elem) { writeln!(f, "{elem}")?; - cur_elem = Some(elem.clone()); + cur_elem = Some(*elem); is_last_at_depth.clear(); } From c7e8f1d9f42a9b9ff6918a5c2130381a9ce82082 Mon Sep 17 00:00:00 2001 From: SamuelBelanger Date: Mon, 6 Jul 2026 15:05:24 -0400 Subject: [PATCH 14/14] fix xtask lint --- crates/cubecl-core/src/codegen/profiler.rs | 4 ++-- crates/cubecl-runtime/src/config/base.rs | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/cubecl-core/src/codegen/profiler.rs b/crates/cubecl-core/src/codegen/profiler.rs index 8b4528de4b..7e79612a2a 100644 --- a/crates/cubecl-core/src/codegen/profiler.rs +++ b/crates/cubecl-core/src/codegen/profiler.rs @@ -59,7 +59,7 @@ impl CompilerProfiler { self.locals.values().for_each(|(local_counter, index)| { declares.extend(self.declare_instructions(local_counter)); - flushes.extend(self.flush_instructions(&local_counter, index, allocator)); + flushes.extend(self.flush_instructions(local_counter, index, allocator)); }); (declares, flushes) @@ -165,7 +165,7 @@ impl CompilerProfiler { let mut entries: Vec<(usize, CountKey)> = self .locals .iter() - .map(|(key, (_, slot))| (*slot, key.clone())) + .map(|(key, (_, slot))| (*slot, *key)) .collect(); entries.sort_by_key(|(slot, _)| *slot); entries.into_iter().map(|(_, key)| key).collect() diff --git a/crates/cubecl-runtime/src/config/base.rs b/crates/cubecl-runtime/src/config/base.rs index c950dece4f..9fabeb3016 100644 --- a/crates/cubecl-runtime/src/config/base.rs +++ b/crates/cubecl-runtime/src/config/base.rs @@ -133,10 +133,10 @@ impl RuntimeConfig for CubeClRuntimeConfig { } } - if let Ok(val) = std::env::var("CUBECL_PROFILE_FLOPS") { - if val == "1" || val.eq_ignore_ascii_case("true") { - self.profiling.hardware_metrics = true; - } + if let Ok(val) = std::env::var("CUBECL_PROFILE_FLOPS") + && (val == "1" || val.eq_ignore_ascii_case("true")) + { + self.profiling.hardware_metrics = true; } self