diff --git a/Cargo.toml b/Cargo.toml index f5f17abf..aa6c20d0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "blade-graphics", "blade-helpers", "blade-macros", + "blade-neural", "blade-particle", "blade-render", "blade-util", @@ -22,6 +23,7 @@ blade-util = { path = "blade-util", version = "0.4.1" } blade-egui = { path = "blade-egui", version = "0.8.1" } blade-particle = { path = "blade-particle", version = "0.1" } blade-macros = { path = "blade-macros", version = "0.3" } +blade-neural = { path = "blade-neural", version = "0.1" } blade-helpers = { path = "blade-helpers", version = "0.2" } blade-asset = { path = "blade-asset", version = "0.2.1" } blade-render = { path = "blade-render", version = "0.5" } diff --git a/blade-neural/Cargo.toml b/blade-neural/Cargo.toml new file mode 100644 index 00000000..7b4be4b4 --- /dev/null +++ b/blade-neural/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "blade-neural" +version = "0.1.0" +edition = "2024" +rust-version = "1.92" +description = "Neural processing (NPU) abstraction for Blade" +keywords = ["neural", "npu", "inference"] +license = "MIT" +repository = "https://github.com/kvark/blade" + +[lib] + +[features] + +[dependencies] +bytemuck = { workspace = true } +log = { workspace = true } +profiling = { workspace = true } + +[package.metadata.cargo_check_external_types] +allowed_external_types = ["bytemuck::*"] + +[lints.rust] +unexpected_cfgs = { level = "allow", check-cfg = [ + 'cfg(neural_cpu)', + 'cfg(neural_coreml)', + 'cfg(neural_directml)', + 'cfg(neural_openvino)', +] } diff --git a/blade-neural/README.md b/blade-neural/README.md new file mode 100644 index 00000000..2d606025 --- /dev/null +++ b/blade-neural/README.md @@ -0,0 +1,72 @@ +# Blade Neural + +[![Crates.io](https://img.shields.io/crates/v/blade-neural.svg?maxAge=2592000)](https://crates.io/crates/blade-neural) + +Blade Neural is a lean, opinionated **NPU abstraction** in the same spirit as +[`blade-graphics`](../blade-graphics) is for GPUs: minimal, ergonomic, and +backed by native platform APIs selected at compile time. + +## Why an op-graph, not a command buffer + +Unlike GPUs, NPUs expose **no portable low-level command-buffer surface**. Every +vendor only exposes the accelerator through a *graph/model execution* API: + +| Platform | Native op-graph surface | +| -------- | ----------------------- | +| Apple | Core ML / BNNSGraph / MPSGraph (direct ANE access is private) | +| Windows | DirectML → Windows ML / ONNX EPs | +| Intel | OpenVINO (+ NPU plugin) | +| Qualcomm | QNN | +| Android | NNAPI | +| Web | WebNN (`MLGraphBuilder`) | + +The portable lowest common denominator that is genuinely programmable across all +of them is an **op-graph builder**: describe a graph of primitive ops, compile +it, bind tensors, and execute. Blade Neural mirrors exactly that, so native +backends can lower onto each vendor API naturally. + +## Usage + +```rust +use blade_neural as bn; + +let ctx = bn::Context::init(bn::ContextDesc::default()).unwrap(); + +let mut b = bn::GraphBuilder::new(); +let x = b.input("x", bn::TensorDesc::new(bn::DataType::F32, &[1, 3])); +let w = b.constant(/* weights */); +let y = b.relu(b.matmul(x, w)); +b.output("y", y); + +let graph = ctx.compile_graph(&b).unwrap(); +ctx.run(&graph, &[("x", &input)], &[("y", &output)]); +``` + +See [`examples/mlp.rs`](examples/mlp.rs) for a complete 2-layer network: + +```bash +cargo run -p blade-neural --example mlp +``` + +## Backends + +The backend is selected automatically by the host platform, like +`blade-graphics`. + +| Backend | Status | +| ------- | ------ | +| CPU reference (pure Rust) | ✅ available everywhere; the correctness oracle | +| MPSGraph (Apple) | planned | +| DirectML (Windows) | planned | +| OpenVINO (Intel NPU) | planned | +| QNN (Qualcomm) | planned | + +The MVP op set is `input`, `constant`, `matmul`, `add`, `mul`, `relu`, +`softmax`, and `reshape` — enough to express an MLP classifier. Convolution, +pooling, normalization, attention, quantized paths, and `blade-graphics` tensor +interop are future work. + +## Scope + +Inference only (no autodiff/training) for now. The crate is `f32`-first; other +data types are reserved in the API and accepted by future native backends. diff --git a/blade-neural/examples/mlp.rs b/blade-neural/examples/mlp.rs new file mode 100644 index 00000000..a145d994 --- /dev/null +++ b/blade-neural/examples/mlp.rs @@ -0,0 +1,44 @@ +//! A tiny 2-layer MLP forward pass with baked-in weights, demonstrating the +//! full builder -> compile -> run flow of `blade-neural`. + +use blade_neural as bn; + +fn desc(shape: &[usize]) -> bn::TensorDesc { + bn::TensorDesc::new(bn::DataType::F32, shape) +} + +fn main() { + let ctx = bn::Context::init(bn::ContextDesc::default()).unwrap(); + + // Network: input[1,4] -> dense(4->3) + relu -> dense(3->2) -> softmax. + let b = bn::GraphBuilder::new(); + let x = b.input("x", desc(&[1, 4])); + + let w1 = b.constant( + desc(&[4, 3]), + bytemuck::cast_slice(&[ + 0.1f32, 0.2, -0.1, 0.0, 0.3, 0.1, -0.2, 0.1, 0.4, 0.5, -0.3, 0.2, + ]), + ); + let b1 = b.constant(desc(&[3]), bytemuck::cast_slice(&[0.1f32, -0.1, 0.0])); + let h = b.relu(b.add(b.matmul(x, w1), b1)); + + let w2 = b.constant( + desc(&[3, 2]), + bytemuck::cast_slice(&[0.3f32, -0.2, 0.1, 0.4, -0.5, 0.2]), + ); + let b2 = b.constant(desc(&[2]), bytemuck::cast_slice(&[0.0f32, 0.1])); + let y = b.softmax(b.add(b.matmul(h, w2), b2), 1); + b.output("y", y); + + let graph = ctx.compile_graph(&b).unwrap(); + + let input = ctx.create_tensor(&desc(&[1, 4])); + input.write_f32(&[1.0, 0.5, -0.5, 2.0]); + let output = ctx.create_tensor(&desc(&[1, 2])); + ctx.run(&graph, &[("x", &input)], &[("y", &output)]); + + let probs = output.read_f32(); + println!("class probabilities: {:?}", probs); + println!("sum = {:.6} (should be ~1.0)", probs.iter().sum::()); +} diff --git a/blade-neural/src/cpu/mod.rs b/blade-neural/src/cpu/mod.rs new file mode 100644 index 00000000..35298fe4 --- /dev/null +++ b/blade-neural/src/cpu/mod.rs @@ -0,0 +1,184 @@ +//! Pure-Rust reference backend. +//! +//! Always available, needs no drivers, and serves as the correctness oracle for +//! future native (NPU) backends. Tensors live in host memory and graphs are +//! evaluated by walking the IR in topological order. + +mod ops; + +use std::sync::{Arc, Mutex}; + +use crate::{ + CompileError, ContextDesc, DataType, GraphBuilder, GraphIr, NotSupportedError, Op, TensorDesc, +}; + +/// A neural context backed by the host CPU. +#[derive(Debug)] +pub struct Context { + validation: bool, +} + +/// Host-resident tensor storage. Cheap to clone (shares the underlying buffer). +#[derive(Clone, Debug)] +pub struct Tensor { + desc: TensorDesc, + data: Arc>>, +} + +impl Tensor { + /// The tensor's element type and shape. + pub fn desc(&self) -> &TensorDesc { + &self.desc + } + /// Overwrite the tensor's raw bytes. + pub fn write(&self, bytes: &[u8]) { + let mut data = self.data.lock().unwrap(); + assert_eq!(bytes.len(), data.len(), "tensor write size mismatch"); + data.copy_from_slice(bytes); + } + /// Overwrite the tensor with `f32` values. + pub fn write_f32(&self, values: &[f32]) { + self.write(bytemuck::cast_slice(values)); + } + /// Read back the tensor's raw bytes. + pub fn read(&self) -> Vec { + self.data.lock().unwrap().clone() + } + /// Read back the tensor as `f32` values. + pub fn read_f32(&self) -> Vec { + bytemuck::cast_slice(&self.read()).to_vec() + } +} + +/// A graph compiled for the reference backend (just the validated IR). +#[derive(Debug)] +pub struct Graph { + ir: GraphIr, +} + +impl Context { + /// Initialize the reference backend. Always succeeds. + pub fn init(desc: ContextDesc) -> Result { + Ok(Self { + validation: desc.validation, + }) + } + + /// Compile a recorded graph. + pub fn compile_graph(&self, builder: &GraphBuilder) -> Result { + profiling::scope!("compile_graph"); + let ir = builder.ir(); + // The reference backend only computes in f32. + for node in &ir.nodes { + if node.desc.dtype != DataType::F32 { + return Err(CompileError::UnsupportedDataType(node.desc.dtype)); + } + } + for (name, operand) in &ir.outputs { + if operand.0 as usize >= ir.nodes.len() { + return Err(CompileError::UnknownOutput(name.clone())); + } + } + Ok(Graph { + ir: GraphIr::clone(&ir), + }) + } + + /// Allocate a zeroed tensor. + pub fn create_tensor(&self, desc: &TensorDesc) -> Tensor { + Tensor { + desc: desc.clone(), + data: Arc::new(Mutex::new(vec![0u8; desc.byte_size()])), + } + } + + /// Execute `graph`, reading `inputs` and writing `outputs`, matched by name. + pub fn run(&self, graph: &Graph, inputs: &[(&str, &Tensor)], outputs: &[(&str, &Tensor)]) { + profiling::scope!("run"); + let nodes = &graph.ir.nodes; + let mut values: Vec> = Vec::with_capacity(nodes.len()); + + for node in nodes { + let value = match node.op { + Op::Input { ref name } => { + let (_, tensor) = inputs + .iter() + .find(|(n, _)| *n == name.as_str()) + .unwrap_or_else(|| panic!("missing input binding {:?}", name)); + if self.validation { + assert_eq!( + &tensor.desc.shape, &node.desc.shape, + "input {:?} shape mismatch", + name + ); + } + tensor.read_f32() + } + Op::Constant { ref data } => bytemuck::cast_slice::(data).to_vec(), + Op::MatMul { a, b } => { + let sa = &nodes[a.0 as usize].desc.shape; + let sb = &nodes[b.0 as usize].desc.shape; + ops::matmul( + &values[a.0 as usize], + &values[b.0 as usize], + sa[0], + sa[1], + sb[1], + ) + } + Op::Add { a, b } => ops::broadcast_binary( + &values[a.0 as usize], + &nodes[a.0 as usize].desc.shape, + &values[b.0 as usize], + &nodes[b.0 as usize].desc.shape, + |x, y| x + y, + ), + Op::Mul { a, b } => ops::broadcast_binary( + &values[a.0 as usize], + &nodes[a.0 as usize].desc.shape, + &values[b.0 as usize], + &nodes[b.0 as usize].desc.shape, + |x, y| x * y, + ), + Op::Relu { x } => { + let mut v = values[x.0 as usize].clone(); + ops::relu(&mut v); + v + } + Op::Softmax { x, axis } => { + ops::softmax(&values[x.0 as usize], &nodes[x.0 as usize].desc.shape, axis) + } + Op::Reshape { x } => values[x.0 as usize].clone(), + }; + values.push(value); + } + + for (name, operand) in &graph.ir.outputs { + let (_, tensor) = outputs + .iter() + .find(|(n, _)| *n == name.as_str()) + .unwrap_or_else(|| panic!("missing output binding {:?}", name)); + tensor.write_f32(&values[operand.0 as usize]); + } + } +} + +impl crate::traits::NeuralDevice for Context { + type Graph = Graph; + type Tensor = Tensor; + + fn compile_graph(&self, builder: &GraphBuilder) -> Result { + Context::compile_graph(self, builder) + } + fn create_tensor(&self, desc: &TensorDesc) -> Self::Tensor { + Context::create_tensor(self, desc) + } + fn run( + &self, + graph: &Self::Graph, + inputs: &[(&str, &Self::Tensor)], + outputs: &[(&str, &Self::Tensor)], + ) { + Context::run(self, graph, inputs, outputs) + } +} diff --git a/blade-neural/src/cpu/ops.rs b/blade-neural/src/cpu/ops.rs new file mode 100644 index 00000000..d3e91d9d --- /dev/null +++ b/blade-neural/src/cpu/ops.rs @@ -0,0 +1,123 @@ +//! Reference kernels for the pure-Rust backend. These operate on flat `f32` +//! slices and prioritize obvious correctness over speed — they are the oracle +//! that hardware backends are checked against. + +use crate::ir::broadcast_shape; + +/// 2D matrix multiply: `a` is `[m, k]`, `b` is `[k, n]`, output is `[m, n]`. +pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + let mut out = vec![0.0f32; m * n]; + for row in 0..m { + for inner in 0..k { + let av = a[row * k + inner]; + if av == 0.0 { + continue; + } + let b_row = &b[inner * n..inner * n + n]; + let out_row = &mut out[row * n..row * n + n]; + for col in 0..n { + out_row[col] += av * b_row[col]; + } + } + } + out +} + +/// Element-wise op with numpy-style broadcasting from shapes `sa`/`sb`. +pub fn broadcast_binary( + a: &[f32], + sa: &[usize], + b: &[f32], + sb: &[usize], + f: impl Fn(f32, f32) -> f32, +) -> Vec { + let out_shape = broadcast_shape(sa, sb); + let count: usize = out_shape.iter().product(); + let stride_a = broadcast_strides(sa, &out_shape); + let stride_b = broadcast_strides(sb, &out_shape); + + let mut out = vec![0.0f32; count]; + let mut index = vec![0usize; out_shape.len()]; + for o in out.iter_mut() { + let mut ia = 0; + let mut ib = 0; + for (d, &i) in index.iter().enumerate() { + ia += i * stride_a[d]; + ib += i * stride_b[d]; + } + *o = f(a[ia], b[ib]); + increment(&mut index, &out_shape); + } + out +} + +/// Strides into a source of shape `src` when iterating over `out`, with 0 stride +/// on broadcast (size-1 or missing) dimensions. +fn broadcast_strides(src: &[usize], out: &[usize]) -> Vec { + let rank = out.len(); + let offset = rank - src.len(); + // Row-major strides of the source. + let mut src_strides = vec![0usize; src.len()]; + let mut acc = 1; + for i in (0..src.len()).rev() { + src_strides[i] = acc; + acc *= src[i]; + } + let mut strides = vec![0usize; rank]; + for (d, stride) in strides.iter_mut().enumerate() { + if d >= offset { + let sd = d - offset; + *stride = if src[sd] == 1 { 0 } else { src_strides[sd] }; + } + } + strides +} + +/// Advance a multi-dimensional index in row-major order. +fn increment(index: &mut [usize], shape: &[usize]) { + for d in (0..shape.len()).rev() { + index[d] += 1; + if index[d] < shape[d] { + return; + } + index[d] = 0; + } +} + +/// Element-wise `max(x, 0)`, in place. +pub fn relu(x: &mut [f32]) { + for v in x.iter_mut() { + if *v < 0.0 { + *v = 0.0; + } + } +} + +/// Numerically-stable softmax along `axis` of a tensor of shape `shape`. +pub fn softmax(x: &[f32], shape: &[usize], axis: usize) -> Vec { + let axis_len = shape[axis]; + // Stride between successive elements along `axis`. + let inner: usize = shape[axis + 1..].iter().product(); + let outer: usize = shape[..axis].iter().product(); + + let mut out = vec![0.0f32; x.len()]; + for o in 0..outer { + for i in 0..inner { + let base = o * axis_len * inner + i; + let mut max = f32::NEG_INFINITY; + for a in 0..axis_len { + max = max.max(x[base + a * inner]); + } + let mut sum = 0.0f32; + for a in 0..axis_len { + let e = (x[base + a * inner] - max).exp(); + out[base + a * inner] = e; + sum += e; + } + for a in 0..axis_len { + out[base + a * inner] /= sum; + } + } + } + out +} diff --git a/blade-neural/src/ir.rs b/blade-neural/src/ir.rs new file mode 100644 index 00000000..40db1752 --- /dev/null +++ b/blade-neural/src/ir.rs @@ -0,0 +1,238 @@ +//! The backend-agnostic op-graph intermediate representation and its builder. +//! +//! This is the shared IR that every backend lowers from, analogous to how +//! `blade-graphics` shares a `naga::Module` across its shader backends. + +use std::cell::{Ref, RefCell}; + +use crate::{Operand, TensorDesc}; + +/// A single operation in the graph. +/// +/// Each variant carries the [`Operand`]s it reads. Operands always refer to +/// *earlier* nodes, so the natural index order of [`GraphIr::nodes`] is a valid +/// topological order. +#[derive(Clone, Debug)] +pub enum Op { + /// A named graph input, bound at execution time. + Input { name: String }, + /// A baked-in constant (e.g. weights). + Constant { data: Box<[u8]> }, + /// 2D matrix multiply: `[M, K] x [K, N] -> [M, N]`. + MatMul { a: Operand, b: Operand }, + /// Element-wise addition with numpy-style broadcasting. + Add { a: Operand, b: Operand }, + /// Element-wise multiplication with numpy-style broadcasting. + Mul { a: Operand, b: Operand }, + /// Element-wise `max(x, 0)`. + Relu { x: Operand }, + /// Softmax along `axis`. + Softmax { x: Operand, axis: usize }, + /// Reinterpret the same data under a new shape. + Reshape { x: Operand }, +} + +impl Op { + /// Stable name for diagnostics. + pub fn name(&self) -> &'static str { + match *self { + Self::Input { .. } => "input", + Self::Constant { .. } => "constant", + Self::MatMul { .. } => "matmul", + Self::Add { .. } => "add", + Self::Mul { .. } => "mul", + Self::Relu { .. } => "relu", + Self::Softmax { .. } => "softmax", + Self::Reshape { .. } => "reshape", + } + } +} + +/// A node produces exactly one [`Operand`], whose id equals the node's index. +#[derive(Clone, Debug)] +pub struct Node { + pub op: Op, + /// Description of the operand this node produces (filled in by shape inference). + pub desc: TensorDesc, +} + +/// The compiled-into-able description of a graph: a topologically ordered list +/// of nodes plus the named outputs to read back. +#[derive(Clone, Debug, Default)] +pub struct GraphIr { + pub nodes: Vec, + pub outputs: Vec<(String, Operand)>, +} + +/// Records operations into a [`GraphIr`]. Backend-agnostic and side-effect free +/// until handed to a backend for compilation. +/// +/// Methods take `&self` (via interior mutability) so graphs read naturally as +/// nested expressions, e.g. `b.relu(b.add(b.matmul(x, w), bias))`. +/// +/// Shape and type mismatches are programming errors and panic with a descriptive +/// message, in keeping with Blade's "you know what you are doing" stance. +#[derive(Debug, Default)] +pub struct GraphBuilder { + ir: RefCell, +} + +impl GraphBuilder { + pub fn new() -> Self { + Self::default() + } + + /// Borrow the recorded IR (used by backends). + pub fn ir(&self) -> Ref<'_, GraphIr> { + self.ir.borrow() + } + + /// The description of an existing operand. + pub fn desc(&self, op: Operand) -> TensorDesc { + self.ir.borrow().nodes[op.0 as usize].desc.clone() + } + + fn push(&self, op: Op, desc: TensorDesc) -> Operand { + let mut ir = self.ir.borrow_mut(); + let id = ir.nodes.len() as u32; + ir.nodes.push(Node { op, desc }); + Operand(id) + } + + /// Declare a named graph input. + pub fn input(&self, name: &str, desc: TensorDesc) -> Operand { + self.push( + Op::Input { + name: name.to_string(), + }, + desc, + ) + } + + /// Bake a constant tensor (e.g. weights) into the graph. + pub fn constant(&self, desc: TensorDesc, data: &[u8]) -> Operand { + assert_eq!( + data.len(), + desc.byte_size(), + "constant data is {} bytes, but its shape needs {}", + data.len(), + desc.byte_size() + ); + self.push(Op::Constant { data: data.into() }, desc) + } + + /// 2D matrix multiply: `[M, K] x [K, N] -> [M, N]`. + pub fn matmul(&self, a: Operand, b: Operand) -> Operand { + let da = self.desc(a); + let db = self.desc(b); + assert!( + da.shape.len() == 2 && db.shape.len() == 2, + "matmul expects 2D operands, got {:?} and {:?}", + da.shape, + db.shape + ); + assert_eq!( + da.shape[1], db.shape[0], + "matmul inner dimensions disagree: {:?} vs {:?}", + da.shape, db.shape + ); + let shape = [da.shape[0], db.shape[1]]; + self.push(Op::MatMul { a, b }, TensorDesc::new(da.dtype, &shape)) + } + + /// Element-wise addition with numpy-style broadcasting. + pub fn add(&self, a: Operand, b: Operand) -> Operand { + let desc = self.broadcast(a, b); + self.push(Op::Add { a, b }, desc) + } + + /// Element-wise multiplication with numpy-style broadcasting. + pub fn mul(&self, a: Operand, b: Operand) -> Operand { + let desc = self.broadcast(a, b); + self.push(Op::Mul { a, b }, desc) + } + + /// Element-wise `max(x, 0)`. + pub fn relu(&self, x: Operand) -> Operand { + let desc = self.desc(x); + self.push(Op::Relu { x }, desc) + } + + /// Softmax normalization along `axis`. + pub fn softmax(&self, x: Operand, axis: usize) -> Operand { + let desc = self.desc(x); + assert!( + axis < desc.shape.len(), + "softmax axis {} out of range for shape {:?}", + axis, + desc.shape + ); + self.push(Op::Softmax { x, axis }, desc) + } + + /// Reinterpret a tensor under a new shape with the same element count. + pub fn reshape(&self, x: Operand, shape: &[usize]) -> Operand { + let src = self.desc(x); + let new_count: usize = shape.iter().product(); + assert_eq!( + src.element_count(), + new_count, + "reshape from {:?} to {:?} changes the element count", + src.shape, + shape + ); + let desc = TensorDesc::new(src.dtype, shape); + self.push(Op::Reshape { x }, desc) + } + + /// Mark an operand as a named output of the graph. + pub fn output(&self, name: &str, value: Operand) { + self.ir.borrow_mut().outputs.push((name.to_string(), value)); + } + + /// Compute the broadcasted output description of an element-wise op. + fn broadcast(&self, a: Operand, b: Operand) -> TensorDesc { + let da = self.desc(a); + let db = self.desc(b); + assert_eq!( + da.dtype, db.dtype, + "element-wise operands disagree on type: {:?} vs {:?}", + da.dtype, db.dtype + ); + let shape = broadcast_shape(&da.shape, &db.shape); + TensorDesc { + dtype: da.dtype, + shape, + } + } +} + +/// Compute the numpy-style broadcasted shape of two shapes, or panic if they are +/// not broadcast-compatible. +pub(crate) fn broadcast_shape(a: &[usize], b: &[usize]) -> Box<[usize]> { + let rank = a.len().max(b.len()); + let mut out = vec![0usize; rank]; + for i in 0..rank { + // Align from the right; missing leading dims act as 1. + let da = if i < rank - a.len() { + 1 + } else { + a[i - (rank - a.len())] + }; + let db = if i < rank - b.len() { + 1 + } else { + b[i - (rank - b.len())] + }; + out[i] = if da == db { + da + } else if da == 1 { + db + } else if db == 1 { + da + } else { + panic!("shapes {:?} and {:?} are not broadcast-compatible", a, b); + }; + } + out.into_boxed_slice() +} diff --git a/blade-neural/src/lib.rs b/blade-neural/src/lib.rs new file mode 100644 index 00000000..4594a0e7 --- /dev/null +++ b/blade-neural/src/lib.rs @@ -0,0 +1,162 @@ +#![allow( + // We don't use syntax sugar where it's not necessary. + clippy::match_like_matches_macro, + // Redundant matching is more explicit. + clippy::redundant_pattern_matching, + // Explicit lifetimes are often easier to reason about. + clippy::needless_lifetimes, + // No need for defaults in the internal types. + clippy::new_without_default, + // Matches are good and extendable, no need to make an exception here. + clippy::single_match, +)] +#![warn(trivial_numeric_casts, unused_extern_crates)] + +//! Blade Neural is a lean NPU abstraction in the spirit of `blade-graphics`. +//! +//! Unlike GPUs, NPUs are not exposed through a portable low-level command +//! buffer. Every vendor surface (Apple MPSGraph/Core ML, Windows DirectML, +//! Intel OpenVINO, Qualcomm QNN, the web's WebNN) is an *op-graph* API: you +//! describe a graph of primitive operations, compile it, bind tensors, and +//! execute. Blade Neural mirrors that lowest common denominator. +//! +//! ``` +//! use blade_neural as bn; +//! let context = bn::Context::init(bn::ContextDesc::default()).unwrap(); +//! +//! let mut builder = bn::GraphBuilder::new(); +//! let x = builder.input("x", bn::TensorDesc::new(bn::DataType::F32, &[1, 3])); +//! let w = builder.constant( +//! bn::TensorDesc::new(bn::DataType::F32, &[3, 2]), +//! bytemuck::cast_slice(&[1.0f32, 0.0, 0.0, 1.0, 1.0, 1.0]), +//! ); +//! let y = builder.relu(builder.matmul(x, w)); +//! builder.output("y", y); +//! +//! let graph = context.compile_graph(&builder).unwrap(); +//! let input = context.create_tensor(&bn::TensorDesc::new(bn::DataType::F32, &[1, 3])); +//! input.write_f32(&[1.0, 2.0, 3.0]); +//! let output = context.create_tensor(&bn::TensorDesc::new(bn::DataType::F32, &[1, 2])); +//! context.run(&graph, &[("x", &input)], &[("y", &output)]); +//! assert_eq!(output.read_f32(), vec![4.0, 5.0]); +//! ``` + +mod ir; +pub mod traits; + +// The backend is selected at compile time, mirroring `blade-graphics`. +// Only the pure-Rust reference backend is wired up today; the commented +// arms below show where native backends slot in as they land. +// +// #[cfg_attr(all(neural_coreml, any(target_os = "macos", target_os = "ios")), path = "coreml/mod.rs")] +// #[cfg_attr(all(neural_directml, target_os = "windows"), path = "directml/mod.rs")] +// #[cfg_attr(neural_openvino, path = "openvino/mod.rs")] +#[cfg_attr(all(), path = "cpu/mod.rs")] +mod hal; + +pub use hal::*; +pub use ir::{GraphBuilder, GraphIr, Node, Op}; + +use std::fmt; + +/// Numeric type of a tensor's elements. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum DataType { + F32, + /// Reserved for native backends; the reference backend does not compute in it yet. + F16, + I32, + I8, + U8, +} + +impl DataType { + /// Size of a single element, in bytes. + pub fn size(&self) -> usize { + match *self { + Self::F32 | Self::I32 => 4, + Self::F16 => 2, + Self::I8 | Self::U8 => 1, + } + } +} + +/// Description of a tensor: its element type and shape. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct TensorDesc { + pub dtype: DataType, + pub shape: Box<[usize]>, +} + +impl TensorDesc { + pub fn new(dtype: DataType, shape: &[usize]) -> Self { + Self { + dtype, + shape: shape.into(), + } + } + /// Number of elements across all dimensions. + pub fn element_count(&self) -> usize { + self.shape.iter().product() + } + /// Total size of the tensor data, in bytes. + pub fn byte_size(&self) -> usize { + self.element_count() * self.dtype.size() + } +} + +/// A handle to a value flowing through the graph, returned by [`GraphBuilder`] +/// methods. It is a small `Copy` token, like a resource handle in `blade-graphics`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct Operand(pub(crate) u32); + +/// Options for initializing a [`Context`]. +#[derive(Clone, Debug, Default)] +pub struct ContextDesc { + /// Enable extra validation of the graph and execution. + pub validation: bool, +} + +/// Error returned when the platform cannot provide a neural context. +#[derive(Debug)] +pub enum NotSupportedError { + PlatformNotSupported, + NoSupportedDeviceFound, +} + +impl fmt::Display for NotSupportedError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + Self::PlatformNotSupported => f.write_str("platform not supported"), + Self::NoSupportedDeviceFound => f.write_str("no supported device found"), + } + } +} + +impl std::error::Error for NotSupportedError {} + +/// Error returned when a graph cannot be compiled by a backend. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CompileError { + /// The graph refers to an output that no operation produces. + UnknownOutput(String), + /// A data type is not supported by this backend. + UnsupportedDataType(DataType), + /// An operation is not supported by this backend. + UnsupportedOp(&'static str), +} + +impl fmt::Display for CompileError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + Self::UnknownOutput(ref name) => write!(f, "unknown graph output {:?}", name), + Self::UnsupportedDataType(dt) => write!(f, "unsupported data type {:?}", dt), + Self::UnsupportedOp(op) => write!(f, "unsupported operation {:?}", op), + } + } +} + +impl std::error::Error for CompileError {} + +#[cfg(test)] +mod tests; diff --git a/blade-neural/src/tests.rs b/blade-neural/src/tests.rs new file mode 100644 index 00000000..aa7b1fdb --- /dev/null +++ b/blade-neural/src/tests.rs @@ -0,0 +1,110 @@ +use crate as bn; + +fn f32_desc(shape: &[usize]) -> bn::TensorDesc { + bn::TensorDesc::new(bn::DataType::F32, shape) +} + +fn context() -> bn::Context { + bn::Context::init(bn::ContextDesc { validation: true }).unwrap() +} + +#[test] +fn matmul_bias_relu() { + let ctx = context(); + let b = bn::GraphBuilder::new(); + let x = b.input("x", f32_desc(&[1, 3])); + // 3x2 weight matrix (column-major view: identity-ish). + let w = b.constant( + f32_desc(&[3, 2]), + bytemuck::cast_slice(&[1.0f32, 0.0, 0.0, 1.0, 1.0, 1.0]), + ); + let bias = b.constant(f32_desc(&[2]), bytemuck::cast_slice(&[-10.0f32, 1.0])); + let y = b.relu(b.add(b.matmul(x, w), bias)); + b.output("y", y); + + let graph = ctx.compile_graph(&b).unwrap(); + let xt = ctx.create_tensor(&f32_desc(&[1, 3])); + xt.write_f32(&[1.0, 2.0, 3.0]); + let yt = ctx.create_tensor(&f32_desc(&[1, 2])); + ctx.run(&graph, &[("x", &xt)], &[("y", &yt)]); + + // matmul: [1*1+2*0+3*1, 1*0+2*1+3*1] = [4, 5]; +bias = [-6, 6]; relu = [0, 6] + assert_eq!(yt.read_f32(), vec![0.0, 6.0]); +} + +#[test] +fn softmax_normalizes() { + let ctx = context(); + let b = bn::GraphBuilder::new(); + let x = b.input("x", f32_desc(&[1, 3])); + let y = b.softmax(x, 1); + b.output("y", y); + + let graph = ctx.compile_graph(&b).unwrap(); + let xt = ctx.create_tensor(&f32_desc(&[1, 3])); + xt.write_f32(&[1.0, 2.0, 3.0]); + let yt = ctx.create_tensor(&f32_desc(&[1, 3])); + ctx.run(&graph, &[("x", &xt)], &[("y", &yt)]); + + let out = yt.read_f32(); + let sum: f32 = out.iter().sum(); + assert!( + (sum - 1.0).abs() < 1e-6, + "softmax should sum to 1, got {sum}" + ); + assert!(out[0] < out[1] && out[1] < out[2], "monotonic: {out:?}"); +} + +#[test] +fn broadcast_add() { + let ctx = context(); + let b = bn::GraphBuilder::new(); + let x = b.input("x", f32_desc(&[2, 2])); + let row = b.constant(f32_desc(&[2]), bytemuck::cast_slice(&[10.0f32, 20.0])); + let y = b.add(x, row); + b.output("y", y); + + let graph = ctx.compile_graph(&b).unwrap(); + let xt = ctx.create_tensor(&f32_desc(&[2, 2])); + xt.write_f32(&[1.0, 2.0, 3.0, 4.0]); + let yt = ctx.create_tensor(&f32_desc(&[2, 2])); + ctx.run(&graph, &[("x", &xt)], &[("y", &yt)]); + assert_eq!(yt.read_f32(), vec![11.0, 22.0, 13.0, 24.0]); +} + +#[test] +fn reshape_preserves_data() { + let ctx = context(); + let b = bn::GraphBuilder::new(); + let x = b.input("x", f32_desc(&[2, 3])); + let y = b.reshape(x, &[3, 2]); + b.output("y", y); + + let graph = ctx.compile_graph(&b).unwrap(); + let xt = ctx.create_tensor(&f32_desc(&[2, 3])); + xt.write_f32(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let yt = ctx.create_tensor(&f32_desc(&[3, 2])); + ctx.run(&graph, &[("x", &xt)], &[("y", &yt)]); + assert_eq!(yt.read_f32(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); +} + +#[test] +fn rejects_non_f32() { + let ctx = context(); + let b = bn::GraphBuilder::new(); + let _ = b.input("x", bn::TensorDesc::new(bn::DataType::I32, &[2])); + let err = ctx.compile_graph(&b).unwrap_err(); + assert_eq!( + err, + bn::CompileError::UnsupportedDataType(bn::DataType::I32) + ); +} + +#[test] +#[should_panic(expected = "broadcast-compatible")] +fn incompatible_shapes_panic() { + let b = bn::GraphBuilder::new(); + let x = b.input("x", f32_desc(&[2, 3])); + let y = b.input("y", f32_desc(&[4, 5])); + let _ = b.add(x, y); +} diff --git a/blade-neural/src/traits.rs b/blade-neural/src/traits.rs new file mode 100644 index 00000000..ab49963e --- /dev/null +++ b/blade-neural/src/traits.rs @@ -0,0 +1,31 @@ +//! Abstract contract that every backend implements, mirroring +//! `blade-graphics::traits`. The concrete `Context`/`Tensor`/`Graph` types are +//! provided by the `cfg`-selected backend module and re-exported from the crate +//! root. + +use std::fmt::Debug; + +use crate::{CompileError, GraphBuilder, TensorDesc}; + +/// A neural device capable of compiling and running op-graphs. +pub trait NeuralDevice { + /// A compiled, executable graph. + type Graph: Send + Sync; + /// Backing storage for a tensor, bound at execution time. + type Tensor: Send + Sync + Clone + Debug; + + /// Compile the recorded graph into an executable form. + fn compile_graph(&self, builder: &GraphBuilder) -> Result; + + /// Allocate a tensor with the given description. + fn create_tensor(&self, desc: &TensorDesc) -> Self::Tensor; + + /// Execute `graph`, reading from `inputs` and writing into `outputs`, + /// matched by name. + fn run( + &self, + graph: &Self::Graph, + inputs: &[(&str, &Self::Tensor)], + outputs: &[(&str, &Self::Tensor)], + ); +}