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