use alloc::sync::Arc;
use core::sync::atomic::{AtomicBool, Ordering};
pub use cubecl_common::profile::{Duration, ProfileDuration, ProfileTicks, TimingMethod};
use cubecl_environment::sync::RwLock;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum TimingRequest {
#[default]
None,
Resolved,
Deferred,
}
pub trait LaunchObserver: Send + Sync {
fn launched(&self, kernel: &'static str);
fn timing(&self) -> TimingRequest {
TimingRequest::None
}
fn profiled(&self, _kernel: &'static str, _profile: ProfileDuration) {}
fn timed(&self, _kernel: &'static str, _duration: Duration, _method: TimingMethod) {}
}
#[must_use = "an observation stops as soon as it is dropped"]
pub struct LaunchObservation {
previous: Option<Arc<dyn LaunchObserver>>,
}
impl LaunchObservation {
pub fn new(observer: Arc<dyn LaunchObserver>) -> Self {
let previous = OBSERVER.write().replace(observer);
OBSERVING.store(true, Ordering::Relaxed);
Self { previous }
}
}
impl Drop for LaunchObservation {
fn drop(&mut self) {
let previous = self.previous.take();
let still_observed = previous.is_some();
*OBSERVER.write() = previous;
OBSERVING.store(still_observed, Ordering::Relaxed);
}
}
impl core::fmt::Debug for LaunchObservation {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("LaunchObservation")
.field("replaced_an_observer", &self.previous.is_some())
.finish()
}
}
pub(crate) fn is_observing() -> bool {
OBSERVING.load(Ordering::Relaxed)
}
pub(crate) fn notify_launch(kernel: &'static str) {
if !OBSERVING.load(Ordering::Relaxed) {
return;
}
if let Some(observer) = OBSERVER.read().as_ref() {
observer.launched(kernel);
}
}
pub(crate) fn timing_wanted() -> bool {
timing_requested() != TimingRequest::None
}
pub(crate) fn warn_logger_takes_deferred_measurements() {
static SAID: AtomicBool = AtomicBool::new(false);
if timing_requested() == TimingRequest::Deferred && !SAID.swap(true, Ordering::Relaxed) {
log::warn!(
"The profiling logger is reading every launch's measurement, so this run's \
launch observer is told durations instead of keeping them, and its kernels \
run one at a time. Turn the profile logging off to measure the pass as it runs."
);
}
}
fn timing_requested() -> TimingRequest {
if !OBSERVING.load(Ordering::Relaxed) {
return TimingRequest::None;
}
OBSERVER
.read()
.as_ref()
.map_or(TimingRequest::None, |observer| observer.timing())
}
pub(crate) fn read_and_notify_timed(
kernel: &'static str,
profile: ProfileDuration,
) -> ProfileDuration {
let method = profile.timing_method();
let observer = installed_observer();
let ticks = cubecl_environment::future::block_on(profile.resolve());
match (&ticks, &observer) {
(Some(ticks), Some(observer)) => deliver_timed(observer, kernel, ticks.duration(), method),
(None, _) => log::warn!(
"Skipped timing a launch of `{kernel}` for its observer: \
the profiled window carried no measurement"
),
(Some(_), None) => {}
}
ProfileDuration::new(alloc::boxed::Box::pin(async move { ticks }), method)
}
fn installed_observer() -> Option<Arc<dyn LaunchObserver>> {
if !OBSERVING.load(Ordering::Relaxed) {
return None;
}
OBSERVER.read().as_ref().map(Arc::clone)
}
fn deliver_timed(
observer: &Arc<dyn LaunchObserver>,
kernel: &'static str,
duration: Duration,
method: TimingMethod,
) {
let slot = OBSERVER.read();
if slot
.as_ref()
.is_some_and(|installed| Arc::ptr_eq(installed, observer))
{
observer.timed(kernel, duration, method);
}
}
pub(crate) fn notify_profiled(kernel: &'static str, profile: ProfileDuration) {
if !OBSERVING.load(Ordering::Relaxed) {
return;
}
let (observer, profile) = {
let slot = OBSERVER.read();
let Some(observer) = slot.as_ref() else {
return;
};
match observer.timing() {
TimingRequest::Deferred => {
observer.profiled(kernel, profile);
return;
}
TimingRequest::None => {
static SAID: AtomicBool = AtomicBool::new(false);
if !SAID.swap(true, Ordering::Relaxed) {
log::warn!(
"Dropped a timing of `{kernel}`: its observer asked for none by the \
time the measurement arrived"
);
}
return;
}
TimingRequest::Resolved => (Arc::clone(observer), profile),
}
};
let method = profile.timing_method();
let Some(ticks) = cubecl_environment::future::block_on(profile.resolve()) else {
log::warn!(
"Skipped timing a launch of `{kernel}` for its observer: \
the profiled window carried no measurement"
);
return;
};
deliver_timed(&observer, kernel, ticks.duration(), method);
}
static OBSERVING: AtomicBool = AtomicBool::new(false);
static OBSERVER: RwLock<Option<Arc<dyn LaunchObserver>>> = RwLock::new(None);
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec::Vec;
use alloc::vec;
use cubecl_environment::sync::Mutex;
#[test]
#[serial_test::serial]
fn launches_arrive_in_issue_order() {
let recorder = Arc::new(Recorder::default());
let watching = LaunchObservation::new(recorder.clone());
notify_launch("first");
notify_launch("second");
assert_eq!(*recorder.0.lock(), ["first", "second"]);
drop(watching);
notify_launch("after");
assert_eq!(
recorder.0.lock().len(),
2,
"an observation that ended must not keep receiving"
);
}
#[test]
#[serial_test::serial]
fn timing_is_off_unless_an_observer_asks() {
assert!(!timing_wanted(), "nothing installed, nothing to time");
let names_only = LaunchObservation::new(Arc::new(Recorder::default()));
assert!(!timing_wanted(), "names only, by default");
drop(names_only);
let timed = Arc::new(Timed::default());
let watching = LaunchObservation::new(timed.clone());
assert!(timing_wanted());
read_and_notify_timed("a_kernel", measured_on_device(7));
assert_eq!(
*timed.0.lock(),
[("a_kernel", Duration::from_micros(7), TimingMethod::Device)]
);
drop(watching);
assert!(!timing_wanted());
}
#[test]
#[serial_test::serial]
fn an_observation_restores_the_one_it_replaced() {
let outer = Arc::new(Recorder::default());
let inner = Arc::new(Recorder::default());
let watching_outer = LaunchObservation::new(outer.clone());
{
let _watching_inner = LaunchObservation::new(inner.clone());
notify_launch("during_the_inner_pass");
}
notify_launch("after_the_inner_pass");
drop(watching_outer);
notify_launch("unobserved");
assert_eq!(*inner.0.lock(), ["during_the_inner_pass"]);
assert_eq!(*outer.0.lock(), ["after_the_inner_pass"]);
}
fn measured(micros: u64) -> ProfileDuration {
let start = cubecl_common::profile::Instant::now();
ProfileDuration::new_system_time(start, start + Duration::from_micros(micros))
}
fn measured_on_device(micros: u64) -> ProfileDuration {
let start = cubecl_common::profile::Instant::now();
let ticks = ProfileTicks::from_start_end(start, start + Duration::from_micros(micros));
ProfileDuration::new(
alloc::boxed::Box::pin(async move { Some(ticks) }),
TimingMethod::Device,
)
}
fn measured_while(
micros: u64,
during_the_read: impl FnOnce() + Send + 'static,
) -> ProfileDuration {
let start = cubecl_common::profile::Instant::now();
ProfileDuration::new(
alloc::boxed::Box::pin(async move {
during_the_read();
Some(ProfileTicks::from_start_end(
start,
start + Duration::from_micros(micros),
))
}),
TimingMethod::System,
)
}
#[test]
#[serial_test::serial]
fn a_measurement_is_read_back_for_an_observer_that_only_wants_durations() {
let timed = Arc::new(Timed::default());
let watching = LaunchObservation::new(timed.clone());
notify_profiled("a_kernel", measured(7));
drop(watching);
assert_eq!(
*timed.0.lock(),
[("a_kernel", Duration::from_micros(7), TimingMethod::System)]
);
}
#[test]
#[serial_test::serial]
fn an_observation_that_ended_mid_read_is_not_told_the_duration() {
let timed = Arc::new(Timed::default());
let watching = Arc::new(Mutex::new(Some(LaunchObservation::new(timed.clone()))));
let ends_the_observation = watching.clone();
notify_profiled(
"a_kernel",
measured_while(7, move || drop(ends_the_observation.lock().take())),
);
assert!(watching.lock().is_none(), "the read-back ended it");
assert!(
timed.0.lock().is_empty(),
"an observation that ended must not keep receiving"
);
}
#[test]
#[serial_test::serial]
fn an_observation_that_ended_mid_read_is_not_told_the_loggers_reading() {
let timed = Arc::new(Timed::default());
let watching = Arc::new(Mutex::new(Some(LaunchObservation::new(timed.clone()))));
let ends_the_observation = watching.clone();
let for_the_logger = read_and_notify_timed(
"a_kernel",
measured_while(7, move || drop(ends_the_observation.lock().take())),
);
assert!(watching.lock().is_none(), "the read-back ended it");
assert!(
timed.0.lock().is_empty(),
"an observation that ended must not keep receiving"
);
let ticks = cubecl_environment::future::block_on(for_the_logger.resolve())
.expect("the logger still gets the reading");
assert_eq!(
ticks.duration(),
Duration::from_micros(7),
"read once, and handed on"
);
}
#[test]
#[serial_test::serial]
fn an_observer_can_keep_a_measurement_unread() {
let kept = Arc::new(Kept::default());
let watching = LaunchObservation::new(kept.clone());
notify_profiled("first", measured(3));
notify_profiled("second", measured(5));
drop(watching);
let read: Vec<(&'static str, Duration)> = core::mem::take(&mut *kept.0.lock())
.into_iter()
.map(|(kernel, profile)| {
let ticks = cubecl_environment::future::block_on(profile.resolve())
.expect("a system measurement always carries its ticks");
(kernel, ticks.duration())
})
.collect();
assert_eq!(
read,
[
("first", Duration::from_micros(3)),
("second", Duration::from_micros(5))
],
"in issue order, each still carrying its own window"
);
}
#[derive(Default)]
struct Kept(Mutex<Vec<(&'static str, ProfileDuration)>>);
impl LaunchObserver for Kept {
fn launched(&self, _kernel: &'static str) {}
fn timing(&self) -> TimingRequest {
TimingRequest::Deferred
}
fn profiled(&self, kernel: &'static str, profile: ProfileDuration) {
self.0.lock().push((kernel, profile));
}
}
#[derive(Default)]
struct Recorder(Mutex<Vec<&'static str>>);
impl LaunchObserver for Recorder {
fn launched(&self, kernel: &'static str) {
self.0.lock().push(kernel);
}
}
#[derive(Default)]
struct Timed(Mutex<Vec<(&'static str, Duration, TimingMethod)>>);
impl LaunchObserver for Timed {
fn launched(&self, _kernel: &'static str) {}
fn timing(&self) -> TimingRequest {
TimingRequest::Resolved
}
fn timed(&self, kernel: &'static str, duration: Duration, method: TimingMethod) {
self.0.lock().push((kernel, duration, method));
}
}
}