cubecl_ir/
debug.rs

1//! Source locations of the ops that kernel expansion inserts.
2//!
3//! The macro-generated code of each `#[cube]` function opens a [frame](DebugState::enter_fn) and
4//! moves its [position](DebugState::set_pos). [`LocationListener`] gives each inserted op the
5//! location of the current frame, nested in the call sites of its callers.
6//!
7//! Expansion inlines every `#[cube]` call, so the call chain is the only record of the source-level
8//! call stack. For an op in `inner`, called from `outer`, called from kernel `k`:
9//!
10//! ```text
11//! CallSite {
12//!     callee: Named("inner", SrcPos(op)),
13//!     caller: CallSite {
14//!         callee: Named("outer", SrcPos(call to inner)),
15//!         caller: Named("k", SrcPos(call to outer)),
16//!     },
17//! }
18//! ```
19
20use alloc::{
21    boxed::Box,
22    collections::{BTreeMap, BTreeSet},
23    string::String,
24    vec::Vec,
25};
26use pliron::{
27    basic_block::BasicBlock,
28    combine::stream::position::SourcePosition,
29    context::{Context, Ptr},
30    irbuild::listener::InsertionListener,
31    location::{Location, Source},
32    operation::Operation,
33    std_deps::path::PathBuf,
34};
35
36use crate::ContextExt;
37
38/// The `#[cube]` call stack during kernel expansion.
39///
40/// A function without debug data, such as a hand-written expand function, opens no frame. Its ops
41/// get the location of the call to it.
42#[derive(Debug, Default)]
43pub struct DebugState {
44    enabled: bool,
45    frames: Vec<Frame>,
46    /// The source text of each file, by the path of its [`Source`]. Only full debug data records
47    /// it.
48    sources: BTreeMap<String, &'static str>,
49    /// The path of each file, at all levels.
50    files: BTreeSet<String>,
51}
52
53/// One `#[cube]` function on the call stack.
54#[derive(Debug)]
55struct Frame {
56    name: &'static str,
57    file: Source,
58    /// The expression the function is at: the call site, while a callee runs.
59    pos: SourcePosition,
60}
61
62impl DebugState {
63    /// An empty call stack that records locations only when `enabled`.
64    #[must_use]
65    pub fn new(enabled: bool) -> Self {
66        Self {
67            enabled,
68            ..Self::default()
69        }
70    }
71
72    /// Whether expansion records locations.
73    #[must_use]
74    pub fn is_enabled(&self) -> bool {
75        self.enabled
76    }
77
78    /// Opens the frame of function `name`, defined at `line` and `column` of `file`.
79    pub fn enter_fn(&mut self, name: &'static str, file: Source, line: u32, column: u32) {
80        self.frames.push(Frame {
81            name,
82            file,
83            pos: position(line, column),
84        });
85    }
86
87    /// Closes the innermost frame.
88    pub fn exit_fn(&mut self) {
89        self.frames.pop();
90    }
91
92    /// Moves the innermost frame to `line` and `column`.
93    pub fn set_pos(&mut self, line: u32, column: u32) {
94        if let Some(frame) = self.frames.last_mut() {
95            frame.pos = position(line, column);
96        }
97    }
98
99    /// Records the file `path`, and `text` as its source text. An empty text records only the
100    /// path.
101    pub fn add_source(&mut self, path: &str, text: &'static str) {
102        if !self.files.contains(path) {
103            self.files.insert(path.into());
104        }
105        if !text.is_empty() && !self.sources.contains_key(path) {
106            self.sources.insert(path.into(), text);
107        }
108    }
109
110    /// The source text of each file, by the path of its [`Source`]. It stays after
111    /// [`finish`](Self::finish), for the compiler.
112    #[must_use]
113    pub fn sources(&self) -> &BTreeMap<String, &'static str> {
114        &self.sources
115    }
116
117    /// The path of each file, with or without its source text. It stays after
118    /// [`finish`](Self::finish), for the compiler.
119    #[must_use]
120    pub fn files(&self) -> &BTreeSet<String> {
121        &self.files
122    }
123
124    /// Stops recording, at the end of kernel expansion. Ops that compiler passes insert later get
125    /// the location of the op they replace.
126    pub fn finish(&mut self) {
127        self.enabled = false;
128        self.frames.clear();
129    }
130
131    /// The location of the current expression, or `None` outside a `#[cube]` function.
132    fn location(&self) -> Option<Location> {
133        let (kernel, callees) = self.frames.split_first()?;
134        let location =
135            callees
136                .iter()
137                .fold(kernel.location(), |caller, callee| Location::CallSite {
138                    callee: Box::new(callee.location()),
139                    caller: Box::new(caller),
140                });
141        Some(location)
142    }
143}
144
145impl Frame {
146    fn location(&self) -> Location {
147        Location::Named {
148            name: self.name.into(),
149            child_loc: Box::new(Location::SrcPos {
150                src: self.file,
151                pos: self.pos,
152            }),
153        }
154    }
155}
156
157/// The position at `line` and `column`. A value past `i32::MAX`, which no source file reaches,
158/// saturates.
159fn position(line: u32, column: u32) -> SourcePosition {
160    SourcePosition {
161        line: i32::try_from(line).unwrap_or(i32::MAX),
162        column: i32::try_from(column).unwrap_or(i32::MAX),
163    }
164}
165
166/// Gives each inserted op without a location the location of the current expression.
167#[derive(Default)]
168pub struct LocationListener;
169
170impl InsertionListener for LocationListener {
171    fn notify_operation_inserted(&mut self, ctx: &Context, operation: Ptr<Operation>) {
172        let Some(debug) = ctx.try_aux_ty::<DebugState>() else {
173            return;
174        };
175        if let Some(loc) = debug.location()
176            && operation.deref(ctx).loc().is_unknown()
177        {
178            operation.deref_mut(ctx).set_loc(loc);
179        }
180    }
181
182    fn notify_block_inserted(&mut self, _ctx: &Context, _block: Ptr<BasicBlock>) {}
183}
184
185/// The source line of an op: the file and the line of the innermost frame of `loc`. A target that
186/// cannot show inlined frames, such as a `#line` directive, uses it. `None` for an unknown
187/// location, or for a position in memory.
188pub fn leaf_line<'c>(ctx: &'c Context, loc: &Location) -> Option<(&'c PathBuf, u32)> {
189    match loc {
190        Location::CallSite { callee, .. } => leaf_line(ctx, callee),
191        Location::Named { child_loc, .. } => leaf_line(ctx, child_loc),
192        Location::SrcPos {
193            src: Source::File(key),
194            pos,
195        } => Some((
196            pliron::uniqued_any::get(ctx, *key),
197            u32::try_from(pos.line).unwrap_or(0),
198        )),
199        Location::Fused { locations, .. } => locations.iter().find_map(|loc| leaf_line(ctx, loc)),
200        Location::SrcPos {
201            src: Source::InMemory,
202            ..
203        }
204        | Location::Unknown => None,
205    }
206}
cubecl_ir/debug.rs

cubecl_ir/
debug.rs
