437de5d651
npu::onnx_limb_model writes the neural-processor member's program (one multiply, one sum over the last axis) as an ONNX opset-17 model with fp16 or int32 tensors, and npu::expand_i32 lays out the int32 inputs. The example npu_onnx writes both models, a batch of inputs and the exact products they must give, for a neural processor reached through a runtime that reads that format. Both models pass the full ONNX checker and give every product of a 16384-pair batch exactly on the reference runtime. Fix: f16_of(±2048) shifted by 10 - 11 and panicked in a debug build. Co-Authored-By: Claude Opus 5.5 <[email protected]>
42 lines
2.4 KiB
Rust
42 lines
2.4 KiB
Rust
//! Writes the limb-product model in the portable format (fp16 and int32), and a
|
|
//! batch of inputs with the products they must give, for a neural-processor
|
|
//! service that runs such models:
|
|
//!
|
|
//! ```text
|
|
//! cargo run --example npu_onnx --features quantum_fabric -- <dir>
|
|
//! ```
|
|
//!
|
|
//! In `<dir>`: `limbs_f16.onnx`, `limbs_i32.onnx`; `x_f16.bin`, `y_f16.bin`
|
|
//! (fp16 bits, little-endian), `x_i32.bin`, `y_i32.bin` (int32, little-endian),
|
|
//! each `[batch, 15, 8]`; `pairs.bin` (the operands, int32 pairs) and
|
|
//! `products.bin` (each exact `a·b`, int64). A service's output `z`
|
|
//! (`[batch, 15]`) is right when `Σ_s z[p, s]·16^s` is `products[p]` for every
|
|
//! `p`.
|
|
use std::io::Write;
|
|
use wai_quantum::quantum_fabric::lanes;
|
|
use wai_quantum::quantum_fabric::npu::{BATCH, Elem, LIMBS, SHIFTS, expand, expand_i32, onnx_limb_model};
|
|
|
|
fn write(path: &std::path::Path, bytes: &[u8]) {
|
|
std::fs::File::create(path).and_then(|mut f| f.write_all(bytes)).unwrap_or_else(|e| panic!("{}: {e}", path.display()));
|
|
}
|
|
|
|
fn main() {
|
|
let dir = std::path::PathBuf::from(std::env::args().nth(1).expect("an output directory"));
|
|
std::fs::create_dir_all(&dir).expect("the directory");
|
|
write(&dir.join("limbs_f16.onnx"), &onnx_limb_model(BATCH as u64, SHIFTS as u64, LIMBS as u64, Elem::F16));
|
|
write(&dir.join("limbs_i32.onnx"), &onnx_limb_model(BATCH as u64, SHIFTS as u64, LIMBS as u64, Elem::I32));
|
|
let mut pairs = lanes::probe_pairs(BATCH, 0x0a1b);
|
|
pairs.truncate(BATCH);
|
|
let n = BATCH * SHIFTS * LIMBS;
|
|
let (mut x, mut y) = (vec![0u16; n], vec![0u16; n]);
|
|
expand(&pairs, &mut x, &mut y);
|
|
write(&dir.join("x_f16.bin"), &x.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
|
|
write(&dir.join("y_f16.bin"), &y.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
|
|
let (mut xi, mut yi) = (vec![0i32; n], vec![0i32; n]);
|
|
expand_i32(&pairs, &mut xi, &mut yi);
|
|
write(&dir.join("x_i32.bin"), &xi.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
|
|
write(&dir.join("y_i32.bin"), &yi.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
|
|
write(&dir.join("pairs.bin"), &pairs.iter().flat_map(|&(a, b)| [a.to_le_bytes(), b.to_le_bytes()].concat()).collect::<Vec<u8>>());
|
|
write(&dir.join("products.bin"), &pairs.iter().flat_map(|&(a, b)| (i64::from(a) * i64::from(b)).to_le_bytes()).collect::<Vec<u8>>());
|
|
println!("wrote {} pairs of a batch of {BATCH} to {}", pairs.len(), dir.display());
|
|
}
|