edgefirst_tensor/protocol.rs
1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! The cross-package tensor protocol.
5//!
6//! Independently-linked EdgeFirst packages cannot share a Rust type: each
7//! Python extension module statically links its own copy of this crate, and
8//! PyO3 caches every `#[pyclass]`'s type object in a static belonging to the
9//! module that defined it. Two modules therefore see two distinct Python
10//! classes even when the Rust types are identical.
11//!
12//! So the contract between packages is a **descriptor**, not a type. A producer
13//! exposes `__edgefirst_tensor__()` returning a `PyCapsule` named
14//! `edgefirst_tensor_v2` whose payload leads with [`TensorDesc`] (followed by
15//! the quantization metadata `TensorDesc` deliberately does not carry -- see
16//! `edgefirst-python-common`'s `QuantDesc`); a consumer reads the descriptor
17//! and never performs an `isinstance` check. That duck typing is
18//! what makes the protocol survive independent release cadences, and it is the
19//! same pattern numpy, pyarrow and DLPack use.
20//!
21//! DLPack cannot serve this role alone: it has no device type for dma-buf,
22//! IOSurface or PBO, which are precisely the zero-copy cases this library
23//! exists for.
24
25/// Protocol version, checked by `TensorDyn::import_descriptor` after the
26/// payload has already been read as this build's `TensorDesc`.
27///
28/// This is the **second line of defense, not the gate**: it cannot catch a
29/// producer whose `TensorDesc` is a different *size*, since the
30/// out-of-bounds/misaligned read has already happened by the time this
31/// field could be inspected. The capsule name (`edgefirst_tensor_v2`, see
32/// `INTEROP.md`'s Versioning section) is what actually gates that --
33/// checked by `PyCapsule::pointer_checked` before any byte of the payload
34/// is read. DLPack reached the same conclusion after shipping an
35/// unversioned `dltensor` capsule and finding no upgrade path: its v1.0 fix
36/// added both a version field *and* a new capsule name, because the field
37/// alone is unreadable until the layout is already known.
38///
39/// This field stays useful for the case the name rule does not cover on its
40/// own: a future change to what a *same-sized, same-layout* `TensorDesc`'s
41/// fields mean. It has not needed to move past `1`.
42pub const ABI_VERSION: u32 = 1;
43
44/// Element type, as reported in [`TensorDesc::dtype`].
45///
46/// These codes cross a package boundary and are read by value, so they are
47/// part of the ABI: append, never reorder. Re-exported from the wire module
48/// [`crate::DType`] itself generates (via `ef_vocabulary!`, see
49/// `vocabulary.rs`), so this and [`crate::DType::code`] are two views of the
50/// same declared literal per variant, not two hand-kept tables that could
51/// drift apart.
52///
53/// `protocol::dtype` is the canonical, documented path -- the source module,
54/// `crate::dtype_wire`, is `#[doc(hidden)]` because it exists only so this
55/// re-export has something public to alias; it is emission plumbing, not a
56/// second API for the same codes.
57pub use crate::dtype_wire as dtype;
58
59/// Map a [`crate::DType`] to its wire code.
60///
61/// Compiled into both backends: [`from_parts`] is reached from
62/// `TensorDyn::descriptor_pinned`, which now exists on the `dynamic`
63/// backend too (task P2a). Nothing here touches storage -- it is the wire
64/// encoding of facts both backends can already answer -- so there was never
65/// a backend-specific reason for the gate this replaces, only a
66/// dead-code-warning one.
67pub(crate) fn dtype_of(d: crate::DType) -> u32 {
68 d.code()
69}
70
71/// Map a wire [`dtype`] code back to a [`crate::DType`]. Inverse of
72/// [`dtype_of`]. `None` for a code this build does not recognise (e.g. a
73/// dtype added by a newer producer this build predates) — the caller
74/// rejects rather than guessing.
75pub fn dtype_to_dtype(code: u32) -> Option<crate::DType> {
76 crate::DType::from_code(code)
77}
78
79/// HAL pixel format, as reported in [`TensorDesc::format`].
80///
81/// Same ABI concern as [`dtype`]: these codes cross a package boundary, so
82/// they are appended, never reordered. [`crate::PixelFormat`] itself
83/// generates this module's per-variant constants (via `ef_vocabulary!`, see
84/// `vocabulary.rs`), so this and [`crate::PixelFormat::code`] are two views
85/// of the same declared literal per variant, not two hand-kept tables that
86/// could drift apart. [`format::NONE`] is the one constant this module adds
87/// beyond what the macro emits -- see its own doc.
88///
89/// This exists because [`TensorDesc::fourcc`] cannot represent every HAL
90/// format: `PlanarRgb`/`PlanarRgba` have no standard FourCC and both encode
91/// as `fourcc = 0`, the same sentinel a non-image tensor uses, so `fourcc`
92/// alone cannot tell "Planar RGB" apart from "no format". `fourcc` stays for
93/// third-party/DRM interop; a HAL-aware consumer prefers `format` and
94/// falls back to `fourcc` only when `format` is `0`.
95///
96/// `protocol::format` is the canonical, documented path -- the source
97/// module, `crate::pixel_format_wire`, is `#[doc(hidden)]` because it exists
98/// only so this re-export has something public to alias; it is emission
99/// plumbing, not a second API for the same codes.
100pub mod format {
101 /// No format (non-image tensor), or a format this build predates. Not a
102 /// [`crate::PixelFormat`] variant -- every code `1..=11` below is one of
103 /// those, generated by the same `ef_vocabulary!` declaration as
104 /// [`crate::PixelFormat::code`]; `0` is a sentinel this module alone
105 /// defines, since "no format" has no corresponding Rust variant.
106 pub const NONE: u32 = 0;
107 pub use crate::format::pixel_format_wire::*;
108}
109
110/// Map a [`crate::PixelFormat`] to its [`mod@format`] wire code.
111///
112/// Compiled into both backends: same reason as [`dtype_of`].
113pub(crate) fn format_of(f: crate::PixelFormat) -> u32 {
114 f.code()
115}
116
117/// Map a [`mod@format`] wire code back to a [`crate::PixelFormat`]. `None` for
118/// [`format::NONE`] or a code this build does not recognise (e.g. a format
119/// added by a newer producer).
120pub fn format_from_code(code: u32) -> Option<crate::PixelFormat> {
121 crate::PixelFormat::from_code(code)
122}
123
124/// Bit flags, as reported in [`TensorDesc::flags`].
125///
126/// A bitfield rather than more discrete fields: the additions this protocol
127/// is most likely to need are booleans, and each one spent as its own `u32`
128/// would grow the descriptor -- which after release costs a capsule-name
129/// bump. Spending one `u32` now buys 32 of them for free.
130pub mod flags {
131 /// [`TensorDesc::sync`] carries a real handle. When clear, `sync` is
132 /// meaningless and must not be waited on.
133 ///
134 /// A flag rather than a sentinel because every plausible sentinel is a
135 /// legal handle somewhere: `0` is a valid `GLsync`-adjacent value and a
136 /// valid fd number is never negative, so `-1` and `0` cannot both be
137 /// spare across the kinds this field spans.
138 ///
139 /// Under [`kind::D3D11_TEXTURE`] `sync` is not a handle at all: it is a
140 /// value on the timeline of the fence whose NT handle
141 /// [`TensorDesc::ptr`] carries, so a producer that sets this flag must
142 /// also fill in `ptr`.
143 pub const SYNC_PRESENT: u32 = 1 << 0;
144}
145
146/// Backing store kind, as reported in [`TensorDesc::kind`].
147pub mod kind {
148 /// Host memory: `Mem` or `Shm`. `ptr` is valid; `handle` is -1.
149 pub const HOST: u32 = 0;
150 /// Linux dma-buf. `handle` is the fd.
151 pub const DMABUF: u32 = 1;
152 /// Apple IOSurface. `handle` is the surface id.
153 pub const IOSURFACE: u32 = 2;
154 /// OpenGL pixel buffer object. `handle` is the buffer id.
155 pub const PBO: u32 = 3;
156 /// CUDA device memory. `ptr` is a device pointer, not host-addressable.
157 pub const CUDA_DEVICE: u32 = 4;
158 /// Windows D3D11 texture. `handle` is the texture's NT shared handle and
159 /// `ptr` the device fence's NT shared handle, both valid in the producing
160 /// process (other modules of the same process included); `sync` is the
161 /// fence value of the last recorded GPU write when `SYNC_PRESENT` is set.
162 /// The `ptr` reuse follows `PBO`'s vtable precedent.
163 pub const D3D11_TEXTURE: u32 = 5;
164}
165
166/// A raw pointer that may cross a `PyCapsule`.
167///
168/// `PyCapsule::new` requires `T: Send`, which `*mut u8` is not. The producer
169/// discharges that obligation by keeping the owning tensor alive for the
170/// capsule's lifetime via the capsule's keepalive; see the module docs.
171#[repr(transparent)]
172#[derive(Clone, Copy, Debug)]
173pub struct SendPtr(pub *mut u8);
174
175// SAFETY: the pointer is only dereferenced by a consumer that received it
176// through a capsule whose keepalive holds the producing tensor alive, so the
177// memory outlives every use.
178unsafe impl Send for SendPtr {}
179unsafe impl Sync for SendPtr {}
180
181impl SendPtr {
182 /// A null pointer, for descriptors whose memory is not host-addressable.
183 pub const fn null() -> Self {
184 SendPtr(std::ptr::null_mut())
185 }
186 /// True when the descriptor carries no host address.
187 pub fn is_null(&self) -> bool {
188 self.0.is_null()
189 }
190}
191
192/// The cross-package descriptor. Layout is part of the contract: it is
193/// `#[repr(C)]` and its size is pinned by test.
194#[repr(C)]
195#[derive(Clone, Copy, Debug)]
196pub struct TensorDesc {
197 /// [`ABI_VERSION`]. Checked, never assumed.
198 pub version: u32,
199 /// One of the [`kind`] constants.
200 pub kind: u32,
201 /// dma-buf fd, IOSurface id, GL buffer id or D3D11 texture NT handle
202 /// value; `-1` when unused.
203 pub handle: i64,
204 /// Host pointer when addressable, else null -- **except under
205 /// [`kind::PBO`]**, where it instead carries a `*const
206 /// crate::pbo::PboOpsVtable`, or null if the producer attached none,
207 /// **and under [`kind::D3D11_TEXTURE`]**, where it carries the process
208 /// device fence's NT handle value, or null when [`Self::sync`] carries
209 /// nothing to wait on.
210 /// Interpretation is keyed by [`Self::kind`], the same established
211 /// pattern [`Self::sync`]'s own doc comment already documents for that
212 /// field: a PBO tensor never takes a `HostPin` (`PboTensor` refuses
213 /// `pin_host` — the address cannot outlive the map), so this field is
214 /// otherwise permanently unused for that kind, which is exactly why it
215 /// is the field that carries the vtable instead of a new one requiring
216 /// a layout change. Valid under the exact same borrow contract as the
217 /// host-pointer case: only while the producer's own keepalive (for the
218 /// capsule protocol, the capsule's) keeps the owning tensor alive.
219 pub ptr: SendPtr,
220 /// Number of significant entries in `shape` and `strides`.
221 pub ndim: u32,
222 /// Element type, one of the [`dtype`] constants.
223 ///
224 /// Deliberately placed here: `ndim` leaves a 4-byte hole before the
225 /// 8-aligned `shape`, so this field costs nothing. Without it `strides`
226 /// — documented in *elements* — cannot be converted to bytes by a
227 /// consumer, and a `Tensor<f32>` is indistinguishable from a
228 /// `Tensor<u8>` of the same shape.
229 pub dtype: u32,
230 /// Logical shape, most-significant dimension first.
231 pub shape: [u64; 8],
232 /// Strides **in bytes**, signed so a consumer can express flips.
233 ///
234 /// Bytes rather than elements so a hardware row pitch is always
235 /// representable: a pitch need not be a whole number of `dtype` elements
236 /// (and a sub-byte dtype has no element size to divide by). Multiply a
237 /// coordinate by the stride directly; do not scale by `dtype` size.
238 pub strides: [i64; 8],
239 /// DRM FourCC of the pixel format, or 0 for a non-image tensor. Kept for
240 /// third-party/DRM interop; a HAL-aware consumer should prefer
241 /// [`Self::format`] and fall back to this only when `format` is 0
242 /// (see [`mod@format`]'s doc for why `fourcc` alone is insufficient).
243 pub fourcc: u32,
244 /// One of the [`mod@format`] constants, or 0 for a non-image tensor.
245 pub format: u32,
246 /// Packed colorimetry, or 0 when unspecified.
247 pub colorimetry: u32,
248 /// Bit flags; see [`mod@flags`]. Zero when nothing applies.
249 ///
250 /// Occupies what was a 4-byte alignment hole before `capacity`, so it
251 /// cost nothing to add.
252 pub flags: u32,
253 /// Bytes available in the producer's underlying allocation, from this
254 /// tensor's window start (its [`TensorTrait::capacity_bytes`
255 /// ](crate::TensorTrait::capacity_bytes)) -- may exceed the byte size
256 /// `shape`/`strides` imply when the producer over-allocated: a pool
257 /// tensor sized for its largest expected image but holding a smaller
258 /// one today, or row-pitch/MCU-alignment padding a decoder needs beyond
259 /// the logical width. Without this, a consumer importing a `HOST`
260 /// descriptor has no way to learn that headroom exists and must treat
261 /// the declared shape as the allocation's exact size -- which starves a
262 /// later `configure_image`/`set_logical_shape` of capacity the producer
263 /// actually has. `DMABUF`/`IOSURFACE` consumers query their real
264 /// capacity from the handle itself and do not depend on this field, but
265 /// it is populated for every kind for consistency.
266 pub capacity: u64,
267 /// Completion handle for work still in flight against this tensor, or 0
268 /// when [`flags::SYNC_PRESENT`] is clear.
269 ///
270 /// Filled in by [`kind::D3D11_TEXTURE`] producers, and left clear by
271 /// every other kind in this build. It was declared before any producer
272 /// set it because the field has to exist before the first release or
273 /// adding it later costs a capsule-name bump, and it is the one gap the
274 /// synchronous-at-the-boundary contract cannot close: a *foreign*
275 /// producer with queued GPU work has no other way to say "wait on this
276 /// first".
277 ///
278 /// Interpretation is keyed by [`Self::kind`], since each backing already
279 /// implies its own fence flavour and a separate discriminant would be
280 /// redundant:
281 ///
282 /// * [`kind::DMABUF`] -- a `sync_file` fd from
283 /// `DMA_BUF_IOCTL_EXPORT_SYNC_FILE`. The consumer owns it and must
284 /// close it.
285 /// * [`kind::PBO`] -- a `GLsync` from `glFenceSync`, valid only in the
286 /// producer's share group.
287 /// * [`kind::CUDA_DEVICE`] -- a `cudaEvent_t`.
288 /// * [`kind::D3D11_TEXTURE`] -- a value on the timeline of the
289 /// `ID3D11Fence` [`Self::ptr`] names, not a handle of its own: the
290 /// handle and the value together are what a consumer waits on.
291 /// * [`kind::HOST`], [`kind::IOSURFACE`] -- no fence flavour defined;
292 /// `SYNC_PRESENT` must be clear.
293 pub sync: u64,
294}
295
296impl TensorDesc {
297 /// Maximum rank the fixed-size `shape`/`strides` arrays can carry.
298 pub const MAX_NDIM: usize = 8;
299
300 /// The logical shape as a slice.
301 pub fn shape(&self) -> &[u64] {
302 &self.shape[..(self.ndim as usize).min(Self::MAX_NDIM)]
303 }
304
305 /// The strides as a slice, in bytes.
306 pub fn strides(&self) -> &[i64] {
307 &self.strides[..(self.ndim as usize).min(Self::MAX_NDIM)]
308 }
309
310 /// Total element count.
311 pub fn len(&self) -> u64 {
312 self.shape().iter().product()
313 }
314
315 /// True when the descriptor addresses no elements.
316 pub fn is_empty(&self) -> bool {
317 self.len() == 0
318 }
319}
320
321// ---------------------------------------------------------------------------
322// Producer side
323// ---------------------------------------------------------------------------
324
325use crate::TensorMemory;
326
327/// Map a backing store to its [`kind`] constant.
328///
329/// `TensorMemory::DmaBuf` is a deliberately shared discriminant covering DMA-BUF,
330/// IOSurface and AHardwareBuffer (it is ABI-stable across platforms), so the
331/// platform decides which protocol kind it reports.
332pub(crate) fn kind_of(memory: TensorMemory) -> u32 {
333 match memory {
334 // Both are a host-addressable pointer with no handle, which is all
335 // `kind` records -- it is a narrower vocabulary than `TensorMemory`
336 // on purpose, not a copy of it that lost a variant.
337 TensorMemory::Mem | TensorMemory::Shm => kind::HOST,
338 TensorMemory::Pbo => kind::PBO,
339 TensorMemory::Cuda => kind::CUDA_DEVICE,
340 // Named specifically rather than through `DmaBuf`. Unreachable
341 // today (no backend reports it) but correct the day one does, which
342 // is cheaper than the panic the obvious `unreachable!()` becomes.
343 TensorMemory::IoSurface => kind::IOSURFACE,
344 TensorMemory::DmaBuf => {
345 #[cfg(any(target_os = "macos", target_os = "ios"))]
346 {
347 kind::IOSURFACE
348 }
349 // Windows has no dma-buf: the shared `DmaBuf` discriminant is a
350 // D3D11 texture there, whose handle field means something else
351 // again (an NT handle, not an fd).
352 #[cfg(target_os = "windows")]
353 {
354 kind::D3D11_TEXTURE
355 }
356 #[cfg(not(any(target_os = "macos", target_os = "ios", target_os = "windows")))]
357 {
358 kind::DMABUF
359 }
360 }
361 }
362}
363
364// ---------------------------------------------------------------------------
365// Consumer side: the per-kind field checks `TensorDyn::import_storage` needs
366// ---------------------------------------------------------------------------
367//
368// Each of these validates one kind's handle/pointer field and produces the
369// value a constructor takes. They live here, next to the descriptor they
370// read, rather than in either backend: a descriptor is an *untrusted*
371// cross-package payload, and two copies of "is this fd representable" or
372// "is this surface id in range" are two places for the refusals to drift.
373
374/// `dup` a `kind::DMABUF` descriptor's fd into one this process owns.
375///
376/// The producer still owns the original and will close it when its own
377/// tensor drops, so an import must have its own -- `from_fd` adopts what it
378/// is given.
379///
380/// # Errors
381///
382/// [`crate::Error::InvalidArgument`] when the descriptor carries no fd, or
383/// one not representable as a file descriptor. `as i32` would wrap a handle
384/// above `i32::MAX` onto an unrelated -- possibly live -- descriptor, which
385/// the `dup` would then adopt as the buffer.
386#[cfg(target_os = "linux")]
387pub(crate) fn dup_descriptor_fd(desc: &TensorDesc) -> crate::Result<std::os::fd::OwnedFd> {
388 use std::os::fd::BorrowedFd;
389 if desc.handle < 0 {
390 return Err(crate::Error::InvalidArgument(
391 "dma-buf descriptor carries no fd".into(),
392 ));
393 }
394 let raw_fd = i32::try_from(desc.handle).map_err(|_| {
395 crate::Error::InvalidArgument(format!(
396 "dma-buf descriptor fd {} is not representable as a file descriptor",
397 desc.handle
398 ))
399 })?;
400 // SAFETY: the producer owns `desc.handle` and the capsule keepalive
401 // holds it open for the duration of this call, so it is a valid fd to
402 // borrow (not adopt) for the dup below.
403 let borrowed = unsafe { BorrowedFd::borrow_raw(raw_fd) };
404 Ok(nix::unistd::dup(borrowed)?)
405}
406
407/// Read a `kind::IOSURFACE` descriptor's surface id.
408///
409/// # Errors
410///
411/// [`crate::Error::InvalidArgument`] when the descriptor carries no id, or
412/// one out of `u32` range. Liveness is *not* checked here -- that belongs
413/// with the lookup, which happens inside the constructor (and, under the
414/// `dynamic` backend, inside `libedgefirst_tensor.so`) so no window exists
415/// between "found it" and "retained it".
416#[cfg(any(target_os = "macos", target_os = "ios"))]
417pub(crate) fn descriptor_surface_id(desc: &TensorDesc) -> crate::Result<u32> {
418 if desc.handle < 0 {
419 return Err(crate::Error::InvalidArgument(
420 "IOSurface descriptor carries no surface id".into(),
421 ));
422 }
423 u32::try_from(desc.handle).map_err(|_| {
424 crate::Error::InvalidArgument(format!(
425 "IOSurface descriptor id {} is out of range",
426 desc.handle
427 ))
428 })
429}
430
431/// Refuse a `kind::HOST` descriptor that carries no address.
432///
433/// # Errors
434///
435/// [`crate::Error::InvalidArgument`], naming the producer-side fix: a
436/// descriptive-only capsule was requested, so no address was pinned.
437pub(crate) fn check_descriptor_host_ptr(desc: &TensorDesc) -> crate::Result<()> {
438 if desc.ptr.is_null() {
439 return Err(crate::Error::InvalidArgument(
440 "host descriptor has no address: the producer was asked \
441 for a descriptive-only capsule (request access=\"read\" \
442 or \"readwrite\")"
443 .into(),
444 ));
445 }
446 Ok(())
447}
448
449/// Read a `kind::PBO` descriptor's GL buffer id, and refuse one carrying no
450/// ops vtable.
451///
452/// Both checks together, unlike the other kinds' one apiece, because the
453/// `PBO` kind is the one that needs *two* of the descriptor's fields to be
454/// present before a constructor can be called at all -- `handle` names the
455/// buffer, `ptr` carries the vtable that can reach it, and either alone is
456/// useless.
457///
458/// # Errors
459///
460/// [`crate::Error::InvalidArgument`] when the buffer id is absent or out of
461/// range, or when no vtable address is carried.
462pub(crate) fn descriptor_pbo_buffer_id(desc: &TensorDesc) -> crate::Result<u32> {
463 if desc.handle < 0 {
464 return Err(crate::Error::InvalidArgument(
465 "PBO descriptor carries no buffer id".into(),
466 ));
467 }
468 let buffer_id = u32::try_from(desc.handle).map_err(|_| {
469 crate::Error::InvalidArgument(format!(
470 "PBO descriptor buffer id {} is out of range",
471 desc.handle
472 ))
473 })?;
474 if desc.ptr.is_null() {
475 return Err(crate::Error::InvalidArgument(
476 "PBO descriptor carries no ops vtable -- the producer either \
477 predates this build's cross-cdylib PBO support or its own \
478 PboTensor could not build one"
479 .into(),
480 ));
481 }
482 Ok(buffer_id)
483}
484
485/// Read a `kind::D3D11_TEXTURE` descriptor's texture handle and, when the
486/// producer recorded a GPU write, the fence handle plus value an import
487/// waits on before reading.
488///
489/// Both checks together, like [`descriptor_pbo_buffer_id`]'s, because this
490/// kind also spends two of the descriptor's fields on one backing:
491/// `handle` names the texture and `ptr` the fence that orders access to it.
492/// The completion is optional -- a producer that has queued no GPU work
493/// leaves `SYNC_PRESENT` clear, and there is then nothing to wait for.
494///
495/// Both handles stay the producer's: they are its own, not duplicates, and
496/// are valid only while the capsule keepalive holds the producing tensor
497/// alive. `from_d3d11_shared_handle` duplicates what it keeps.
498///
499/// # Errors
500///
501/// [`crate::Error::InvalidArgument`] when the descriptor carries no texture
502/// handle.
503#[cfg(target_os = "windows")]
504pub(crate) fn descriptor_d3d11_handles(
505 desc: &TensorDesc,
506) -> crate::Result<(
507 std::os::windows::io::RawHandle,
508 Option<(std::os::windows::io::RawHandle, u64)>,
509)> {
510 if desc.handle <= 0 {
511 return Err(crate::Error::InvalidArgument(
512 "D3D11_TEXTURE descriptor without a texture handle".into(),
513 ));
514 }
515 let tex = desc.handle as usize as std::os::windows::io::RawHandle;
516 let completion = (desc.flags & flags::SYNC_PRESENT != 0 && !desc.ptr.is_null())
517 .then_some((desc.ptr.0 as std::os::windows::io::RawHandle, desc.sync));
518 Ok((tex, completion))
519}
520
521/// The pixel format and image dimensions a `kind::D3D11_TEXTURE` import is
522/// built at, with the descriptor's shape checked against them.
523///
524/// The dimensions come from the texture, not from `shape`: a shape cannot
525/// tell an allocation grid from an addressing grid, and a producer may
526/// legitimately send either (`descriptor_pinned` emits `shape()`, which is
527/// the allocation shape until someone calls `set_logical_shape`). The one
528/// thing the texture cannot supply is a semi-planar image's width -- such a
529/// texture is as wide as the driver's row pitch -- which is why `shape` is
530/// passed down to [`crate::d3d11_shared_handle_geometry`] rather than only
531/// checked here.
532///
533/// A third spelling is accepted here and nowhere else: a packed *window* of
534/// the texture, which is what a `view()` describes (`[h', w', c]` with `h'`
535/// and `w'` no larger than the texture's, and no smaller than one). The
536/// import still opens the whole texture; `restore_d3d11_logical_shape`
537/// narrows it to the window and the producer's plane offset places it. The
538/// blob transport does not take this spelling because nothing narrows on
539/// that path, so a window shape would be accepted and then ignored.
540///
541/// # Errors
542///
543/// [`crate::Error::InvalidArgument`] when the descriptor names no pixel
544/// format, or when its shape is none of the three spellings of the geometry
545/// the texture reports. Propagates whatever opening and describing the
546/// texture reports.
547///
548/// # Safety
549///
550/// `texture` must be an NT shared handle of a D3D11 texture, valid in this
551/// process. It stays owned by the caller.
552#[cfg(target_os = "windows")]
553pub(crate) unsafe fn descriptor_d3d11_geometry(
554 desc: &TensorDesc,
555 texture: std::os::windows::io::RawHandle,
556 shape: &[usize],
557) -> crate::Result<(crate::PixelFormat, usize, usize)> {
558 let format = crate::PixelFormat::from_code(desc.format).ok_or_else(|| {
559 crate::Error::InvalidArgument(
560 "D3D11_TEXTURE descriptor without a pixel format: the texture's own \
561 description cannot say which format it was allocated for"
562 .into(),
563 )
564 })?;
565 // SAFETY: the caller guarantees `texture` is a shared NT handle valid in
566 // this process.
567 let (width, height) =
568 unsafe { d3d11_geometry_with_rule(format, texture, shape, ShapeRule::OrWindow) }?;
569 Ok((format, width, height))
570}
571
572/// Which spellings of a texture's geometry a `shape` may take.
573#[cfg(target_os = "windows")]
574#[derive(Clone, Copy, PartialEq, Eq)]
575enum ShapeRule {
576 /// The allocation shape or the addressing shape, exactly.
577 Exact,
578 /// `Exact`, or a packed window of the texture (a `view()`).
579 OrWindow,
580}
581
582/// Whether `shape` is a packed sub-rectangle of a `width` x `height` image
583/// in `format`: the allocation spelling with a smaller height or width. A
584/// window is never a semi-planar or planar image (`Tensor::view` makes
585/// packed views only) and never empty.
586#[cfg(target_os = "windows")]
587fn is_packed_window_of(
588 format: crate::PixelFormat,
589 shape: &[usize],
590 width: usize,
591 height: usize,
592) -> bool {
593 if format.layout() != crate::PixelLayout::Packed {
594 return false;
595 }
596 let [h, w, c] = shape else {
597 return false;
598 };
599 *c == format.channels() && (1..=height).contains(h) && (1..=width).contains(w)
600}
601
602/// The image dimensions a shared D3D11 texture reports, with `shape` checked
603/// against them.
604///
605/// Split out of [`descriptor_d3d11_geometry`] so the blob's reference-mode
606/// import ([`crate::blob::import`]) makes the same refusal rather than a
607/// second copy of it: the two transports differ in how they name the format,
608/// not in which shapes describe a given texture.
609///
610/// # Errors
611///
612/// [`crate::Error::InvalidArgument`] when `shape` is neither spelling of the
613/// geometry the texture reports. Propagates whatever opening and describing
614/// the texture reports.
615///
616/// # Safety
617///
618/// `texture` must be an NT shared handle of a D3D11 texture, valid in this
619/// process. It stays owned by the caller.
620#[cfg(target_os = "windows")]
621pub(crate) unsafe fn d3d11_geometry_checked(
622 format: crate::PixelFormat,
623 texture: std::os::windows::io::RawHandle,
624 shape: &[usize],
625) -> crate::Result<(usize, usize)> {
626 // SAFETY: forwarded unchanged.
627 unsafe { d3d11_geometry_with_rule(format, texture, shape, ShapeRule::Exact) }
628}
629
630#[cfg(target_os = "windows")]
631unsafe fn d3d11_geometry_with_rule(
632 format: crate::PixelFormat,
633 texture: std::os::windows::io::RawHandle,
634 shape: &[usize],
635 rule: ShapeRule,
636) -> crate::Result<(usize, usize)> {
637 // SAFETY: the caller guarantees `texture` is a shared NT handle valid in
638 // this process; the helper opens its own texture and drops it on return.
639 let (width, height) =
640 unsafe { crate::d3d11_shared_handle_geometry(texture, format, Some(shape)) }?;
641 let allocation = format.allocation_shape(width, height);
642 let addressing = format.addressing_shape(width, height);
643 if allocation.as_deref() == Some(shape) || addressing.as_deref() == Some(shape) {
644 return Ok((width, height));
645 }
646 if rule == ShapeRule::OrWindow && is_packed_window_of(format, shape, width, height) {
647 return Ok((width, height));
648 }
649 let window = match rule {
650 ShapeRule::OrWindow => " nor a packed window",
651 ShapeRule::Exact => "",
652 };
653 Err(crate::Error::InvalidArgument(format!(
654 "D3D11 texture shape {shape:?} is neither the {format:?} allocation \
655 shape {allocation:?} nor the addressing shape {addressing:?}{window} of the \
656 {width}x{height} texture it names"
657 )))
658}
659
660/// Inputs to [`from_parts`], grouped into one named-field struct rather than
661/// a positional argument list.
662///
663/// This crossed nine parameters once the pixel-format and capacity inputs
664/// were added, several of them same-typed (`fourcc: u32`, a `u32`
665/// handle-adjacent value, `colorimetry: u32`) — a positional call at that
666/// width compiles happily with two arguments transposed and silently
667/// corrupts the descriptor. Named fields make that a compile error instead. Internal to
668/// this crate: `TensorDyn::descriptor_pinned` is presently the only caller --
669/// on both backends now, not just `static` (task P2a).
670pub(crate) struct DescParts<'a, 'p> {
671 pub dims: &'a [usize],
672 pub memory: TensorMemory,
673 pub dtype: crate::DType,
674 /// DRM FourCC for [`TensorDesc::fourcc`]; 0 when the format has none.
675 pub fourcc: u32,
676 /// The HAL format, source for both the row-dimension decision below and
677 /// [`TensorDesc::format`] (via [`format_of`]).
678 pub format: Option<crate::PixelFormat>,
679 /// Row pitch in bytes; see [`from_parts`]'s doc for how it maps to a
680 /// stride index.
681 pub row_stride: Option<usize>,
682 pub handle: i64,
683 pub colorimetry: u32,
684 /// Producer's [`crate::TensorTrait::capacity_bytes`] for
685 /// [`TensorDesc::capacity`]; see that field's docs.
686 pub capacity: u64,
687 pub pin: Option<&'a crate::HostPin<'p>>,
688 /// The producer's `PboOpsVtable` address, for [`TensorDesc::ptr`] under
689 /// [`kind::PBO`] — see that field's own doc comment. `None` for every
690 /// other kind (mutually exclusive with `pin`: a PBO tensor never has
691 /// one, since `PboTensor` refuses `pin_host`).
692 pub pbo_vtable_ptr: Option<*const std::ffi::c_void>,
693 /// The process device fence's NT handle value, for [`TensorDesc::ptr`]
694 /// under [`kind::D3D11_TEXTURE`] — the third meaning that one field
695 /// carries, keyed by kind exactly as the PBO vtable is. `0` for every
696 /// other kind and for a device with no fence.
697 pub fence_handle: usize,
698 /// The fence value of the producer's last recorded GPU write, for
699 /// [`TensorDesc::sync`] under [`kind::D3D11_TEXTURE`]. `None` for every
700 /// other kind and when nothing has been recorded.
701 pub sync: Option<u64>,
702}
703
704/// Build a descriptor from its parts.
705///
706/// The single construction site: `TensorDyn::descriptor` and any future
707/// `Tensor<T>::descriptor` both route here, so the layout and the stride
708/// convention are defined once.
709///
710/// `ptr` is null unless a [`HostPin`](crate::HostPin) is supplied, or the
711/// tensor is PBO-backed and `pbo_vtable_ptr` is supplied instead — see
712/// [`TensorDesc::ptr`]'s own doc comment for why that one field carries two
713/// different meanings by [`TensorDesc::kind`]. A raw address is only
714/// meaningful while something guarantees the memory (or, for PBO, the
715/// `PboHandle` the vtable addresses) stays put, so the descriptor carries
716/// one only when the producer has pinned or built a vtable — and the
717/// capsule's keepalive then holds the relevant tensor alive for the
718/// consumer's benefit either way.
719///
720/// `row_stride` is the producer's [`TensorDyn::row_stride`](crate::TensorDyn::row_stride)
721/// **in bytes**, or `None` for a tightly-packed tensor. A pitch-aligned image
722/// carries padding between rows that the shape alone cannot express, so when
723/// present it overrides the stride of the dimension that is actually the row
724/// dimension for `format`'s layout. `format` is also the source of
725/// [`TensorDesc::format`] (via [`format_of`]) and is taken separately
726/// from `fourcc` rather than derived from it: `PlanarRgb`/`PlanarRgba` both
727/// encode as `fourcc = 0` (no standard FourCC exists for them), the same
728/// sentinel a non-image tensor uses, so `fourcc` alone cannot tell Planar
729/// apart from "no format" — for either purpose.
730///
731/// * Packed (`[H, W, C]`) / SemiPlanar (`[combined_H, W]`): dimension 0.
732/// * Planar (`[C, H, W]`, e.g. `PlanarRgb`): dimension 1 — dimension 0 there
733/// is the *plane* count, and clobbering it with the row pitch would
734/// corrupt the plane stride while leaving the true per-row stride
735/// (dimension 1) unfixed. `Tensor` tracks only one row-pitch scalar (no
736/// separate per-plane stride), so dimension 0 keeps its packed value here
737/// exactly as it did before any padding was known — unchanged, not
738/// recomputed.
739///
740/// `row_stride` is a byte pitch and is carried verbatim. It need not be a
741/// whole number of `dtype` elements: `set_row_stride` validates only
742/// `stride >= minimum`, not element alignment, and since strides are bytes
743/// there is no division to lose a remainder to. The element-stride
744/// representation this replaced could not express such a pitch at all — it
745/// fell back to the packed stride and reported a pitch the buffer does not
746/// have, across a package/ABI boundary.
747///
748/// `handle`, `colorimetry` and `capacity` are carried through verbatim; `-1`
749/// and `0` are the "unused"/"undefined" sentinels for the first two.
750pub(crate) fn from_parts(parts: DescParts) -> TensorDesc {
751 let DescParts {
752 dims,
753 memory,
754 dtype,
755 fourcc,
756 format,
757 row_stride,
758 handle,
759 colorimetry,
760 capacity,
761 pin,
762 pbo_vtable_ptr,
763 fence_handle,
764 sync,
765 } = parts;
766 let desc_kind = kind_of(memory);
767 // Both fields are meaningful only under `D3D11_TEXTURE`, and only when
768 // the device has a fence to name: a value with no handle behind it would
769 // set `SYNC_PRESENT` on a descriptor no consumer can wait on.
770 let completion = (desc_kind == kind::D3D11_TEXTURE && fence_handle != 0)
771 .then_some(sync)
772 .flatten()
773 .map(|value| (fence_handle, value));
774 // `ndim` carries the TRUE rank even when it exceeds the eight slots the
775 // descriptor can hold, so a consumer can detect the case and refuse.
776 // Reporting the clamped value instead would ship a descriptor whose
777 // `ndim` disagrees with its own `dims`, and the consumer would address
778 // `product(dims[..8])` elements of a larger allocation -- silently wrong
779 // data across a package boundary. This mirrors the `row_stride` override
780 // below, which likewise refuses to truncate quietly.
781 let ndim = dims.len();
782 let filled = ndim.min(TensorDesc::MAX_NDIM);
783 if ndim > TensorDesc::MAX_NDIM {
784 log::warn!(
785 "tensor of rank {ndim} exceeds the descriptor's {} shape slots; \
786 the descriptor will be rejected on import",
787 TensorDesc::MAX_NDIM
788 );
789 }
790 let mut shape = [0u64; TensorDesc::MAX_NDIM];
791 let mut strides = [0i64; TensorDesc::MAX_NDIM];
792 for (i, d) in dims.iter().take(filled).enumerate() {
793 shape[i] = *d as u64;
794 }
795 // Row-major and contiguous, in BYTES: stride[i] = dtype size * product of
796 // trailing dims. Bytes, not elements, because a hardware pitch need not be
797 // a whole number of elements (and sub-byte dtypes have no element size at
798 // all) -- see the `row_stride` override below.
799 let mut acc: i64 = dtype.size() as i64;
800 // `filled`, not `ndim`: the arrays hold MAX_NDIM slots and `ndim` may now
801 // legitimately exceed that.
802 for i in (0..filled).rev() {
803 strides[i] = acc;
804 acc *= shape[i] as i64;
805 }
806 if ndim >= 2 {
807 if let Some(rs) = row_stride {
808 // Planar stacks channels ahead of rows ([C, H, W]): the row
809 // dimension is 1, not 0. Every other layout this protocol
810 // carries (Packed [H, W, C], SemiPlanar [combined_H, W]) has the
811 // row dimension at 0. `format` is `None` for a non-image tensor,
812 // in which case `row_stride` should not have been supplied in
813 // the first place; dimension 0 is a harmless default for that
814 // case since it matches the pre-padding behaviour.
815 let row_dim = match format.map(|f| f.layout()) {
816 Some(crate::PixelLayout::Planar) if ndim >= 3 => 1,
817 _ => 0,
818 };
819 // Bytes in, bytes out. The element-stride representation had to
820 // divide here, and a remainder made the pitch unrepresentable --
821 // it fell back to the packed stride and reported a pitch the
822 // buffer does not have. There is nothing left to get wrong.
823 strides[row_dim] = rs as i64;
824 }
825 }
826 TensorDesc {
827 version: ABI_VERSION,
828 kind: desc_kind,
829 handle,
830 // Keyed by kind, because that is what this field's meaning is keyed
831 // by. Under `D3D11_TEXTURE` it is the fence handle or nothing, even
832 // when the producer also pinned a host address -- a Python capsule
833 // asked for `access="read"` does pin a texture tensor, and writing
834 // that address here would hand the consumer a host pointer to open
835 // as a fence. The pin stays the producer's own keepalive; a D3D11
836 // import reopens the texture and takes its own CPU mapping, so it
837 // never wanted the address. Otherwise `pin` and `pbo_vtable_ptr`
838 // are mutually exclusive by construction (see
839 // `DescParts::pbo_vtable_ptr`'s own doc comment) and `pin` takes
840 // precedence purely because it is the older, more general case.
841 ptr: if desc_kind == kind::D3D11_TEXTURE {
842 completion.map_or_else(SendPtr::null, |(h, _)| SendPtr(h as *mut u8))
843 } else {
844 pin.map(|p| SendPtr(p.as_mut_ptr()))
845 .or_else(|| pbo_vtable_ptr.map(|p| SendPtr(p as *mut u8)))
846 .unwrap_or_else(SendPtr::null)
847 },
848 ndim: ndim as u32,
849 dtype: dtype_of(dtype),
850 shape,
851 strides,
852 fourcc,
853 format: format.map(format_of).unwrap_or(format::NONE),
854 colorimetry,
855 // `D3D11_TEXTURE` is the one kind that fills these in; every other
856 // producer in this build leaves them clear. See the field docs.
857 flags: if completion.is_some() {
858 flags::SYNC_PRESENT
859 } else {
860 0
861 },
862 capacity,
863 sync: completion.map_or(0, |(_, value)| value),
864 }
865}
866
867// ---------------------------------------------------------------------------
868// Quantization, as it crosses the capsule.
869//
870// This lives here rather than in `edgefirst-python-common` beside
871// `TensorCapsulePayload`, which is the only thing that embeds it, for one
872// reason: it is unsafe pointer/length decoding of a versioned wire format,
873// and `edgefirst-python-common` cannot host a test that runs. That crate
874// pins `pyo3/extension-module`, which suppresses linking libpython, so its
875// test binary cannot link at all -- and `make test-rust` excludes it by
876// name regardless. Here, the round trip is covered by `quant_desc_tests`
877// below, which `make test` runs.
878//
879// It does NOT go in `TensorDesc`: that descriptor is also the C ABI's, and
880// the C libraries pass real `ef_tensor` handles into one shared
881// `libedgefirst_tensor.so` and never rebuild a tensor from a descriptor, so
882// they never lose the metadata. See `python-common/INTEROP.md`.
883// ---------------------------------------------------------------------------
884
885/// A `*const T` that may cross a `PyCapsule`.
886///
887/// The same obligation [`crate::SendPtr`] discharges for
888/// [`crate::TensorDesc::ptr`], for the two arrays [`QuantDesc`] points at: the
889/// producer keeps them alive for the capsule's lifetime, here via
890/// `TensorCapsulePayload::quant_keepalive`.
891#[repr(transparent)]
892#[derive(Clone, Copy, Debug)]
893pub struct SendConstPtr<T>(pub *const T);
894
895// SAFETY: the pointer is only dereferenced by a consumer that received it
896// through a capsule whose `quant_keepalive` owns the pointed-to arrays, so
897// the memory outlives every use. See `QuantDesc`.
898unsafe impl<T> Send for SendConstPtr<T> {}
899unsafe impl<T> Sync for SendConstPtr<T> {}
900
901impl<T> SendConstPtr<T> {
902 /// A null pointer, for a descriptor that carries no array.
903 pub const fn null() -> Self {
904 SendConstPtr(std::ptr::null())
905 }
906}
907
908/// Quantization metadata, in a form that can cross an `.so` boundary.
909///
910/// [`crate::TensorDesc`] deliberately does not carry this: it is the *shared*
911/// descriptor, pinned by `ABI_VERSION` and consumed by the C API, where
912/// tensors travel as real `ef_tensor` handles into one shared
913/// `libedgefirst_tensor.so` and no reconstruction ever happens. The Python
914/// packages are the only consumers that rebuild a tensor from a descriptor,
915/// so this rides in the capsule payload rather than widening a C ABI that
916/// does not need it.
917///
918/// Variable-length by nature -- per-channel quantization carries one scale
919/// (and optionally one zero-point) per channel -- so the arrays are
920/// *borrowed*, not inlined: `scales`/`zero_points` point into the
921/// `Quantization` that `TensorCapsulePayload::quant_keepalive` owns, on
922/// exactly the terms [`crate::TensorDesc::ptr`] borrows the producer's host
923/// address. A consumer copies the values out during import (see
924/// [`QuantDesc::to_quantization`]) while the capsule is still alive, and
925/// never retains the pointers.
926#[repr(C)]
927#[derive(Clone, Copy, Debug)]
928pub struct QuantDesc {
929 /// Number of entries in `scales` (and in `zero_points` when non-null).
930 /// **Zero means the tensor carries no quantization** -- the one field a
931 /// consumer must check before reading either pointer.
932 pub len: u64,
933 /// Channel axis for per-channel quantization, **plus one**; `0` means
934 /// per-tensor. Encoded this way so the whole struct zeroes to "absent"
935 /// rather than needing a separate presence flag for an `Option<usize>`.
936 pub axis_plus_one: u64,
937 /// `len` scale factors. Null iff `len == 0`.
938 pub scales: SendConstPtr<f32>,
939 /// `len` zero-points, or null for symmetric quantization.
940 pub zero_points: SendConstPtr<i32>,
941}
942
943impl QuantDesc {
944 /// The descriptor for a tensor with no quantization metadata.
945 pub const fn absent() -> Self {
946 Self {
947 len: 0,
948 axis_plus_one: 0,
949 scales: SendConstPtr::null(),
950 zero_points: SendConstPtr::null(),
951 }
952 }
953
954 /// Describe `quant`, borrowing its arrays.
955 ///
956 /// The caller must keep `quant` alive for as long as the descriptor is
957 /// readable -- in the capsule protocol, by storing the same `Arc` in
958 /// `TensorCapsulePayload::quant_keepalive`.
959 pub fn borrowing(quant: &crate::Quantization) -> Self {
960 Self {
961 len: quant.scale().len() as u64,
962 axis_plus_one: quant.axis().map_or(0, |a| a as u64 + 1),
963 scales: SendConstPtr(quant.scale().as_ptr()),
964 zero_points: quant
965 .zero_point()
966 .map_or(SendConstPtr::null(), |zp| SendConstPtr(zp.as_ptr())),
967 }
968 }
969
970 /// Copy the borrowed arrays into an owned [`crate::Quantization`], or `None`
971 /// when the producer declared no quantization.
972 ///
973 /// `Ok(None)` means "the producer has none"; a descriptor that declares
974 /// quantization but cannot be read back as a valid [`crate::Quantization`] is
975 /// an `Err`, not a second spelling of `None`. Conflating the two is
976 /// what made the bug this whole capsule version exists to fix so hard
977 /// to place: a tensor that quietly arrives unquantized surfaces much
978 /// later as the consumer's own "requires quantization metadata", which
979 /// names the symptom and not the cause.
980 ///
981 /// # Safety
982 ///
983 /// `self` must have come from a capsule that is still alive, so the
984 /// arrays `scales`/`zero_points` point at are still owned by the
985 /// producer's `quant_keepalive`.
986 pub unsafe fn to_quantization(&self) -> crate::Result<Option<crate::Quantization>> {
987 if self.len == 0 {
988 return Ok(None);
989 }
990 if self.scales.0.is_null() {
991 return Err(crate::Error::InvalidArgument(format!(
992 "malformed tensor capsule: quantization declares {} scale(s) but \
993 carries a null scale array",
994 self.len
995 )));
996 }
997 let len = self.len as usize;
998 // SAFETY: the caller's obligation, above: the producer's
999 // `quant_keepalive` owns these arrays and the capsule holding it is
1000 // alive, and `len` came from the same producer as the pointers.
1001 let scales = unsafe { std::slice::from_raw_parts(self.scales.0, len) }.to_vec();
1002 let zero_points = if self.zero_points.0.is_null() {
1003 None
1004 } else {
1005 Some(unsafe { std::slice::from_raw_parts(self.zero_points.0, len) }.to_vec())
1006 };
1007 // Rebuilt through the named constructors so an ill-formed
1008 // descriptor cannot produce a `Quantization` that the type's own
1009 // invariants say is impossible -- `crate::Quantization::validate` rejects
1010 // per-channel scales with no axis and a per-tensor scale that
1011 // carries one, so those shapes cannot be forced into existence here
1012 // either. A descriptor that names one of them is reported, not
1013 // silently downgraded to "unquantized".
1014 let axis = self.axis_plus_one.checked_sub(1).map(|a| a as usize);
1015 let quant = match (axis, zero_points) {
1016 (None, None) if len == 1 => Ok(crate::Quantization::per_tensor_symmetric(scales[0])),
1017 (None, Some(zp)) if len == 1 => Ok(crate::Quantization::per_tensor(scales[0], zp[0])),
1018 (Some(a), None) => crate::Quantization::per_channel_symmetric(scales, a),
1019 (Some(a), Some(zp)) => crate::Quantization::per_channel(scales, zp, a),
1020 // `len > 1` with no axis: per-channel scales that name no
1021 // channel dimension. `Quantization` has no such shape.
1022 (None, _) => {
1023 return Err(crate::Error::InvalidArgument(format!(
1024 "malformed tensor capsule: {len} quantization scales but no \
1025 channel axis; per-channel quantization must name one"
1026 )));
1027 }
1028 };
1029 quant.map(Some).map_err(|e| {
1030 crate::Error::InvalidArgument(format!(
1031 "malformed tensor capsule: quantization metadata is not valid: {e}"
1032 ))
1033 })
1034 }
1035}
1036
1037#[cfg(test)]
1038mod quant_desc_tests {
1039 use super::QuantDesc;
1040 use crate::Quantization;
1041
1042 /// Borrow `q`, then copy back out of the descriptor while `q` is still
1043 /// alive -- exactly the producer/consumer ordering the capsule enforces
1044 /// via `TensorCapsulePayload::quant_keepalive`.
1045 fn round_trip(q: &Quantization) -> Option<Quantization> {
1046 let desc = QuantDesc::borrowing(q);
1047 // SAFETY: `q` outlives `desc` here, which is the whole contract.
1048 unsafe { desc.to_quantization() }.expect("a well-formed descriptor must import")
1049 }
1050
1051 #[test]
1052 fn absent_descriptor_is_none_not_an_error() {
1053 // SAFETY: an all-zero descriptor names no arrays to read.
1054 let imported = unsafe { QuantDesc::absent().to_quantization() };
1055 assert!(imported
1056 .expect("absent is well-formed, not an error")
1057 .is_none());
1058 }
1059
1060 #[test]
1061 fn per_tensor_round_trips() {
1062 for q in [
1063 Quantization::per_tensor_symmetric(0.5),
1064 Quantization::per_tensor(0.125, -7),
1065 ] {
1066 assert_eq!(round_trip(&q).as_ref(), Some(&q));
1067 }
1068 }
1069
1070 /// The arms this type exists for. A per-tensor descriptor carries one
1071 /// scale and no axis, so it would round-trip correctly even if the
1072 /// length and the axis were both ignored -- only a per-channel case
1073 /// exercises the borrowed variable-length arrays.
1074 #[test]
1075 fn per_channel_symmetric_round_trips_with_its_length_and_axis() {
1076 let q = Quantization::per_channel_symmetric(vec![0.5, 0.25, 0.125], 0)
1077 .expect("per-channel symmetric");
1078 let imported = round_trip(&q).expect("per-channel must survive the descriptor");
1079 assert_eq!(imported, q);
1080 // Spelled out rather than left to `PartialEq`: dropping the axis or
1081 // truncating the array to one entry is the specific failure the
1082 // borrowed-array encoding risks, and both would still be a valid
1083 // `Quantization`.
1084 assert_eq!(imported.scale(), &[0.5, 0.25, 0.125]);
1085 assert_eq!(imported.axis(), Some(0));
1086 assert_eq!(imported.zero_point(), None);
1087 }
1088
1089 #[test]
1090 fn per_channel_asymmetric_round_trips_with_its_zero_points() {
1091 let q = Quantization::per_channel(vec![0.5, 0.25, 0.125], vec![1, 2, 3], 2)
1092 .expect("per-channel asymmetric");
1093 let imported = round_trip(&q).expect("per-channel must survive the descriptor");
1094 assert_eq!(imported, q);
1095 assert_eq!(imported.scale().len(), 3);
1096 assert_eq!(imported.axis(), Some(2));
1097 assert_eq!(imported.zero_point(), Some(&[1, 2, 3][..]));
1098 }
1099
1100 /// A non-zero axis must survive as itself, not be clamped to 0 -- the
1101 /// `+1` encoding is the one place an off-by-one would hide.
1102 #[test]
1103 fn the_axis_plus_one_encoding_is_reversible() {
1104 for axis in 0..4usize {
1105 let q =
1106 Quantization::per_channel_symmetric(vec![0.5, 0.25], axis).expect("per-channel");
1107 assert_eq!(QuantDesc::borrowing(&q).axis_plus_one, axis as u64 + 1);
1108 assert_eq!(round_trip(&q).and_then(|i| i.axis()), Some(axis));
1109 }
1110 }
1111
1112 /// A descriptor that survives the capsule-name check but disagrees with
1113 /// itself is *reported*, not forced into a `Quantization` whose own
1114 /// invariants say it cannot exist -- and not quietly downgraded to
1115 /// "unquantized" either, which would resurface as the consumer's own
1116 /// "requires quantization metadata" and name the symptom, not the cause.
1117 #[test]
1118 fn per_channel_without_an_axis_is_an_error_not_a_silent_none() {
1119 let q = Quantization::per_channel_symmetric(vec![0.5, 0.25], 0).expect("per-channel");
1120 let mut desc = QuantDesc::borrowing(&q);
1121 desc.axis_plus_one = 0;
1122 // SAFETY: `q` still owns the arrays `desc` points at.
1123 let err = unsafe { desc.to_quantization() }
1124 .expect_err("per-channel scales naming no axis is malformed");
1125 assert!(
1126 err.to_string().contains("no channel axis"),
1127 "the error must say what is wrong with the descriptor, got: {err}"
1128 );
1129 }
1130
1131 /// The other malformed shape the length field admits: a non-zero count
1132 /// with nothing to read it from. Checked before any dereference.
1133 #[test]
1134 fn a_null_scale_array_with_a_non_zero_length_is_an_error() {
1135 let mut desc = QuantDesc::absent();
1136 desc.len = 3;
1137 // SAFETY: nothing is dereferenced -- the null check precedes it.
1138 let err =
1139 unsafe { desc.to_quantization() }.expect_err("a length with a null array is malformed");
1140 assert!(err.to_string().contains("null scale array"), "got: {err}");
1141 }
1142}