use crate::Error;
use super::arch::{self, DnnArch};
use super::weights::WeightList;
use crate::in_place::write_fields;
#[cfg(feature = "deep-plc")]
use super::{builtin, models};
#[cfg(feature = "heap")]
use alloc::boxed::Box;
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
use alloc::sync::Arc;
use super::models::ERR_NO_MODEL;
#[derive(Clone)]
pub struct DnnBlob {
inner: Inner,
}
#[derive(Clone)]
enum Inner {
#[cfg(feature = "deep-plc")]
Builtin,
Static(WeightList<'static>),
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
Shared(Arc<SharedBlob>),
}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
struct SharedBlob {
list: WeightList<'static>,
bytes: *mut [Block],
}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
unsafe impl Send for SharedBlob {}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
unsafe impl Sync for SharedBlob {}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
impl Drop for SharedBlob {
fn drop(&mut self) {
drop(unsafe { Box::from_raw(self.bytes) });
}
}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
#[derive(Clone, Copy)]
#[repr(C, align(64))]
struct Block([u8; 64]);
impl DnnBlob {
pub fn from_static(bytes: &'static [u8]) -> Result<DnnBlob, Error> {
Ok(DnnBlob { inner: Inner::Static(WeightList::parse(bytes)?) })
}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
pub fn from_bytes(bytes: &[u8]) -> Result<DnnBlob, Error> {
let blocks = alloc::vec![Block([0; 64]); bytes.len().div_ceil(64)].into_boxed_slice();
let raw: *mut [Block] = Box::into_raw(blocks);
let copy: &mut [u8] = unsafe { core::slice::from_raw_parts_mut(raw.cast::<u8>(), bytes.len()) };
copy.copy_from_slice(bytes);
let view: &'static [u8] = unsafe { core::slice::from_raw_parts(raw.cast::<u8>(), bytes.len()) };
let list = match WeightList::parse(view) {
Ok(list) => list,
Err(e) => {
drop(unsafe { Box::from_raw(raw) });
return Err(e);
}
};
Ok(DnnBlob { inner: Inner::Shared(Arc::new(SharedBlob { list, bytes: raw })) })
}
#[cfg(feature = "deep-plc")]
pub fn builtin() -> DnnBlob {
DnnBlob { inner: Inner::Builtin }
}
#[cfg(feature = "__internals")]
#[doc(hidden)]
pub fn weights(&self) -> &WeightList<'_> {
self.list()
}
fn list(&self) -> &WeightList<'static> {
match &self.inner {
#[cfg(feature = "deep-plc")]
Inner::Builtin => builtin::weights(),
Inner::Static(list) => list,
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
Inner::Shared(shared) => &shared.list,
}
}
pub fn prefault(&self) {
for blob in self.list().blobs() {
prefault(blob);
}
}
}
pub(crate) fn prefault(bytes: &[u8]) {
for i in (0..bytes.len()).step_by(4096) {
unsafe { core::ptr::read_volatile(bytes.as_ptr().add(i)) };
}
}
impl core::fmt::Debug for DnnBlob {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let kind = match &self.inner {
#[cfg(feature = "deep-plc")]
Inner::Builtin => "built-in",
Inner::Static(_) => "static",
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
Inner::Shared(_) => "copied",
};
f.debug_struct("DnnBlob").field("kind", &kind).field("weights", self.list()).finish()
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub struct DnnModels(u8);
impl DnnModels {
pub const NONE: DnnModels = DnnModels(0);
pub const DEEP_PLC: DnnModels = DnnModels(1 << 0);
pub const LACE: DnnModels = DnnModels(1 << 1);
pub const NOLACE: DnnModels = DnnModels(1 << 2);
pub const BBWENET: DnnModels = DnnModels(1 << 3);
pub const DRED_ENCODER: DnnModels = DnnModels(1 << 4);
pub const DRED_DECODER: DnnModels = DnnModels(1 << 5);
pub const fn bits(self) -> u8 {
self.0
}
pub const fn contains(self, other: DnnModels) -> bool {
self.0 & other.0 == other.0
}
pub const fn union(self, other: DnnModels) -> DnnModels {
DnnModels(self.0 | other.0)
}
pub const fn is_empty(self) -> bool {
self.0 == 0
}
}
#[cfg(feature = "deep-plc")]
enum Slot<T: 'static> {
Empty,
Builtin(&'static T),
#[cfg(feature = "heap")]
Loaded { model: Box<T>, _blob: DnnBlob },
#[cfg(not(feature = "heap"))]
Loaded(T),
}
#[cfg(feature = "deep-plc")]
impl<T: 'static> Slot<T> {
fn builtin(model: Option<&'static T>) -> Self {
model.map_or(Slot::Empty, Slot::Builtin)
}
#[cfg(feature = "heap")]
fn loaded(model: T, blob: &DnnBlob) -> Self {
Slot::Loaded { model: Box::new(model), _blob: blob.clone() }
}
#[cfg(not(feature = "heap"))]
fn loaded(model: T, _blob: &DnnBlob) -> Self {
Slot::Loaded(model)
}
fn get(&self) -> Option<&T> {
match self {
Slot::Empty => None,
Slot::Builtin(m) => Some(m),
#[cfg(feature = "heap")]
Slot::Loaded { model, .. } => Some(model),
#[cfg(not(feature = "heap"))]
Slot::Loaded(m) => Some(m),
}
}
}
pub struct DecoderDnn {
arch: DnnArch,
#[cfg(feature = "deep-plc")]
deep_plc: Slot<models::DeepPlcModels<'static>>,
#[cfg(feature = "osce")]
lace: Slot<models::lace::LaceLayers<'static>>,
#[cfg(feature = "osce")]
nolace: Slot<models::nolace::NolaceLayers<'static>>,
#[cfg(feature = "osce-bwe")]
bbwenet: Slot<models::bbwenet::BbwenetLayers<'static>>,
}
impl DecoderDnn {
pub fn new() -> DecoderDnn {
let mut d = core::mem::MaybeUninit::<DecoderDnn>::uninit();
unsafe {
DecoderDnn::write_new(d.as_mut_ptr());
d.assume_init()
}
}
pub(crate) unsafe fn write_new(this: *mut DecoderDnn) {
#[cfg(feature = "deep-plc")]
let set = builtin::models();
unsafe {
write_fields!(this, {
arch: arch::detect(),
#[cfg(feature = "deep-plc")]
deep_plc: Slot::builtin(set.deep_plc.as_ref()),
#[cfg(feature = "osce")]
lace: Slot::builtin(set.lace.as_ref()),
#[cfg(feature = "osce")]
nolace: Slot::builtin(set.nolace.as_ref()),
#[cfg(feature = "osce-bwe")]
bbwenet: Slot::builtin(set.bbwenet.as_ref()),
});
}
}
pub fn arch(&self) -> DnnArch {
self.arch
}
pub fn models(&self) -> DnnModels {
#[allow(unused_mut)]
let mut m = DnnModels::NONE;
#[cfg(feature = "deep-plc")]
if self.deep_plc.get().is_some() {
m = m.union(DnnModels::DEEP_PLC);
}
#[cfg(feature = "osce")]
if self.lace.get().is_some() {
m = m.union(DnnModels::LACE);
}
#[cfg(feature = "osce")]
if self.nolace.get().is_some() {
m = m.union(DnnModels::NOLACE);
}
#[cfg(feature = "osce-bwe")]
if self.bbwenet.get().is_some() {
m = m.union(DnnModels::BBWENET);
}
m
}
pub fn load(&mut self, blob: &DnnBlob) -> Result<DnnModels, Error> {
#[allow(unused_mut)]
let mut loaded = DnnModels::NONE;
#[cfg(feature = "deep-plc")]
{
if let Inner::Builtin = blob.inner {
let set = builtin::models();
if let Some(m) = set.deep_plc.as_ref() {
self.deep_plc = Slot::Builtin(m);
loaded = loaded.union(DnnModels::DEEP_PLC);
}
#[cfg(feature = "osce")]
if let Some(m) = set.lace.as_ref() {
self.lace = Slot::Builtin(m);
loaded = loaded.union(DnnModels::LACE);
}
#[cfg(feature = "osce")]
if let Some(m) = set.nolace.as_ref() {
self.nolace = Slot::Builtin(m);
loaded = loaded.union(DnnModels::NOLACE);
}
#[cfg(feature = "osce-bwe")]
if let Some(m) = set.bbwenet.as_ref() {
self.bbwenet = Slot::Builtin(m);
loaded = loaded.union(DnnModels::BBWENET);
}
} else {
let list = blob.list();
if let Ok(m) = models::DeepPlcModels::init(list) {
self.deep_plc = Slot::loaded(m, blob);
loaded = loaded.union(DnnModels::DEEP_PLC);
}
#[cfg(feature = "osce")]
if let Ok(m) = models::lace::LaceLayers::init(list) {
self.lace = Slot::loaded(m, blob);
loaded = loaded.union(DnnModels::LACE);
}
#[cfg(feature = "osce")]
if let Ok(m) = models::nolace::NolaceLayers::init(list) {
self.nolace = Slot::loaded(m, blob);
loaded = loaded.union(DnnModels::NOLACE);
}
#[cfg(feature = "osce-bwe")]
if let Ok(m) = models::bbwenet::BbwenetLayers::init(list) {
self.bbwenet = Slot::loaded(m, blob);
loaded = loaded.union(DnnModels::BBWENET);
}
}
}
if loaded.is_empty() {
return Err(ERR_NO_MODEL);
}
blob.prefault();
Ok(loaded)
}
#[cfg(feature = "deep-plc")]
pub fn deep_plc(&self) -> Option<&models::DeepPlcModels<'_>> {
self.deep_plc.get()
}
#[cfg(feature = "osce")]
pub fn lace(&self) -> Option<&models::lace::LaceLayers<'_>> {
self.lace.get()
}
#[cfg(feature = "osce")]
pub fn nolace(&self) -> Option<&models::nolace::NolaceLayers<'_>> {
self.nolace.get()
}
#[cfg(feature = "osce-bwe")]
pub fn bbwenet(&self) -> Option<&models::bbwenet::BbwenetLayers<'_>> {
self.bbwenet.get()
}
}
impl Default for DecoderDnn {
fn default() -> DecoderDnn {
DecoderDnn::new()
}
}
#[cfg(feature = "dred")]
pub(crate) struct DredDecoderDnn {
arch: DnnArch,
model: Slot<models::dred_rdovae_dec::RdovaeDec<'static>>,
}
#[cfg(feature = "dred")]
impl DredDecoderDnn {
pub(crate) fn new() -> Self {
DredDecoderDnn { arch: arch::detect(), model: Slot::builtin(builtin::rdovae_dec()) }
}
pub(crate) fn load(&mut self, blob: &DnnBlob) -> Result<DnnModels, Error> {
if let Inner::Builtin = blob.inner {
self.model = Slot::Builtin(builtin::rdovae_dec().ok_or(ERR_NO_MODEL)?);
} else {
let model = models::dred_rdovae_dec::RdovaeDec::init(blob.list()).map_err(|_| ERR_NO_MODEL)?;
self.model = Slot::loaded(model, blob);
}
blob.prefault();
Ok(DnnModels::DRED_DECODER)
}
pub(crate) fn model(&self) -> Option<&models::dred_rdovae_dec::RdovaeDec<'_>> {
self.model.get()
}
pub(crate) fn arch(&self) -> DnnArch {
self.arch
}
}
#[cfg(feature = "dred")]
pub(crate) struct DredEncoderDnn {
arch: DnnArch,
model: Slot<models::DredEncModels<'static>>,
}
#[cfg(feature = "dred")]
impl DredEncoderDnn {
pub(crate) fn new() -> Self {
DredEncoderDnn { arch: arch::detect(), model: Slot::builtin(builtin::dred_enc()) }
}
pub(crate) fn load(&mut self, blob: &DnnBlob) -> Result<DnnModels, Error> {
if let Inner::Builtin = blob.inner {
self.model = Slot::Builtin(builtin::dred_enc().ok_or(ERR_NO_MODEL)?);
} else {
let model = models::DredEncModels::init(blob.list()).map_err(|_| ERR_NO_MODEL)?;
self.model = Slot::loaded(model, blob);
}
blob.prefault();
Ok(DnnModels::DRED_ENCODER)
}
pub(crate) fn model(&self) -> Option<&models::DredEncModels<'_>> {
self.model.get()
}
pub(crate) fn arch(&self) -> DnnArch {
self.arch
}
}
#[cfg(all(test, feature = "deep-plc"))]
mod tests {
use super::*;
#[test]
fn a_new_holder_has_every_builtin_model() {
let d = DecoderDnn::new();
assert_eq!(d.models(), builtin::models().models());
assert!(d.models().contains(DnnModels::DEEP_PLC));
assert!(d.deep_plc().is_some());
}
#[test]
fn a_failed_load_changes_nothing() {
let mut d = DecoderDnn::new();
let before = d.models();
let empty = DnnBlob::from_static(&[]).unwrap();
assert_eq!(d.load(&empty), Err(ERR_NO_MODEL));
assert_eq!(d.models(), before);
let shared = builtin::models().deep_plc.as_ref().unwrap();
assert!(core::ptr::eq(d.deep_plc().unwrap(), shared));
}
#[cfg(all(feature = "heap", target_has_atomic = "ptr"))]
#[test]
fn a_copied_blob_outlives_its_source_and_is_kept_by_the_decoder() {
let mut d = DecoderDnn::new();
let mut bytes = std::vec::Vec::new();
for &(_, blob) in builtin::blobs() {
bytes.extend_from_slice(blob);
}
let mut shifted = std::vec![0u8; 1];
shifted.extend_from_slice(&bytes);
let blob = DnnBlob::from_bytes(&shifted[1..]).unwrap();
drop(shifted);
let got = d.load(&blob).unwrap();
assert_eq!(got, builtin::models().models());
drop(blob);
let plc = &d.deep_plc().unwrap().plc;
let builtin_plc = &builtin::models().deep_plc.as_ref().unwrap().plc;
assert_eq!(plc.plc_dense_in.bias, builtin_plc.plc_dense_in.bias);
assert!(!core::ptr::eq(plc, builtin_plc));
}
#[test]
fn the_builtin_blob_brings_back_the_shared_models() {
let mut d = DecoderDnn::new();
let blob = DnnBlob::from_static(builtin::plc()).unwrap();
assert_eq!(d.load(&blob), Err(ERR_NO_MODEL));
assert_eq!(d.load(&DnnBlob::builtin()), Ok(builtin::models().models()));
}
}