use super::BrickProfiler;
use crate::brick::exec_graph::{BrickId, ExecutionNode, ExecutionNodeId};
impl BrickProfiler {
pub fn enable_graph(&mut self) {
self.graph_enabled = true;
}
pub fn disable_graph(&mut self) {
self.graph_enabled = false;
}
#[must_use]
pub fn is_graph_enabled(&self) -> bool {
self.graph_enabled
}
#[must_use]
pub fn execution_graph(&self) -> &crate::brick::exec_graph::ExecutionGraph {
&self.execution_graph
}
pub fn execution_graph_mut(&mut self) -> &mut crate::brick::exec_graph::ExecutionGraph {
&mut self.execution_graph
}
pub fn graph_push_scope(&mut self, node: ExecutionNode) -> Option<ExecutionNodeId> {
if !self.graph_enabled {
return None;
}
debug_assert!(
self.execution_graph.num_nodes() < 100_000,
"CB-BUDGET: execution graph has {} nodes, exceeds 100k budget",
self.execution_graph.num_nodes()
);
Some(self.execution_graph.push_scope(node))
}
pub fn graph_pop_scope(&mut self) -> Option<ExecutionNodeId> {
if !self.graph_enabled {
return None;
}
self.execution_graph.pop_scope()
}
pub fn graph_record_brick(
&mut self,
brick_id: BrickId,
timing_ns: u64,
elements: u64,
) -> Option<ExecutionNodeId> {
if !self.graph_enabled {
return None;
}
let node = ExecutionNode::Brick { id: brick_id, timing_ns, elements };
Some(self.execution_graph.add_node_in_scope(node))
}
pub fn graph_record_kernel(
&mut self,
name: &str,
ptx_hash: u64,
grid: (u32, u32, u32),
block: (u32, u32, u32),
shared_mem: u32,
) -> Option<ExecutionNodeId> {
if !self.graph_enabled {
return None;
}
Some(self.execution_graph.record_kernel_launch(name, ptx_hash, grid, block, shared_mem))
}
#[must_use]
pub fn graph_to_dot(&self) -> String {
self.execution_graph.to_dot()
}
#[cfg(feature = "execution-graph")]
#[must_use]
pub fn graph_to_csr(&self) -> trueno_graph::CsrGraph {
self.execution_graph.to_csr()
}
pub fn graph_clear(&mut self) {
self.execution_graph.clear();
}
#[must_use]
pub fn graph_is_scope_balanced(&self) -> bool {
self.execution_graph.is_scope_balanced()
}
}