diff --git a/crates/ptwm-core/src/extension/mod.rs b/crates/ptwm-core/src/extension/mod.rs index 3399469..d6a2be5 100644 --- a/crates/ptwm-core/src/extension/mod.rs +++ b/crates/ptwm-core/src/extension/mod.rs @@ -11,6 +11,7 @@ pub mod id; pub mod kind; pub mod lifecycle; pub mod manifest; +pub mod resolve; pub mod table; pub use builder::ExtensionTableBuilder; @@ -22,6 +23,7 @@ pub use id::{CanonicalId, ContributionRef}; pub use kind::Kind; pub use lifecycle::Lifecycle; pub use manifest::{ContributionDecl, Manifest}; +pub use resolve::{ResolveError, resolve_codec_selector}; pub use table::{Attestation, ExtensionTable, ExtensionTableEntry}; /// Sentinel public key for in-tree built-ins. Built-ins are trusted by diff --git a/crates/ptwm-core/src/extension/resolve.rs b/crates/ptwm-core/src/extension/resolve.rs new file mode 100644 index 0000000..bc4d4ba --- /dev/null +++ b/crates/ptwm-core/src/extension/resolve.rs @@ -0,0 +1,194 @@ +//! Resolve a user-supplied codec selector to a `CanonicalId`. +//! +//! Selectors are a convenience for humans at the call site. The container +//! records the resolved `CanonicalId`, so renaming a contribution can +//! never invalidate an existing archive. + +use thiserror::Error; + +use crate::discovery::DiscoveredContribution; +use crate::extension::{CanonicalId, builtin_canonical_id}; + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum ResolveError { + #[error("no codec named or identified by {selector:?}")] + NotFound { selector: String }, + + #[error( + "codec name {selector:?} is ambiguous across {} contributions: {}; \ + select by canonical id instead", + candidates.len(), + candidates.join(", ") + )] + Ambiguous { + selector: String, + candidates: Vec, + }, +} + +fn looks_like_canonical_id(s: &str) -> bool { + let body = s.strip_prefix("blake3:").unwrap_or(s); + body.len() == 64 && body.chars().all(|c| c.is_ascii_hexdigit()) +} + +pub fn resolve_codec_selector( + selector: &str, + installed: &[DiscoveredContribution], +) -> Result { + if looks_like_canonical_id(selector) { + let to_parse = if selector.starts_with("blake3:") { + selector.to_string() + } else { + format!("blake3:{}", selector) + }; + return CanonicalId::parse(&to_parse).map_err(|_| ResolveError::NotFound { + selector: selector.to_string(), + }); + } + + let mut hits: Vec = Vec::new(); + if crate::codec::REGISTRY + .iter() + .any(|(name, _)| *name == selector) + { + hits.push(builtin_canonical_id(selector)); + } + for d in installed { + for c in &d.manifest.contributions { + if c.label == selector { + if let Ok(id) = CanonicalId::parse(&c.id) { + hits.push(id); + } + } + } + } + + match hits.len() { + 0 => Err(ResolveError::NotFound { + selector: selector.to_string(), + }), + 1 => Ok(hits[0]), + _ => Err(ResolveError::Ambiguous { + selector: selector.to_string(), + candidates: hits.iter().map(|id| id.to_string()).collect(), + }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn two_contributions_labelled(label: &str) -> Vec { + use crate::extension::manifest::{BundleHeader, ContributionDecl, Manifest}; + use crate::extension::{Kind, Lifecycle}; + use std::path::PathBuf; + + let manifest1 = Manifest { + bundle: BundleHeader { + name: "bundle1".to_string(), + version: "1.0.0".to_string(), + author_pubkey: + "ed25519:0000000000000000000000000000000000000000000000000000000000000000" + .to_string(), + description: None, + }, + contributions: vec![ContributionDecl { + id: crate::extension::builtin_canonical_id("huffman").to_string(), + label: label.to_string(), + kind: Kind::PlaneCodec, + abi_version: 1, + lifecycle: Lifecycle::Thread, + flavors: vec![], + capabilities: Default::default(), + install_hint: None, + }], + }; + + let manifest2 = Manifest { + bundle: BundleHeader { + name: "bundle2".to_string(), + version: "1.0.0".to_string(), + author_pubkey: + "ed25519:0000000000000000000000000000000000000000000000000000000000000000" + .to_string(), + description: None, + }, + contributions: vec![ContributionDecl { + id: crate::extension::builtin_canonical_id("zstd").to_string(), + label: label.to_string(), + kind: Kind::PlaneCodec, + abi_version: 1, + lifecycle: Lifecycle::Thread, + flavors: vec![], + capabilities: Default::default(), + install_hint: None, + }], + }; + + vec![ + DiscoveredContribution { + manifest: manifest1, + manifest_path: PathBuf::from("/tmp/manifest1.toml"), + bundle_dir: PathBuf::from("/tmp/bundle1"), + installed_flavors: 0, + }, + DiscoveredContribution { + manifest: manifest2, + manifest_path: PathBuf::from("/tmp/manifest2.toml"), + bundle_dir: PathBuf::from("/tmp/bundle2"), + installed_flavors: 0, + }, + ] + } + + #[test] + fn resolves_a_builtin_name() { + let got = resolve_codec_selector("huffman", &[] as &[DiscoveredContribution]) + .expect("builtin must resolve"); + assert_eq!(got, crate::extension::builtin_canonical_id("huffman")); + } + + #[test] + fn resolves_a_canonical_id_verbatim() { + let id = crate::extension::builtin_canonical_id("zstd"); + let got = resolve_codec_selector(&id.to_string(), &[] as &[DiscoveredContribution]) + .expect("id must resolve"); + assert_eq!(got, id); + } + + #[test] + fn resolves_bare_hex_canonical_id() { + let id = crate::extension::builtin_canonical_id("huffman"); + let id_str = id.to_string(); + let bare_hex = &id_str[7..]; // Skip "blake3:" + let got = resolve_codec_selector(bare_hex, &[] as &[DiscoveredContribution]) + .expect("bare hex must resolve"); + assert_eq!(got, id); + } + + #[test] + fn unknown_name_reports_not_found() { + match resolve_codec_selector("no_such_codec", &[] as &[DiscoveredContribution]) { + Err(ResolveError::NotFound { selector }) => assert_eq!(selector, "no_such_codec"), + other => panic!("expected NotFound, got {other:?}"), + } + } + + #[test] + fn ambiguous_name_names_every_candidate() { + // Two installed contributions sharing a label must not silently + // resolve to whichever was discovered first. + let installed = two_contributions_labelled("gpu_codec"); + match resolve_codec_selector("gpu_codec", &installed) { + Err(ResolveError::Ambiguous { + selector, + candidates, + }) => { + assert_eq!(selector, "gpu_codec"); + assert_eq!(candidates.len(), 2, "both candidates must be named"); + } + other => panic!("expected Ambiguous, got {other:?}"), + } + } +} diff --git a/crates/ptwm-core/src/flavor/mod.rs b/crates/ptwm-core/src/flavor/mod.rs index 27d7ea9..ca45d0a 100644 --- a/crates/ptwm-core/src/flavor/mod.rs +++ b/crates/ptwm-core/src/flavor/mod.rs @@ -22,7 +22,8 @@ pub use native::{ }; pub use router::{ DeltaSchemeRouter, DispatchedDeltaScheme, DispatchedHardwareBackendCuda, DispatchedPlaneCodec, - HardwareBackendRouter, PlaneCodecRouter, + DispatchedPlaneCodecCuda, HardwareBackendRouter, NativePlaneCodecCudaAdapter, + PlaneCodecCudaRouter, PlaneCodecRouter, }; pub use third_party::ThirdPartyPlaneCodec; pub use wasm::{ diff --git a/crates/ptwm-core/src/flavor/native.rs b/crates/ptwm-core/src/flavor/native.rs index 47a355e..a6f5b58 100644 --- a/crates/ptwm-core/src/flavor/native.rs +++ b/crates/ptwm-core/src/flavor/native.rs @@ -43,6 +43,11 @@ pub struct NativeSymbols { /// absent, the dispatcher falls back to the stateless /// `plane_codec_v1_decode` symbol. pub plane_codec_v1_decode_stateful: Option, + /// Optional CUDA device-pointer path. A codec that only implements + /// the host path leaves these absent. + pub plane_codec_v1_encode_cuda: Option, + pub plane_codec_v1_decode_cuda: Option, + pub plane_codec_v1_cuda_stream_handle: Option, // Transform pub transform_v1_forward: Option, pub transform_v1_inverse: Option, @@ -135,6 +140,62 @@ pub type HardwareBackendCudaDispatchDecodeFn = unsafe extern "C" fn( device_ordinal: u32, ) -> i64; +pub const PLANE_CODEC_CUDA_ENCODE_SYMBOL: &[u8] = b"ptwm_plane_codec_v1_encode_cuda\0"; +pub const PLANE_CODEC_CUDA_DECODE_SYMBOL: &[u8] = b"ptwm_plane_codec_v1_decode_cuda\0"; +pub const PLANE_CODEC_CUDA_STREAM_HANDLE_SYMBOL: &[u8] = + b"ptwm_plane_codec_v1_cuda_stream_handle\0"; + +/// Returns the extension's CUDA stream handle for `device_ordinal`, or 0 +/// if unavailable. Same contract as the hardware-backend equivalent. +pub type PlaneCodecCudaStreamHandleFn = unsafe extern "C" fn(device_ordinal: u32) -> u64; + +/// `plane_codec_v1_decode_cuda`: decode `in_dev_ptr` into `out_dev_ptr`, +/// two distinct device buffers. `state_bytes` carries the codec's small +/// per-tensor state in the same `(state_format_version, state_bytes)` +/// shape the existing `decode_stateful` path uses. `codec_id` +/// self-describes the wire format for the extension to validate. +/// +/// Completion contract: by the time this returns, the kernel has FULLY +/// COMPLETED (the extension synchronizes its own stream). `out_dev_ptr` +/// is valid and visible to any subsequent CUDA operation on any stream, +/// with no further caller-side synchronization. Launch-time and +/// execution-time errors are both visible in the return code. +pub type PlaneCodecCudaDecodeFn = unsafe extern "C" fn( + state_format_version: u8, + state_ptr: *const u8, + state_len: usize, + codec_id_ptr: *const u8, + codec_id_len: usize, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, +) -> i64; + +/// `plane_codec_v1_encode_cuda`: encode `in_dev_ptr` into `out_dev_ptr`. +/// Same completion contract as the decode symbol. +/// +/// `needed_out` exists because an encoder cannot size its output before +/// running: when `out_cap` is insufficient the extension writes the +/// required byte count to `*needed_out` and returns `-2`. The generic +/// return-code decoder cannot carry that number, so it is passed +/// explicitly. On success `*needed_out` is left untouched. +#[allow(clippy::too_many_arguments)] +pub type PlaneCodecCudaEncodeFn = unsafe extern "C" fn( + state_format_version: u8, + state_ptr: *const u8, + state_len: usize, + codec_id_ptr: *const u8, + codec_id_len: usize, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + needed_out: *mut u64, + device_ordinal: u32, +) -> i64; + pub struct NativeExtension { // Held for its drop-time side effect: keeping the dlopen handle alive // so all fn pointers in `symbols` remain valid. @@ -163,6 +224,9 @@ impl NativeExtension { plane_codec_v1_encode: None, plane_codec_v1_decode: None, plane_codec_v1_decode_stateful: None, + plane_codec_v1_encode_cuda: None, + plane_codec_v1_decode_cuda: None, + plane_codec_v1_cuda_stream_handle: None, transform_v1_forward: None, transform_v1_inverse: None, delta_scheme_v1_encode: None, @@ -185,6 +249,23 @@ impl NativeExtension { &library, b"ptwm_plane_codec_v1_decode_stateful\0", ); + // Optional CUDA device-pointer path. A codec that only + // implements the host path leaves these absent; a codec + // that implements CUDA must export all three, so a + // partial set is a packaging error worth failing on. + symbols.plane_codec_v1_encode_cuda = + resolve::(&library, PLANE_CODEC_CUDA_ENCODE_SYMBOL); + symbols.plane_codec_v1_decode_cuda = + resolve::(&library, PLANE_CODEC_CUDA_DECODE_SYMBOL); + symbols.plane_codec_v1_cuda_stream_handle = resolve::( + &library, + PLANE_CODEC_CUDA_STREAM_HANDLE_SYMBOL, + ); + validate_cuda_symbol_set([ + symbols.plane_codec_v1_encode_cuda.is_some(), + symbols.plane_codec_v1_decode_cuda.is_some(), + symbols.plane_codec_v1_cuda_stream_handle.is_some(), + ])?; if symbols.plane_codec_v1_encode.is_none() || symbols.plane_codec_v1_decode.is_none() { @@ -455,6 +536,99 @@ impl NativeExtension { }; decode_rc(rc, out_len) } + + /// Returns the extension's CUDA stream handle for `device_ordinal`, + /// or 0 when the extension has none. + pub fn plane_codec_cuda_stream_handle(&self, device_ordinal: u32) -> u64 { + match self.symbols.plane_codec_v1_cuda_stream_handle { + Some(f) => unsafe { f(device_ordinal) }, + None => 0, + } + } + + /// Invoke `ptwm_plane_codec_v1_decode_cuda`. `in_dev_ptr` and + /// `out_dev_ptr` are distinct device buffers; see the type's doc + /// comment for the completion contract. + #[allow(clippy::too_many_arguments)] + pub fn invoke_plane_codec_decode_cuda( + &self, + state_bytes: &[u8], + codec_id: &CanonicalId, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, + ) -> Result { + let f = self + .symbols + .plane_codec_v1_decode_cuda + .ok_or(CodecError::Unsupported { + feature: "plane_codec_v1_decode_cuda not present".into(), + })?; + let codec_id_bytes = codec_id.as_bytes(); + let rc = unsafe { + f( + 1, + state_bytes.as_ptr(), + state_bytes.len(), + codec_id_bytes.as_ptr(), + codec_id_bytes.len(), + in_dev_ptr, + in_len, + out_dev_ptr, + out_cap, + device_ordinal, + ) + }; + decode_rc(rc, out_cap) + } + + /// Invoke `ptwm_plane_codec_v1_encode_cuda`. + /// + /// On `-2` (buffer too small) the extension has written the required + /// byte count through `needed_out`; that value is returned in + /// `CodecError::BufferTooSmall { needed }` rather than the zero the + /// generic return-code decoder would produce. + #[allow(clippy::too_many_arguments)] + pub fn invoke_plane_codec_encode_cuda( + &self, + state_bytes: &[u8], + codec_id: &CanonicalId, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, + ) -> Result { + let f = self + .symbols + .plane_codec_v1_encode_cuda + .ok_or(CodecError::Unsupported { + feature: "plane_codec_v1_encode_cuda not present".into(), + })?; + let codec_id_bytes = codec_id.as_bytes(); + let mut needed: u64 = 0; + let rc = unsafe { + f( + 1, + state_bytes.as_ptr(), + state_bytes.len(), + codec_id_bytes.as_ptr(), + codec_id_bytes.len(), + in_dev_ptr, + in_len, + out_dev_ptr, + out_cap, + &mut needed as *mut u64, + device_ordinal, + ) + }; + match decode_rc(rc, out_cap) { + Err(CodecError::BufferTooSmall { .. }) => Err(CodecError::BufferTooSmall { needed }), + other => other, + } + } } fn decode_rc(rc: i64, out_capacity: usize) -> Result { @@ -482,6 +656,22 @@ fn decode_rc(rc: i64, out_capacity: usize) -> Result { Ok(n) } +/// A `plane_codec` may implement the CUDA path or not, but not halfway: +/// exporting `decode_cuda` without `encode_cuda` would resolve, then fail +/// only when something tried to encode. Reject the partial set at load. +fn validate_cuda_symbol_set(present: [bool; 3]) -> Result<(), CodecError> { + let any = present.iter().any(|p| *p); + let all = present.iter().all(|p| *p); + if any && !all { + return Err(CodecError::Unsupported { + feature: "partial plane_codec_v1 CUDA symbol set: a codec exporting any \ + of encode_cuda/decode_cuda/cuda_stream_handle must export all three" + .into(), + }); + } + Ok(()) +} + fn resolve(library: &Library, name: &[u8]) -> Option { unsafe { library.get::(name).ok().map(|sym: Symbol| *sym) } } @@ -511,13 +701,14 @@ mod tests { /// Build a bare `NativeExtension` with every `NativeSymbols` field set /// to `None`, for exercising "symbol absent" invoke paths without - /// dlopen-ing a real contribution artifact from disk. + /// dlopen-ing a real contribution artifact from disk. Shared across + /// the hardware-backend and plane-codec CUDA absent-symbol tests below. /// /// `Library::this()` wraps a handle to the already-loaded host process /// instead of opening a new shared object, which gives a real, valid /// `Library` (fn pointers can be safely dropped/never resolved against /// it) without depending on any particular file existing on disk. - fn native_extension_with_no_hardware_backend_symbols() -> NativeExtension { + fn native_extension_with_no_optional_symbols() -> NativeExtension { let library: Library = libloading::os::unix::Library::this().into(); NativeExtension { library, @@ -528,6 +719,9 @@ mod tests { plane_codec_v1_encode: None, plane_codec_v1_decode: None, plane_codec_v1_decode_stateful: None, + plane_codec_v1_encode_cuda: None, + plane_codec_v1_decode_cuda: None, + plane_codec_v1_cuda_stream_handle: None, transform_v1_forward: None, transform_v1_inverse: None, delta_scheme_v1_encode: None, @@ -552,11 +746,75 @@ mod tests { fn hardware_backend_invoke_stream_handle_errors_when_symbol_absent() { // A NativeExtension whose hardware_backend_v1_cuda_stream_handle symbol // was never resolved (None) must return Unsupported, not panic/UB. - let ext = native_extension_with_no_hardware_backend_symbols(); // test-only constructor, see Step 3 + let ext = native_extension_with_no_optional_symbols(); let res = ext.invoke_hardware_backend_cuda_stream_handle(0); assert!(matches!(res, Err(CodecError::Unsupported { .. }))); } + #[test] + fn plane_codec_cuda_invocations_report_absent_symbols() { + // With no library loaded the CUDA symbols are None; the invoke + // methods must surface Unsupported (or 0 for the stream-handle + // probe) rather than panic on unwrap. + let ext = native_extension_with_no_optional_symbols(); + let codec_id = CanonicalId::from_bytes([0xAB; 32]); + + assert_eq!(ext.plane_codec_cuda_stream_handle(0), 0); + assert!(matches!( + ext.invoke_plane_codec_encode_cuda(&[], &codec_id, 0, 0, 0, 0, 0), + Err(CodecError::Unsupported { .. }) + )); + assert!(matches!( + ext.invoke_plane_codec_decode_cuda(&[], &codec_id, 0, 0, 0, 0, 0), + Err(CodecError::Unsupported { .. }) + )); + } + + #[test] + fn cuda_plane_codec_symbol_names_are_the_documented_ones() { + // Guards against a rename drifting from the plugin-facing contract: + // these exact strings are what a third-party .so must export. + assert_eq!( + PLANE_CODEC_CUDA_ENCODE_SYMBOL, + b"ptwm_plane_codec_v1_encode_cuda\0" + ); + assert_eq!( + PLANE_CODEC_CUDA_DECODE_SYMBOL, + b"ptwm_plane_codec_v1_decode_cuda\0" + ); + assert_eq!( + PLANE_CODEC_CUDA_STREAM_HANDLE_SYMBOL, + b"ptwm_plane_codec_v1_cuda_stream_handle\0" + ); + } + + #[test] + fn a_complete_cuda_symbol_set_is_accepted() { + assert!(validate_cuda_symbol_set([true, true, true]).is_ok()); + } + + #[test] + fn no_cuda_symbols_is_accepted_as_a_host_only_codec() { + assert!(validate_cuda_symbol_set([false, false, false]).is_ok()); + } + + #[test] + fn a_partial_cuda_symbol_set_is_rejected() { + for present in [ + [true, false, false], + [false, true, false], + [true, true, false], + ] { + assert!( + matches!( + validate_cuda_symbol_set(present), + Err(CodecError::Unsupported { .. }) + ), + "partial set {present:?} must be rejected" + ); + } + } + #[test] fn hardware_backend_missing_required_symbols_is_unsupported() { // A library with neither ptwm_hardware_backend_v1_cuda_stream_handle diff --git a/crates/ptwm-core/src/flavor/router.rs b/crates/ptwm-core/src/flavor/router.rs index 8e7406d..6d357ae 100644 --- a/crates/ptwm-core/src/flavor/router.rs +++ b/crates/ptwm-core/src/flavor/router.rs @@ -707,6 +707,244 @@ impl HardwareBackendRouter { } } +// ── PlaneCodec CUDA adapter ───────────────────────────────────────────────── + +/// A plane codec that operates on CUDA device pointers. +/// +/// Distinct from [`DispatchedPlaneCodec`], whose buffers are host memory. +/// Implementors move no data across the host boundary: both input and +/// output stay device-resident. +pub trait DispatchedPlaneCodecCuda: Send + Sync { + fn cuda_stream_handle(&self, device_ordinal: u32) -> u64; + + #[allow(clippy::too_many_arguments)] + fn encode_cuda( + &self, + state_bytes: &[u8], + codec_id: &CanonicalId, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, + ) -> Result; + + #[allow(clippy::too_many_arguments)] + fn decode_cuda( + &self, + state_bytes: &[u8], + codec_id: &CanonicalId, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, + ) -> Result; +} + +pub struct NativePlaneCodecCudaAdapter { + inner: Arc, +} + +impl NativePlaneCodecCudaAdapter { + pub fn new(inner: Arc) -> Self { + Self { inner } + } +} + +impl DispatchedPlaneCodecCuda for NativePlaneCodecCudaAdapter { + fn cuda_stream_handle(&self, device_ordinal: u32) -> u64 { + self.inner.plane_codec_cuda_stream_handle(device_ordinal) + } + + #[allow(clippy::too_many_arguments)] + fn encode_cuda( + &self, + state_bytes: &[u8], + codec_id: &CanonicalId, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, + ) -> Result { + self.inner.invoke_plane_codec_encode_cuda( + state_bytes, + codec_id, + in_dev_ptr, + in_len, + out_dev_ptr, + out_cap, + device_ordinal, + ) + } + + #[allow(clippy::too_many_arguments)] + fn decode_cuda( + &self, + state_bytes: &[u8], + codec_id: &CanonicalId, + in_dev_ptr: u64, + in_len: usize, + out_dev_ptr: u64, + out_cap: usize, + device_ordinal: u32, + ) -> Result { + self.inner.invoke_plane_codec_decode_cuda( + state_bytes, + codec_id, + in_dev_ptr, + in_len, + out_dev_ptr, + out_cap, + device_ordinal, + ) + } +} + +/// Resolves a `CanonicalId` to a `Box`. +/// +/// Same shape as [`HardwareBackendRouter`]: a `Mutex>` cache, the +/// same `new` / `new_with_policy` constructor split, and the same +/// `hardware_class`-presence precondition gated through [`check`] against +/// `self.policy` before any native load is attempted. See +/// [`HardwareBackendRouter::resolve`]'s doc comment for the full reasoning +/// behind that precondition; it applies here unchanged. +/// +/// It differs from [`HardwareBackendRouter`] in two ways. First, it caches +/// `Arc` and constructs +/// [`NativePlaneCodecCudaAdapter`] rather than +/// `NativeHardwareBackendCudaAdapter`. Second, a `plane_codec` contribution +/// that resolves to a real installed manifest entry but has no loadable +/// native artifact reports a fixed, non-formatted message +/// ("plane_codec CUDA dispatch requires the native flavor") rather than one +/// naming the canonical id and bundle directory: a WASM contribution cannot +/// hold device pointers (WASM's linear address space is 32-bit), so this +/// router is native-only by construction, exactly like +/// [`HardwareBackendRouter`]. +pub struct PlaneCodecCudaRouter { + installed: Vec, + policy: HostPolicy, + cache: Mutex>>, +} + +impl std::fmt::Debug for PlaneCodecCudaRouter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PlaneCodecCudaRouter") + .field("installed_count", &self.installed.len()) + .finish_non_exhaustive() + } +} + +impl PlaneCodecCudaRouter { + /// Create a new router with the given set of discovered installed + /// extensions, gated by the default (empty-`available_hardware`) + /// [`HostPolicy`]. Pass an empty `Vec` to get `Unsupported` for every + /// id (there are no built-in `plane_codec` CUDA contributions). + /// + /// As with [`HardwareBackendRouter::new`], this constructor's + /// default-deny behavior stays fixed for every existing caller: it + /// never admits a `hardware_class`-declaring contribution. A caller + /// that has independently confirmed CUDA availability on the host + /// uses [`Self::new_with_policy`] instead. + pub fn new(installed: Vec) -> Self { + Self { + installed, + policy: HostPolicy::default(), + cache: Mutex::new(HashMap::new()), + } + } + + /// Create a new router with the given set of discovered installed + /// extensions, gated by a caller-supplied [`HostPolicy`] instead of + /// the default. + pub fn new_with_policy(installed: Vec, policy: HostPolicy) -> Self { + Self { + installed, + policy, + cache: Mutex::new(HashMap::new()), + } + } + + /// Resolve `canonical_id` to a `DispatchedPlaneCodecCuda`. + pub fn get( + &self, + canonical_id: &CanonicalId, + ) -> Result, CodecError> { + { + let cache = self.cache.lock().unwrap(); + if let Some(codec) = cache.get(canonical_id) { + return Ok(Arc::clone(codec)); + } + } + + let codec = self.resolve(canonical_id)?; + let arc = Arc::from(codec); + { + let mut cache = self.cache.lock().unwrap(); + cache.insert(*canonical_id, Arc::clone(&arc)); + } + Ok(arc) + } + + /// Look up `canonical_id`, gate it through a `hardware_class` + /// precondition and [`check`], and dlopen the native artifact if + /// admitted. Mirrors [`HardwareBackendRouter::resolve`] exactly up to + /// the native-load step; see that method's doc comment for why the + /// `hardware_class`-presence check runs before [`check`] itself. + fn resolve( + &self, + canonical_id: &CanonicalId, + ) -> Result, CodecError> { + let Some(install) = find_install(canonical_id, &self.installed) else { + return Err(CodecError::Unsupported { + feature: format!( + "no plane_codec CUDA contribution registered for canonical id {canonical_id}" + ), + }); + }; + let bundle_dir = &install.bundle_dir; + let entry = manifest_entry_for(install, canonical_id)?; + + // Same precondition as HardwareBackendRouter::resolve: check()'s + // hardware_class rule only denies a *declared* mismatched class, + // so omission must be enforced here, before check() runs. + if !matches!( + entry.capabilities.get("hardware_class"), + Some(CapabilityValue::Text(_)) + ) { + return Err(CodecError::Unsupported { + feature: "plane_codec CUDA contribution must declare hardware_class as a \ + capability" + .into(), + }); + } + + // Deny before dlopen, not after: see self.policy in + // Self::new / Self::new_with_policy. + match check(&entry, &self.policy, &VendorTable::default()) { + CapabilityVerdict::Denied { reason } => { + return Err(CodecError::Unsupported { + feature: format!("plane_codec CUDA capability check denied: {reason}"), + }); + } + CapabilityVerdict::Admitted => {} + } + + if let Ok(native_path) = find_native_path(bundle_dir) { + let token = VerifiedToken::new_unchecked(); + let ext = NativeExtension::load(&native_path, &entry, token)?; + return Ok(Box::new(NativePlaneCodecCudaAdapter::new(Arc::new(ext)))); + } + + // A WASM contribution cannot hold device pointers: this router is + // native-only by construction, same as HardwareBackendRouter. + Err(CodecError::Unsupported { + feature: "plane_codec CUDA dispatch requires the native flavor".into(), + }) + } +} + // ── Helpers ─────────────────────────────────────────────────────────────────── /// Find the `DiscoveredContribution` whose manifest declares a contribution @@ -1123,6 +1361,134 @@ mod tests { } } + #[test] + fn native_adapter_implements_the_cuda_plane_codec_trait() { + fn assert_impl() {} + assert_impl::(); + } + + #[test] + fn cuda_plane_codec_router_denies_when_policy_lacks_cuda() { + // Default policy declares no available hardware, so a contribution + // requiring hardware_class = "cuda" must not resolve. + let router = PlaneCodecCudaRouter::new(Vec::new()); + let id = crate::extension::builtin_canonical_id("identity"); + assert!(router.get(&id).is_err()); + } + + /// Build a `DiscoveredContribution` declaring a single `plane_codec` + /// contribution with the given `hardware_class` capability, rooted at + /// `bundle_dir` (which need not exist on disk for capability-check + /// tests: the check must run, and deny, before any filesystem access + /// is attempted). + fn discovered_plane_codec_cuda_contribution( + hardware_class: &str, + bundle_dir: &str, + ) -> (DiscoveredContribution, CanonicalId) { + let id_bytes = [0xDEu8; 32]; + let id = CanonicalId::from_bytes(id_bytes); + let id_hex: String = id_bytes.iter().map(|b| format!("{:02x}", b)).collect(); + + let mut caps = CapabilityMap::new(); + caps.set( + "hardware_class", + CapabilityValue::Text(hardware_class.to_string()), + ); + + let manifest = Manifest { + bundle: BundleHeader { + name: "test-plane-codec-cuda".into(), + version: "0.1.0".into(), + author_pubkey: "ed25519:00".into(), + description: None, + }, + contributions: vec![ContributionDecl { + id: format!("blake3:{id_hex}"), + label: "io.example.plane_codec_cuda".into(), + kind: Kind::PlaneCodec, + abi_version: 1, + lifecycle: Lifecycle::Process, + flavors: vec!["native".into()], + capabilities: caps, + install_hint: None, + }], + }; + + let install = DiscoveredContribution { + manifest, + manifest_path: PathBuf::from(bundle_dir).join("manifest.toml"), + bundle_dir: PathBuf::from(bundle_dir), + installed_flavors: crate::discovery::FLAVOR_NATIVE, + }; + + (install, id) + } + + #[test] + fn cuda_plane_codec_router_denies_by_capability_not_by_absence() { + // The preceding test passes for the wrong reason: with an empty + // `installed` list the identity codec is never registered, so it + // fails via the "not found" branch without ever reaching the + // capability-policy check. This test closes that gap: an *installed* + // contribution that declares hardware_class = "cuda" must be denied + // by the capability check specifically (against HostPolicy::default(), + // whose available_hardware is empty), not by "not found". A nonexistent + // bundle_dir is used so a filesystem error cannot masquerade as the + // intended denial. + let (install, id) = + discovered_plane_codec_cuda_contribution("cuda", "/this/bundle/dir/does/not/exist"); + let router = PlaneCodecCudaRouter::new(vec![install]); + let res = router.get(&id); + match res { + Ok(_) => panic!("expected capability denial, got Ok"), + Err(CodecError::Unsupported { feature }) => { + assert!( + feature.contains("capability check denied"), + "expected a capability-check denial, got: {feature}" + ); + assert!( + !feature.contains("no plane_codec CUDA contribution registered"), + "denial should not be the 'not found' branch, got: {feature}" + ); + } + Err(other) => panic!("expected capability denial, got {other:?}"), + } + } + + #[test] + fn cuda_plane_codec_router_new_with_policy_admits_the_case_default_denies() { + // Same fixture, but with hardware_class = "cpu" and a policy that + // explicitly lists "cpu" in available_hardware: admission must + // clear the capability check and fail only afterward, at the + // native-artifact-not-found step (there is nothing to dlopen at + // the nonexistent bundle_dir), and the fixed native-flavor-only + // message is what should surface — proving the WASM/native + // fallback branch, not the capability gate, produced the denial. + let (install, id) = + discovered_plane_codec_cuda_contribution("cpu", "/this/bundle/dir/does/not/exist"); + let policy = HostPolicy { + available_hardware: vec!["cpu".into()], + ..HostPolicy::default() + }; + let router = PlaneCodecCudaRouter::new_with_policy(vec![install], policy); + let res = router.get(&id); + match res { + Ok(_) => panic!("expected a native-flavor-required error, got Ok"), + Err(CodecError::Unsupported { feature }) => { + assert!( + !feature.contains("capability check denied"), + "capability check should have admitted this contribution, got: {feature}" + ); + assert_eq!( + feature, "plane_codec CUDA dispatch requires the native flavor", + "expected the fixed native-flavor-required message once past capability \ + admission" + ); + } + Err(other) => panic!("expected Unsupported, got {other:?}"), + } + } + #[test] fn hardware_backend_router_new_with_policy_admits_the_case_default_denies() { // Same "cuda" fixture shape reused with hardware_class = "cpu" to diff --git a/crates/ptwm-core/src/flavor/third_party.rs b/crates/ptwm-core/src/flavor/third_party.rs index c3a4385..9e2d032 100644 --- a/crates/ptwm-core/src/flavor/third_party.rs +++ b/crates/ptwm-core/src/flavor/third_party.rs @@ -3,13 +3,17 @@ //! `crate::codec::PlaneCodec`. The container decode loop can then treat //! every plane uniformly through a single `Box`. //! -//! Scope: decode-only. The flat ABI used by third-party codecs is -//! input-bytes → output-bytes (with an optional `state_bytes` blob on -//! decode), which doesn't expose the structured `Encoded` return shape -//! the in-tree trial-encode loop expects. `encode` therefore returns -//! `Err(InvalidContainer)` — v1 third-party codecs participate only at -//! decode time. Reaching `encode` here would mean the trial-encode loop -//! reached a third-party codec, which is itself a bug. +//! Scope: decode is always reachable; encode depends on how the adapter +//! was built. The flat ABI used by third-party codecs is input-bytes → +//! output-bytes, which doesn't expose the structured `Encoded` return +//! shape the in-tree trial-encode loop expects. An adapter built through +//! [`ThirdPartyPlaneCodec::new`] represents a codec the trial-encode loop +//! is considering among several candidates, so `encode` returns +//! `Err(InvalidContainer)` there: reaching `encode` on that path would +//! mean the trial loop reached a third-party codec, which is itself a +//! bug. An adapter built through [`ThirdPartyPlaneCodec::new_explicit`] +//! represents a codec the caller named directly; there is no trial loop +//! to satisfy, so `encode` is reachable. use std::sync::Arc; @@ -21,11 +25,30 @@ use crate::layout::PlaneLayout; pub struct ThirdPartyPlaneCodec { inner: Arc, + explicitly_selected: bool, } impl ThirdPartyPlaneCodec { + /// Trial-encode path: the dispatcher is considering this codec among + /// several candidates, so `encode` stays refused (see module docs). pub fn new(inner: Arc) -> Self { - Self { inner } + Self { + inner, + explicitly_selected: false, + } + } + + /// Explicit-selection path: the caller named this codec directly, so + /// there is no trial-encode loop to satisfy and `encode` is reachable. + pub fn new_explicit(inner: Arc) -> Self { + Self { + inner, + explicitly_selected: true, + } + } + + pub(crate) fn allows_encode(explicitly_selected: bool) -> bool { + explicitly_selected } } @@ -40,15 +63,38 @@ impl PlaneCodec for ThirdPartyPlaneCodec { fn encode( &self, - _plane: &[u8], + plane: &[u8], _shared_state: Option<&[u8]>, _layout: &PlaneLayout, ) -> Result { - Err(PtwmCoreError::InvalidContainer( - "third-party codecs cannot participate in the trial-encode loop \ - (v1 dispatches third-party codecs only at decode time)" - .into(), - )) + if !Self::allows_encode(self.explicitly_selected) { + return Err(PtwmCoreError::InvalidContainer( + "third-party codecs do not participate in the trial-encode loop; \ + select this codec explicitly to encode with it" + .into(), + )); + } + let mut out = vec![0u8; plane.len() * 2 + 1024]; + let written = self.inner.encode(plane, &mut out).map_err(|e| { + map_codec_error(e, CodecDirection::Encode, "third-party encode", out.len()) + })?; + // `written` crosses the flat ABI from the third-party side, so it + // is untrusted: a codec that reports more bytes than the buffer + // holds must fail loudly here rather than let `truncate` silently + // no-op and hand the caller a zero-padded buffer as if it were + // real encoded output. + if written > out.len() { + return Err(PtwmCoreError::InvalidContainer(format!( + "third-party encode reported {written} bytes written into a {}-byte buffer", + out.len() + ))); + } + out.truncate(written); + Ok(Encoded { + state_bytes: Vec::new(), + state_format_version: 0, + payload: out, + }) } fn decode( @@ -63,7 +109,9 @@ impl PlaneCodec for ThirdPartyPlaneCodec { let written = self .inner .decode_with_state(state_format_version, state_bytes, payload, &mut out) - .map_err(|e| map_codec_error(e, "third-party decode", out.len()))?; + .map_err(|e| { + map_codec_error(e, CodecDirection::Decode, "third-party decode", out.len()) + })?; // The container decode loop for non-chunked planes doesn't // separately verify the byte count, so enforce it here: // a third-party codec that writes the wrong number of bytes @@ -78,15 +126,100 @@ impl PlaneCodec for ThirdPartyPlaneCodec { } } -fn map_codec_error(e: CodecError, ctx: &'static str, got_cap: usize) -> PtwmCoreError { +/// Which call site is mapping a [`CodecError`], so the non-`BufferTooSmall` +/// fallback below can pick the matching [`PtwmCoreError`] variant instead of +/// always reporting a decode failure. +enum CodecDirection { + Encode, + Decode, +} + +fn map_codec_error( + e: CodecError, + direction: CodecDirection, + ctx: &'static str, + got_cap: usize, +) -> PtwmCoreError { match e { CodecError::BufferTooSmall { needed } => PtwmCoreError::BufferTooSmall { expected: needed as usize, got: got_cap, }, - other => PtwmCoreError::CodecDecode { - codec: ctx, - msg: format!("{other:?}"), - }, + other => { + let msg = format!("{other:?}"); + match direction { + CodecDirection::Encode => PtwmCoreError::CodecEncode { codec: ctx, msg }, + CodecDirection::Decode => PtwmCoreError::CodecDecode { codec: ctx, msg }, + } + } + } +} + +#[cfg(test)] +mod explicit_selection_tests { + use super::*; + + /// Minimal `DispatchedPlaneCodec` that echoes its input, so the tests + /// exercise the gate rather than any real codec's behavior. + struct EchoCodec; + + impl DispatchedPlaneCodec for EchoCodec { + fn encode(&self, input: &[u8], output: &mut [u8]) -> Result { + output[..input.len()].copy_from_slice(input); + Ok(input.len()) + } + fn decode(&self, input: &[u8], output: &mut [u8]) -> Result { + output[..input.len()].copy_from_slice(input); + Ok(input.len()) + } + } + + #[test] + fn encode_is_refused_for_trial_selected_codecs() { + let codec = ThirdPartyPlaneCodec::new(Arc::new(EchoCodec)); + let err = codec + .encode(&[1, 2, 3], None, &PlaneLayout::default()) + .expect_err("the trial-encode loop must not reach a third-party encoder"); + assert!( + format!("{err}").contains("select this codec explicitly"), + "the error must tell the caller how to proceed, got: {err}" + ); + } + + #[test] + fn encode_is_permitted_for_explicitly_selected_codecs() { + let codec = ThirdPartyPlaneCodec::new_explicit(Arc::new(EchoCodec)); + let encoded = codec + .encode(&[1, 2, 3], None, &PlaneLayout::default()) + .expect("explicit selection must reach the encoder"); + assert_eq!(encoded.payload.len(), 3, "echo codec returns its input"); + } + + /// `DispatchedPlaneCodec` whose `encode` reports writing more bytes + /// than the output buffer holds, so the tests can exercise the + /// out-of-bounds guard without a codec that actually overruns the + /// buffer (which would panic in the stub instead of in the guard). + struct OverclaimingCodec; + + impl DispatchedPlaneCodec for OverclaimingCodec { + fn encode(&self, _input: &[u8], output: &mut [u8]) -> Result { + Ok(output.len() + 1) + } + fn decode(&self, input: &[u8], output: &mut [u8]) -> Result { + output[..input.len()].copy_from_slice(input); + Ok(input.len()) + } + } + + #[test] + fn encode_fails_loudly_when_third_party_overreports_bytes_written() { + let codec = ThirdPartyPlaneCodec::new_explicit(Arc::new(OverclaimingCodec)); + let err = codec + .encode(&[1, 2, 3], None, &PlaneLayout::default()) + .expect_err("an over-large byte count must not become a zero-padded payload"); + assert!( + format!("{err}").contains("bytes written"), + "the error should describe the mismatched byte count, got: {err}" + ); } } diff --git a/crates/ptwm-core/tests/cuda_plane_codec_rust_only.rs b/crates/ptwm-core/tests/cuda_plane_codec_rust_only.rs new file mode 100644 index 0000000..00350db --- /dev/null +++ b/crates/ptwm-core/tests/cuda_plane_codec_rust_only.rs @@ -0,0 +1,16 @@ +//! Decode through the CUDA plane-codec router from Rust alone. +//! +//! This test exists to pin a layering property: nothing on the decode +//! path may require the Python bindings. It links `ptwm-core` only. + +use ptwm_core::flavor::PlaneCodecCudaRouter; + +#[test] +fn cuda_plane_codec_router_is_constructible_without_python() { + // Constructing and querying the router must not need an interpreter. + let router = PlaneCodecCudaRouter::new(Vec::new()); + let id = ptwm_core::extension::builtin_canonical_id("identity"); + // Default policy denies CUDA, so this resolves to an error rather + // than a panic -- the point is that it runs at all. + assert!(router.get(&id).is_err()); +} diff --git a/crates/ptwm-py/src/device_buffer.rs b/crates/ptwm-py/src/device_buffer.rs new file mode 100644 index 0000000..5e27032 --- /dev/null +++ b/crates/ptwm-py/src/device_buffer.rs @@ -0,0 +1,466 @@ +//! Shared DLPack capsule parsing and buffer-protocol fallback helpers. +//! +//! Extracted from `hardware.rs` so `plane_codec_cuda.rs` can reuse the same +//! device-pointer extraction and input-resolution logic rather than +//! duplicating a hand-rolled DLPack parser. +//! +//! # DLPack capsule parsing +//! +//! `torch.Tensor.__dlpack__()` returns a `PyCapsule` wrapping a +//! `DLManagedTensor` (frozen C ABI, defined by `dlpack.h`). No `dlpack` +//! crate dependency is used here: the pinned `pyo3 = "0.24"` in this +//! workspace conflicts with the version range the `dlpark` crate's `pyo3` +//! feature requires, so the relevant structs are reproduced below directly +//! against the stable DLPack C layout instead. + +use std::os::raw::{c_int, c_void}; + +use pyo3::buffer::PyBuffer; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyCapsuleMethods, PyDict}; + +pub(crate) fn to_pyerr(e: E) -> PyErr { + pyo3::exceptions::PyValueError::new_err(e.to_string()) +} + +// --------------------------------------------------------------------------- +// DLPack: frozen C structs (dlpack.h) +// --------------------------------------------------------------------------- + +/// The capsule name mandated by the DLPack Python spec for an +/// unconsumed `__dlpack__()` capsule. A producer that has already +/// consumed the capsule renames it to `"used_dltensor"`; this module only +/// ever sees freshly produced capsules, so `"dltensor"` is the only name +/// accepted here. +pub(crate) const DLPACK_CAPSULE_NAME: &str = "dltensor"; + +/// `DLDeviceType::kDLCUDA`, per `dlpack.h`. +pub(crate) const DL_CUDA: c_int = 2; + +#[repr(C)] +#[derive(Clone, Copy)] +pub(crate) struct DLDevice { + pub(crate) device_type: c_int, // kDLCUDA = 2 + pub(crate) device_id: c_int, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub(crate) struct DLDataType { + pub(crate) code: u8, // kDLBfloat = 4, kDLUInt = 1, etc. + pub(crate) bits: u8, + pub(crate) lanes: u16, +} + +#[repr(C)] +pub(crate) struct DLTensor { + pub(crate) data: *mut c_void, + pub(crate) device: DLDevice, + pub(crate) ndim: c_int, + pub(crate) dtype: DLDataType, + pub(crate) shape: *mut i64, + pub(crate) strides: *mut i64, // may be null (implies row-major contiguous) + pub(crate) byte_offset: u64, +} + +#[repr(C)] +pub(crate) struct DLManagedTensor { + pub(crate) dl_tensor: DLTensor, + pub(crate) manager_ctx: *mut c_void, + pub(crate) deleter: Option, +} + +/// True when `strides` (in elements, per the DLPack spec) matches the +/// row-major-contiguous layout implied by `shape`. +/// +/// A dimension of size 0 or 1 is skipped: its stride is a don't-care for +/// contiguity purposes (this mirrors the convention PyTorch's own +/// `is_contiguous()` uses for singleton dimensions), which keeps this +/// check from rejecting perfectly usable tensors that merely have an +/// arbitrary stride recorded on a size-1 axis. +pub(crate) fn is_row_major_contiguous(shape: &[i64], strides: &[i64]) -> bool { + debug_assert_eq!(shape.len(), strides.len()); + let mut expected: i64 = 1; + for i in (0..shape.len()).rev() { + let dim = shape[i]; + if dim < 0 { + return false; + } + if dim > 1 && strides[i] != expected { + return false; + } + expected = expected.saturating_mul(dim.max(1)); + } + true +} + +/// Extract (device pointer as u64, device ordinal, byte length) from a +/// PyCapsule returned by `tensor.__dlpack__(stream=...)`. Rejects +/// non-contiguous tensors (a non-null `strides` pointer whose values do +/// not match the row-major-contiguous stride for `shape`), matching the +/// existing `.contiguous()` convention already used on the CPU path in +/// `python/ptwm/core/_compressor.py::_to_raw_bytes`. +/// +/// `role` names the argument this capsule came from (e.g. `"compressed"`, +/// `"out"`, `"src"`, `"dst"`), for the contiguity error message only — it +/// mirrors the `role` prefix `resolve_input`'s own buffer-protocol errors +/// already use, so an error from either branch names the actual argument +/// rather than a hardcoded caller. +pub(crate) fn extract_device_ptr( + capsule: &Bound<'_, PyAny>, + role: &str, +) -> PyResult<(u64, i32, usize)> { + let capsule: &Bound<'_, PyCapsule> = capsule + .downcast::() + .map_err(|e| to_pyerr(e.to_string()))?; + + // `PyCapsule::pointer()` fetches the capsule's own stored name and + // passes it straight back into `PyCapsule_GetPointer`, so it always + // "succeeds" for any valid capsule regardless of what that name is. + // The DLPack Python spec pins the name to a specific string; check it + // explicitly here rather than trusting an unnamed or wrongly-named + // capsule to actually contain a `DLManagedTensor`. + let name = capsule.name().map_err(to_pyerr)?; + match name { + Some(n) if n.to_str().map_err(to_pyerr)? == DLPACK_CAPSULE_NAME => {} + Some(n) => { + return Err(to_pyerr(format!( + "expected a DLPack capsule named '{DLPACK_CAPSULE_NAME}', got '{}' \ + (has this capsule already been consumed?)", + n.to_string_lossy() + ))); + } + None => { + return Err(to_pyerr(format!( + "expected a DLPack capsule named '{DLPACK_CAPSULE_NAME}', capsule has no name" + ))); + } + } + + let raw_ptr = capsule.pointer(); + if raw_ptr.is_null() { + return Err(to_pyerr("DLPack capsule pointer is null")); + } + + // SAFETY: the name check above confirms this is a "dltensor" capsule + // per the DLPack Python spec, which guarantees the capsule's opaque + // pointer references a `DLManagedTensor` laid out per the frozen + // dlpack.h C ABI. The tensor memory is kept alive by the capsule + // object itself, which the caller holds for the duration of this + // call (it is not dropped until the enclosing `#[pyfunction]` returns). + let managed: &DLManagedTensor = unsafe { &*raw_ptr.cast::() }; + let dl_tensor = &managed.dl_tensor; + + if dl_tensor.device.device_type != DL_CUDA { + return Err(to_pyerr(format!( + "expected a CUDA DLPack tensor (device_type={DL_CUDA}), got device_type={}", + dl_tensor.device.device_type + ))); + } + + if dl_tensor.ndim < 0 { + return Err(to_pyerr("DLPack tensor has negative ndim")); + } + let ndim = dl_tensor.ndim as usize; + + let shape: &[i64] = if ndim == 0 { + &[] + } else { + if dl_tensor.shape.is_null() { + return Err(to_pyerr("DLPack tensor shape pointer is null")); + } + // SAFETY: shape is non-null (checked above) and the + // `DLManagedTensor` contract guarantees `ndim` valid `i64` + // entries at that pointer. + unsafe { std::slice::from_raw_parts(dl_tensor.shape, ndim) } + }; + + if !dl_tensor.strides.is_null() { + // SAFETY: same contract as `shape` above; non-null `strides` + // points to `ndim` valid `i64` entries per the DLPack spec. + let strides = unsafe { std::slice::from_raw_parts(dl_tensor.strides, ndim) }; + if !is_row_major_contiguous(shape, strides) { + return Err(to_pyerr(format!( + "{role}: non-contiguous tensor passed to device-buffer dispatch; \ + call .contiguous() on the tensor before passing it in" + ))); + } + } + + let elem_bits = dl_tensor.dtype.bits as u64 * dl_tensor.dtype.lanes as u64; + if elem_bits == 0 || !elem_bits.is_multiple_of(8) { + return Err(to_pyerr(format!( + "unsupported DLPack dtype: bits={} lanes={} does not divide evenly into bytes", + dl_tensor.dtype.bits, dl_tensor.dtype.lanes + ))); + } + let elem_bytes = elem_bits / 8; + + let num_elements: u64 = shape + .iter() + .try_fold(1u64, |acc, &d| { + if d < 0 { + None + } else { + acc.checked_mul(d as u64) + } + }) + .ok_or_else(|| to_pyerr("DLPack tensor shape has a negative dimension or overflows"))?; + + let byte_len = num_elements + .checked_mul(elem_bytes) + .ok_or_else(|| to_pyerr("DLPack tensor byte length overflows usize"))?; + + if dl_tensor.data.is_null() { + return Err(to_pyerr("DLPack tensor data pointer is null")); + } + let device_ptr = (dl_tensor.data as u64) + .checked_add(dl_tensor.byte_offset) + .ok_or_else(|| to_pyerr("DLPack tensor data pointer + byte_offset overflows u64"))?; + + Ok((device_ptr, dl_tensor.device.device_id, byte_len as usize)) +} + +// --------------------------------------------------------------------------- +// Buffer-protocol fallback: plain `bytes` / `bytearray` inputs +// --------------------------------------------------------------------------- + +/// Either a DLPack capsule (produced by `tensor.__dlpack__()`) or a plain +/// buffer-protocol object (`bytes`, `bytearray`, ...), resolved once and +/// kept alive for the duration of the dispatch call. +/// +/// `resolve_input` picks the variant based on whether the Python object +/// exposes `__dlpack__`; real `torch.Tensor` arguments always do, so the +/// DLPack path is unchanged for them. Plain `bytes`/`bytearray` do not, so +/// they fall back to `Buffer`, which reads/writes the host memory the +/// buffer protocol exposes directly. This lets a CPU-only caller (this +/// crate's own interop tests, and any other host-memory caller) exercise a +/// dispatch contribution without constructing a fake CUDA-shaped DLPack +/// tensor; a contribution that treats its "device" pointers as ordinary +/// host pointers (e.g. a CPU-flavor `hardware_backend`) can be driven +/// through this path exactly as it expects. Both `hardware.rs`'s +/// `compressed`/`out` argument pair and `plane_codec_cuda.rs`'s `src`/`dst` +/// pair resolve through this same enum; the argument names live in the +/// `role` string carried alongside each variant, not in the type itself. +pub(crate) enum InputHandle<'py> { + Dlpack(Bound<'py, PyAny>, &'static str), + Buffer(PyBuffer), +} + +impl InputHandle<'_> { + /// Returns `(pointer, device_ordinal, byte_len)`. `device_ordinal` is + /// `None` for the buffer-protocol fallback, which carries no device + /// information; callers must skip the device-ordinal-agreement check + /// in that case rather than treat `None` as a mismatch. + pub(crate) fn ptr_len_ordinal(&self) -> PyResult<(u64, Option, usize)> { + match self { + InputHandle::Dlpack(capsule, role) => { + let (ptr, ordinal, len) = extract_device_ptr(capsule, role)?; + Ok((ptr, Some(ordinal), len)) + } + InputHandle::Buffer(buf) => Ok((buf.buf_ptr() as u64, None, buf.len_bytes())), + } + } +} + +/// Resolve one dispatch argument (named by `role`, e.g. `compressed`/`out` +/// or `src`/`dst` depending on the caller) to an [`InputHandle`]. +/// +/// `dlpack_kwargs` is only used on the DLPack branch (`stream=` has no +/// meaning for a plain host buffer). `require_writable` rejects a +/// read-only buffer on the fallback branch; the DLPack branch has no +/// equivalent read-only concept at this layer; a decode into a read-only +/// tensor's backing memory is between the caller and whatever `torch` +/// enforces, unchanged from before this fallback was added. +pub(crate) fn resolve_input<'py>( + obj: &Bound<'py, PyAny>, + dlpack_kwargs: &Bound<'py, PyDict>, + role: &'static str, + require_writable: bool, +) -> PyResult> { + if obj.hasattr("__dlpack__")? { + let capsule = obj.call_method("__dlpack__", (), Some(dlpack_kwargs))?; + return Ok(InputHandle::Dlpack(capsule, role)); + } + + let buf = PyBuffer::::get(obj)?; + if !buf.is_c_contiguous() { + return Err(to_pyerr(format!("{role}: buffer must be C-contiguous"))); + } + if require_writable && buf.readonly() { + return Err(to_pyerr(format!( + "{role}: buffer-protocol fallback requires a writable buffer \ + (e.g. bytearray), got a read-only buffer" + ))); + } + Ok(InputHandle::Buffer(buf)) +} + +// --------------------------------------------------------------------------- +// Device-ordinal agreement check +// --------------------------------------------------------------------------- + +/// Pure comparison behind [`check_device_ordinal_agreement`]: `Some(message)` +/// when the two resolved ordinals disagree with `device_ordinal`, `None` when +/// they agree or either side is unknown (the buffer-protocol fallback, +/// which carries no device information to check against). +/// +/// Kept separate from `check_device_ordinal_agreement` itself, and free of +/// any `pyo3` type, so it can be unit-tested directly here (matching this +/// module's existing `is_row_major_contiguous` convention): building or +/// dropping a real `PyErr` requires the CPython C API, which this crate's +/// plain `cargo test` binary cannot link (the `extension-module` pyo3 +/// feature deliberately omits linking libpython, since a real build of this +/// crate is loaded into an already-running Python process instead). +fn device_ordinal_mismatch_message( + device_ordinal: u32, + in_ordinal: Option, + out_ordinal: Option, + in_role: &str, + out_role: &str, +) -> Option { + let (in_ordinal, out_ordinal) = match (in_ordinal, out_ordinal) { + (Some(a), Some(b)) => (a, b), + _ => return None, + }; + if in_ordinal != device_ordinal as i32 || out_ordinal != device_ordinal as i32 { + Some(format!( + "device_ordinal mismatch: caller supplied {device_ordinal}, but {in_role} \ + tensor reports device {in_ordinal} and {out_role} tensor reports device {out_ordinal}" + )) + } else { + None + } +} + +/// Verify that a caller-supplied `device_ordinal` (used to acquire a CUDA +/// stream before any tensor was inspected) agrees with both resolved +/// tensors' own DLPack-reported device ordinals. +/// +/// A mismatch means the stream and the tensor memory belong to different +/// devices, which would silently corrupt a decode/encode rather than fail +/// loudly, so this returns an error naming all three ordinals. Either +/// ordinal is `None` on the buffer-protocol fallback branch (see +/// [`InputHandle::ptr_len_ordinal`]), which carries no device information; +/// the check is a no-op in that case rather than treated as a mismatch. +/// +/// `in_role`/`out_role` name the two arguments in the error message (e.g. +/// `"compressed"`/`"out"` for `hardware.rs`, `"src"`/`"dst"` for +/// `plane_codec_cuda.rs`), so the message names the actual caller +/// arguments rather than a hardcoded pair. +pub(crate) fn check_device_ordinal_agreement( + device_ordinal: u32, + in_ordinal: Option, + out_ordinal: Option, + in_role: &str, + out_role: &str, +) -> PyResult<()> { + match device_ordinal_mismatch_message( + device_ordinal, + in_ordinal, + out_ordinal, + in_role, + out_role, + ) { + Some(msg) => Err(to_pyerr(msg)), + None => Ok(()), + } +} + +#[cfg(test)] +mod dlpack_tests { + use super::*; + + #[test] + fn dl_data_type_bf16_matches_dlpack_spec_code() { + // DLPack's DLDataTypeCode for bfloat16 is kDLBfloat = 4, per the + // frozen DLPack C header (dlpack.h, DLDataTypeCode enum). This + // guards against a transcription error in the hand-rolled struct + // above; it is not a full DLPack conformance test. + let dt = DLDataType { + code: 4, + bits: 16, + lanes: 1, + }; + assert_eq!(dt.code, 4); + assert_eq!(dt.bits, 16); + } + + #[test] + fn row_major_contiguous_accepts_standard_layout() { + // A [2, 3] row-major tensor has strides [3, 1] (in elements). + assert!(is_row_major_contiguous(&[2, 3], &[3, 1])); + } + + #[test] + fn row_major_contiguous_rejects_transposed_layout() { + // The same [2, 3] tensor transposed (a view, not a copy) has + // strides [1, 2], which is not row-major-contiguous. + assert!(!is_row_major_contiguous(&[2, 3], &[1, 2])); + } + + #[test] + fn row_major_contiguous_ignores_singleton_dimension_stride() { + // A [1, 3] tensor's size-1 leading dimension carries an + // arbitrary/don't-care stride in many producers; only the size-3 + // trailing dimension's stride must be 1. + assert!(is_row_major_contiguous(&[1, 3], &[999, 1])); + } +} + +#[cfg(test)] +mod device_ordinal_agreement_tests { + use super::*; + + // These exercise `device_ordinal_mismatch_message` directly rather than + // `check_device_ordinal_agreement` itself: the latter's compiled body + // unconditionally contains a call to `to_pyerr`/`PyErr::new_err`, so + // even a test run that only takes the `Ok` branch at runtime would make + // that call statically reachable from this test binary and fail to + // link (see the doc comment on `device_ordinal_mismatch_message`). + + #[test] + fn matching_ordinals_pass() { + assert!(device_ordinal_mismatch_message(0, Some(0), Some(0), "src", "dst").is_none()); + } + + #[test] + fn mismatched_ordinal_produces_descriptive_message() { + let msg = device_ordinal_mismatch_message(0, Some(0), Some(1), "src", "dst") + .expect("mismatched ordinals must be reported"); + assert!(msg.contains("device_ordinal mismatch"), "{msg}"); + assert!(msg.contains("src"), "{msg}"); + assert!(msg.contains("dst"), "{msg}"); + } + + #[test] + fn buffer_protocol_fallback_both_none_passes() { + // Neither side carries a device ordinal on the buffer-protocol + // fallback branch, so there is nothing to check and this must not + // be treated as a mismatch. + assert!(device_ordinal_mismatch_message(0, None, None, "src", "dst").is_none()); + } + + // The next two pin the exact wording for each call site's argument + // names: `hardware.rs` uses "compressed"/"out", `plane_codec_cuda.rs` + // uses "src"/"dst" for both its encode and decode functions. A silent + // change to either message would fail one of these. + + #[test] + fn message_names_the_hardware_backend_argument_pair() { + assert_eq!( + device_ordinal_mismatch_message(0, Some(1), Some(1), "compressed", "out").unwrap(), + "device_ordinal mismatch: caller supplied 0, but compressed tensor reports \ + device 1 and out tensor reports device 1" + ); + } + + #[test] + fn message_names_the_plane_codec_cuda_argument_pair() { + assert_eq!( + device_ordinal_mismatch_message(0, Some(1), Some(1), "src", "dst").unwrap(), + "device_ordinal mismatch: caller supplied 0, but src tensor reports \ + device 1 and dst tensor reports device 1" + ); + } +} diff --git a/crates/ptwm-py/src/hardware.rs b/crates/ptwm-py/src/hardware.rs index f973f75..2b1914c 100644 --- a/crates/ptwm-py/src/hardware.rs +++ b/crates/ptwm-py/src/hardware.rs @@ -1,5 +1,5 @@ //! PyO3 bridge for invoking a `hardware_backend` contribution by canonical -//! id, plus a hand-rolled DLPack capsule parser. +//! id. //! //! There is no in-tree consumer of `hardware_backend` (unlike `plane_codec`, //! which the container encode/decode loop calls internally): this module @@ -10,31 +10,19 @@ //! `delta_scheme.rs` and `ext.rs` rather than introducing a persistent //! router object. //! -//! # DLPack capsule parsing -//! -//! `torch.Tensor.__dlpack__()` returns a `PyCapsule` wrapping a -//! `DLManagedTensor` (frozen C ABI, defined by `dlpack.h`). No `dlpack` -//! crate dependency is used here: the pinned `pyo3 = "0.24"` in this -//! workspace conflicts with the version range the `dlpark` crate's `pyo3` -//! feature requires, so the relevant structs are reproduced below directly -//! against the stable DLPack C layout instead. - -use std::os::raw::{c_int, c_void}; +//! DLPack capsule parsing and the buffer-protocol fallback live in +//! `crate::device_buffer`, shared with `plane_codec_cuda.rs`. -use pyo3::buffer::PyBuffer; use pyo3::prelude::*; -use pyo3::types::{PyCapsule, PyCapsuleMethods, PyDict}; +use pyo3::types::PyDict; use ptwm_core::discovery::scan_all_cached; use ptwm_core::extension::CanonicalId; use ptwm_core::flavor::HardwareBackendRouter; +use crate::device_buffer::{check_device_ordinal_agreement, resolve_input, to_pyerr}; use crate::policy::PyResolvedPolicy; -fn to_pyerr(e: E) -> PyErr { - pyo3::exceptions::PyValueError::new_err(e.to_string()) -} - /// Build a `HardwareBackendRouter` for one call. /// /// `policy` is optional. When `None`, the router is built via @@ -57,267 +45,6 @@ fn hardware_backend_router(policy: Option<&PyResolvedPolicy>) -> PyResult, -} - -/// True when `strides` (in elements, per the DLPack spec) matches the -/// row-major-contiguous layout implied by `shape`. -/// -/// A dimension of size 0 or 1 is skipped: its stride is a don't-care for -/// contiguity purposes (this mirrors the convention PyTorch's own -/// `is_contiguous()` uses for singleton dimensions), which keeps this -/// check from rejecting perfectly usable tensors that merely have an -/// arbitrary stride recorded on a size-1 axis. -fn is_row_major_contiguous(shape: &[i64], strides: &[i64]) -> bool { - debug_assert_eq!(shape.len(), strides.len()); - let mut expected: i64 = 1; - for i in (0..shape.len()).rev() { - let dim = shape[i]; - if dim < 0 { - return false; - } - if dim > 1 && strides[i] != expected { - return false; - } - expected = expected.saturating_mul(dim.max(1)); - } - true -} - -/// Extract (device pointer as u64, device ordinal, byte length) from a -/// PyCapsule returned by `tensor.__dlpack__(stream=...)`. Rejects -/// non-contiguous tensors (a non-null `strides` pointer whose values do -/// not match the row-major-contiguous stride for `shape`), matching the -/// existing `.contiguous()` convention already used on the CPU path in -/// `python/ptwm/core/_compressor.py::_to_raw_bytes`. -fn extract_device_ptr(capsule: &Bound<'_, PyAny>) -> PyResult<(u64, i32, usize)> { - let capsule: &Bound<'_, PyCapsule> = capsule - .downcast::() - .map_err(|e| to_pyerr(e.to_string()))?; - - // `PyCapsule::pointer()` fetches the capsule's own stored name and - // passes it straight back into `PyCapsule_GetPointer`, so it always - // "succeeds" for any valid capsule regardless of what that name is. - // The DLPack Python spec pins the name to a specific string; check it - // explicitly here rather than trusting an unnamed or wrongly-named - // capsule to actually contain a `DLManagedTensor`. - let name = capsule.name().map_err(to_pyerr)?; - match name { - Some(n) if n.to_str().map_err(to_pyerr)? == DLPACK_CAPSULE_NAME => {} - Some(n) => { - return Err(to_pyerr(format!( - "expected a DLPack capsule named '{DLPACK_CAPSULE_NAME}', got '{}' \ - (has this capsule already been consumed?)", - n.to_string_lossy() - ))); - } - None => { - return Err(to_pyerr(format!( - "expected a DLPack capsule named '{DLPACK_CAPSULE_NAME}', capsule has no name" - ))); - } - } - - let raw_ptr = capsule.pointer(); - if raw_ptr.is_null() { - return Err(to_pyerr("DLPack capsule pointer is null")); - } - - // SAFETY: the name check above confirms this is a "dltensor" capsule - // per the DLPack Python spec, which guarantees the capsule's opaque - // pointer references a `DLManagedTensor` laid out per the frozen - // dlpack.h C ABI. The tensor memory is kept alive by the capsule - // object itself, which the caller holds for the duration of this - // call (it is not dropped until the enclosing `#[pyfunction]` returns). - let managed: &DLManagedTensor = unsafe { &*raw_ptr.cast::() }; - let dl_tensor = &managed.dl_tensor; - - if dl_tensor.device.device_type != DL_CUDA { - return Err(to_pyerr(format!( - "expected a CUDA DLPack tensor (device_type={DL_CUDA}), got device_type={}", - dl_tensor.device.device_type - ))); - } - - if dl_tensor.ndim < 0 { - return Err(to_pyerr("DLPack tensor has negative ndim")); - } - let ndim = dl_tensor.ndim as usize; - - let shape: &[i64] = if ndim == 0 { - &[] - } else { - if dl_tensor.shape.is_null() { - return Err(to_pyerr("DLPack tensor shape pointer is null")); - } - // SAFETY: shape is non-null (checked above) and the - // `DLManagedTensor` contract guarantees `ndim` valid `i64` - // entries at that pointer. - unsafe { std::slice::from_raw_parts(dl_tensor.shape, ndim) } - }; - - if !dl_tensor.strides.is_null() { - // SAFETY: same contract as `shape` above; non-null `strides` - // points to `ndim` valid `i64` entries per the DLPack spec. - let strides = unsafe { std::slice::from_raw_parts(dl_tensor.strides, ndim) }; - if !is_row_major_contiguous(shape, strides) { - return Err(to_pyerr( - "non-contiguous tensor passed to hardware_backend dispatch; \ - call .contiguous() on the tensor before passing it in", - )); - } - } - - let elem_bits = dl_tensor.dtype.bits as u64 * dl_tensor.dtype.lanes as u64; - if elem_bits == 0 || !elem_bits.is_multiple_of(8) { - return Err(to_pyerr(format!( - "unsupported DLPack dtype: bits={} lanes={} does not divide evenly into bytes", - dl_tensor.dtype.bits, dl_tensor.dtype.lanes - ))); - } - let elem_bytes = elem_bits / 8; - - let num_elements: u64 = shape - .iter() - .try_fold(1u64, |acc, &d| { - if d < 0 { - None - } else { - acc.checked_mul(d as u64) - } - }) - .ok_or_else(|| to_pyerr("DLPack tensor shape has a negative dimension or overflows"))?; - - let byte_len = num_elements - .checked_mul(elem_bytes) - .ok_or_else(|| to_pyerr("DLPack tensor byte length overflows usize"))?; - - if dl_tensor.data.is_null() { - return Err(to_pyerr("DLPack tensor data pointer is null")); - } - let device_ptr = (dl_tensor.data as u64) - .checked_add(dl_tensor.byte_offset) - .ok_or_else(|| to_pyerr("DLPack tensor data pointer + byte_offset overflows u64"))?; - - Ok((device_ptr, dl_tensor.device.device_id, byte_len as usize)) -} - -// --------------------------------------------------------------------------- -// Buffer-protocol fallback: plain `bytes` / `bytearray` inputs -// --------------------------------------------------------------------------- - -/// Either a DLPack capsule (produced by `tensor.__dlpack__()`) or a plain -/// buffer-protocol object (`bytes`, `bytearray`, ...), resolved once and -/// kept alive for the duration of the dispatch call. -/// -/// `resolve_input` picks the variant based on whether the Python object -/// exposes `__dlpack__`; real `torch.Tensor` arguments always do, so the -/// DLPack path is unchanged for them. Plain `bytes`/`bytearray` do not, so -/// they fall back to `Buffer`, which reads/writes the host memory the -/// buffer protocol exposes directly. This lets a CPU-only caller (this -/// crate's own interop tests, and any other host-memory caller) exercise a -/// `hardware_backend` contribution without constructing a fake CUDA-shaped -/// DLPack tensor; a `hardware_class = "cpu"` contribution such as -/// `ref_hardware_backend` treats its "device" pointers as ordinary host -/// pointers already (see that crate's own doc comments), so passing a host -/// pointer through this path is exactly what such a contribution expects. -enum InputHandle<'py> { - Dlpack(Bound<'py, PyAny>), - Buffer(PyBuffer), -} - -impl InputHandle<'_> { - /// Returns `(pointer, device_ordinal, byte_len)`. `device_ordinal` is - /// `None` for the buffer-protocol fallback, which carries no device - /// information; callers must skip the device-ordinal-agreement check - /// in that case rather than treat `None` as a mismatch. - fn ptr_len_ordinal(&self) -> PyResult<(u64, Option, usize)> { - match self { - InputHandle::Dlpack(capsule) => { - let (ptr, ordinal, len) = extract_device_ptr(capsule)?; - Ok((ptr, Some(ordinal), len)) - } - InputHandle::Buffer(buf) => Ok((buf.buf_ptr() as u64, None, buf.len_bytes())), - } - } -} - -/// Resolve one `compressed`/`out` argument to an [`InputHandle`]. -/// -/// `dlpack_kwargs` is only used on the DLPack branch (`stream=` has no -/// meaning for a plain host buffer). `require_writable` rejects a -/// read-only buffer on the fallback branch; the DLPack branch has no -/// equivalent read-only concept at this layer; a decode into a read-only -/// tensor's backing memory is between the caller and whatever `torch` -/// enforces, unchanged from before this fallback was added. -fn resolve_input<'py>( - obj: &Bound<'py, PyAny>, - dlpack_kwargs: &Bound<'py, PyDict>, - role: &str, - require_writable: bool, -) -> PyResult> { - if obj.hasattr("__dlpack__")? { - let capsule = obj.call_method("__dlpack__", (), Some(dlpack_kwargs))?; - return Ok(InputHandle::Dlpack(capsule)); - } - - let buf = PyBuffer::::get(obj)?; - if !buf.is_c_contiguous() { - return Err(to_pyerr(format!("{role}: buffer must be C-contiguous"))); - } - if require_writable && buf.readonly() { - return Err(to_pyerr(format!( - "{role}: buffer-protocol fallback requires a writable buffer \ - (e.g. bytearray), got a read-only buffer" - ))); - } - Ok(InputHandle::Buffer(buf)) -} - // --------------------------------------------------------------------------- // PyO3-visible functions // --------------------------------------------------------------------------- @@ -407,15 +134,9 @@ pub fn hardware_backend_dispatch_decode_cuda<'py>( // different devices, which would silently corrupt the decode rather // than fail loudly. Neither ordinal is known on the buffer-protocol // fallback branch (`None`), so the check is skipped there rather than - // treated as a mismatch. - if let (Some(in_ordinal), Some(out_ordinal)) = (in_ordinal, out_ordinal) { - if in_ordinal != device_ordinal as i32 || out_ordinal != device_ordinal as i32 { - return Err(to_pyerr(format!( - "device_ordinal mismatch: caller supplied {device_ordinal}, but compressed \ - tensor reports device {in_ordinal} and out tensor reports device {out_ordinal}" - ))); - } - } + // treated as a mismatch. See `check_device_ordinal_agreement`, shared + // with `plane_codec_cuda.rs`. + check_device_ordinal_agreement(device_ordinal, in_ordinal, out_ordinal, "compressed", "out")?; py.allow_threads(|| { backend @@ -443,44 +164,3 @@ pub fn register(m: &Bound<'_, pyo3::types::PyModule>) -> PyResult<()> { )?)?; Ok(()) } - -#[cfg(test)] -mod dlpack_tests { - use super::*; - - #[test] - fn dl_data_type_bf16_matches_dlpack_spec_code() { - // DLPack's DLDataTypeCode for bfloat16 is kDLBfloat = 4, per the - // frozen DLPack C header (dlpack.h, DLDataTypeCode enum). This - // guards against a transcription error in the hand-rolled struct - // above; it is not a full DLPack conformance test. - let dt = DLDataType { - code: 4, - bits: 16, - lanes: 1, - }; - assert_eq!(dt.code, 4); - assert_eq!(dt.bits, 16); - } - - #[test] - fn row_major_contiguous_accepts_standard_layout() { - // A [2, 3] row-major tensor has strides [3, 1] (in elements). - assert!(is_row_major_contiguous(&[2, 3], &[3, 1])); - } - - #[test] - fn row_major_contiguous_rejects_transposed_layout() { - // The same [2, 3] tensor transposed (a view, not a copy) has - // strides [1, 2], which is not row-major-contiguous. - assert!(!is_row_major_contiguous(&[2, 3], &[1, 2])); - } - - #[test] - fn row_major_contiguous_ignores_singleton_dimension_stride() { - // A [1, 3] tensor's size-1 leading dimension carries an - // arbitrary/don't-care stride in many producers; only the size-3 - // trailing dimension's stride must be 1. - assert!(is_row_major_contiguous(&[1, 3], &[999, 1])); - } -} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index b4abf2b..7a56c71 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -17,10 +17,12 @@ use ptwm_core::container::ContainerReader; use ptwm_core::{PtwmCoreError, codec_tagged, delta, entropy}; mod delta_scheme; +mod device_buffer; mod ext; mod hardware; mod host_flavor; mod inspect; +mod plane_codec_cuda; mod policy; mod trust; @@ -63,6 +65,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { inspect::register(m)?; delta_scheme::register(m)?; hardware::register(m)?; + plane_codec_cuda::register(m)?; Ok(()) } diff --git a/crates/ptwm-py/src/plane_codec_cuda.rs b/crates/ptwm-py/src/plane_codec_cuda.rs new file mode 100644 index 0000000..6bc1788 --- /dev/null +++ b/crates/ptwm-py/src/plane_codec_cuda.rs @@ -0,0 +1,196 @@ +//! PyO3 bridge for dispatching a `plane_codec` CUDA (device-resident) +//! encode/decode by selector. +//! +//! Mirrors `hardware.rs`'s `hardware_backend_dispatch_decode_cuda` in +//! structure: build a `PlaneCodecCudaRouter` for the call (gated by an +//! optional caller-supplied policy, same as `hardware_backend_router`), +//! resolve device pointers through the shared DLPack/buffer-protocol +//! helpers in `crate::device_buffer`, enforce device-ordinal agreement, +//! then dispatch with the GIL released. +//! +//! Unlike `hardware_backend_dispatch_decode_cuda`, both `selector` and +//! `codec_id` are resolved through `resolve_codec_selector` rather than +//! parsed directly as canonical ids: this lets a caller pass a +//! human-typed name (a builtin like `"identity"`, or an installed +//! contribution's label) as well as a canonical id, matching the +//! selector-resolution convenience `resolve_codec_selector` provides +//! elsewhere. Both selectors are resolved against the same `installed` +//! scan used to build the router, so name resolution and router lookup +//! agree on what is actually installed for this call. + +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use ptwm_core::discovery::DiscoveredContribution; +use ptwm_core::discovery::scan_all_cached; +use ptwm_core::extension::resolve_codec_selector; +use ptwm_core::flavor::PlaneCodecCudaRouter; + +use crate::device_buffer::{check_device_ordinal_agreement, resolve_input, to_pyerr}; +use crate::policy::PyResolvedPolicy; + +/// Build a `PlaneCodecCudaRouter` for one call, plus the `installed` scan +/// it was built from (also needed by the caller to resolve selectors +/// against the same install snapshot). +/// +/// `policy` follows exactly the same optional-gating convention as +/// `hardware.rs`'s `hardware_backend_router`: `None` keeps the router's +/// own default-deny `HostPolicy` (today's behavior, unchanged); `Some` +/// passes the caller-supplied `ResolvedPolicy`'s `HostPolicy` through to +/// `PlaneCodecCudaRouter::new_with_policy` instead. +fn plane_codec_cuda_router( + policy: Option<&PyResolvedPolicy>, +) -> PyResult<(PlaneCodecCudaRouter, Vec)> { + let installed = scan_all_cached().map_err(to_pyerr)?; + let router = match policy { + Some(p) => PlaneCodecCudaRouter::new_with_policy(installed.clone(), p.host_policy()), + None => PlaneCodecCudaRouter::new(installed.clone()), + }; + Ok((router, installed)) +} + +// --------------------------------------------------------------------------- +// PyO3-visible functions +// --------------------------------------------------------------------------- + +/// Dispatch a CUDA plane-codec encode through `selector`'s `plane_codec` +/// CUDA contribution. +/// +/// `selector` names the `plane_codec` contribution to dispatch through +/// (a builtin name, an installed contribution's label, or a canonical +/// id); `codec_id` names the specific codec variant passed through to the +/// contribution's `encode_cuda`. Both are resolved via +/// `resolve_codec_selector` against the same installed-extension snapshot +/// used to build the router. +/// +/// `src`/`dst` are ordinarily `torch.Tensor` objects already resident on +/// the CUDA device identified by `device_ordinal`, read via +/// `__dlpack__()`; a plain buffer-protocol object (`bytes`/`bytearray`) +/// is accepted as a fallback too. See `hardware.rs`'s +/// `hardware_backend_dispatch_decode_cuda` doc comment for the full +/// reasoning behind the buffer-protocol fallback, the explicit +/// `device_ordinal` parameter, and the device-ordinal agreement check — +/// all identical here. +#[pyfunction] +#[pyo3(signature = ( + selector, + codec_id, + state_bytes, + src, + dst, + device_ordinal, + policy=None, +))] +#[allow(clippy::too_many_arguments)] +pub fn plane_codec_encode_cuda<'py>( + py: Python<'py>, + selector: String, + codec_id: String, + state_bytes: &[u8], + src: &Bound<'py, PyAny>, + dst: &Bound<'py, PyAny>, + device_ordinal: u32, + policy: Option<&PyResolvedPolicy>, +) -> PyResult { + let (router, installed) = plane_codec_cuda_router(policy)?; + let id = resolve_codec_selector(&selector, &installed).map_err(to_pyerr)?; + let codec = resolve_codec_selector(&codec_id, &installed).map_err(to_pyerr)?; + let backend = router.get(&id).map_err(to_pyerr)?; + + // `DispatchedPlaneCodecCuda::cuda_stream_handle` is infallible (`u64`, + // not `Result`): unlike `hardware.rs`'s + // `HardwareBackendCuda::cuda_stream_handle`, there is no `?`/`map_err` + // here. + let stream = backend.cuda_stream_handle(device_ordinal); + + let dlpack_kwargs = PyDict::new(py); + dlpack_kwargs.set_item("stream", stream)?; + let src_handle = resolve_input(src, &dlpack_kwargs, "src", false)?; + let dst_handle = resolve_input(dst, &dlpack_kwargs, "dst", true)?; + + let (in_dev_ptr, in_ordinal, in_len) = src_handle.ptr_len_ordinal()?; + let (out_dev_ptr, out_ordinal, out_len) = dst_handle.ptr_len_ordinal()?; + + // Same device-ordinal agreement check as + // `hardware_backend_dispatch_decode_cuda`: skipped on the + // buffer-protocol fallback branch, which carries no device ordinal. + check_device_ordinal_agreement(device_ordinal, in_ordinal, out_ordinal, "src", "dst")?; + + py.allow_threads(|| { + backend + .encode_cuda( + state_bytes, + &codec, + in_dev_ptr, + in_len, + out_dev_ptr, + out_len, + device_ordinal, + ) + .map_err(to_pyerr) + }) +} + +/// Dispatch a CUDA plane-codec decode through `selector`'s `plane_codec` +/// CUDA contribution. See [`plane_codec_encode_cuda`] for the full +/// argument and behavior documentation; this differs only in calling +/// `decode_cuda` instead of `encode_cuda`. +#[pyfunction] +#[pyo3(signature = ( + selector, + codec_id, + state_bytes, + src, + dst, + device_ordinal, + policy=None, +))] +#[allow(clippy::too_many_arguments)] +pub fn plane_codec_decode_cuda<'py>( + py: Python<'py>, + selector: String, + codec_id: String, + state_bytes: &[u8], + src: &Bound<'py, PyAny>, + dst: &Bound<'py, PyAny>, + device_ordinal: u32, + policy: Option<&PyResolvedPolicy>, +) -> PyResult { + let (router, installed) = plane_codec_cuda_router(policy)?; + let id = resolve_codec_selector(&selector, &installed).map_err(to_pyerr)?; + let codec = resolve_codec_selector(&codec_id, &installed).map_err(to_pyerr)?; + let backend = router.get(&id).map_err(to_pyerr)?; + + // Infallible: see the identical comment in `plane_codec_encode_cuda`. + let stream = backend.cuda_stream_handle(device_ordinal); + + let dlpack_kwargs = PyDict::new(py); + dlpack_kwargs.set_item("stream", stream)?; + let src_handle = resolve_input(src, &dlpack_kwargs, "src", false)?; + let dst_handle = resolve_input(dst, &dlpack_kwargs, "dst", true)?; + + let (in_dev_ptr, in_ordinal, in_len) = src_handle.ptr_len_ordinal()?; + let (out_dev_ptr, out_ordinal, out_len) = dst_handle.ptr_len_ordinal()?; + + check_device_ordinal_agreement(device_ordinal, in_ordinal, out_ordinal, "src", "dst")?; + + py.allow_threads(|| { + backend + .decode_cuda( + state_bytes, + &codec, + in_dev_ptr, + in_len, + out_dev_ptr, + out_len, + device_ordinal, + ) + .map_err(to_pyerr) + }) +} + +pub fn register(m: &Bound<'_, pyo3::types::PyModule>) -> PyResult<()> { + m.add_function(pyo3::wrap_pyfunction!(plane_codec_encode_cuda, m)?)?; + m.add_function(pyo3::wrap_pyfunction!(plane_codec_decode_cuda, m)?)?; + Ok(()) +} diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 92431fe..3b824be 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -85,6 +85,8 @@ class CompressionConfig: falsify ablation measurements. Ignored when ``method`` selects a forced- codec path (ZSTD / RANS / IDENTITY). """ + codec: str | None = None + device: int | None = None @classmethod def from_resolved_policy( @@ -160,3 +162,5 @@ class DecompressionConfig: interrogated (e.g. listing tensor names). When ``False`` (the default), opening such a container raises ``ValueError``. """ + codec: str | None = None + device: int | None = None diff --git a/python/ptwm/_rust/plane_codec_cuda.py b/python/ptwm/_rust/plane_codec_cuda.py new file mode 100644 index 0000000..69f91ea --- /dev/null +++ b/python/ptwm/_rust/plane_codec_cuda.py @@ -0,0 +1,17 @@ +"""Thin shim exposing the plane_codec CUDA dispatch surface. + +Re-exports the `ptwm._core` plane_codec CUDA functions as +`ptwm._rust.plane_codec_cuda`. +""" + +from __future__ import annotations + +from ptwm._core import ( # type: ignore[import] + plane_codec_decode_cuda, + plane_codec_encode_cuda, +) + +__all__ = [ + "plane_codec_decode_cuda", + "plane_codec_encode_cuda", +] diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index cdb7366..c40daf5 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -129,6 +129,8 @@ def compress_file( threads: int | None = None, quiet: bool = False, codec_menu: list[CodecId] | None = None, + codec: str | None = None, + device: int | None = None, ) -> None: """Compress a single file.""" from ptwm import ( # noqa: PLC0415 @@ -165,6 +167,8 @@ def compress_file( method=Method(method), threads=threads, codec_menu=codec_menu, + codec=codec, + device=device, ) ) @@ -190,6 +194,10 @@ def compress_file( test_buffer += compressed_chunk if verification: + # `device` is the compression-side ordinal already passed to + # CompressionConfig above; this round-trip check always decodes on + # CPU, so it is left unset on DecompressionConfig rather than + # passed through. decompressor = Decompressor(DecompressionConfig(threads=threads)) if test: with full_path.open("rb") as f: @@ -236,6 +244,8 @@ def compress_file_delta( is_streaming: bool = False, threads: int | None = None, codec_menu: list[CodecId] | None = None, + codec: str | None = None, + device: int | None = None, ) -> None: """Compress a file using delta compression.""" from ptwm import ( # noqa: PLC0415 @@ -281,6 +291,8 @@ def compress_file_delta( method=Method(method), threads=threads, codec_menu=codec_menu, + codec=codec, + device=device, ) ) @@ -289,6 +301,8 @@ def compress_file_delta( compressed_data = compressor.compress(file_data, delta_second_data=delta_file) if verification: + # `device` is not forwarded here either; see the note in + # `compress_file`. decompressor = Decompressor( DecompressionConfig( delta_second_data=delta_path.read_bytes(), @@ -328,6 +342,8 @@ def compress_safetensors_file( threads: int | None = None, quiet: bool = False, codec_menu: list[CodecId] | None = None, + codec: str | None = None, + device: int | None = None, ) -> None: """Compress a safetensors file.""" import torch # noqa: PLC0415 @@ -391,6 +407,8 @@ def compress_safetensors_file( method=compression_method, threads=threads, codec_menu=codec_menu, + codec=codec, + device=device, ) ) compressor = compressor_cache[dtype_str] @@ -447,6 +465,8 @@ def compress_path( threads: int | None = None, file_compression: bool = False, codec_menu: list[CodecId] | None = None, + codec: str | None = None, + device: int | None = None, ) -> None: """Compress all files with the given suffix in the specified path.""" overwrite_first = True @@ -541,10 +561,22 @@ def compress_path( threads, True, # quiet codec_menu, + codec, + device, ) else: compression_func = compress_safetensors_file - comp_args = (delete, True, hf_cache, method, threads, True, codec_menu) + comp_args = ( + delete, + True, + hf_cache, + method, + threads, + True, + codec_menu, + codec, + device, + ) failures: list[tuple[str, BaseException]] = [] with ProcessPoolExecutor(max_workers=max_processes) as executor: @@ -809,6 +841,19 @@ def add_compress_parser(subparsers): "order1-scale-ac. Ignored when --method selects a forced-codec path." ), ) + parser.add_argument( + "--codec", + default=None, + help="Select one codec explicitly, by name (e.g. zstd) or canonical id. " + "Skips the automatic per-plane codec search.", + ) + parser.add_argument( + "--device", + type=int, + default=None, + help="CUDA device ordinal for GPU-resident codecs. Defaults to 0. " + "Must match the input tensor's device when that tensor is already on GPU.", + ) parser.add_argument( "--explore", action="store_true", @@ -1109,6 +1154,8 @@ def handle_compress(args): threads=args.threads, file_compression=args.file_compression, codec_menu=codec_menu, + codec=args.codec, + device=args.device, ) elif path.is_file(): if args.delta: @@ -1126,6 +1173,8 @@ def handle_compress(args): is_streaming=args.is_streaming, threads=args.threads, codec_menu=codec_menu, + codec=args.codec, + device=args.device, ) elif path.suffix == ".safetensors" and not args.file_compression: compress_safetensors_file( @@ -1136,6 +1185,8 @@ def handle_compress(args): method=args.method, threads=args.threads, codec_menu=codec_menu, + codec=args.codec, + device=args.device, ) else: compress_file( @@ -1151,6 +1202,8 @@ def handle_compress(args): is_streaming=args.is_streaming, threads=args.threads, codec_menu=codec_menu, + codec=args.codec, + device=args.device, ) elif args.hf_cache and args.model: compress_path( @@ -1172,6 +1225,8 @@ def handle_compress(args): threads=args.threads, file_compression=args.file_compression, codec_menu=codec_menu, + codec=args.codec, + device=args.device, ) else: print( # noqa: T201 diff --git a/python/ptwm/cli/decompress.py b/python/ptwm/cli/decompress.py index f20875e..e432303 100644 --- a/python/ptwm/cli/decompress.py +++ b/python/ptwm/cli/decompress.py @@ -19,6 +19,7 @@ def decompress_file( hf_cache: bool = False, threads: int | None = None, quiet: bool = False, + device: int | None = None, ) -> None: """Decompress a single file.""" from ptwm import DecompressionConfig, Decompressor # noqa: PLC0415 @@ -44,7 +45,7 @@ def decompress_file( return output_file = decompressed_path - decompressor = Decompressor(DecompressionConfig(threads=threads)) + decompressor = Decompressor(DecompressionConfig(threads=threads, device=device)) with full_path.open("rb") as infile, output_file.open("wb") as outfile: chunk = infile.read() @@ -74,6 +75,7 @@ def decompress_file_delta( force: bool = False, hf_cache: bool = False, threads: int | None = None, + device: int | None = None, ) -> None: """Decompress a file using delta compression.""" from ptwm import DecompressionConfig, Decompressor # noqa: PLC0415 @@ -113,6 +115,7 @@ def decompress_file_delta( DecompressionConfig( delta_second_data=delta_path.read_bytes(), threads=threads, + device=device, ) ) @@ -142,6 +145,7 @@ def decompress_safetensors_file( hf_cache: bool = False, threads: int | None = None, quiet: bool = False, + device: int | None = None, ) -> None: """Decompress a safetensors file.""" from safetensors import safe_open # noqa: PLC0415 @@ -174,7 +178,7 @@ def decompress_safetensors_file( return tensors = {} - decompressor = Decompressor(DecompressionConfig(threads=threads)) + decompressor = Decompressor(DecompressionConfig(threads=threads, device=device)) with safe_open(filename, "pt", "cpu") as f: metadata_raw = f.metadata() compressed_metadata = get_compressed_tensors_metadata(metadata_raw) @@ -218,6 +222,7 @@ def decompress_path( model: str = "", branch: str = "main", threads: int | None = None, + device: int | None = None, ) -> None: """Decompress every .ptwm file under ``path``.""" overwrite_first = True @@ -320,7 +325,9 @@ def decompress_path( else decompress_file ) future_to_file[ - executor.submit(func, file, delete, True, hf_cache, threads, True) + executor.submit( + func, file, delete, True, hf_cache, threads, True, device + ) ] = file remaining_files = collections.deque(file_list[max_processes:]) @@ -338,7 +345,14 @@ def decompress_path( ) future_to_file[ executor.submit( - func, next_file, delete, True, hf_cache, threads, True + func, + next_file, + delete, + True, + hf_cache, + threads, + True, + device, ) ] = next_file @@ -370,6 +384,13 @@ def add_decompress_parser(subparsers): default=None, help="The amount of threads to be used.", ) + parser.add_argument( + "--device", + type=int, + default=None, + help="CUDA device ordinal for GPU-resident codecs. Defaults to 0. " + "Must match the input tensor's device when that tensor is already on GPU.", + ) parser.add_argument( "--max_processes", type=int, @@ -434,6 +455,7 @@ def handle_decompress(args): model=args.model, branch=args.model_branch, threads=args.threads, + device=args.device, ) elif path.is_file(): if args.delta: @@ -444,6 +466,7 @@ def handle_decompress(args): force=args.force, hf_cache=args.hf_cache, threads=args.threads, + device=args.device, ) elif path.name.endswith(".ptwm.safetensors"): decompress_safetensors_file( @@ -452,6 +475,7 @@ def handle_decompress(args): force=args.force, hf_cache=args.hf_cache, threads=args.threads, + device=args.device, ) else: decompress_file( @@ -460,6 +484,7 @@ def handle_decompress(args): force=args.force, hf_cache=args.hf_cache, threads=args.threads, + device=args.device, ) else: print(f"{RED}Error: Path '{path_str}' not found.{RESET}", file=sys.stderr) # noqa: T201 diff --git a/python/ptwm/core/_compressor.py b/python/ptwm/core/_compressor.py index 14962c9..56e98fc 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -83,6 +83,29 @@ def compress( once; the reference's BLAKE3 digest is recorded as a Dependency so the decoder can verify that the caller supplied the matching reference. """ + if isinstance(data, torch.Tensor) and data.device.type == "cuda": + if ( + self.config.device is not None + and self.config.device != data.device.index + ): + msg = ( + f"device mismatch: config.device={self.config.device} but the " + f"input tensor is on cuda:{data.device.index}. Move the tensor " + f"or change config.device; ptwm will not relocate it implicitly." + ) + raise ValueError(msg) + + if self.config.codec is not None: + msg = ( + f"CompressionConfig.codec={self.config.codec!r} is set, but " + "explicit per-codec dispatch is not wired through " + "Compressor.compress() yet: the automatic per-plane codec " + "search always runs regardless of this value, so setting it " + "would silently do nothing. Leave codec unset until explicit " + "dispatch lands." + ) + raise CompressionMethodNotSupportedError(msg) + if ( self.config.delta_compressed_type is not None and self.config.delta_compressed_type != 0 diff --git a/python/ptwm/core/_decompressor.py b/python/ptwm/core/_decompressor.py index e1ebe5d..474312f 100644 --- a/python/ptwm/core/_decompressor.py +++ b/python/ptwm/core/_decompressor.py @@ -52,6 +52,26 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten match, and the decoded bytes are XOR-ed with it to recover the original tensor before any format reconstruction. """ + if self.config.codec is not None: + msg = ( + f"DecompressionConfig.codec={self.config.codec!r} is set, but " + "explicit per-codec dispatch is not wired through " + "Decompressor.decompress() yet: decoding always dispatches on " + "the codec recorded in the container header regardless of " + "this value, so setting it would silently do nothing. Leave " + "codec unset until explicit dispatch lands." + ) + raise ValueError(msg) + if self.config.device is not None: + msg = ( + f"DecompressionConfig.device={self.config.device} is set, but " + "GPU-resident decode dispatch is not wired through " + "Decompressor.decompress() yet: decoding always runs on CPU " + "regardless of this value, so setting it would silently do " + "nothing. Leave device unset until explicit dispatch lands." + ) + raise ValueError(msg) + mv_data = memoryview(data) magic = b"\x89PTWM" if len(mv_data) < len(magic): diff --git a/tests/test_cli_codec_flag.py b/tests/test_cli_codec_flag.py new file mode 100644 index 0000000..d972c46 --- /dev/null +++ b/tests/test_cli_codec_flag.py @@ -0,0 +1,17 @@ +import pytest +from ptwm.cli.main import main + + +def test_compress_help_lists_codec_and_device(capsys): + with pytest.raises(SystemExit): + main(["compress", "--help"]) + out, _ = capsys.readouterr() + assert "--codec" in out + assert "--device" in out + + +def test_decompress_help_lists_device(capsys): + with pytest.raises(SystemExit): + main(["decompress", "--help"]) + out, _ = capsys.readouterr() + assert "--device" in out diff --git a/tests/test_codec_selection.py b/tests/test_codec_selection.py new file mode 100644 index 0000000..87b4021 --- /dev/null +++ b/tests/test_codec_selection.py @@ -0,0 +1,20 @@ +import pytest +import torch + +from ptwm import CompressionConfig, Compressor + + +def test_codec_defaults_to_none(): + assert CompressionConfig().codec is None + + +def test_device_defaults_to_none(): + assert CompressionConfig().device is None + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +def test_device_mismatch_with_tensor_device_raises(): + t = torch.zeros(64, dtype=torch.bfloat16, device="cuda:0") + cfg = CompressionConfig(codec="example_codec", device=1) + with pytest.raises(ValueError, match="device"): + Compressor(cfg).compress(t) diff --git a/tests/test_plane_codec_cuda.py b/tests/test_plane_codec_cuda.py new file mode 100644 index 0000000..b43c000 --- /dev/null +++ b/tests/test_plane_codec_cuda.py @@ -0,0 +1,42 @@ +"""Test for the CUDA plane-codec PyO3 bridge. + +`plane_codec_encode_cuda`/`plane_codec_decode_cuda` resolve a selector and a +codec id, then dispatch through `PlaneCodecCudaRouter`. No installed +`plane_codec` CUDA contribution exists yet, so with an empty installed +list, resolving even a real builtin canonical id (such as "identity") +still fails at the router lookup (no native contribution is registered for +it) before any device pointer or stream is ever touched. That makes this +assertion reachable without CUDA, without torch, and without any native +extension build: it exercises the router-miss branch of the bridge using +the buffer-protocol fallback (plain `bytes`/`bytearray`) for both `src` +and `dst`. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + + +@pytest.fixture(autouse=True) +def isolated_extensions(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + # Hermetic: keep this test from ever picking up a real installed + # extension on the host running the suite. + monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "data")) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + return tmp_path + + +def test_unregistered_codec_raises() -> None: + from ptwm._rust.plane_codec_cuda import plane_codec_decode_cuda + + with pytest.raises(ValueError, match="no plane_codec CUDA contribution registered"): + plane_codec_decode_cuda("identity", "identity", b"", b"", bytearray(16), 0) + + +def test_encode_unregistered_codec_raises() -> None: + from ptwm._rust.plane_codec_cuda import plane_codec_encode_cuda + + with pytest.raises(ValueError, match="no plane_codec CUDA contribution registered"): + plane_codec_encode_cuda("identity", "identity", b"", b"", bytearray(16), 0)