use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use knok_core::{type_check, Expr, Graph, Input, Let, TensorType, TypedExpr, TypedGraph};
use crate::Result;
use super::{expr::next_tuple_id, tensor::TraceTensor};
#[derive(Default)]
pub struct TraceContext {
inputs: Vec<Input>,
}
impl TraceContext {
pub fn input<T: TraceTensor>(&mut self, name: &str) -> T {
let ty = T::tensor_type();
self.inputs.push(Input {
name: name.into(),
ty,
});
T::from_expr(Expr::Var(name.into()))
}
pub(crate) fn finish<O: TraceOutput>(
self,
name: &str,
backend: &str,
output: O,
) -> Result<TypedGraph> {
let body = output.exprs();
let reserved_names = self
.inputs
.iter()
.map(|input| input.name.clone())
.collect::<BTreeSet<_>>();
let (lets, body) = let_bind_nodes(body, reserved_names);
let outputs = O::types();
let graph = Graph {
name: name.into(),
backend: backend.into(),
inputs: self.inputs,
outputs,
lets,
body,
};
type_check(graph, &[]).map_err(Into::into)
}
}
fn let_bind_nodes(body: Vec<Expr>, mut reserved_names: BTreeSet<String>) -> (Vec<Let>, Vec<Expr>) {
let mut values = BTreeMap::new();
let mut visited = BTreeSet::new();
let mut order = Vec::new();
for expr in &body {
collect_nodes(expr, &mut visited, &mut values, &mut order);
}
let mut names = BTreeMap::new();
for node_id in &order {
let name = unique_node_name(*node_id, &mut reserved_names);
names.insert(*node_id, name);
}
let lets = order
.into_iter()
.map(|node_id| Let {
names: vec![names
.get(&node_id)
.expect("node traversal recorded an id without a generated name")
.clone()],
value: rewrite_node_refs(
values
.get(&node_id)
.expect("node traversal recorded an id without a value")
.as_ref()
.clone(),
&names,
),
})
.collect();
let body = body
.into_iter()
.map(|expr| rewrite_node_refs(expr, &names))
.collect();
(lets, body)
}
fn collect_nodes(
expr: &Expr,
visited: &mut BTreeSet<u64>,
values: &mut BTreeMap<u64, Arc<Expr>>,
order: &mut Vec<u64>,
) {
match expr {
Expr::Unary { value, .. } => collect_nodes(value, visited, values, order),
Expr::Binary { lhs, rhs, .. } => {
collect_nodes(lhs, visited, values, order);
collect_nodes(rhs, visited, values, order);
}
Expr::Node { node_id, value } => {
if visited.insert(*node_id) {
collect_nodes(value, visited, values, order);
values.insert(*node_id, value.clone());
order.push(*node_id);
}
}
Expr::TupleGet { value, .. } => collect_nodes(value, visited, values, order),
Expr::Call { args, .. } => {
for arg in args {
collect_nodes(arg, visited, values, order);
}
}
Expr::Var(_) | Expr::Const { .. } => {}
}
}
fn rewrite_node_refs(expr: Expr, names: &BTreeMap<u64, String>) -> Expr {
match expr {
Expr::Unary { op, value } => Expr::Unary {
op,
value: Box::new(rewrite_node_refs(*value, names)),
},
Expr::Binary { op, lhs, rhs } => Expr::Binary {
op,
lhs: Box::new(rewrite_node_refs(*lhs, names)),
rhs: Box::new(rewrite_node_refs(*rhs, names)),
},
Expr::Node { node_id, .. } => Expr::Var(
names
.get(&node_id)
.expect("node traversal missed a referenced node")
.clone(),
),
Expr::TupleGet {
tuple_id,
value,
index,
} => Expr::TupleGet {
tuple_id,
value: Arc::new(rewrite_node_refs(value.as_ref().clone(), names)),
index,
},
Expr::Call { op, args } => Expr::Call {
op,
args: args
.into_iter()
.map(|arg| rewrite_node_refs(arg, names))
.collect(),
},
Expr::Var(_) | Expr::Const { .. } => expr,
}
}
fn unique_node_name(node_id: u64, reserved_names: &mut BTreeSet<String>) -> String {
let base = node_name(node_id);
let mut candidate = base.clone();
let mut suffix = 1;
while reserved_names.contains(&candidate) {
candidate = format!("{base}_{suffix}");
suffix += 1;
}
reserved_names.insert(candidate.clone());
candidate
}
fn node_name(node_id: u64) -> String {
format!("__knok_node_{node_id}")
}
pub trait TraceOutput {
fn exprs(self) -> Vec<Expr>;
fn types() -> Vec<TensorType>;
}
pub trait TraceVars: TraceOutput {
fn var_count() -> usize;
fn from_tuple_expr(tuple_id: u64, value: Expr) -> Self;
}
impl<T: TraceTensor> TraceOutput for T {
fn exprs(self) -> Vec<Expr> {
vec![self.into_expr()]
}
fn types() -> Vec<TensorType> {
vec![T::tensor_type()]
}
}
impl<T: TraceTensor> TraceVars for T {
fn var_count() -> usize {
1
}
fn from_tuple_expr(tuple_id: u64, value: Expr) -> Self {
T::from_expr(Expr::TupleGet {
tuple_id,
value: Arc::new(value),
index: 0,
})
}
}
macro_rules! impl_tuple_output {
($($name:ident: $index:tt),+) => {
impl<$($name: TraceTensor),+> TraceOutput for ($($name,)+) {
fn exprs(self) -> Vec<Expr> {
#[allow(non_snake_case)]
let ($($name,)+) = self;
vec![$($name.into_expr()),+]
}
fn types() -> Vec<TensorType> {
vec![$($name::tensor_type()),+]
}
}
impl<$($name: TraceTensor),+> TraceVars for ($($name,)+) {
fn var_count() -> usize {
[$($name::tensor_type()),+].len()
}
fn from_tuple_expr(tuple_id: u64, value: Expr) -> Self {
(
$(
$name::from_expr(Expr::TupleGet {
tuple_id,
value: Arc::new(value.clone()),
index: $index,
}),
)+
)
}
}
};
}
impl_tuple_output!(A: 0, B: 1);
impl_tuple_output!(A: 0, B: 1, C: 2);
impl_tuple_output!(A: 0, B: 1, C: 2, D: 3);
impl_tuple_output!(A: 0, B: 1, C: 2, D: 3, E: 4);
impl_tuple_output!(A: 0, B: 1, C: 2, D: 3, E: 4, F: 5);
pub fn typed_expr<T: TraceTensor>(value: T) -> TypedExpr {
TypedExpr {
kind: value.expr().clone(),
ty: T::tensor_type(),
}
}
pub(crate) fn tuple_id() -> u64 {
next_tuple_id()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::trace::*;
use knok_core::ElementType;
#[test]
fn traces_helper_loop_to_expression_tree() {
fn block(x: T2<f32, 4, 4>) -> T2<f32, 4, 4> {
relu(matmul(x.clone(), x) + 1.0)
}
let mut context = TraceContext::default();
let x = context.input::<T2<f32, 4, 4>>("x");
let mut y = x;
for _ in 0..2 {
y = block(y);
}
let graph = context.finish("forward", "llvm-cpu", y).unwrap();
assert_eq!(graph.inputs.len(), 1);
assert_eq!(graph.outputs[0].shape, vec![4, 4]);
assert_eq!(graph.outputs[0].elem, ElementType::F32);
}
#[test]
fn split_traces_as_tuple_projections_without_context_side_effects() {
let mut context = TraceContext::default();
let x = context.input::<T1<f32, 4>>("x");
let output: (T1<f32, 2>, T1<f32, 2>) = split(x, 0, [2, 2]);
let graph = context.finish("split", "llvm-cpu", output).unwrap();
assert!(graph.lets.is_empty());
assert_eq!(graph.outputs, vec![T1::<f32, 2>::tensor_type(); 2]);
}
#[test]
fn generated_node_bindings_do_not_shadow_inputs() {
let input_name = "__knok_node_7".to_string();
let body = vec![Expr::Node {
node_id: 7,
value: std::sync::Arc::new(Expr::Var(input_name.clone())),
}];
let reserved_names = BTreeSet::from([input_name.clone()]);
let (lets, body) = let_bind_nodes(body, reserved_names);
assert_eq!(lets[0].names, vec!["__knok_node_7_1"]);
assert_eq!(lets[0].value, Expr::Var(input_name));
assert_eq!(body, vec![Expr::Var("__knok_node_7_1".into())]);
}
#[test]
fn traces_elementwise_creation_and_predicate_ops() {
let mut context = TraceContext::default();
let x = context.input::<T1<f32, 4>>("x");
let y = context.input::<T1<f32, 4>>("y");
let mask = context.input::<T1<bool, 4>>("mask");
let _: T1<f32, 4> = abs(x.clone());
let _: T1<f32, 4> = ceil(x.clone());
let _: T1<f32, 4> = exp(x.clone());
let _: T1<f32, 4> = exp2(x.clone());
let _: T1<f32, 4> = expm1(x.clone());
let _: T1<f32, 4> = floor(x.clone());
let _: T1<f32, 4> = log(x.clone());
let _: T1<f32, 4> = log1p(x.clone());
let _: T1<f32, 4> = log2(x.clone());
let _: T1<f32, 4> = log10(x.clone());
let _: T1<f32, 4> = rint(x.clone());
let _: T1<f32, 4> = round(x.clone());
let _: T1<f32, 4> = sigmoid(x.clone());
let _: T1<f32, 4> = sin(x.clone());
let _: T1<f32, 4> = cos(x.clone());
let _: T1<f32, 4> = sqrt(x.clone());
let _: T1<f32, 4> = tan(x.clone());
let _: T1<f32, 4> = tanh(x.clone());
let _: T1<f32, 4> = square(x.clone());
let _: T1<f32, 4> = reciprocal(x.clone());
let _: T1<bool, 4> = isnan(x.clone());
let _: T1<f32, 4> = zeros_like(x.clone());
let _: T1<f32, 4> = ones_like(x.clone());
let _: T1<f32, 4> = minimum(x.clone(), y.clone());
let _: T1<f32, 4> = pow(x.clone(), 2.0);
let _: T1<bool, 4> = greater_equal(x.clone(), y.clone());
let _: T1<bool, 4> = less(x.clone(), y.clone());
let _: T1<bool, 4> = equal(x.clone(), y.clone());
let _: T1<bool, 4> = not_equal(x.clone(), y.clone());
let _: T1<bool, 4> = logical_or(mask.clone(), false);
let _: T1<bool, 4> = logical_xor(mask.clone(), true);
let _: T1<i32, 4> = arange_to(4);
let _: T1<i32, 4> = arange(0, 4);
let _: T2<f32, 2, 2> = identity();
let predicate = logical_and(greater(x.clone(), 0.0), less_equal(y.clone(), 6.0));
let selected: T1<f32, 4> = r#where(
predicate.clone(),
clip(maximum(x.clone(), y.clone()), 0.0, 6.0),
full_like(x.clone(), 1.0),
);
let ints: T1<i32, 4> = arange_step(0, 8, 2);
let line: T1<f32, 4> = linspace(0.0, 1.0);
let eye_matrix: T2<f32, 2, 2> = eye();
let graph = context
.finish(
"elementwise_creation",
"llvm-cpu",
(
relu(selected),
logical_not(predicate),
ints,
line,
eye_matrix,
ones_like(x),
),
)
.unwrap();
assert_eq!(graph.outputs.len(), 6);
assert_eq!(graph.outputs[0], T1::<f32, 4>::tensor_type());
assert_eq!(graph.outputs[1], T1::<bool, 4>::tensor_type());
assert_eq!(graph.outputs[2].elem, ElementType::I32);
}
#[test]
fn traces_reduction_linalg_and_conv_ops() {
let mut context = TraceContext::default();
let matrix = context.input::<T2<f32, 2, 3>>("matrix");
let rhs = context.input::<T2<f32, 3, 2>>("rhs");
let vector = context.input::<T1<f32, 3>>("vector");
let square = context.input::<T2<f32, 3, 3>>("square");
let flags = context.input::<T2<bool, 2, 3>>("flags");
let image = context.input::<T4<f32, 1, 4, 4, 2>>("image");
let kernel = context.input::<T4<f32, 3, 3, 2, 1>>("kernel");
let _: T0<f32> = sum(matrix.clone());
let _: T0<f32> = prod(matrix.clone());
let _: T0<f32> = mean(matrix.clone());
let _: T0<f32> = max(matrix.clone());
let _: T0<f32> = amax(matrix.clone());
let _: T0<f32> = min(matrix.clone());
let _: T0<f32> = amin(matrix.clone());
let _: T0<i64> = argmax(matrix.clone());
let _: T0<i64> = argmin(matrix.clone());
let _: T0<f32> = var(matrix.clone());
let _: T0<f32> = std(matrix.clone());
let _: T0<f32> = ptp(matrix.clone());
let _: T1<f32, 2> = prod_axis(matrix.clone(), 1);
let _: T1<f32, 2> = mean_axis(matrix.clone(), 1);
let _: T1<f32, 2> = max_axis(matrix.clone(), 1);
let _: T1<f32, 2> = amax_axis(matrix.clone(), 1);
let _: T1<f32, 2> = min_axis(matrix.clone(), 1);
let _: T1<f32, 2> = amin_axis(matrix.clone(), 1);
let _: T1<f32, 2> = var_axis(matrix.clone(), 1);
let _: T1<f32, 2> = std_axis(matrix.clone(), 1);
let _: T1<f32, 2> = ptp_axis(matrix.clone(), 1);
let _: T0<bool> = all(flags.clone());
let _: T0<bool> = any(flags.clone());
let _: T1<bool, 2> = all_axis(flags.clone(), 1);
let _: T1<bool, 2> = any_axis(flags.clone(), 1);
let _: T1<i64, 2> = argmin_axis(matrix.clone(), 1);
let _: T2<f32, 2, 3> = softmax(matrix.clone());
let _: T0<f32> = dot(vector.clone(), vector.clone());
let _: T0<f32> = inner(vector.clone(), vector.clone());
let _: T1<f32, 2> = vecdot_axis(matrix.clone(), matrix.clone(), 1);
let _: T0<f32> = trace(square.clone());
let _: T0<f32> = trace_axes(square.clone(), 0, 1);
let _: T1<f32, 3> = diagonal_axes(square.clone(), 0, 1);
let _: T4<f32, 1, 4, 4, 1> = conv2d_options(
image.clone(),
kernel.clone(),
Conv2dOptions::new()
.padding(1, 1, 1, 1)
.stride(1, 1)
.dilation(1, 1)
.groups(1),
);
let product: T2<f32, 2, 2> = matmul(matrix.clone(), rhs);
let row_sum: T1<f32, 2> = sum_axis(matrix.clone(), 1);
let row_argmax: T1<i64, 2> = argmax_axis(matrix.clone(), 1);
let outer_product: T2<f32, 3, 3> = outer(vector.clone(), vector);
let diagonal_values: T1<f32, 3> = diagonal(square);
let conv: T4<f32, 1, 2, 2, 1> = conv2d(image, kernel);
let graph = context
.finish(
"reduction_linalg_conv",
"llvm-cpu",
(
product,
row_sum,
row_argmax,
outer_product,
diagonal_values,
conv,
),
)
.unwrap();
assert_eq!(graph.outputs.len(), 6);
assert_eq!(graph.outputs[0], T2::<f32, 2, 2>::tensor_type());
assert_eq!(graph.outputs[2], T1::<i64, 2>::tensor_type());
assert_eq!(graph.outputs[5], T4::<f32, 1, 2, 2, 1>::tensor_type());
}
#[test]
fn traces_shape_and_indexing_ops() {
let mut context = TraceContext::default();
let x = context.input::<T2<f32, 2, 3>>("x");
let y = context.input::<T2<f32, 2, 3>>("y");
let row = context.input::<T1<f32, 3>>("row");
let indices = context.input::<T2<i64, 2, 2>>("indices");
let cube = context.input::<T3<f32, 2, 3, 4>>("cube");
let _: T2<f32, 2, 3> = broadcast(row.clone());
let _: T1<f32, 3> = squeeze(unsqueeze::<T2<f32, 1, 3>>(row.clone()));
let _: T2<f32, 4, 5> = pad(x.clone(), [1, 1]);
let _: T1<f32, 3> = take(x.clone(), 0, 0);
let _: T2<f32, 2, 2> = take_along_axis(x.clone(), indices.clone(), 1);
let _: T2<f32, 2, 6> = repeat(x.clone(), 1, 2);
let _: T2<f32, 2, 3> = flip(x.clone());
let _: T2<f32, 2, 3> = flip_axes(x.clone(), [1]);
let _: T2<f32, 2, 3> = roll(x.clone(), 1, 1);
let _: T2<f32, 3, 2> = transpose_axes(x.clone(), [1, 0]);
let _: T3<f32, 4, 2, 3> = permute(cube.clone(), [2, 0, 1]);
let _: T3<f32, 4, 2, 3> = permute_dims(cube.clone(), [2, 0, 1]);
let _: T3<f32, 2, 4, 3> = swapaxes(cube.clone(), 1, 2);
let _: T3<f32, 3, 2, 4> = moveaxis(cube, 1, 0);
let reshaped: T1<f32, 6> = reshape(x.clone());
let transposed: T2<f32, 3, 2> = transpose(x.clone());
let gathered: T3<f32, 2, 2, 2> = gather(x.clone(), indices, 1);
let sliced: T2<f32, 2, 2> = slice(x.clone(), [0, 1]);
let concatenated: T2<f32, 4, 3> = concat(x.clone(), y.clone(), 0);
let stacked: T3<f32, 2, 2, 3> = stack(x, y, 0);
let tiled: T2<f32, 4, 6> = tile(row, [4, 2]);
let graph = context
.finish(
"shape_indexing",
"llvm-cpu",
(
reshaped,
transposed,
gathered,
sliced,
concatenated,
stacked,
),
)
.unwrap();
assert_eq!(graph.outputs.len(), 6);
assert_eq!(graph.outputs[0], T1::<f32, 6>::tensor_type());
assert_eq!(graph.outputs[2], T3::<f32, 2, 2, 2>::tensor_type());
assert_eq!(tiled.expr().clone(), tiled.into_expr());
}
}