Skip to main content

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    ///
144    /// Bit 0.
145    pub const SYNC_PRESENT: u32 = 1;
146}
147
148/// Backing store kind, as reported in [`TensorDesc::kind`].
149pub mod kind {
150    /// Host memory: `Mem` or `Shm`. `ptr` is valid; `handle` is -1.
151    pub const HOST: u32 = 0;
152    /// Linux dma-buf. `handle` is the fd.
153    pub const DMABUF: u32 = 1;
154    /// Apple IOSurface. `handle` is the surface id.
155    pub const IOSURFACE: u32 = 2;
156    /// OpenGL pixel buffer object. `handle` is the buffer id.
157    pub const PBO: u32 = 3;
158    /// CUDA device memory. `ptr` is a device pointer, not host-addressable.
159    pub const CUDA_DEVICE: u32 = 4;
160    /// Windows D3D11 texture. `handle` is the texture's NT shared handle and
161    /// `ptr` the device fence's NT shared handle, both valid in the producing
162    /// process (other modules of the same process included); `sync` is the
163    /// fence value of the last recorded GPU write when `SYNC_PRESENT` is set.
164    /// The `ptr` reuse follows `PBO`'s vtable precedent.
165    pub const D3D11_TEXTURE: u32 = 5;
166}
167
168/// A raw pointer that may cross a `PyCapsule`.
169///
170/// `PyCapsule::new` requires `T: Send`, which `*mut u8` is not. The producer
171/// discharges that obligation by keeping the owning tensor alive for the
172/// capsule's lifetime via the capsule's keepalive; see the module docs.
173#[repr(transparent)]
174#[derive(Clone, Copy, Debug)]
175pub struct SendPtr(pub *mut u8);
176
177// SAFETY: the pointer is only dereferenced by a consumer that received it
178// through a capsule whose keepalive holds the producing tensor alive, so the
179// memory outlives every use.
180unsafe impl Send for SendPtr {}
181unsafe impl Sync for SendPtr {}
182
183impl SendPtr {
184    /// A null pointer, for descriptors whose memory is not host-addressable.
185    pub const fn null() -> Self {
186        SendPtr(std::ptr::null_mut())
187    }
188    /// True when the descriptor carries no host address.
189    pub fn is_null(&self) -> bool {
190        self.0.is_null()
191    }
192}
193
194/// The cross-package descriptor. Layout is part of the contract: it is
195/// `#[repr(C)]` and its size is pinned by test.
196#[repr(C)]
197#[derive(Clone, Copy, Debug)]
198pub struct TensorDesc {
199    /// [`ABI_VERSION`]. Checked, never assumed.
200    pub version: u32,
201    /// One of the [`kind`] constants.
202    pub kind: u32,
203    /// dma-buf fd, IOSurface id, GL buffer id or D3D11 texture NT handle
204    /// value; `-1` when unused.
205    pub handle: i64,
206    /// Host pointer when addressable, else null -- **except under
207    /// [`kind::PBO`]**, where it instead carries a `*const
208    /// crate::pbo::PboOpsVtable`, or null if the producer attached none,
209    /// **and under [`kind::D3D11_TEXTURE`]**, where it carries the process
210    /// device fence's NT handle value, or null when [`Self::sync`] carries
211    /// nothing to wait on.
212    /// Interpretation is keyed by [`Self::kind`], the same established
213    /// pattern [`Self::sync`]'s own doc comment already documents for that
214    /// field: a PBO tensor never takes a `HostPin` (`PboTensor` refuses
215    /// `pin_host` — the address cannot outlive the map), so this field is
216    /// otherwise permanently unused for that kind, which is exactly why it
217    /// is the field that carries the vtable instead of a new one requiring
218    /// a layout change. Valid under the exact same borrow contract as the
219    /// host-pointer case: only while the producer's own keepalive (for the
220    /// capsule protocol, the capsule's) keeps the owning tensor alive.
221    pub ptr: SendPtr,
222    /// Number of significant entries in `shape` and `strides`.
223    pub ndim: u32,
224    /// Element type, one of the [`dtype`] constants.
225    ///
226    /// Deliberately placed here: `ndim` leaves a 4-byte hole before the
227    /// 8-aligned `shape`, so this field costs nothing. Without it `strides`
228    /// — documented in *elements* — cannot be converted to bytes by a
229    /// consumer, and a `Tensor<f32>` is indistinguishable from a
230    /// `Tensor<u8>` of the same shape.
231    pub dtype: u32,
232    /// Logical shape, most-significant dimension first.
233    pub shape: [u64; 8],
234    /// Strides **in bytes**, signed so a consumer can express flips.
235    ///
236    /// Bytes rather than elements so a hardware row pitch is always
237    /// representable: a pitch need not be a whole number of `dtype` elements
238    /// (and a sub-byte dtype has no element size to divide by). Multiply a
239    /// coordinate by the stride directly; do not scale by `dtype` size.
240    pub strides: [i64; 8],
241    /// DRM FourCC of the pixel format, or 0 for a non-image tensor. Kept for
242    /// third-party/DRM interop; a HAL-aware consumer should prefer
243    /// [`Self::format`] and fall back to this only when `format` is 0
244    /// (see [`mod@format`]'s doc for why `fourcc` alone is insufficient).
245    pub fourcc: u32,
246    /// One of the [`mod@format`] constants, or 0 for a non-image tensor.
247    pub format: u32,
248    /// Packed colorimetry, or 0 when unspecified.
249    pub colorimetry: u32,
250    /// Bit flags; see [`mod@flags`]. Zero when nothing applies.
251    ///
252    /// Occupies what was a 4-byte alignment hole before `capacity`, so it
253    /// cost nothing to add.
254    pub flags: u32,
255    /// Bytes available in the producer's underlying allocation, from this
256    /// tensor's window start (its [`TensorTrait::capacity_bytes`
257    /// ](crate::TensorTrait::capacity_bytes)) -- may exceed the byte size
258    /// `shape`/`strides` imply when the producer over-allocated: a pool
259    /// tensor sized for its largest expected image but holding a smaller
260    /// one today, or row-pitch/MCU-alignment padding a decoder needs beyond
261    /// the logical width. Without this, a consumer importing a `HOST`
262    /// descriptor has no way to learn that headroom exists and must treat
263    /// the declared shape as the allocation's exact size -- which starves a
264    /// later `configure_image`/`set_logical_shape` of capacity the producer
265    /// actually has. `DMABUF`/`IOSURFACE` consumers query their real
266    /// capacity from the handle itself and do not depend on this field, but
267    /// it is populated for every kind for consistency.
268    pub capacity: u64,
269    /// Completion handle for work still in flight against this tensor, or 0
270    /// when [`flags::SYNC_PRESENT`] is clear.
271    ///
272    /// Filled in by [`kind::D3D11_TEXTURE`] producers, and left clear by
273    /// every other kind in this build. It was declared before any producer
274    /// set it because the field has to exist before the first release or
275    /// adding it later costs a capsule-name bump, and it is the one gap the
276    /// synchronous-at-the-boundary contract cannot close: a *foreign*
277    /// producer with queued GPU work has no other way to say "wait on this
278    /// first".
279    ///
280    /// Interpretation is keyed by [`Self::kind`], since each backing already
281    /// implies its own fence flavour and a separate discriminant would be
282    /// redundant:
283    ///
284    /// * [`kind::DMABUF`] -- a `sync_file` fd from
285    ///   `DMA_BUF_IOCTL_EXPORT_SYNC_FILE`. The consumer owns it and must
286    ///   close it.
287    /// * [`kind::PBO`] -- a `GLsync` from `glFenceSync`, valid only in the
288    ///   producer's share group.
289    /// * [`kind::CUDA_DEVICE`] -- a `cudaEvent_t`.
290    /// * [`kind::D3D11_TEXTURE`] -- a value on the timeline of the
291    ///   `ID3D11Fence` [`Self::ptr`] names, not a handle of its own: the
292    ///   handle and the value together are what a consumer waits on.
293    /// * [`kind::HOST`], [`kind::IOSURFACE`] -- no fence flavour defined;
294    ///   `SYNC_PRESENT` must be clear.
295    pub sync: u64,
296}
297
298impl TensorDesc {
299    /// Maximum rank the fixed-size `shape`/`strides` arrays can carry.
300    pub const MAX_NDIM: usize = 8;
301
302    /// The logical shape as a slice.
303    pub fn shape(&self) -> &[u64] {
304        &self.shape[..(self.ndim as usize).min(Self::MAX_NDIM)]
305    }
306
307    /// The strides as a slice, in bytes.
308    pub fn strides(&self) -> &[i64] {
309        &self.strides[..(self.ndim as usize).min(Self::MAX_NDIM)]
310    }
311
312    /// Total element count.
313    pub fn len(&self) -> u64 {
314        self.shape().iter().product()
315    }
316
317    /// True when the descriptor addresses no elements.
318    pub fn is_empty(&self) -> bool {
319        self.len() == 0
320    }
321}
322
323// ---------------------------------------------------------------------------
324// Producer side
325// ---------------------------------------------------------------------------
326
327use crate::TensorMemory;
328
329/// Map a backing store to its [`kind`] constant.
330///
331/// `TensorMemory::DmaBuf` is a deliberately shared discriminant covering DMA-BUF,
332/// IOSurface and AHardwareBuffer (it is ABI-stable across platforms), so the
333/// platform decides which protocol kind it reports.
334pub(crate) fn kind_of(memory: TensorMemory) -> u32 {
335    match memory {
336        // Both are a host-addressable pointer with no handle, which is all
337        // `kind` records -- it is a narrower vocabulary than `TensorMemory`
338        // on purpose, not a copy of it that lost a variant.
339        TensorMemory::Mem | TensorMemory::Shm => kind::HOST,
340        TensorMemory::Pbo => kind::PBO,
341        TensorMemory::Cuda => kind::CUDA_DEVICE,
342        // Named specifically rather than through `DmaBuf`. Unreachable
343        // today (no backend reports it) but correct the day one does, which
344        // is cheaper than the panic the obvious `unreachable!()` becomes.
345        TensorMemory::IoSurface => kind::IOSURFACE,
346        TensorMemory::DmaBuf => {
347            #[cfg(any(target_os = "macos", target_os = "ios"))]
348            {
349                kind::IOSURFACE
350            }
351            // Windows has no dma-buf: the shared `DmaBuf` discriminant is a
352            // D3D11 texture there, whose handle field means something else
353            // again (an NT handle, not an fd).
354            #[cfg(target_os = "windows")]
355            {
356                kind::D3D11_TEXTURE
357            }
358            #[cfg(not(any(target_os = "macos", target_os = "ios", target_os = "windows")))]
359            {
360                kind::DMABUF
361            }
362        }
363    }
364}
365
366// ---------------------------------------------------------------------------
367// Consumer side: the per-kind field checks `TensorDyn::import_storage` needs
368// ---------------------------------------------------------------------------
369//
370// Each of these validates one kind's handle/pointer field and produces the
371// value a constructor takes. They live here, next to the descriptor they
372// read, rather than in either backend: a descriptor is an *untrusted*
373// cross-package payload, and two copies of "is this fd representable" or
374// "is this surface id in range" are two places for the refusals to drift.
375
376/// `dup` a `kind::DMABUF` descriptor's fd into one this process owns.
377///
378/// The producer still owns the original and will close it when its own
379/// tensor drops, so an import must have its own -- `from_fd` adopts what it
380/// is given.
381///
382/// # Errors
383///
384/// [`crate::Error::InvalidArgument`] when the descriptor carries no fd, or
385/// one not representable as a file descriptor. `as i32` would wrap a handle
386/// above `i32::MAX` onto an unrelated -- possibly live -- descriptor, which
387/// the `dup` would then adopt as the buffer.
388#[cfg(target_os = "linux")]
389pub(crate) fn dup_descriptor_fd(desc: &TensorDesc) -> crate::Result<std::os::fd::OwnedFd> {
390    use std::os::fd::BorrowedFd;
391    if desc.handle < 0 {
392        return Err(crate::Error::InvalidArgument(
393            "dma-buf descriptor carries no fd".into(),
394        ));
395    }
396    let raw_fd = i32::try_from(desc.handle).map_err(|_| {
397        crate::Error::InvalidArgument(format!(
398            "dma-buf descriptor fd {} is not representable as a file descriptor",
399            desc.handle
400        ))
401    })?;
402    // SAFETY: the producer owns `desc.handle` and the capsule keepalive
403    // holds it open for the duration of this call, so it is a valid fd to
404    // borrow (not adopt) for the dup below.
405    let borrowed = unsafe { BorrowedFd::borrow_raw(raw_fd) };
406    Ok(nix::unistd::dup(borrowed)?)
407}
408
409/// Read a `kind::IOSURFACE` descriptor's surface id.
410///
411/// # Errors
412///
413/// [`crate::Error::InvalidArgument`] when the descriptor carries no id, or
414/// one out of `u32` range. Liveness is *not* checked here -- that belongs
415/// with the lookup, which happens inside the constructor (and, under the
416/// `dynamic` backend, inside `libedgefirst_tensor.so`) so no window exists
417/// between "found it" and "retained it".
418#[cfg(any(target_os = "macos", target_os = "ios"))]
419pub(crate) fn descriptor_surface_id(desc: &TensorDesc) -> crate::Result<u32> {
420    if desc.handle < 0 {
421        return Err(crate::Error::InvalidArgument(
422            "IOSurface descriptor carries no surface id".into(),
423        ));
424    }
425    u32::try_from(desc.handle).map_err(|_| {
426        crate::Error::InvalidArgument(format!(
427            "IOSurface descriptor id {} is out of range",
428            desc.handle
429        ))
430    })
431}
432
433/// Refuse a `kind::HOST` descriptor that carries no address.
434///
435/// # Errors
436///
437/// [`crate::Error::InvalidArgument`], naming the producer-side fix: a
438/// descriptive-only capsule was requested, so no address was pinned.
439pub(crate) fn check_descriptor_host_ptr(desc: &TensorDesc) -> crate::Result<()> {
440    if desc.ptr.is_null() {
441        return Err(crate::Error::InvalidArgument(
442            "host descriptor has no address: the producer was asked \
443             for a descriptive-only capsule (request access=\"read\" \
444             or \"readwrite\")"
445                .into(),
446        ));
447    }
448    Ok(())
449}
450
451/// Read a `kind::PBO` descriptor's GL buffer id, and refuse one carrying no
452/// ops vtable.
453///
454/// Both checks together, unlike the other kinds' one apiece, because the
455/// `PBO` kind is the one that needs *two* of the descriptor's fields to be
456/// present before a constructor can be called at all -- `handle` names the
457/// buffer, `ptr` carries the vtable that can reach it, and either alone is
458/// useless.
459///
460/// # Errors
461///
462/// [`crate::Error::InvalidArgument`] when the buffer id is absent or out of
463/// range, or when no vtable address is carried.
464pub(crate) fn descriptor_pbo_buffer_id(desc: &TensorDesc) -> crate::Result<u32> {
465    if desc.handle < 0 {
466        return Err(crate::Error::InvalidArgument(
467            "PBO descriptor carries no buffer id".into(),
468        ));
469    }
470    let buffer_id = u32::try_from(desc.handle).map_err(|_| {
471        crate::Error::InvalidArgument(format!(
472            "PBO descriptor buffer id {} is out of range",
473            desc.handle
474        ))
475    })?;
476    if desc.ptr.is_null() {
477        return Err(crate::Error::InvalidArgument(
478            "PBO descriptor carries no ops vtable -- the producer either \
479             predates this build's cross-cdylib PBO support or its own \
480             PboTensor could not build one"
481                .into(),
482        ));
483    }
484    Ok(buffer_id)
485}
486
487/// Read a `kind::D3D11_TEXTURE` descriptor's texture handle and, when the
488/// producer recorded a GPU write, the fence handle plus value an import
489/// waits on before reading.
490///
491/// Both checks together, like [`descriptor_pbo_buffer_id`]'s, because this
492/// kind also spends two of the descriptor's fields on one backing:
493/// `handle` names the texture and `ptr` the fence that orders access to it.
494/// The completion is optional -- a producer that has queued no GPU work
495/// leaves `SYNC_PRESENT` clear, and there is then nothing to wait for.
496///
497/// Both handles stay the producer's: they are its own, not duplicates, and
498/// are valid only while the capsule keepalive holds the producing tensor
499/// alive. `from_d3d11_shared_handle` duplicates what it keeps.
500///
501/// # Errors
502///
503/// [`crate::Error::InvalidArgument`] when the descriptor carries no texture
504/// handle.
505#[cfg(target_os = "windows")]
506pub(crate) fn descriptor_d3d11_handles(
507    desc: &TensorDesc,
508) -> crate::Result<(
509    std::os::windows::io::RawHandle,
510    Option<(std::os::windows::io::RawHandle, u64)>,
511)> {
512    if desc.handle <= 0 {
513        return Err(crate::Error::InvalidArgument(
514            "D3D11_TEXTURE descriptor without a texture handle".into(),
515        ));
516    }
517    let tex = desc.handle as usize as std::os::windows::io::RawHandle;
518    let completion = (desc.flags & flags::SYNC_PRESENT != 0 && !desc.ptr.is_null())
519        .then_some((desc.ptr.0 as std::os::windows::io::RawHandle, desc.sync));
520    Ok((tex, completion))
521}
522
523/// The pixel format and image dimensions a `kind::D3D11_TEXTURE` import is
524/// built at, with the descriptor's shape checked against them.
525///
526/// The dimensions come from the texture, not from `shape`: a shape cannot
527/// tell an allocation grid from an addressing grid, and a producer may
528/// legitimately send either (`descriptor_pinned` emits `shape()`, which is
529/// the allocation shape until someone calls `set_logical_shape`). The one
530/// thing the texture cannot supply is a semi-planar image's width -- such a
531/// texture is as wide as the driver's row pitch -- which is why `shape` is
532/// passed down to [`crate::d3d11_shared_handle_geometry`] rather than only
533/// checked here.
534///
535/// A third spelling is accepted here and nowhere else: a packed *window* of
536/// the texture, which is what a `view()` describes (`[h', w', c]` with `h'`
537/// and `w'` no larger than the texture's, and no smaller than one). The
538/// import still opens the whole texture; `restore_d3d11_logical_shape`
539/// narrows it to the window and the producer's plane offset places it. The
540/// blob transport does not take this spelling because nothing narrows on
541/// that path, so a window shape would be accepted and then ignored.
542///
543/// # Errors
544///
545/// [`crate::Error::InvalidArgument`] when the descriptor names no pixel
546/// format, or when its shape is none of the three spellings of the geometry
547/// the texture reports. Propagates whatever opening and describing the
548/// texture reports.
549///
550/// # Safety
551///
552/// `texture` must be an NT shared handle of a D3D11 texture, valid in this
553/// process. It stays owned by the caller.
554#[cfg(target_os = "windows")]
555pub(crate) unsafe fn descriptor_d3d11_geometry(
556    desc: &TensorDesc,
557    texture: std::os::windows::io::RawHandle,
558    shape: &[usize],
559) -> crate::Result<(crate::PixelFormat, usize, usize)> {
560    let format = crate::PixelFormat::from_code(desc.format).ok_or_else(|| {
561        crate::Error::InvalidArgument(
562            "D3D11_TEXTURE descriptor without a pixel format: the texture's own \
563             description cannot say which format it was allocated for"
564                .into(),
565        )
566    })?;
567    // SAFETY: the caller guarantees `texture` is a shared NT handle valid in
568    // this process.
569    let (width, height) =
570        unsafe { d3d11_geometry_with_rule(format, texture, shape, ShapeRule::OrWindow) }?;
571    Ok((format, width, height))
572}
573
574/// Which spellings of a texture's geometry a `shape` may take.
575#[cfg(target_os = "windows")]
576#[derive(Clone, Copy, PartialEq, Eq)]
577enum ShapeRule {
578    /// The allocation shape or the addressing shape, exactly.
579    #[cfg_attr(feature = "dynamic", allow(dead_code))]
580    Exact,
581    /// `Exact`, or a packed window of the texture (a `view()`).
582    OrWindow,
583}
584
585/// Whether `shape` is a packed sub-rectangle of a `width` x `height` image
586/// in `format`: the allocation spelling with a smaller height or width. A
587/// window is never a semi-planar or planar image (`Tensor::view` makes
588/// packed views only) and never empty.
589#[cfg(target_os = "windows")]
590fn is_packed_window_of(
591    format: crate::PixelFormat,
592    shape: &[usize],
593    width: usize,
594    height: usize,
595) -> bool {
596    if format.layout() != crate::PixelLayout::Packed {
597        return false;
598    }
599    let [h, w, c] = shape else {
600        return false;
601    };
602    *c == format.channels() && (1..=height).contains(h) && (1..=width).contains(w)
603}
604
605/// The image dimensions a shared D3D11 texture reports, with `shape` checked
606/// against them.
607///
608/// Split out of [`descriptor_d3d11_geometry`] so the blob's reference-mode
609/// import ([`crate::blob::import`]) makes the same refusal rather than a
610/// second copy of it: the two transports differ in how they name the format,
611/// not in which shapes describe a given texture.
612///
613/// # Errors
614///
615/// [`crate::Error::InvalidArgument`] when `shape` is neither spelling of the
616/// geometry the texture reports. Propagates whatever opening and describing
617/// the texture reports.
618///
619/// # Safety
620///
621/// `texture` must be an NT shared handle of a D3D11 texture, valid in this
622/// process. It stays owned by the caller.
623#[cfg(target_os = "windows")]
624#[cfg_attr(feature = "dynamic", allow(dead_code))]
625pub(crate) unsafe fn d3d11_geometry_checked(
626    format: crate::PixelFormat,
627    texture: std::os::windows::io::RawHandle,
628    shape: &[usize],
629) -> crate::Result<(usize, usize)> {
630    // SAFETY: forwarded unchanged.
631    unsafe { d3d11_geometry_with_rule(format, texture, shape, ShapeRule::Exact) }
632}
633
634#[cfg(target_os = "windows")]
635unsafe fn d3d11_geometry_with_rule(
636    format: crate::PixelFormat,
637    texture: std::os::windows::io::RawHandle,
638    shape: &[usize],
639    rule: ShapeRule,
640) -> crate::Result<(usize, usize)> {
641    // SAFETY: the caller guarantees `texture` is a shared NT handle valid in
642    // this process; the helper opens its own texture and drops it on return.
643    let (width, height) =
644        unsafe { crate::d3d11_shared_handle_geometry(texture, format, Some(shape)) }?;
645    let allocation = format.allocation_shape(width, height);
646    let addressing = format.addressing_shape(width, height);
647    if allocation.as_deref() == Some(shape) || addressing.as_deref() == Some(shape) {
648        return Ok((width, height));
649    }
650    if rule == ShapeRule::OrWindow && is_packed_window_of(format, shape, width, height) {
651        return Ok((width, height));
652    }
653    let window = match rule {
654        ShapeRule::OrWindow => " nor a packed window",
655        ShapeRule::Exact => "",
656    };
657    Err(crate::Error::InvalidArgument(format!(
658        "D3D11 texture shape {shape:?} is neither the {format:?} allocation \
659         shape {allocation:?} nor the addressing shape {addressing:?}{window} of the \
660         {width}x{height} texture it names"
661    )))
662}
663
664/// Inputs to [`from_parts`], grouped into one named-field struct rather than
665/// a positional argument list.
666///
667/// This crossed nine parameters once the pixel-format and capacity inputs
668/// were added, several of them same-typed (`fourcc: u32`, a `u32`
669/// handle-adjacent value, `colorimetry: u32`) — a positional call at that
670/// width compiles happily with two arguments transposed and silently
671/// corrupts the descriptor. Named fields make that a compile error instead. Internal to
672/// this crate: `TensorDyn::descriptor_pinned` is presently the only caller --
673/// on both backends now, not just `static` (task P2a).
674pub(crate) struct DescParts<'a, 'p> {
675    pub dims: &'a [usize],
676    pub memory: TensorMemory,
677    pub dtype: crate::DType,
678    /// DRM FourCC for [`TensorDesc::fourcc`]; 0 when the format has none.
679    pub fourcc: u32,
680    /// The HAL format, source for both the row-dimension decision below and
681    /// [`TensorDesc::format`] (via [`format_of`]).
682    pub format: Option<crate::PixelFormat>,
683    /// Row pitch in bytes; see [`from_parts`]'s doc for how it maps to a
684    /// stride index.
685    pub row_stride: Option<usize>,
686    pub handle: i64,
687    pub colorimetry: u32,
688    /// Producer's [`crate::TensorTrait::capacity_bytes`] for
689    /// [`TensorDesc::capacity`]; see that field's docs.
690    pub capacity: u64,
691    pub pin: Option<&'a crate::HostPin<'p>>,
692    /// The producer's `PboOpsVtable` address, for [`TensorDesc::ptr`] under
693    /// [`kind::PBO`] — see that field's own doc comment. `None` for every
694    /// other kind (mutually exclusive with `pin`: a PBO tensor never has
695    /// one, since `PboTensor` refuses `pin_host`).
696    pub pbo_vtable_ptr: Option<*const std::ffi::c_void>,
697    /// The process device fence's NT handle value, for [`TensorDesc::ptr`]
698    /// under [`kind::D3D11_TEXTURE`] — the third meaning that one field
699    /// carries, keyed by kind exactly as the PBO vtable is. `0` for every
700    /// other kind and for a device with no fence.
701    pub fence_handle: usize,
702    /// The fence value of the producer's last recorded GPU write, for
703    /// [`TensorDesc::sync`] under [`kind::D3D11_TEXTURE`]. `None` for every
704    /// other kind and when nothing has been recorded.
705    pub sync: Option<u64>,
706}
707
708/// Row-major byte strides for `shape`, with the row dimension replaced by
709/// `row_stride` when one is recorded.
710///
711/// The one definition of the stride convention both wire formats carry: the
712/// descriptor ([`from_parts`]) and the blob (`blob::export`). See
713/// [`from_parts`] for which dimension is the row dimension per layout, and
714/// why a planar image's plane stride follows the pitch. Also the source of
715/// the C API's `ef_tensor_strides`, so every surface reports one layout.
716pub fn c_byte_strides(
717    shape: &[u64],
718    esz: i64,
719    layout: Option<crate::PixelLayout>,
720    row_stride: Option<usize>,
721) -> Vec<i64> {
722    let mut strides = vec![0i64; shape.len()];
723    let mut acc = esz;
724    for (stride, dim) in strides.iter_mut().zip(shape).rev() {
725        *stride = acc;
726        acc *= *dim as i64;
727    }
728    let Some(rs) = row_stride.map(|rs| rs as i64) else {
729        return strides;
730    };
731    // A formatted tensor one rank above its image rank is a batch `[N, ...]`
732    // of images stacked at the padded pitch, so the row override applies to
733    // the per-image dims and `N` steps over one padded image.
734    let image_rank = match layout {
735        Some(crate::PixelLayout::SemiPlanar) => Some(2),
736        Some(crate::PixelLayout::Packed | crate::PixelLayout::Planar) => Some(3),
737        None => None,
738    };
739    let batch_dims = match image_rank {
740        Some(rank) if shape.len() == rank + 1 => 1,
741        _ => 0,
742    };
743    let (lead, image) = strides.split_at_mut(batch_dims);
744    let image_bytes = match (layout, image, &shape[batch_dims..]) {
745        (Some(crate::PixelLayout::Planar), [plane, row, _, ..], [planes, rows, ..]) => {
746            *row = rs;
747            *plane = rs * *rows as i64;
748            *plane * *planes as i64
749        }
750        // A rank-1 tensor has no row to pad.
751        (_, [row, _, ..], [rows, ..]) => {
752            *row = rs;
753            rs * *rows as i64
754        }
755        _ => return strides,
756    };
757    if let [n] = lead {
758        *n = image_bytes;
759    }
760    strides
761}
762
763/// Build a descriptor from its parts.
764///
765/// The single construction site: `TensorDyn::descriptor` and any future
766/// `Tensor<T>::descriptor` both route here, so the layout and the stride
767/// convention are defined once.
768///
769/// `ptr` is null unless a [`HostPin`](crate::HostPin) is supplied, or the
770/// tensor is PBO-backed and `pbo_vtable_ptr` is supplied instead — see
771/// [`TensorDesc::ptr`]'s own doc comment for why that one field carries two
772/// different meanings by [`TensorDesc::kind`]. A raw address is only
773/// meaningful while something guarantees the memory (or, for PBO, the
774/// `PboHandle` the vtable addresses) stays put, so the descriptor carries
775/// one only when the producer has pinned or built a vtable — and the
776/// capsule's keepalive then holds the relevant tensor alive for the
777/// consumer's benefit either way.
778///
779/// `row_stride` is the producer's [`TensorDyn::row_stride`](crate::TensorDyn::row_stride)
780/// **in bytes**, or `None` for a tightly-packed tensor. A pitch-aligned image
781/// carries padding between rows that the shape alone cannot express, so when
782/// present it overrides the stride of the dimension that is actually the row
783/// dimension for `format`'s layout. `format` is also the source of
784/// [`TensorDesc::format`] (via [`format_of`]) and is taken separately
785/// from `fourcc` rather than derived from it: `PlanarRgb`/`PlanarRgba` both
786/// encode as `fourcc = 0` (no standard FourCC exists for them), the same
787/// sentinel a non-image tensor uses, so `fourcc` alone cannot tell Planar
788/// apart from "no format" — for either purpose.
789///
790/// * Packed (`[H, W, C]`) / SemiPlanar (`[combined_H, W]`): dimension 0.
791/// * Planar (`[C, H, W]`, e.g. `PlanarRgb`): dimension 1 — dimension 0 there
792///   is the *plane* count, and clobbering it with the row pitch would
793///   corrupt the plane stride while leaving the true per-row stride
794///   (dimension 1) unfixed. Dimension 0 becomes `row_stride * H` instead: a
795///   strided planar tensor stores its planes as `C * H` rows of one pitch,
796///   back to back (the layout `Tensor::map`, `copy_to_flat` and
797///   [`PixelFormat::plane_table`](crate::PixelFormat::plane_table) all
798///   assume), so the tight `H * W * size` would place plane 1 inside plane
799///   0's row padding.
800///
801/// `row_stride` is a byte pitch and is carried verbatim. It need not be a
802/// whole number of `dtype` elements: `set_row_stride` validates only
803/// `stride >= minimum`, not element alignment, and since strides are bytes
804/// there is no division to lose a remainder to. The element-stride
805/// representation this replaced could not express such a pitch at all — it
806/// fell back to the packed stride and reported a pitch the buffer does not
807/// have, across a package/ABI boundary.
808///
809/// `handle`, `colorimetry` and `capacity` are carried through verbatim; `-1`
810/// and `0` are the "unused"/"undefined" sentinels for the first two.
811pub(crate) fn from_parts(parts: DescParts) -> TensorDesc {
812    let DescParts {
813        dims,
814        memory,
815        dtype,
816        fourcc,
817        format,
818        row_stride,
819        handle,
820        colorimetry,
821        capacity,
822        pin,
823        pbo_vtable_ptr,
824        fence_handle,
825        sync,
826    } = parts;
827    let desc_kind = kind_of(memory);
828    // Both fields are meaningful only under `D3D11_TEXTURE`, and only when
829    // the device has a fence to name: a value with no handle behind it would
830    // set `SYNC_PRESENT` on a descriptor no consumer can wait on.
831    let completion = (desc_kind == kind::D3D11_TEXTURE && fence_handle != 0)
832        .then_some(sync)
833        .flatten()
834        .map(|value| (fence_handle, value));
835    // `ndim` carries the TRUE rank even when it exceeds the eight slots the
836    // descriptor can hold, so a consumer can detect the case and refuse.
837    // Reporting the clamped value instead would ship a descriptor whose
838    // `ndim` disagrees with its own `dims`, and the consumer would address
839    // `product(dims[..8])` elements of a larger allocation -- silently wrong
840    // data across a package boundary. This mirrors the `row_stride` override
841    // below, which likewise refuses to truncate quietly.
842    let ndim = dims.len();
843    let filled = ndim.min(TensorDesc::MAX_NDIM);
844    if ndim > TensorDesc::MAX_NDIM {
845        log::warn!(
846            "tensor of rank {ndim} exceeds the descriptor's {} shape slots; \
847             the descriptor will be rejected on import",
848            TensorDesc::MAX_NDIM
849        );
850    }
851    let mut shape = [0u64; TensorDesc::MAX_NDIM];
852    let mut strides = [0i64; TensorDesc::MAX_NDIM];
853    for (i, d) in dims.iter().take(filled).enumerate() {
854        shape[i] = *d as u64;
855    }
856    // `filled`, not `ndim`: the arrays hold MAX_NDIM slots and `ndim` may
857    // legitimately exceed that.
858    let byte_strides = c_byte_strides(
859        &shape[..filled],
860        dtype.size() as i64,
861        format.map(|f| f.layout()),
862        row_stride,
863    );
864    strides[..filled].copy_from_slice(&byte_strides);
865    TensorDesc {
866        version: ABI_VERSION,
867        kind: desc_kind,
868        handle,
869        // Keyed by kind, because that is what this field's meaning is keyed
870        // by. Under `D3D11_TEXTURE` it is the fence handle or nothing, even
871        // when the producer also pinned a host address -- a Python capsule
872        // asked for `access="read"` does pin a texture tensor, and writing
873        // that address here would hand the consumer a host pointer to open
874        // as a fence. The pin stays the producer's own keepalive; a D3D11
875        // import reopens the texture and takes its own CPU mapping, so it
876        // never wanted the address. Otherwise `pin` and `pbo_vtable_ptr`
877        // are mutually exclusive by construction (see
878        // `DescParts::pbo_vtable_ptr`'s own doc comment) and `pin` takes
879        // precedence purely because it is the older, more general case.
880        ptr: if desc_kind == kind::D3D11_TEXTURE {
881            completion.map_or_else(SendPtr::null, |(h, _)| SendPtr(h as *mut u8))
882        } else {
883            pin.map(|p| SendPtr(p.as_mut_ptr()))
884                .or_else(|| pbo_vtable_ptr.map(|p| SendPtr(p as *mut u8)))
885                .unwrap_or_else(SendPtr::null)
886        },
887        ndim: ndim as u32,
888        dtype: dtype_of(dtype),
889        shape,
890        strides,
891        fourcc,
892        format: format.map(format_of).unwrap_or(format::NONE),
893        colorimetry,
894        // `D3D11_TEXTURE` is the one kind that fills these in; every other
895        // producer in this build leaves them clear. See the field docs.
896        flags: if completion.is_some() {
897            flags::SYNC_PRESENT
898        } else {
899            0
900        },
901        capacity,
902        sync: completion.map_or(0, |(_, value)| value),
903    }
904}
905
906// ---------------------------------------------------------------------------
907// Quantization, as it crosses the capsule.
908//
909// This lives here rather than in `edgefirst-python-common` beside
910// `TensorCapsulePayload`, which is the only thing that embeds it, for one
911// reason: it is unsafe pointer/length decoding of a versioned wire format,
912// and `edgefirst-python-common` cannot host a test that runs. That crate
913// pins `pyo3/extension-module`, which suppresses linking libpython, so its
914// test binary cannot link at all -- and `make test-rust` excludes it by
915// name regardless. Here, the round trip is covered by `quant_desc_tests`
916// below, which `make test` runs.
917//
918// It does NOT go in `TensorDesc`: that descriptor is also the C ABI's, and
919// the C libraries pass real `ef_tensor` handles into one shared
920// `libedgefirst_tensor.so` and never rebuild a tensor from a descriptor, so
921// they never lose the metadata. See `python-common/INTEROP.md`.
922// ---------------------------------------------------------------------------
923
924/// A `*const T` that may cross a `PyCapsule`.
925///
926/// The same obligation [`crate::SendPtr`] discharges for
927/// [`crate::TensorDesc::ptr`], for the two arrays [`QuantDesc`] points at: the
928/// producer keeps them alive for the capsule's lifetime, here via
929/// `TensorCapsulePayload::quant_keepalive`.
930#[repr(transparent)]
931#[derive(Clone, Copy, Debug)]
932pub struct SendConstPtr<T>(pub *const T);
933
934// SAFETY: the pointer is only dereferenced by a consumer that received it
935// through a capsule whose `quant_keepalive` owns the pointed-to arrays, so
936// the memory outlives every use. See `QuantDesc`.
937unsafe impl<T> Send for SendConstPtr<T> {}
938unsafe impl<T> Sync for SendConstPtr<T> {}
939
940impl<T> SendConstPtr<T> {
941    /// A null pointer, for a descriptor that carries no array.
942    pub const fn null() -> Self {
943        SendConstPtr(std::ptr::null())
944    }
945}
946
947/// Quantization metadata, in a form that can cross an `.so` boundary.
948///
949/// [`crate::TensorDesc`] deliberately does not carry this: it is the *shared*
950/// descriptor, pinned by `ABI_VERSION` and consumed by the C API, where
951/// tensors travel as real `ef_tensor` handles into one shared
952/// `libedgefirst_tensor.so` and no reconstruction ever happens. The Python
953/// packages are the only consumers that rebuild a tensor from a descriptor,
954/// so this rides in the capsule payload rather than widening a C ABI that
955/// does not need it.
956///
957/// Variable-length by nature -- per-channel quantization carries one scale
958/// (and optionally one zero-point) per channel -- so the arrays are
959/// *borrowed*, not inlined: `scales`/`zero_points` point into the
960/// `Quantization` that `TensorCapsulePayload::quant_keepalive` owns, on
961/// exactly the terms [`crate::TensorDesc::ptr`] borrows the producer's host
962/// address. A consumer copies the values out during import (see
963/// [`QuantDesc::to_quantization`]) while the capsule is still alive, and
964/// never retains the pointers.
965#[repr(C)]
966#[derive(Clone, Copy, Debug)]
967pub struct QuantDesc {
968    /// Number of entries in `scales` (and in `zero_points` when non-null).
969    /// **Zero means the tensor carries no quantization** -- the one field a
970    /// consumer must check before reading either pointer.
971    pub len: u64,
972    /// Channel axis for per-channel quantization, **plus one**; `0` means
973    /// per-tensor. Encoded this way so the whole struct zeroes to "absent"
974    /// rather than needing a separate presence flag for an `Option<usize>`.
975    pub axis_plus_one: u64,
976    /// `len` scale factors. Null iff `len == 0`.
977    pub scales: SendConstPtr<f32>,
978    /// `len` zero-points, or null for symmetric quantization.
979    pub zero_points: SendConstPtr<i32>,
980}
981
982impl QuantDesc {
983    /// The descriptor for a tensor with no quantization metadata.
984    pub const fn absent() -> Self {
985        Self {
986            len: 0,
987            axis_plus_one: 0,
988            scales: SendConstPtr::null(),
989            zero_points: SendConstPtr::null(),
990        }
991    }
992
993    /// Describe `quant`, borrowing its arrays.
994    ///
995    /// The caller must keep `quant` alive for as long as the descriptor is
996    /// readable -- in the capsule protocol, by storing the same `Arc` in
997    /// `TensorCapsulePayload::quant_keepalive`.
998    pub fn borrowing(quant: &crate::Quantization) -> Self {
999        Self {
1000            len: quant.scale().len() as u64,
1001            axis_plus_one: quant.axis().map_or(0, |a| a as u64 + 1),
1002            scales: SendConstPtr(quant.scale().as_ptr()),
1003            zero_points: quant
1004                .zero_point()
1005                .map_or(SendConstPtr::null(), |zp| SendConstPtr(zp.as_ptr())),
1006        }
1007    }
1008
1009    /// Copy the borrowed arrays into an owned [`crate::Quantization`], or `None`
1010    /// when the producer declared no quantization.
1011    ///
1012    /// `Ok(None)` means "the producer has none"; a descriptor that declares
1013    /// quantization but cannot be read back as a valid [`crate::Quantization`] is
1014    /// an `Err`, not a second spelling of `None`. Conflating the two is
1015    /// what made the bug this whole capsule version exists to fix so hard
1016    /// to place: a tensor that quietly arrives unquantized surfaces much
1017    /// later as the consumer's own "requires quantization metadata", which
1018    /// names the symptom and not the cause.
1019    ///
1020    /// # Safety
1021    ///
1022    /// `self` must have come from a capsule that is still alive, so the
1023    /// arrays `scales`/`zero_points` point at are still owned by the
1024    /// producer's `quant_keepalive`.
1025    pub unsafe fn to_quantization(&self) -> crate::Result<Option<crate::Quantization>> {
1026        if self.len == 0 {
1027            return Ok(None);
1028        }
1029        if self.scales.0.is_null() {
1030            return Err(crate::Error::InvalidArgument(format!(
1031                "malformed tensor capsule: quantization declares {} scale(s) but \
1032                 carries a null scale array",
1033                self.len
1034            )));
1035        }
1036        let len = self.len as usize;
1037        // SAFETY: the caller's obligation, above: the producer's
1038        // `quant_keepalive` owns these arrays and the capsule holding it is
1039        // alive, and `len` came from the same producer as the pointers.
1040        let scales = unsafe { std::slice::from_raw_parts(self.scales.0, len) }.to_vec();
1041        let zero_points = if self.zero_points.0.is_null() {
1042            None
1043        } else {
1044            Some(unsafe { std::slice::from_raw_parts(self.zero_points.0, len) }.to_vec())
1045        };
1046        // Rebuilt through the named constructors so an ill-formed
1047        // descriptor cannot produce a `Quantization` that the type's own
1048        // invariants say is impossible -- `crate::Quantization::validate` rejects
1049        // per-channel scales with no axis and a per-tensor scale that
1050        // carries one, so those shapes cannot be forced into existence here
1051        // either. A descriptor that names one of them is reported, not
1052        // silently downgraded to "unquantized".
1053        let axis = self.axis_plus_one.checked_sub(1).map(|a| a as usize);
1054        let quant = match (axis, zero_points) {
1055            (None, None) if len == 1 => Ok(crate::Quantization::per_tensor_symmetric(scales[0])),
1056            (None, Some(zp)) if len == 1 => Ok(crate::Quantization::per_tensor(scales[0], zp[0])),
1057            (Some(a), None) => crate::Quantization::per_channel_symmetric(scales, a),
1058            (Some(a), Some(zp)) => crate::Quantization::per_channel(scales, zp, a),
1059            // `len > 1` with no axis: per-channel scales that name no
1060            // channel dimension. `Quantization` has no such shape.
1061            (None, _) => {
1062                return Err(crate::Error::InvalidArgument(format!(
1063                    "malformed tensor capsule: {len} quantization scales but no \
1064                     channel axis; per-channel quantization must name one"
1065                )));
1066            }
1067        };
1068        quant.map(Some).map_err(|e| {
1069            crate::Error::InvalidArgument(format!(
1070                "malformed tensor capsule: quantization metadata is not valid: {e}"
1071            ))
1072        })
1073    }
1074}
1075
1076#[cfg(test)]
1077mod quant_desc_tests {
1078    use super::QuantDesc;
1079    use crate::Quantization;
1080
1081    /// Borrow `q`, then copy back out of the descriptor while `q` is still
1082    /// alive -- exactly the producer/consumer ordering the capsule enforces
1083    /// via `TensorCapsulePayload::quant_keepalive`.
1084    fn round_trip(q: &Quantization) -> Option<Quantization> {
1085        let desc = QuantDesc::borrowing(q);
1086        // SAFETY: `q` outlives `desc` here, which is the whole contract.
1087        unsafe { desc.to_quantization() }.expect("a well-formed descriptor must import")
1088    }
1089
1090    #[test]
1091    fn absent_descriptor_is_none_not_an_error() {
1092        // SAFETY: an all-zero descriptor names no arrays to read.
1093        let imported = unsafe { QuantDesc::absent().to_quantization() };
1094        assert!(imported
1095            .expect("absent is well-formed, not an error")
1096            .is_none());
1097    }
1098
1099    #[test]
1100    fn per_tensor_round_trips() {
1101        for q in [
1102            Quantization::per_tensor_symmetric(0.5),
1103            Quantization::per_tensor(0.125, -7),
1104        ] {
1105            assert_eq!(round_trip(&q).as_ref(), Some(&q));
1106        }
1107    }
1108
1109    /// The arms this type exists for. A per-tensor descriptor carries one
1110    /// scale and no axis, so it would round-trip correctly even if the
1111    /// length and the axis were both ignored -- only a per-channel case
1112    /// exercises the borrowed variable-length arrays.
1113    #[test]
1114    fn per_channel_symmetric_round_trips_with_its_length_and_axis() {
1115        let q = Quantization::per_channel_symmetric(vec![0.5, 0.25, 0.125], 0)
1116            .expect("per-channel symmetric");
1117        let imported = round_trip(&q).expect("per-channel must survive the descriptor");
1118        assert_eq!(imported, q);
1119        // Spelled out rather than left to `PartialEq`: dropping the axis or
1120        // truncating the array to one entry is the specific failure the
1121        // borrowed-array encoding risks, and both would still be a valid
1122        // `Quantization`.
1123        assert_eq!(imported.scale(), &[0.5, 0.25, 0.125]);
1124        assert_eq!(imported.axis(), Some(0));
1125        assert_eq!(imported.zero_point(), None);
1126    }
1127
1128    #[test]
1129    fn per_channel_asymmetric_round_trips_with_its_zero_points() {
1130        let q = Quantization::per_channel(vec![0.5, 0.25, 0.125], vec![1, 2, 3], 2)
1131            .expect("per-channel asymmetric");
1132        let imported = round_trip(&q).expect("per-channel must survive the descriptor");
1133        assert_eq!(imported, q);
1134        assert_eq!(imported.scale().len(), 3);
1135        assert_eq!(imported.axis(), Some(2));
1136        assert_eq!(imported.zero_point(), Some(&[1, 2, 3][..]));
1137    }
1138
1139    /// A non-zero axis must survive as itself, not be clamped to 0 -- the
1140    /// `+1` encoding is the one place an off-by-one would hide.
1141    #[test]
1142    fn the_axis_plus_one_encoding_is_reversible() {
1143        for axis in 0..4usize {
1144            let q =
1145                Quantization::per_channel_symmetric(vec![0.5, 0.25], axis).expect("per-channel");
1146            assert_eq!(QuantDesc::borrowing(&q).axis_plus_one, axis as u64 + 1);
1147            assert_eq!(round_trip(&q).and_then(|i| i.axis()), Some(axis));
1148        }
1149    }
1150
1151    /// A descriptor that survives the capsule-name check but disagrees with
1152    /// itself is *reported*, not forced into a `Quantization` whose own
1153    /// invariants say it cannot exist -- and not quietly downgraded to
1154    /// "unquantized" either, which would resurface as the consumer's own
1155    /// "requires quantization metadata" and name the symptom, not the cause.
1156    #[test]
1157    fn per_channel_without_an_axis_is_an_error_not_a_silent_none() {
1158        let q = Quantization::per_channel_symmetric(vec![0.5, 0.25], 0).expect("per-channel");
1159        let mut desc = QuantDesc::borrowing(&q);
1160        desc.axis_plus_one = 0;
1161        // SAFETY: `q` still owns the arrays `desc` points at.
1162        let err = unsafe { desc.to_quantization() }
1163            .expect_err("per-channel scales naming no axis is malformed");
1164        assert!(
1165            err.to_string().contains("no channel axis"),
1166            "the error must say what is wrong with the descriptor, got: {err}"
1167        );
1168    }
1169
1170    /// The other malformed shape the length field admits: a non-zero count
1171    /// with nothing to read it from. Checked before any dereference.
1172    #[test]
1173    fn a_null_scale_array_with_a_non_zero_length_is_an_error() {
1174        let mut desc = QuantDesc::absent();
1175        desc.len = 3;
1176        // SAFETY: nothing is dereferenced -- the null check precedes it.
1177        let err =
1178            unsafe { desc.to_quantization() }.expect_err("a length with a null array is malformed");
1179        assert!(err.to_string().contains("null scale array"), "got: {err}");
1180    }
1181}
1182
1183#[cfg(test)]
1184mod desc_tests {
1185    use super::{c_byte_strides, from_parts, DescParts};
1186    use crate::{DType, PixelFormat, PixelLayout, TensorMemory};
1187
1188    #[test]
1189    fn a_padded_batch_steps_n_over_one_padded_image() {
1190        // [N, H, W, C] = [2, 2, 8, 3] at a 32 B pitch: each tile is 2 rows.
1191        assert_eq!(
1192            c_byte_strides(&[2, 2, 8, 3], 1, Some(PixelLayout::Packed), Some(32)),
1193            [64, 32, 3, 1]
1194        );
1195        // [N, C, H, W] = [2, 3, 4, 8] at 16 B: 3 planes of 4 rows per image.
1196        assert_eq!(
1197            c_byte_strides(&[2, 3, 4, 8], 1, Some(PixelLayout::Planar), Some(16)),
1198            [192, 64, 16, 1]
1199        );
1200        // [N, H*k, W] = [2, 6, 8] NV12 at 32 B: 6 combined rows per image.
1201        assert_eq!(
1202            c_byte_strides(&[2, 6, 8], 1, Some(PixelLayout::SemiPlanar), Some(32)),
1203            [192, 32, 1]
1204        );
1205        // Unformatted rank 4 has no batch reading: the pitch lands on dim 0.
1206        assert_eq!(
1207            c_byte_strides(&[2, 2, 8, 3], 1, None, Some(64)),
1208            [64, 24, 3, 1]
1209        );
1210    }
1211
1212    #[test]
1213    fn a_tight_tensor_gets_row_major_byte_strides() {
1214        assert_eq!(c_byte_strides(&[3, 8], 4, None, None), [32, 4]);
1215        assert_eq!(
1216            c_byte_strides(&[48, 64, 3], 1, Some(PixelLayout::Packed), None),
1217            [192, 3, 1]
1218        );
1219    }
1220
1221    #[test]
1222    fn a_recorded_pitch_replaces_the_row_stride_of_a_packed_or_semi_planar_image() {
1223        assert_eq!(
1224            c_byte_strides(&[48, 64, 3], 1, Some(PixelLayout::Packed), Some(256)),
1225            [256, 3, 1]
1226        );
1227        assert_eq!(
1228            c_byte_strides(&[48, 64], 1, Some(PixelLayout::SemiPlanar), Some(128)),
1229            [128, 1]
1230        );
1231    }
1232
1233    /// Planes of a strided planar image are `H` rows of the pitch apart, so
1234    /// the plane stride follows the pitch rather than staying `H * W`.
1235    #[test]
1236    fn a_recorded_pitch_moves_both_the_row_and_the_plane_stride_of_a_planar_image() {
1237        assert_eq!(
1238            c_byte_strides(&[3, 48, 64], 1, Some(PixelLayout::Planar), Some(128)),
1239            [128 * 48, 128, 1]
1240        );
1241        // Bytes, not elements: an f32 planar image at a 320-byte pitch.
1242        assert_eq!(
1243            c_byte_strides(&[3, 48, 64], 4, Some(PixelLayout::Planar), Some(320)),
1244            [320 * 48, 320, 4]
1245        );
1246    }
1247
1248    #[test]
1249    fn a_rank_one_tensor_has_no_row_to_pad() {
1250        assert_eq!(c_byte_strides(&[16], 1, None, Some(64)), [1]);
1251    }
1252
1253    fn parts(
1254        memory: TensorMemory,
1255        fence_handle: usize,
1256        sync: Option<u64>,
1257    ) -> DescParts<'static, 'static> {
1258        DescParts {
1259            dims: &[48, 64, 3],
1260            memory,
1261            dtype: DType::U8,
1262            fourcc: 0,
1263            format: Some(PixelFormat::Rgb),
1264            row_stride: None,
1265            handle: -1,
1266            colorimetry: 0,
1267            capacity: 48 * 64 * 3,
1268            pin: None,
1269            pbo_vtable_ptr: None,
1270            fence_handle,
1271            sync,
1272        }
1273    }
1274
1275    /// Only a D3D11 texture has a fence to name. A completion handed to any
1276    /// other kind is dropped, never advertised as `SYNC_PRESENT`.
1277    #[test]
1278    fn a_completion_on_a_kind_with_no_fence_is_not_advertised() {
1279        let d = from_parts(parts(TensorMemory::Mem, 5, Some(9)));
1280        assert_eq!(d.flags, 0);
1281        assert_eq!(d.sync, 0);
1282        assert!(d.ptr.is_null());
1283    }
1284
1285    /// On Windows the shared `DmaBuf` discriminant is a D3D11 texture: its
1286    /// completion needs both the fence handle and a recorded value.
1287    #[cfg(target_os = "windows")]
1288    #[test]
1289    fn a_texture_completion_needs_a_fence_handle() {
1290        use super::flags;
1291        let d = from_parts(parts(TensorMemory::DmaBuf, 5, Some(9)));
1292        assert_eq!(d.flags, flags::SYNC_PRESENT);
1293        assert_eq!(d.sync, 9);
1294        assert_eq!(d.ptr.0 as usize, 5);
1295
1296        let d = from_parts(parts(TensorMemory::DmaBuf, 0, Some(9)));
1297        assert_eq!(
1298            d.flags, 0,
1299            "a value with no fence handle cannot be waited on"
1300        );
1301        assert_eq!(d.sync, 0);
1302        assert!(d.ptr.is_null());
1303    }
1304}