#![allow(missing_docs)]
use super::BrickId;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ExecutionNodeId(pub u32);
impl ExecutionNodeId {
pub const MAX_BUDGET: u32 = 100_000;
#[inline]
pub fn validate(self) -> bool {
debug_assert!(
self.0 < Self::MAX_BUDGET,
"CB-BUDGET: node id {} exceeds max budget {}",
self.0,
Self::MAX_BUDGET
);
self.0 < Self::MAX_BUDGET
}
}
#[derive(Debug, Clone)]
pub enum ExecutionNode {
Brick { id: BrickId, timing_ns: u64, elements: u64 },
Kernel {
name: String,
ptx_hash: u64,
grid: (u32, u32, u32),
block: (u32, u32, u32),
shared_mem: u32,
timing_ns: Option<u64>,
arithmetic_intensity: Option<f32>,
achieved_tflops: Option<f32>,
},
Transfer {
src: String,
dst: String,
bytes: u64,
direction: TransferDirection,
timing_ns: Option<u64>,
},
Function { name: String, file: Option<String>, line: Option<u32> },
Layer { index: u32 },
AsyncTask {
name: String,
poll_count: u64,
yield_count: u64,
total_poll_ns: u64,
},
}
impl ExecutionNode {
pub fn name(&self) -> String {
match self {
Self::Brick { id, .. } => id.name().to_string(),
Self::Kernel { name, .. } => name.clone(),
Self::Function { name, .. } => name.clone(),
Self::Layer { index } => format!("Layer{}", index),
Self::Transfer { src, dst, direction, .. } => {
let dir = match direction {
TransferDirection::H2D => "H2D",
TransferDirection::D2H => "D2H",
TransferDirection::D2D => "D2D",
};
format!("{}:{}->{}", dir, src, dst)
}
Self::AsyncTask { name, .. } => name.clone(),
}
}
pub fn is_kernel(&self) -> bool {
matches!(self, Self::Kernel { .. })
}
pub fn is_brick(&self) -> bool {
matches!(self, Self::Brick { .. })
}
pub fn is_transfer(&self) -> bool {
matches!(self, Self::Transfer { .. })
}
pub fn timing_ns(&self) -> Option<u64> {
match self {
Self::Brick { timing_ns, .. } => Some(*timing_ns),
Self::Kernel { timing_ns, .. } => *timing_ns,
Self::Transfer { timing_ns, .. } => *timing_ns,
_ => None,
}
}
pub fn ptx_hash(&self) -> Option<u64> {
match self {
Self::Kernel { ptx_hash, .. } => Some(*ptx_hash),
_ => None,
}
}
pub fn arithmetic_intensity(&self) -> Option<f32> {
match self {
Self::Kernel { arithmetic_intensity, .. } => *arithmetic_intensity,
_ => None,
}
}
pub fn achieved_tflops(&self) -> Option<f32> {
match self {
Self::Kernel { achieved_tflops, .. } => *achieved_tflops,
_ => None,
}
}
pub fn transfer_bytes(&self) -> Option<u64> {
match self {
Self::Transfer { bytes, .. } => Some(*bytes),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum EdgeType {
Calls,
Contains,
Launches,
Sequence,
DependsOn,
Transfer {
bytes: u64,
direction: TransferDirection,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransferDirection {
H2D,
D2H,
D2D,
}
#[derive(Debug, Clone)]
pub struct ExecutionEdge {
pub src: ExecutionNodeId,
pub dst: ExecutionNodeId,
pub edge_type: EdgeType,
pub weight: f32,
}