fuel_cpu_backend/
quantized.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! CPU-side quantized storage adapter — bridges `Box<dyn QuantizedType>`
3//! (per-block-format scalar ops from fuel-quantized) to the
4//! backend-agnostic `DynQuantizedStorage` trait that fuel-core dispatches
5//! against.
6//!
7//! Owns the `QuantizedDeviceKernels` impl on `CpuBackendDevice` so
8//! `Device::qzeros` / `Device::load_quantized` reach CPU through the
9//! standard `DynBackendDevice::as_quantized_kernels` accessor.
10
11use crate::CpuStorage;
12use crate::dyn_impl::CpuBackendDevice;
13use fuel_backend_contract::dyn_backend::DynBackendStorage;
14use fuel_backend_contract::quantized::{DynQuantizedStorage, QuantizedDeviceKernels};
15use fuel_ir::quantized::GgmlDType;
16use fuel_ir::{DType, Error, HostBuffer, Layout, Result, Shape};
17use fuel_quantized::{QuantizedType, cpu_from_data, cpu_zeros};
18use half::f16;
19use std::any::Any;
20use std::borrow::Cow;
21
22pub struct CpuQStorage(pub Box<dyn QuantizedType>);
23
24impl std::fmt::Debug for CpuQStorage {
25    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
26        write!(f, "CpuQStorage[{:?}]", self.0.dtype())
27    }
28}
29
30fn cpu_src_as_slice(src: &dyn DynBackendStorage) -> Result<&[f32]> {
31    let cpu = src
32        .as_any()
33        .downcast_ref::<CpuStorage>()
34        .ok_or_else(|| Error::Msg("expected cpu storage for quantize".into()).bt())?;
35    cpu.0.as_slice::<f32>()
36}
37
38impl DynQuantizedStorage for CpuQStorage {
39    fn dtype(&self) -> GgmlDType {
40        self.0.dtype()
41    }
42    fn block_size(&self) -> usize {
43        self.0.block_size()
44    }
45    fn storage_size_in_bytes(&self) -> usize {
46        self.0.storage_size_in_bytes()
47    }
48    fn quantize(&mut self, src: &dyn DynBackendStorage) -> Result<()> {
49        self.0.dtype().check_quantize()?;
50        self.0.from_float(cpu_src_as_slice(src)?);
51        Ok(())
52    }
53    fn quantize_imatrix(
54        &mut self,
55        src: &dyn DynBackendStorage,
56        imatrix_weights: &[f32],
57        n_per_row: usize,
58    ) -> Result<()> {
59        self.0.dtype().check_quantize()?;
60        self.0
61            .from_float_imatrix(cpu_src_as_slice(src)?, imatrix_weights, n_per_row);
62        Ok(())
63    }
64    fn quantize_onto(&mut self, src: &dyn DynBackendStorage) -> Result<()> {
65        self.0.dtype().check_quantize()?;
66        self.0.from_float(cpu_src_as_slice(src)?);
67        Ok(())
68    }
69    fn quantize_imatrix_onto(
70        &mut self,
71        src: &dyn DynBackendStorage,
72        imatrix_weights: &[f32],
73        n_per_row: usize,
74    ) -> Result<()> {
75        self.0.dtype().check_quantize()?;
76        self.0
77            .from_float_imatrix(cpu_src_as_slice(src)?, imatrix_weights, n_per_row);
78        Ok(())
79    }
80    fn dequantize(&self, elem_count: usize) -> Result<Box<dyn DynBackendStorage>> {
81        let buf = self.0.dequantize(elem_count)?;
82        Ok(Box::new(CpuStorage(buf)))
83    }
84    fn data(&self) -> Result<Cow<'_, [u8]>> {
85        let data_ptr = self.0.as_ptr();
86        let size_in_bytes = self.0.storage_size_in_bytes();
87        let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
88        Ok(Cow::from(data))
89    }
90    fn fwd(
91        &self,
92        self_shape: &Shape,
93        input: &dyn DynBackendStorage,
94        layout: &Layout,
95    ) -> Result<(Box<dyn DynBackendStorage>, Shape)> {
96        let cpu = input
97            .as_any()
98            .downcast_ref::<CpuStorage>()
99            .ok_or_else(|| Error::Msg("qmatmul: expected cpu storage".into()).bt())?;
100        let storage = &cpu.0;
101        if !layout.is_contiguous() {
102            return Err(Error::Msg(format!("input tensor is not contiguous {layout:?}")).bt());
103        }
104        let src_shape = layout.shape();
105        let (n, k) = self_shape.dims2()?;
106        if src_shape.rank() < 2 {
107            return Err(Error::Msg(format!("input tensor has only one dimension {layout:?}")).bt());
108        }
109        let mut dst_shape = src_shape.dims().to_vec();
110        let last_k = dst_shape.pop().unwrap();
111        if last_k != k {
112            return Err(Error::Msg(format!(
113                "input tensor {layout:?} incompatible with {self_shape:?}"
114            ))
115            .bt());
116        }
117        dst_shape.push(n);
118        let dst_shape = Shape::from(dst_shape);
119        let out = match storage.dtype() {
120            DType::F32 => {
121                let slice = storage.as_slice::<f32>()?;
122                let slice =
123                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
124                let mut dst_storage = vec![0f32; dst_shape.elem_count()];
125                self.0
126                    .matmul_t((dst_shape.elem_count() / n, k, n), slice, &mut dst_storage)?;
127                HostBuffer::F32(dst_storage)
128            }
129            DType::F16 => {
130                let slice = storage.as_slice::<f16>()?;
131                let slice =
132                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
133                let mut dst_storage = vec![f16::ZERO; dst_shape.elem_count()];
134                self.0
135                    .matmul_t_f16((dst_shape.elem_count() / n, k, n), slice, &mut dst_storage)?;
136                HostBuffer::F16(dst_storage)
137            }
138            _ => return Err(Error::Msg("Expected f32/f16".into()).bt()),
139        };
140        Ok((Box::new(CpuStorage(out)), dst_shape))
141    }
142    fn as_any(&self) -> &dyn Any {
143        self
144    }
145    fn device_arc_dyn(
146        &self,
147    ) -> std::sync::Arc<dyn fuel_backend_contract::dyn_backend::DynBackendDevice> {
148        std::sync::Arc::new(CpuBackendDevice)
149    }
150}
151
152impl QuantizedDeviceKernels for CpuBackendDevice {
153    fn qzeros(&self, elem_count: usize, dtype: GgmlDType) -> Result<Box<dyn DynQuantizedStorage>> {
154        Ok(Box::new(CpuQStorage(cpu_zeros(dtype, elem_count))))
155    }
156    fn load_quantized(
157        &self,
158        dtype: GgmlDType,
159        data: Cow<'_, [u8]>,
160    ) -> Result<Box<dyn DynQuantizedStorage>> {
161        Ok(Box::new(CpuQStorage(cpu_from_data(dtype, data)?)))
162    }
163}
fuel_cpu_backend/quantized.rs

fuel_cpu_backend/
quantized.rs
