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_v1` wrapping [`TensorDesc`]; a consumer reads the
15//! descriptor and never performs an `isinstance` check. That duck typing is
16//! what makes the protocol survive independent release cadences, and it is the
17//! same pattern numpy, pyarrow and DLPack use.
18//!
19//! DLPack cannot serve this role alone: it has no device type for dma-buf,
20//! IOSurface or PBO, which are precisely the zero-copy cases this library
21//! exists for.
22
23/// Protocol version, checked by `TensorDyn::import_descriptor` after the
24/// payload has already been read as this build's `TensorDesc`.
25///
26/// This is the **second line of defense, not the gate**: it cannot catch a
27/// producer whose `TensorDesc` is a different *size*, since the
28/// out-of-bounds/misaligned read has already happened by the time this
29/// field could be inspected. The capsule name (`edgefirst_tensor_v1`, see
30/// `INTEROP.md`'s Versioning section) is what actually gates that --
31/// checked by `PyCapsule::pointer_checked` before any byte of the payload
32/// is read. DLPack reached the same conclusion after shipping an
33/// unversioned `dltensor` capsule and finding no upgrade path: its v1.0 fix
34/// added both a version field *and* a new capsule name, because the field
35/// alone is unreadable until the layout is already known.
36///
37/// This field stays useful for the case the name rule does not cover on its
38/// own: a future change to what a *same-sized, same-layout* `TensorDesc`'s
39/// fields mean. It has not needed to move past `1`.
40pub const ABI_VERSION: u32 = 1;
41
42/// Element type, as reported in [`TensorDesc::dtype`].
43///
44/// These codes cross a package boundary and are read by value, so they are
45/// part of the ABI: append, never reorder. Re-exported from the wire module
46/// [`crate::DType`] itself generates (via `ef_vocabulary!`, see
47/// `vocabulary.rs`), so this and [`crate::DType::code`] are two views of the
48/// same declared literal per variant, not two hand-kept tables that could
49/// drift apart.
50///
51/// `protocol::dtype` is the canonical, documented path -- the source module,
52/// `crate::dtype_wire`, is `#[doc(hidden)]` because it exists only so this
53/// re-export has something public to alias; it is emission plumbing, not a
54/// second API for the same codes.
55pub use crate::dtype_wire as dtype;
56
57/// Map a [`crate::DType`] to its wire code.
58///
59/// Compiled into both backends: [`from_parts`] is reached from
60/// `TensorDyn::descriptor_pinned`, which now exists on the `dynamic`
61/// backend too (task P2a). Nothing here touches storage -- it is the wire
62/// encoding of facts both backends can already answer -- so there was never
63/// a backend-specific reason for the gate this replaces, only a
64/// dead-code-warning one.
65pub(crate) fn dtype_of(d: crate::DType) -> u32 {
66 d.code()
67}
68
69/// Map a wire [`dtype`] code back to a [`crate::DType`]. Inverse of
70/// [`dtype_of`]. `None` for a code this build does not recognise (e.g. a
71/// dtype added by a newer producer this build predates) — the caller
72/// rejects rather than guessing.
73pub fn dtype_to_dtype(code: u32) -> Option<crate::DType> {
74 crate::DType::from_code(code)
75}
76
77/// HAL pixel format, as reported in [`TensorDesc::format`].
78///
79/// Same ABI concern as [`dtype`]: these codes cross a package boundary, so
80/// they are appended, never reordered. [`crate::PixelFormat`] itself
81/// generates this module's per-variant constants (via `ef_vocabulary!`, see
82/// `vocabulary.rs`), so this and [`crate::PixelFormat::code`] are two views
83/// of the same declared literal per variant, not two hand-kept tables that
84/// could drift apart. [`format::NONE`] is the one constant this module adds
85/// beyond what the macro emits -- see its own doc.
86///
87/// This exists because [`TensorDesc::fourcc`] cannot represent every HAL
88/// format: `PlanarRgb`/`PlanarRgba` have no standard FourCC and both encode
89/// as `fourcc = 0`, the same sentinel a non-image tensor uses, so `fourcc`
90/// alone cannot tell "Planar RGB" apart from "no format". `fourcc` stays for
91/// third-party/DRM interop; a HAL-aware consumer prefers `format` and
92/// falls back to `fourcc` only when `format` is `0`.
93///
94/// `protocol::format` is the canonical, documented path -- the source
95/// module, `crate::pixel_format_wire`, is `#[doc(hidden)]` because it exists
96/// only so this re-export has something public to alias; it is emission
97/// plumbing, not a second API for the same codes.
98pub mod format {
99 /// No format (non-image tensor), or a format this build predates. Not a
100 /// [`crate::PixelFormat`] variant -- every code `1..=11` below is one of
101 /// those, generated by the same `ef_vocabulary!` declaration as
102 /// [`crate::PixelFormat::code`]; `0` is a sentinel this module alone
103 /// defines, since "no format" has no corresponding Rust variant.
104 pub const NONE: u32 = 0;
105 pub use crate::format::pixel_format_wire::*;
106}
107
108/// Map a [`crate::PixelFormat`] to its [`mod@format`] wire code.
109///
110/// Compiled into both backends: same reason as [`dtype_of`].
111pub(crate) fn format_of(f: crate::PixelFormat) -> u32 {
112 f.code()
113}
114
115/// Map a [`mod@format`] wire code back to a [`crate::PixelFormat`]. `None` for
116/// [`format::NONE`] or a code this build does not recognise (e.g. a format
117/// added by a newer producer).
118pub fn format_from_code(code: u32) -> Option<crate::PixelFormat> {
119 crate::PixelFormat::from_code(code)
120}
121
122/// Bit flags, as reported in [`TensorDesc::flags`].
123///
124/// A bitfield rather than more discrete fields: the additions this protocol
125/// is most likely to need are booleans, and each one spent as its own `u32`
126/// would grow the descriptor -- which after release costs a capsule-name
127/// bump. Spending one `u32` now buys 32 of them for free.
128pub mod flags {
129 /// [`TensorDesc::sync`] carries a real handle. When clear, `sync` is
130 /// meaningless and must not be waited on.
131 ///
132 /// A flag rather than a sentinel because every plausible sentinel is a
133 /// legal handle somewhere: `0` is a valid `GLsync`-adjacent value and a
134 /// valid fd number is never negative, so `-1` and `0` cannot both be
135 /// spare across the kinds this field spans.
136 ///
137 /// Under [`kind::D3D11_TEXTURE`] `sync` is not a handle at all: it is a
138 /// value on the timeline of the fence whose NT handle
139 /// [`TensorDesc::ptr`] carries, so a producer that sets this flag must
140 /// also fill in `ptr`.
141 pub const SYNC_PRESENT: u32 = 1 << 0;
142}
143
144/// Backing store kind, as reported in [`TensorDesc::kind`].
145pub mod kind {
146 /// Host memory: `Mem` or `Shm`. `ptr` is valid; `handle` is -1.
147 pub const HOST: u32 = 0;
148 /// Linux dma-buf. `handle` is the fd.
149 pub const DMABUF: u32 = 1;
150 /// Apple IOSurface. `handle` is the surface id.
151 pub const IOSURFACE: u32 = 2;
152 /// OpenGL pixel buffer object. `handle` is the buffer id.
153 pub const PBO: u32 = 3;
154 /// CUDA device memory. `ptr` is a device pointer, not host-addressable.
155 pub const CUDA_DEVICE: u32 = 4;
156 /// Windows D3D11 texture. `handle` is the texture's NT shared handle and
157 /// `ptr` the device fence's NT shared handle, both valid in the producing
158 /// process (other modules of the same process included); `sync` is the
159 /// fence value of the last recorded GPU write when `SYNC_PRESENT` is set.
160 /// The `ptr` reuse follows `PBO`'s vtable precedent.
161 pub const D3D11_TEXTURE: u32 = 5;
162}
163
164/// A raw pointer that may cross a `PyCapsule`.
165///
166/// `PyCapsule::new` requires `T: Send`, which `*mut u8` is not. The producer
167/// discharges that obligation by keeping the owning tensor alive for the
168/// capsule's lifetime via the capsule's keepalive; see the module docs.
169#[repr(transparent)]
170#[derive(Clone, Copy, Debug)]
171pub struct SendPtr(pub *mut u8);
172
173// SAFETY: the pointer is only dereferenced by a consumer that received it
174// through a capsule whose keepalive holds the producing tensor alive, so the
175// memory outlives every use.
176unsafe impl Send for SendPtr {}
177unsafe impl Sync for SendPtr {}
178
179impl SendPtr {
180 /// A null pointer, for descriptors whose memory is not host-addressable.
181 pub const fn null() -> Self {
182 SendPtr(std::ptr::null_mut())
183 }
184 /// True when the descriptor carries no host address.
185 pub fn is_null(&self) -> bool {
186 self.0.is_null()
187 }
188}
189
190/// The cross-package descriptor. Layout is part of the contract: it is
191/// `#[repr(C)]` and its size is pinned by test.
192#[repr(C)]
193#[derive(Clone, Copy, Debug)]
194pub struct TensorDesc {
195 /// [`ABI_VERSION`]. Checked, never assumed.
196 pub version: u32,
197 /// One of the [`kind`] constants.
198 pub kind: u32,
199 /// dma-buf fd, IOSurface id, GL buffer id or D3D11 texture NT handle
200 /// value; `-1` when unused.
201 pub handle: i64,
202 /// Host pointer when addressable, else null -- **except under
203 /// [`kind::PBO`]**, where it instead carries a `*const
204 /// crate::pbo::PboOpsVtable`, or null if the producer attached none,
205 /// **and under [`kind::D3D11_TEXTURE`]**, where it carries the process
206 /// device fence's NT handle value, or null when [`Self::sync`] carries
207 /// nothing to wait on.
208 /// Interpretation is keyed by [`Self::kind`], the same established
209 /// pattern [`Self::sync`]'s own doc comment already documents for that
210 /// field: a PBO tensor never takes a `HostPin` (`PboTensor` refuses
211 /// `pin_host` — the address cannot outlive the map), so this field is
212 /// otherwise permanently unused for that kind, which is exactly why it
213 /// is the field that carries the vtable instead of a new one requiring
214 /// a layout change. Valid under the exact same borrow contract as the
215 /// host-pointer case: only while the producer's own keepalive (for the
216 /// capsule protocol, the capsule's) keeps the owning tensor alive.
217 pub ptr: SendPtr,
218 /// Number of significant entries in `shape` and `strides`.
219 pub ndim: u32,
220 /// Element type, one of the [`dtype`] constants.
221 ///
222 /// Deliberately placed here: `ndim` leaves a 4-byte hole before the
223 /// 8-aligned `shape`, so this field costs nothing. Without it `strides`
224 /// — documented in *elements* — cannot be converted to bytes by a
225 /// consumer, and a `Tensor<f32>` is indistinguishable from a
226 /// `Tensor<u8>` of the same shape.
227 pub dtype: u32,
228 /// Logical shape, most-significant dimension first.
229 pub shape: [u64; 8],
230 /// Strides **in bytes**, signed so a consumer can express flips.
231 ///
232 /// Bytes rather than elements so a hardware row pitch is always
233 /// representable: a pitch need not be a whole number of `dtype` elements
234 /// (and a sub-byte dtype has no element size to divide by). Multiply a
235 /// coordinate by the stride directly; do not scale by `dtype` size.
236 pub strides: [i64; 8],
237 /// DRM FourCC of the pixel format, or 0 for a non-image tensor. Kept for
238 /// third-party/DRM interop; a HAL-aware consumer should prefer
239 /// [`Self::format`] and fall back to this only when `format` is 0
240 /// (see [`mod@format`]'s doc for why `fourcc` alone is insufficient).
241 pub fourcc: u32,
242 /// One of the [`mod@format`] constants, or 0 for a non-image tensor.
243 pub format: u32,
244 /// Packed colorimetry, or 0 when unspecified.
245 pub colorimetry: u32,
246 /// Bit flags; see [`mod@flags`]. Zero when nothing applies.
247 ///
248 /// Occupies what was a 4-byte alignment hole before `capacity`, so it
249 /// cost nothing to add.
250 pub flags: u32,
251 /// Bytes available in the producer's underlying allocation, from this
252 /// tensor's window start (its [`TensorTrait::capacity_bytes`
253 /// ](crate::TensorTrait::capacity_bytes)) -- may exceed the byte size
254 /// `shape`/`strides` imply when the producer over-allocated: a pool
255 /// tensor sized for its largest expected image but holding a smaller
256 /// one today, or row-pitch/MCU-alignment padding a decoder needs beyond
257 /// the logical width. Without this, a consumer importing a `HOST`
258 /// descriptor has no way to learn that headroom exists and must treat
259 /// the declared shape as the allocation's exact size -- which starves a
260 /// later `configure_image`/`set_logical_shape` of capacity the producer
261 /// actually has. `DMABUF`/`IOSURFACE` consumers query their real
262 /// capacity from the handle itself and do not depend on this field, but
263 /// it is populated for every kind for consistency.
264 pub capacity: u64,
265 /// Completion handle for work still in flight against this tensor, or 0
266 /// when [`flags::SYNC_PRESENT`] is clear.
267 ///
268 /// Filled in by [`kind::D3D11_TEXTURE`] producers, and left clear by
269 /// every other kind in this build. It was declared before any producer
270 /// set it because the field has to exist before the first release or
271 /// adding it later costs a capsule-name bump, and it is the one gap the
272 /// synchronous-at-the-boundary contract cannot close: a *foreign*
273 /// producer with queued GPU work has no other way to say "wait on this
274 /// first".
275 ///
276 /// Interpretation is keyed by [`Self::kind`], since each backing already
277 /// implies its own fence flavour and a separate discriminant would be
278 /// redundant:
279 ///
280 /// * [`kind::DMABUF`] -- a `sync_file` fd from
281 /// `DMA_BUF_IOCTL_EXPORT_SYNC_FILE`. The consumer owns it and must
282 /// close it.
283 /// * [`kind::PBO`] -- a `GLsync` from `glFenceSync`, valid only in the
284 /// producer's share group.
285 /// * [`kind::CUDA_DEVICE`] -- a `cudaEvent_t`.
286 /// * [`kind::D3D11_TEXTURE`] -- a value on the timeline of the
287 /// `ID3D11Fence` [`Self::ptr`] names, not a handle of its own: the
288 /// handle and the value together are what a consumer waits on.
289 /// * [`kind::HOST`], [`kind::IOSURFACE`] -- no fence flavour defined;
290 /// `SYNC_PRESENT` must be clear.
291 pub sync: u64,
292}
293
294impl TensorDesc {
295 /// Maximum rank the fixed-size `shape`/`strides` arrays can carry.
296 pub const MAX_NDIM: usize = 8;
297
298 /// The logical shape as a slice.
299 pub fn shape(&self) -> &[u64] {
300 &self.shape[..(self.ndim as usize).min(Self::MAX_NDIM)]
301 }
302
303 /// The strides as a slice, in bytes.
304 pub fn strides(&self) -> &[i64] {
305 &self.strides[..(self.ndim as usize).min(Self::MAX_NDIM)]
306 }
307
308 /// Total element count.
309 pub fn len(&self) -> u64 {
310 self.shape().iter().product()
311 }
312
313 /// True when the descriptor addresses no elements.
314 pub fn is_empty(&self) -> bool {
315 self.len() == 0
316 }
317}
318
319// ---------------------------------------------------------------------------
320// Producer side
321// ---------------------------------------------------------------------------
322
323use crate::TensorMemory;
324
325/// Map a backing store to its [`kind`] constant.
326///
327/// `TensorMemory::DmaBuf` is a deliberately shared discriminant covering DMA-BUF,
328/// IOSurface and AHardwareBuffer (it is ABI-stable across platforms), so the
329/// platform decides which protocol kind it reports.
330pub(crate) fn kind_of(memory: TensorMemory) -> u32 {
331 match memory {
332 // Both are a host-addressable pointer with no handle, which is all
333 // `kind` records -- it is a narrower vocabulary than `TensorMemory`
334 // on purpose, not a copy of it that lost a variant.
335 TensorMemory::Mem | TensorMemory::Shm => kind::HOST,
336 TensorMemory::Pbo => kind::PBO,
337 TensorMemory::Cuda => kind::CUDA_DEVICE,
338 // Named specifically rather than through `DmaBuf`. Unreachable
339 // today (no backend reports it) but correct the day one does, which
340 // is cheaper than the panic the obvious `unreachable!()` becomes.
341 TensorMemory::IoSurface => kind::IOSURFACE,
342 TensorMemory::DmaBuf => {
343 #[cfg(any(target_os = "macos", target_os = "ios"))]
344 {
345 kind::IOSURFACE
346 }
347 // Windows has no dma-buf: the shared `DmaBuf` discriminant is a
348 // D3D11 texture there, whose handle field means something else
349 // again (an NT handle, not an fd).
350 #[cfg(target_os = "windows")]
351 {
352 kind::D3D11_TEXTURE
353 }
354 #[cfg(not(any(target_os = "macos", target_os = "ios", target_os = "windows")))]
355 {
356 kind::DMABUF
357 }
358 }
359 }
360}
361
362// ---------------------------------------------------------------------------
363// Consumer side: the per-kind field checks `TensorDyn::import_storage` needs
364// ---------------------------------------------------------------------------
365//
366// Each of these validates one kind's handle/pointer field and produces the
367// value a constructor takes. They live here, next to the descriptor they
368// read, rather than in either backend: a descriptor is an *untrusted*
369// cross-package payload, and two copies of "is this fd representable" or
370// "is this surface id in range" are two places for the refusals to drift.
371
372/// `dup` a `kind::DMABUF` descriptor's fd into one this process owns.
373///
374/// The producer still owns the original and will close it when its own
375/// tensor drops, so an import must have its own -- `from_fd` adopts what it
376/// is given.
377///
378/// # Errors
379///
380/// [`crate::Error::InvalidArgument`] when the descriptor carries no fd, or
381/// one not representable as a file descriptor. `as i32` would wrap a handle
382/// above `i32::MAX` onto an unrelated -- possibly live -- descriptor, which
383/// the `dup` would then adopt as the buffer.
384#[cfg(target_os = "linux")]
385pub(crate) fn dup_descriptor_fd(desc: &TensorDesc) -> crate::Result<std::os::fd::OwnedFd> {
386 use std::os::fd::BorrowedFd;
387 if desc.handle < 0 {
388 return Err(crate::Error::InvalidArgument(
389 "dma-buf descriptor carries no fd".into(),
390 ));
391 }
392 let raw_fd = i32::try_from(desc.handle).map_err(|_| {
393 crate::Error::InvalidArgument(format!(
394 "dma-buf descriptor fd {} is not representable as a file descriptor",
395 desc.handle
396 ))
397 })?;
398 // SAFETY: the producer owns `desc.handle` and the capsule keepalive
399 // holds it open for the duration of this call, so it is a valid fd to
400 // borrow (not adopt) for the dup below.
401 let borrowed = unsafe { BorrowedFd::borrow_raw(raw_fd) };
402 Ok(nix::unistd::dup(borrowed)?)
403}
404
405/// Read a `kind::IOSURFACE` descriptor's surface id.
406///
407/// # Errors
408///
409/// [`crate::Error::InvalidArgument`] when the descriptor carries no id, or
410/// one out of `u32` range. Liveness is *not* checked here -- that belongs
411/// with the lookup, which happens inside the constructor (and, under the
412/// `dynamic` backend, inside `libedgefirst_tensor.so`) so no window exists
413/// between "found it" and "retained it".
414#[cfg(any(target_os = "macos", target_os = "ios"))]
415pub(crate) fn descriptor_surface_id(desc: &TensorDesc) -> crate::Result<u32> {
416 if desc.handle < 0 {
417 return Err(crate::Error::InvalidArgument(
418 "IOSurface descriptor carries no surface id".into(),
419 ));
420 }
421 u32::try_from(desc.handle).map_err(|_| {
422 crate::Error::InvalidArgument(format!(
423 "IOSurface descriptor id {} is out of range",
424 desc.handle
425 ))
426 })
427}
428
429/// Refuse a `kind::HOST` descriptor that carries no address.
430///
431/// # Errors
432///
433/// [`crate::Error::InvalidArgument`], naming the producer-side fix: a
434/// descriptive-only capsule was requested, so no address was pinned.
435pub(crate) fn check_descriptor_host_ptr(desc: &TensorDesc) -> crate::Result<()> {
436 if desc.ptr.is_null() {
437 return Err(crate::Error::InvalidArgument(
438 "host descriptor has no address: the producer was asked \
439 for a descriptive-only capsule (request access=\"read\" \
440 or \"readwrite\")"
441 .into(),
442 ));
443 }
444 Ok(())
445}
446
447/// Read a `kind::PBO` descriptor's GL buffer id, and refuse one carrying no
448/// ops vtable.
449///
450/// Both checks together, unlike the other kinds' one apiece, because the
451/// `PBO` kind is the one that needs *two* of the descriptor's fields to be
452/// present before a constructor can be called at all -- `handle` names the
453/// buffer, `ptr` carries the vtable that can reach it, and either alone is
454/// useless.
455///
456/// # Errors
457///
458/// [`crate::Error::InvalidArgument`] when the buffer id is absent or out of
459/// range, or when no vtable address is carried.
460pub(crate) fn descriptor_pbo_buffer_id(desc: &TensorDesc) -> crate::Result<u32> {
461 if desc.handle < 0 {
462 return Err(crate::Error::InvalidArgument(
463 "PBO descriptor carries no buffer id".into(),
464 ));
465 }
466 let buffer_id = u32::try_from(desc.handle).map_err(|_| {
467 crate::Error::InvalidArgument(format!(
468 "PBO descriptor buffer id {} is out of range",
469 desc.handle
470 ))
471 })?;
472 if desc.ptr.is_null() {
473 return Err(crate::Error::InvalidArgument(
474 "PBO descriptor carries no ops vtable -- the producer either \
475 predates this build's cross-cdylib PBO support or its own \
476 PboTensor could not build one"
477 .into(),
478 ));
479 }
480 Ok(buffer_id)
481}
482
483/// Read a `kind::D3D11_TEXTURE` descriptor's texture handle and, when the
484/// producer recorded a GPU write, the fence handle plus value an import
485/// waits on before reading.
486///
487/// Both checks together, like [`descriptor_pbo_buffer_id`]'s, because this
488/// kind also spends two of the descriptor's fields on one backing:
489/// `handle` names the texture and `ptr` the fence that orders access to it.
490/// The completion is optional -- a producer that has queued no GPU work
491/// leaves `SYNC_PRESENT` clear, and there is then nothing to wait for.
492///
493/// Both handles stay the producer's: they are its own, not duplicates, and
494/// are valid only while the capsule keepalive holds the producing tensor
495/// alive. `from_d3d11_shared_handle` duplicates what it keeps.
496///
497/// # Errors
498///
499/// [`crate::Error::InvalidArgument`] when the descriptor carries no texture
500/// handle.
501#[cfg(target_os = "windows")]
502pub(crate) fn descriptor_d3d11_handles(
503 desc: &TensorDesc,
504) -> crate::Result<(
505 std::os::windows::io::RawHandle,
506 Option<(std::os::windows::io::RawHandle, u64)>,
507)> {
508 if desc.handle <= 0 {
509 return Err(crate::Error::InvalidArgument(
510 "D3D11_TEXTURE descriptor without a texture handle".into(),
511 ));
512 }
513 let tex = desc.handle as usize as std::os::windows::io::RawHandle;
514 let completion = (desc.flags & flags::SYNC_PRESENT != 0 && !desc.ptr.is_null())
515 .then_some((desc.ptr.0 as std::os::windows::io::RawHandle, desc.sync));
516 Ok((tex, completion))
517}
518
519/// The pixel format and image dimensions a `kind::D3D11_TEXTURE` import is
520/// built at, with the descriptor's shape checked against them.
521///
522/// The dimensions come from the texture, not from `shape`: a shape cannot
523/// tell an allocation grid from an addressing grid, and a producer may
524/// legitimately send either (`descriptor_pinned` emits `shape()`, which is
525/// the allocation shape until someone calls `set_logical_shape`). The one
526/// thing the texture cannot supply is a semi-planar image's width -- such a
527/// texture is as wide as the driver's row pitch -- which is why `shape` is
528/// passed down to [`crate::d3d11_shared_handle_geometry`] rather than only
529/// checked here.
530///
531/// # Errors
532///
533/// [`crate::Error::InvalidArgument`] when the descriptor names no pixel
534/// format, or when its shape is neither spelling of the geometry the
535/// texture reports. Propagates whatever opening and describing the texture
536/// reports.
537///
538/// # Safety
539///
540/// `texture` must be an NT shared handle of a D3D11 texture, valid in this
541/// process. It stays owned by the caller.
542#[cfg(target_os = "windows")]
543pub(crate) unsafe fn descriptor_d3d11_geometry(
544 desc: &TensorDesc,
545 texture: std::os::windows::io::RawHandle,
546 shape: &[usize],
547) -> crate::Result<(crate::PixelFormat, usize, usize)> {
548 let format = crate::PixelFormat::from_code(desc.format).ok_or_else(|| {
549 crate::Error::InvalidArgument(
550 "D3D11_TEXTURE descriptor without a pixel format: the texture's own \
551 description cannot say which format it was allocated for"
552 .into(),
553 )
554 })?;
555 // SAFETY: the caller guarantees `texture` is a shared NT handle valid in
556 // this process.
557 let (width, height) = unsafe { d3d11_geometry_checked(format, texture, shape) }?;
558 Ok((format, width, height))
559}
560
561/// The image dimensions a shared D3D11 texture reports, with `shape` checked
562/// against them.
563///
564/// Split out of [`descriptor_d3d11_geometry`] so the blob's reference-mode
565/// import ([`crate::blob::import`]) makes the same refusal rather than a
566/// second copy of it: the two transports differ in how they name the format,
567/// not in which shapes describe a given texture.
568///
569/// # Errors
570///
571/// [`crate::Error::InvalidArgument`] when `shape` is neither spelling of the
572/// geometry the texture reports. Propagates whatever opening and describing
573/// the texture reports.
574///
575/// # Safety
576///
577/// `texture` must be an NT shared handle of a D3D11 texture, valid in this
578/// process. It stays owned by the caller.
579#[cfg(target_os = "windows")]
580pub(crate) unsafe fn d3d11_geometry_checked(
581 format: crate::PixelFormat,
582 texture: std::os::windows::io::RawHandle,
583 shape: &[usize],
584) -> crate::Result<(usize, usize)> {
585 // SAFETY: the caller guarantees `texture` is a shared NT handle valid in
586 // this process; the helper opens its own texture and drops it on return.
587 let (width, height) =
588 unsafe { crate::d3d11_shared_handle_geometry(texture, format, Some(shape)) }?;
589 let allocation = format.allocation_shape(width, height);
590 let addressing = format.addressing_shape(width, height);
591 if allocation.as_deref() != Some(shape) && addressing.as_deref() != Some(shape) {
592 return Err(crate::Error::InvalidArgument(format!(
593 "D3D11 texture shape {shape:?} is neither the {format:?} allocation \
594 shape {allocation:?} nor the addressing shape {addressing:?} of the \
595 {width}x{height} texture it names"
596 )));
597 }
598 Ok((width, height))
599}
600
601/// Inputs to [`from_parts`], grouped into one named-field struct rather than
602/// a positional argument list.
603///
604/// This crossed nine parameters once the pixel-format and capacity inputs
605/// were added, several of them same-typed (`fourcc: u32`, a `u32`
606/// handle-adjacent value, `colorimetry: u32`) — a positional call at that
607/// width compiles happily with two arguments transposed and silently
608/// corrupts the descriptor. Named fields make that a compile error instead. Internal to
609/// this crate: `TensorDyn::descriptor_pinned` is presently the only caller --
610/// on both backends now, not just `static` (task P2a).
611pub(crate) struct DescParts<'a, 'p> {
612 pub dims: &'a [usize],
613 pub memory: TensorMemory,
614 pub dtype: crate::DType,
615 /// DRM FourCC for [`TensorDesc::fourcc`]; 0 when the format has none.
616 pub fourcc: u32,
617 /// The HAL format, source for both the row-dimension decision below and
618 /// [`TensorDesc::format`] (via [`format_of`]).
619 pub format: Option<crate::PixelFormat>,
620 /// Row pitch in bytes; see [`from_parts`]'s doc for how it maps to a
621 /// stride index.
622 pub row_stride: Option<usize>,
623 pub handle: i64,
624 pub colorimetry: u32,
625 /// Producer's [`crate::TensorTrait::capacity_bytes`] for
626 /// [`TensorDesc::capacity`]; see that field's docs.
627 pub capacity: u64,
628 pub pin: Option<&'a crate::HostPin<'p>>,
629 /// The producer's `PboOpsVtable` address, for [`TensorDesc::ptr`] under
630 /// [`kind::PBO`] — see that field's own doc comment. `None` for every
631 /// other kind (mutually exclusive with `pin`: a PBO tensor never has
632 /// one, since `PboTensor` refuses `pin_host`).
633 pub pbo_vtable_ptr: Option<*const std::ffi::c_void>,
634 /// The process device fence's NT handle value, for [`TensorDesc::ptr`]
635 /// under [`kind::D3D11_TEXTURE`] — the third meaning that one field
636 /// carries, keyed by kind exactly as the PBO vtable is. `0` for every
637 /// other kind and for a device with no fence.
638 pub fence_handle: usize,
639 /// The fence value of the producer's last recorded GPU write, for
640 /// [`TensorDesc::sync`] under [`kind::D3D11_TEXTURE`]. `None` for every
641 /// other kind and when nothing has been recorded.
642 pub sync: Option<u64>,
643}
644
645/// Build a descriptor from its parts.
646///
647/// The single construction site: `TensorDyn::descriptor` and any future
648/// `Tensor<T>::descriptor` both route here, so the layout and the stride
649/// convention are defined once.
650///
651/// `ptr` is null unless a [`HostPin`](crate::HostPin) is supplied, or the
652/// tensor is PBO-backed and `pbo_vtable_ptr` is supplied instead — see
653/// [`TensorDesc::ptr`]'s own doc comment for why that one field carries two
654/// different meanings by [`TensorDesc::kind`]. A raw address is only
655/// meaningful while something guarantees the memory (or, for PBO, the
656/// `PboHandle` the vtable addresses) stays put, so the descriptor carries
657/// one only when the producer has pinned or built a vtable — and the
658/// capsule's keepalive then holds the relevant tensor alive for the
659/// consumer's benefit either way.
660///
661/// `row_stride` is the producer's [`TensorDyn::row_stride`](crate::TensorDyn::row_stride)
662/// **in bytes**, or `None` for a tightly-packed tensor. A pitch-aligned image
663/// carries padding between rows that the shape alone cannot express, so when
664/// present it overrides the stride of the dimension that is actually the row
665/// dimension for `format`'s layout. `format` is also the source of
666/// [`TensorDesc::format`] (via [`format_of`]) and is taken separately
667/// from `fourcc` rather than derived from it: `PlanarRgb`/`PlanarRgba` both
668/// encode as `fourcc = 0` (no standard FourCC exists for them), the same
669/// sentinel a non-image tensor uses, so `fourcc` alone cannot tell Planar
670/// apart from "no format" — for either purpose.
671///
672/// * Packed (`[H, W, C]`) / SemiPlanar (`[combined_H, W]`): dimension 0.
673/// * Planar (`[C, H, W]`, e.g. `PlanarRgb`): dimension 1 — dimension 0 there
674/// is the *plane* count, and clobbering it with the row pitch would
675/// corrupt the plane stride while leaving the true per-row stride
676/// (dimension 1) unfixed. `Tensor` tracks only one row-pitch scalar (no
677/// separate per-plane stride), so dimension 0 keeps its packed value here
678/// exactly as it did before any padding was known — unchanged, not
679/// recomputed.
680///
681/// `row_stride` is a byte pitch and is carried verbatim. It need not be a
682/// whole number of `dtype` elements: `set_row_stride` validates only
683/// `stride >= minimum`, not element alignment, and since strides are bytes
684/// there is no division to lose a remainder to. The element-stride
685/// representation this replaced could not express such a pitch at all — it
686/// fell back to the packed stride and reported a pitch the buffer does not
687/// have, across a package/ABI boundary.
688///
689/// `handle`, `colorimetry` and `capacity` are carried through verbatim; `-1`
690/// and `0` are the "unused"/"undefined" sentinels for the first two.
691pub(crate) fn from_parts(parts: DescParts) -> TensorDesc {
692 let DescParts {
693 dims,
694 memory,
695 dtype,
696 fourcc,
697 format,
698 row_stride,
699 handle,
700 colorimetry,
701 capacity,
702 pin,
703 pbo_vtable_ptr,
704 fence_handle,
705 sync,
706 } = parts;
707 let desc_kind = kind_of(memory);
708 // Both fields are meaningful only under `D3D11_TEXTURE`, and only when
709 // the device has a fence to name: a value with no handle behind it would
710 // set `SYNC_PRESENT` on a descriptor no consumer can wait on.
711 let completion = (desc_kind == kind::D3D11_TEXTURE && fence_handle != 0)
712 .then_some(sync)
713 .flatten()
714 .map(|value| (fence_handle, value));
715 // `ndim` carries the TRUE rank even when it exceeds the eight slots the
716 // descriptor can hold, so a consumer can detect the case and refuse.
717 // Reporting the clamped value instead would ship a descriptor whose
718 // `ndim` disagrees with its own `dims`, and the consumer would address
719 // `product(dims[..8])` elements of a larger allocation -- silently wrong
720 // data across a package boundary. This mirrors the `row_stride` override
721 // below, which likewise refuses to truncate quietly.
722 let ndim = dims.len();
723 let filled = ndim.min(TensorDesc::MAX_NDIM);
724 if ndim > TensorDesc::MAX_NDIM {
725 log::warn!(
726 "tensor of rank {ndim} exceeds the descriptor's {} shape slots; \
727 the descriptor will be rejected on import",
728 TensorDesc::MAX_NDIM
729 );
730 }
731 let mut shape = [0u64; TensorDesc::MAX_NDIM];
732 let mut strides = [0i64; TensorDesc::MAX_NDIM];
733 for (i, d) in dims.iter().take(filled).enumerate() {
734 shape[i] = *d as u64;
735 }
736 // Row-major and contiguous, in BYTES: stride[i] = dtype size * product of
737 // trailing dims. Bytes, not elements, because a hardware pitch need not be
738 // a whole number of elements (and sub-byte dtypes have no element size at
739 // all) -- see the `row_stride` override below.
740 let mut acc: i64 = dtype.size() as i64;
741 // `filled`, not `ndim`: the arrays hold MAX_NDIM slots and `ndim` may now
742 // legitimately exceed that.
743 for i in (0..filled).rev() {
744 strides[i] = acc;
745 acc *= shape[i] as i64;
746 }
747 if ndim >= 2 {
748 if let Some(rs) = row_stride {
749 // Planar stacks channels ahead of rows ([C, H, W]): the row
750 // dimension is 1, not 0. Every other layout this protocol
751 // carries (Packed [H, W, C], SemiPlanar [combined_H, W]) has the
752 // row dimension at 0. `format` is `None` for a non-image tensor,
753 // in which case `row_stride` should not have been supplied in
754 // the first place; dimension 0 is a harmless default for that
755 // case since it matches the pre-padding behaviour.
756 let row_dim = match format.map(|f| f.layout()) {
757 Some(crate::PixelLayout::Planar) if ndim >= 3 => 1,
758 _ => 0,
759 };
760 // Bytes in, bytes out. The element-stride representation had to
761 // divide here, and a remainder made the pitch unrepresentable --
762 // it fell back to the packed stride and reported a pitch the
763 // buffer does not have. There is nothing left to get wrong.
764 strides[row_dim] = rs as i64;
765 }
766 }
767 TensorDesc {
768 version: ABI_VERSION,
769 kind: desc_kind,
770 handle,
771 // Keyed by kind, because that is what this field's meaning is keyed
772 // by. Under `D3D11_TEXTURE` it is the fence handle or nothing, even
773 // when the producer also pinned a host address -- a Python capsule
774 // asked for `access="read"` does pin a texture tensor, and writing
775 // that address here would hand the consumer a host pointer to open
776 // as a fence. The pin stays the producer's own keepalive; a D3D11
777 // import reopens the texture and takes its own CPU mapping, so it
778 // never wanted the address. Otherwise `pin` and `pbo_vtable_ptr`
779 // are mutually exclusive by construction (see
780 // `DescParts::pbo_vtable_ptr`'s own doc comment) and `pin` takes
781 // precedence purely because it is the older, more general case.
782 ptr: if desc_kind == kind::D3D11_TEXTURE {
783 completion.map_or_else(SendPtr::null, |(h, _)| SendPtr(h as *mut u8))
784 } else {
785 pin.map(|p| SendPtr(p.as_mut_ptr()))
786 .or_else(|| pbo_vtable_ptr.map(|p| SendPtr(p as *mut u8)))
787 .unwrap_or_else(SendPtr::null)
788 },
789 ndim: ndim as u32,
790 dtype: dtype_of(dtype),
791 shape,
792 strides,
793 fourcc,
794 format: format.map(format_of).unwrap_or(format::NONE),
795 colorimetry,
796 // `D3D11_TEXTURE` is the one kind that fills these in; every other
797 // producer in this build leaves them clear. See the field docs.
798 flags: if completion.is_some() {
799 flags::SYNC_PRESENT
800 } else {
801 0
802 },
803 capacity,
804 sync: completion.map_or(0, |(_, value)| value),
805 }
806}