use crate::Error;
use super::MAX_INPUTS;
use super::nnet::{Conv2dLayer, LinearLayer};
pub const WEIGHT_BLOB_VERSION: i32 = 0;
pub const WEIGHT_BLOCK_SIZE: usize = 64;
pub const WEIGHT_NAME_SIZE: usize = 44;
pub const WEIGHT_TYPE_FLOAT: i32 = 0;
pub const WEIGHT_TYPE_INT: i32 = 1;
pub const WEIGHT_TYPE_QWEIGHT: i32 = 2;
pub const WEIGHT_TYPE_INT8: i32 = 3;
pub const SPARSE_BLOCK_SIZE: usize = 32;
pub const MAX_BLOBS: usize = 8;
pub const ERR_INVALID_BLOB: Error = Error::bad_arg("invalid DNN blob");
pub const ERR_NOT_ALIGNED: Error = Error::bad_arg("DNN blob is not aligned");
pub const ERR_TOO_MANY_BLOBS: Error = Error::bad_arg("too many DNN blobs");
pub const ERR_MISSING_ARRAY: Error = Error::bad_arg("DNN weight array is missing");
pub const ERR_ARRAY_SIZE: Error = Error::bad_arg("DNN weight array has the wrong size");
pub const ERR_SPARSE_INDEX: Error = Error::bad_arg("invalid DNN sparse index table");
pub const ERR_NO_WEIGHTS: Error = Error::bad_arg("DNN layer has no weights");
pub const ERR_INT8_MINUS_128: Error = Error::bad_arg("DNN int8 weight is -128");
pub const ERR_LAYER_SHAPE: Error = Error::bad_arg("invalid DNN layer shape");
#[derive(Clone, Copy)]
pub struct WeightArray<'w> {
pub name: &'w [u8],
pub kind: i32,
pub data: &'w [u8],
}
impl core::fmt::Debug for WeightArray<'_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("WeightArray")
.field("name", &core::str::from_utf8(self.name).unwrap_or("(not UTF-8)"))
.field("kind", &self.kind)
.field("size", &self.data.len())
.finish()
}
}
fn read_i32(b: &[u8], at: usize) -> Option<i32> {
let bytes = b.get(at..at + 4)?;
Some(i32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
}
fn read_record(blob: &[u8], off: usize) -> Result<(WeightArray<'_>, usize), Error> {
let remaining = blob.len().checked_sub(off).ok_or(ERR_INVALID_BLOB)?;
if remaining < WEIGHT_BLOCK_SIZE {
return Err(ERR_INVALID_BLOB);
}
let head = &blob[off..off + WEIGHT_BLOCK_SIZE];
let kind = read_i32(head, 8).ok_or(ERR_INVALID_BLOB)?;
let size = read_i32(head, 12).ok_or(ERR_INVALID_BLOB)?;
let block_size = read_i32(head, 16).ok_or(ERR_INVALID_BLOB)?;
let name_field = &head[20..WEIGHT_BLOCK_SIZE];
if block_size < size
|| i64::from(block_size) > (remaining - WEIGHT_BLOCK_SIZE) as i64
|| name_field[WEIGHT_NAME_SIZE - 1] != 0
|| size <= 0
{
return Err(ERR_INVALID_BLOB);
}
let (size, block_size) = (size as usize, block_size as usize);
let start = off + WEIGHT_BLOCK_SIZE;
let data = &blob[start..start + size];
if !(data.as_ptr() as usize).is_multiple_of(4) {
return Err(ERR_NOT_ALIGNED);
}
let name_len = name_field.iter().position(|&b| b == 0).unwrap_or(WEIGHT_NAME_SIZE);
let array = WeightArray { name: &name_field[..name_len], kind, data };
Ok((array, start + block_size))
}
fn validate(blob: &[u8]) -> Result<(), Error> {
let mut off = 0;
while off < blob.len() {
off = read_record(blob, off)?.1;
}
Ok(())
}
#[derive(Clone)]
struct Records<'w> {
blob: &'w [u8],
off: usize,
}
impl<'w> Iterator for Records<'w> {
type Item = WeightArray<'w>;
fn next(&mut self) -> Option<WeightArray<'w>> {
if self.off >= self.blob.len() {
return None;
}
let (array, next) = read_record(self.blob, self.off).ok()?;
self.off = next;
Some(array)
}
}
#[derive(Clone, Copy, Default)]
pub struct WeightList<'w> {
blobs: [&'w [u8]; MAX_BLOBS],
count: usize,
}
impl core::fmt::Debug for WeightList<'_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("WeightList")
.field("blobs", &self.count)
.field("arrays", &self.len())
.field("bytes", &self.blobs().iter().map(|b| b.len()).sum::<usize>())
.finish()
}
}
impl<'w> WeightList<'w> {
pub const fn empty() -> Self {
WeightList { blobs: [&[]; MAX_BLOBS], count: 0 }
}
pub fn parse(blob: &'w [u8]) -> Result<Self, Error> {
Self::parse_all(&[blob])
}
pub fn parse_all(blobs: &[&'w [u8]]) -> Result<Self, Error> {
let mut list = Self::empty();
for &blob in blobs {
list.push(blob)?;
}
Ok(list)
}
pub fn push(&mut self, blob: &'w [u8]) -> Result<(), Error> {
if self.count == MAX_BLOBS {
return Err(ERR_TOO_MANY_BLOBS);
}
validate(blob)?;
self.blobs[self.count] = blob;
self.count += 1;
Ok(())
}
pub fn blobs(&self) -> &[&'w [u8]] {
&self.blobs[..self.count]
}
pub fn find(&self, name: &str) -> Option<WeightArray<'w>> {
self.iter().find(|a| a.name == name.as_bytes())
}
pub fn iter(&self) -> impl Iterator<Item = WeightArray<'w>> + use<'w> {
let (blobs, count) = (self.blobs, self.count);
blobs.into_iter().take(count).flat_map(|blob| Records { blob, off: 0 })
}
pub fn len(&self) -> usize {
self.iter().count()
}
pub fn is_empty(&self) -> bool {
self.iter().next().is_none()
}
}
fn cast<T: Pod>(data: &[u8]) -> Option<&[T]> {
let (head, mid, tail) = unsafe { data.align_to::<T>() };
(head.is_empty() && tail.is_empty()).then_some(mid)
}
trait Pod: Copy {}
impl Pod for f32 {}
impl Pod for i32 {}
impl Pod for i8 {}
fn find_sized<'w, T: Pod>(list: &WeightList<'w>, name: &str, size: usize) -> Result<&'w [T], Error> {
let a = list.find(name).ok_or(ERR_MISSING_ARRAY)?;
if a.data.len() != size {
return Err(ERR_ARRAY_SIZE);
}
cast(a.data).ok_or(ERR_NOT_ALIGNED)
}
fn find_optional<'w, T: Pod>(list: &WeightList<'w>, name: &str, size: usize) -> Result<Option<&'w [T]>, Error> {
match find_sized(list, name, size) {
Ok(a) => Ok(Some(a)),
Err(ERR_MISSING_ARRAY) => Ok(None),
Err(e) => Err(e),
}
}
fn find_idx<'w>(list: &WeightList<'w>, name: &str, nb_inputs: usize, nb_outputs: usize) -> Result<(&'w [i32], usize), Error> {
let a = list.find(name).ok_or(ERR_MISSING_ARRAY)?;
let whole = a.data.len() / 4 * 4;
let idx: &[i32] = cast(&a.data[..whole]).ok_or(ERR_NOT_ALIGNED)?;
let mut rest = idx;
let mut rows_left = nb_outputs;
let mut total_blocks = 0usize;
while let Some((&count, tail)) = rest.split_first() {
let count = usize::try_from(count).map_err(|_| ERR_SPARSE_INDEX)?;
if tail.len() < count {
return Err(ERR_SPARSE_INDEX);
}
let (positions, tail) = tail.split_at(count);
for &pos in positions {
let pos = usize::try_from(pos).map_err(|_| ERR_SPARSE_INDEX)?;
if !pos.is_multiple_of(4) || pos.checked_add(3).is_none_or(|end| end >= nb_inputs) {
return Err(ERR_SPARSE_INDEX);
}
}
rows_left = rows_left.checked_sub(8).ok_or(ERR_SPARSE_INDEX)?;
total_blocks += count;
rest = tail;
}
if rows_left != 0 {
return Err(ERR_SPARSE_INDEX);
}
Ok((idx, total_blocks))
}
pub fn linear_init<'w>(
list: &WeightList<'w>,
bias: Option<&str>,
subias: Option<&str>,
weights: Option<&str>,
float_weights: Option<&str>,
weights_idx: Option<&str>,
diag: Option<&str>,
scale: Option<&str>,
nb_inputs: usize,
nb_outputs: usize,
) -> Result<LinearLayer<'w>, Error> {
let out_bytes = nb_outputs.checked_mul(4).ok_or(ERR_LAYER_SHAPE)?;
let mut layer = LinearLayer { nb_inputs, nb_outputs, ..LinearLayer::default() };
if let Some(name) = bias {
layer.bias = Some(find_sized(list, name, out_bytes)?);
}
if let Some(name) = subias {
layer.subias = Some(find_sized(list, name, out_bytes)?);
}
let nb_weights = if let Some(name) = weights_idx {
let (idx, blocks) = find_idx(list, name, nb_inputs, nb_outputs)?;
layer.weights_idx = Some(idx);
blocks.checked_mul(SPARSE_BLOCK_SIZE)
} else {
nb_inputs.checked_mul(nb_outputs)
}
.ok_or(ERR_LAYER_SHAPE)?;
if let Some(name) = weights {
layer.weights = Some(find_sized(list, name, nb_weights)?);
}
if let Some(name) = float_weights {
let bytes = nb_weights.checked_mul(4).ok_or(ERR_LAYER_SHAPE)?;
layer.float_weights = find_optional(list, name, bytes)?;
}
if let Some(name) = diag {
layer.diag = Some(find_sized(list, name, out_bytes)?);
if nb_inputs.checked_mul(3) != Some(nb_outputs) {
return Err(ERR_LAYER_SHAPE);
}
}
if weights.is_some() {
layer.scale = Some(find_sized(list, scale.ok_or(ERR_MISSING_ARRAY)?, out_bytes)?);
}
if let Some(w) = layer.weights {
if w.contains(&-128) {
return Err(ERR_INT8_MINUS_128);
}
if !nb_outputs.is_multiple_of(8) || !nb_inputs.is_multiple_of(4) || nb_inputs > MAX_INPUTS {
return Err(ERR_LAYER_SHAPE);
}
}
if layer.weights.is_none() && layer.float_weights.is_none() {
return Err(ERR_NO_WEIGHTS);
}
Ok(layer)
}
pub fn conv2d_init<'w>(
list: &WeightList<'w>,
bias: Option<&str>,
float_weights: Option<&str>,
in_channels: usize,
out_channels: usize,
ktime: usize,
kheight: usize,
) -> Result<Conv2dLayer<'w>, Error> {
let mut layer = Conv2dLayer { in_channels, out_channels, ktime, kheight, ..Conv2dLayer::default() };
if let Some(name) = bias {
let bytes = out_channels.checked_mul(4).ok_or(ERR_LAYER_SHAPE)?;
layer.bias = Some(find_sized(list, name, bytes)?);
}
if let Some(name) = float_weights {
let bytes = in_channels
.checked_mul(out_channels)
.and_then(|n| n.checked_mul(ktime))
.and_then(|n| n.checked_mul(kheight))
.and_then(|n| n.checked_mul(4))
.ok_or(ERR_LAYER_SHAPE)?;
layer.float_weights = find_optional(list, name, bytes)?;
}
if layer.float_weights.is_none() {
return Err(ERR_NO_WEIGHTS);
}
Ok(layer)
}
#[cfg(test)]
mod tests {
use super::*;
#[repr(C, align(64))]
struct Buf([u8; 1024]);
fn record(b: &mut [u8], off: usize, name: &str, kind: i32, data: &[u8]) -> usize {
let block = data.len().div_ceil(64) * 64;
b[off..off + 4].copy_from_slice(b"DNNw");
b[off + 4..off + 8].copy_from_slice(&0i32.to_le_bytes());
b[off + 8..off + 12].copy_from_slice(&kind.to_le_bytes());
b[off + 12..off + 16].copy_from_slice(&(data.len() as i32).to_le_bytes());
b[off + 16..off + 20].copy_from_slice(&(block as i32).to_le_bytes());
b[off + 20..off + 64].fill(0);
b[off + 20..off + 20 + name.len()].copy_from_slice(name.as_bytes());
b[off + 64..off + 64 + data.len()].copy_from_slice(data);
b[off + 64 + data.len()..off + 64 + block].fill(0);
off + 64 + block
}
fn floats(v: &[f32], out: &mut [u8]) -> usize {
for (i, x) in v.iter().enumerate() {
out[4 * i..4 * i + 4].copy_from_slice(&x.to_le_bytes());
}
4 * v.len()
}
#[test]
fn parse_find_and_iterate() {
let mut buf = Buf([0; 1024]);
let mut tmp = [0u8; 64];
let n = floats(&[1.0, 2.0], &mut tmp);
let mut off = record(&mut buf.0, 0, "l_bias", WEIGHT_TYPE_FLOAT, &tmp[..n]);
off = record(&mut buf.0, off, "l_bias", WEIGHT_TYPE_FLOAT, &tmp[..4]);
let list = WeightList::parse(&buf.0[..off]).unwrap();
assert_eq!(list.len(), 2);
assert_eq!(list.find("l_bias").unwrap().data.len(), 8);
assert!(list.find("l_bia").is_none());
assert!(list.find("l_bias_").is_none());
let names = list.iter().filter(|a| a.name == b"l_bias").count();
assert_eq!(names, 2);
assert!(WeightList::parse(&[]).unwrap().is_empty());
}
#[test]
fn libopus_refusals() {
let mut buf = Buf([0; 1024]);
let off = record(&mut buf.0, 0, "a", WEIGHT_TYPE_INT8, &[1, 2, 3]);
assert_eq!(off, 128);
let good = buf.0;
let check = |edit: &dyn Fn(&mut [u8]), len: usize| {
let mut b = Buf(good);
edit(&mut b.0);
WeightList::parse(&b.0[..len]).map(|_| ())
};
assert_eq!(check(&|_| {}, 128), Ok(()));
for len in [1, 63, 64, 66, 127] {
assert_eq!(check(&|_| {}, len), Err(ERR_INVALID_BLOB), "length {len}");
}
assert_eq!(check(&|_| {}, 129), Err(ERR_INVALID_BLOB));
assert_eq!(check(&|_| {}, 192), Err(ERR_INVALID_BLOB));
let set = |at: usize, v: i32| move |b: &mut [u8]| b[at..at + 4].copy_from_slice(&v.to_le_bytes());
assert_eq!(check(&set(12, 0), 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&set(12, -1), 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&set(16, 2), 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&set(16, 128), 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&set(16, i32::MAX), 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&set(16, -64), 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&|b: &mut [u8]| b[63] = b'x', 128), Err(ERR_INVALID_BLOB));
assert_eq!(check(&|b: &mut [u8]| b[0] = b'x', 128), Ok(()));
assert_eq!(check(&set(4, 7), 128), Ok(()));
assert_eq!(check(&set(8, 99), 128), Ok(()));
assert_eq!(check(&set(16, 4), 68), Ok(()));
}
#[test]
fn alignment() {
let mut buf = Buf([0; 1024]);
let mut off = record(&mut buf.0, 0, "a", WEIGHT_TYPE_INT8, &[1, 2, 3, 4, 5]);
buf.0[16..20].copy_from_slice(&5i32.to_le_bytes());
off = off - 64 + 5;
let end = record(&mut buf.0, off, "b", WEIGHT_TYPE_INT8, &[1]);
assert_eq!(WeightList::parse(&buf.0[..end]).map(|_| ()), Err(ERR_NOT_ALIGNED));
assert!(WeightList::parse(&buf.0[..off]).is_ok());
let mut buf2 = Buf([0; 1024]);
let end = record(&mut buf2.0, 1, "a", WEIGHT_TYPE_INT8, &[1]);
assert_eq!(WeightList::parse(&buf2.0[1..end]).map(|_| ()), Err(ERR_NOT_ALIGNED));
}
#[test]
fn several_blobs() {
let mut a = Buf([0; 1024]);
let mut b = Buf([0; 1024]);
let ea = record(&mut a.0, 0, "x", WEIGHT_TYPE_INT8, &[1]);
let eb = record(&mut b.0, 0, "x", WEIGHT_TYPE_INT8, &[2, 2]);
let eb = record(&mut b.0, eb, "y", WEIGHT_TYPE_INT8, &[3]);
let list = WeightList::parse_all(&[&a.0[..ea], &b.0[..eb]]).unwrap();
assert_eq!(list.blobs().len(), 2);
assert_eq!(list.len(), 3);
assert_eq!(list.find("x").unwrap().data, &[1]);
assert_eq!(list.find("y").unwrap().data, &[3]);
let mut full = WeightList::empty();
for _ in 0..MAX_BLOBS {
full.push(&a.0[..ea]).unwrap();
}
assert_eq!(full.push(&a.0[..ea]), Err(ERR_TOO_MANY_BLOBS));
assert_eq!(WeightList::parse_all(&[&a.0[..ea], &b.0[..eb - 1]]).map(|_| ()), Err(ERR_INVALID_BLOB));
}
#[test]
fn layers() {
let mut buf = Buf([0; 1024]);
let mut tmp = [0u8; 256];
let mut off = 0;
let n = floats(&[0.5; 8], &mut tmp);
off = record(&mut buf.0, off, "i_bias", WEIGHT_TYPE_FLOAT, &tmp[..n]);
off = record(&mut buf.0, off, "i_scale", WEIGHT_TYPE_FLOAT, &tmp[..n]);
off = record(&mut buf.0, off, "i_weights_int8", WEIGHT_TYPE_INT8, &[1u8; 32]);
let n = floats(&[0.25; 6], &mut tmp);
off = record(&mut buf.0, off, "f_weights_float", WEIGHT_TYPE_FLOAT, &tmp[..n]);
let n = floats(&[0.25; 3], &mut tmp);
off = record(&mut buf.0, off, "f_bias", WEIGHT_TYPE_FLOAT, &tmp[..n]);
let list = WeightList::parse(&buf.0[..off]).unwrap();
let i8l = |list: &WeightList<'_>, nin, nout| {
linear_init(list, Some("i_bias"), None, Some("i_weights_int8"), Some("i_weights_float"), None, None, Some("i_scale"), nin, nout)
.map(|_| ())
};
let l = linear_init(&list, Some("i_bias"), None, Some("i_weights_int8"), Some("i_weights_float"), None, None, Some("i_scale"), 4, 8).unwrap();
assert_eq!((l.nb_inputs, l.nb_outputs), (4, 8));
assert_eq!(l.weights.unwrap(), &[1i8; 32]);
assert_eq!(l.scale.unwrap(), &[0.5f32; 8]);
assert!(l.float_weights.is_none() && l.subias.is_none() && l.weights_idx.is_none());
assert_eq!(i8l(&list, 8, 8), Err(ERR_ARRAY_SIZE));
assert_eq!(i8l(&list, 2, 16), Err(ERR_ARRAY_SIZE));
assert_eq!(i8l(&list, 2, usize::MAX), Err(ERR_LAYER_SHAPE));
let r = linear_init(&list, None, None, Some("i_weights_int8"), None, None, None, Some("i_scale"), usize::MAX, 2);
assert_eq!(r.map(|_| ()), Err(ERR_LAYER_SHAPE));
let r = linear_init(&list, None, None, Some("i_weights_int8"), None, None, None, None, 4, 8);
assert_eq!(r.map(|_| ()), Err(ERR_MISSING_ARRAY));
let f = linear_init(&list, Some("f_bias"), None, None, Some("f_weights_float"), None, None, None, 2, 3).unwrap();
assert_eq!(f.float_weights.unwrap().len(), 6);
let r = linear_init(&list, Some("f_bias"), None, None, Some("none_weights_float"), None, None, None, 2, 3);
assert_eq!(r.map(|_| ()), Err(ERR_NO_WEIGHTS));
let r = linear_init(&list, Some("f_bias"), None, None, Some("f_weights_float"), None, None, None, 3, 3);
assert_eq!(r.map(|_| ()), Err(ERR_ARRAY_SIZE));
let c = conv2d_init(&list, Some("i_scale"), Some("f_weights_float"), 1, 2, 1, 3);
assert_eq!(c.map(|_| ()), Err(ERR_ARRAY_SIZE));
let c = conv2d_init(&list, None, Some("f_weights_float"), 1, 2, 1, 3).unwrap();
assert_eq!(c.float_weights.unwrap().len(), 6);
assert_eq!(conv2d_init(&list, None, None, 1, 2, 1, 3).map(|_| ()), Err(ERR_NO_WEIGHTS));
assert_eq!(conv2d_init(&list, None, Some("x"), 1, 2, 1, 3).map(|_| ()), Err(ERR_NO_WEIGHTS));
assert_eq!(conv2d_init(&list, None, Some("x"), usize::MAX, 2, 1, 3).map(|_| ()), Err(ERR_LAYER_SHAPE));
}
}