PoC Neural abstraction library

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" }

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)',
+] }

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.

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::<f32>());
+}

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<Mutex<Vec<u8>>>,
+}
+
+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<u8> {
+        self.data.lock().unwrap().clone()
+    }
+    /// Read back the tensor as `f32` values.
+    pub fn read_f32(&self) -> Vec<f32> {
+        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<Self, NotSupportedError> {
+        Ok(Self {
+            validation: desc.validation,
+        })
+    }
+
+    /// Compile a recorded graph.
+    pub fn compile_graph(&self, builder: &GraphBuilder) -> Result<Graph, CompileError> {
+        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<f32>> = 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::<u8, f32>(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<Self::Graph, CompileError> {
+        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)
+    }
+}