grafos_std/

gpu.rs

//! GPU resource module.
//!
//! Provides typed access to fabric GPU resources via the persistent
//! v1 session API ([`GpuSession`]). Programs acquire a [`GpuLease`],
//! build a [`GpuSession`] on top of it, and drive device memory,
//! module loads, launches, and synchronisation via the
//! `fabricbios_gpu_v1` host ABI.
//!
//! The legacy one-shot v0 `gpu_submit` SDK builder has been removed;
//! see `docs/grafos/gpu-bridge-deferred-waves.md` §3 for context.

extern crate alloc;
use alloc::vec::Vec;

use crate::error::Result;
use crate::host;
use crate::lease::{self, LeaseInfo, LeaseStatus, SharedLeaseState};

#[cfg(feature = "observe")]
const LEASE_TAG_GPU: u8 = 0x05;

/// A GPU lease that auto-frees on drop.
///
/// Anchors session-based GPU work: a [`GpuSession`] built from this
/// lease drives persistent device memory, module loads, launches, and
/// synchronisation via the `fabricbios_gpu_v1` ABI. Created via
/// [`GpuBuilder::acquire`].
pub struct GpuLease {
    state: SharedLeaseState,
}

impl GpuLease {
    /// Lease metadata snapshot (id, creation time, expiry, and status).
    pub fn info(&self) -> LeaseInfo {
        lease::info(&self.state)
    }

    /// Unique lease identifier.
    pub fn lease_id(&self) -> u128 {
        lease::lease_id(&self.state)
    }

    /// Lease creation timestamp (unix seconds).
    pub fn created_at_unix_secs(&self) -> u64 {
        lease::created_at_unix_secs(&self.state)
    }

    /// Lease expiry timestamp (unix seconds).
    pub fn expires_at_unix_secs(&self) -> u64 {
        lease::expires_at_unix_secs(&self.state)
    }

    /// Current lease status.
    ///
    /// Queries the host for the authoritative lease status, falling
    /// back to local tracking on error.
    pub fn status(&self) -> LeaseStatus {
        if let Ok(s) = host::gpu_lease_query(lease::lease_id(&self.state)) {
            match s {
                0 => LeaseStatus::Active,
                1 => LeaseStatus::Expired,
                2 => LeaseStatus::Revoked,
                _ => LeaseStatus::Revoked,
            }
        } else {
            lease::status(&self.state)
        }
    }

    /// Renew the lease TTL by `duration_secs`.
    pub fn renew(&self, duration_secs: u64) -> Result<()> {
        host::gpu_lease_renew(lease::lease_id(&self.state), duration_secs as u32)?;
        lease::renew(&self.state, duration_secs)
    }

    /// Explicitly revoke/free this lease.
    pub fn free(&self) {
        let _ = host::gpu_lease_free(lease::lease_id(&self.state));
        lease::free(&self.state);
    }
}

impl Drop for GpuLease {
    fn drop(&mut self) {
        #[cfg(feature = "observe")]
        crate::observe_hooks::on_lease_dropped(LEASE_TAG_GPU, lease::lease_id(&self.state));
        let _ = host::gpu_lease_free(lease::lease_id(&self.state));
        lease::free(&self.state);
    }
}

/// Builder for acquiring a fabric GPU lease.
///
/// Allows specifying minimum VRAM requirements.
///
/// # Examples
///
/// ```rust
/// use grafos_std::gpu::GpuBuilder;
///
/// # grafos_std::host::reset_mock();
/// let lease = GpuBuilder::new().min_vram(1024).acquire()?;
/// // lease.gpu() is available for submit() calls
/// # Ok::<(), grafos_std::FabricError>(())
/// ```
pub struct GpuBuilder {
    _min_vram: u64,
    _lease_secs: u32,
    _exclusivity: Option<GpuExclusivityClass>,
    _affinities: Vec<crate::affinity::Affinity>,
}

/// Per-lease GPU exclusivity class.
///
/// Mirrors `fabricbios_core::lease_gpu_exclusivity::GpuExclusivityClass` at
/// the SDK layer.  See `docs/spec/gpu-exclusivity-wire-format.md` for the
/// wire format and `docs/spec/resource-isolation-and-exclusivity.md` §5.3
/// for the semantics.
///
/// When passed to [`GpuBuilder::exclusivity`], the runtime emits
/// `TLV_LEASE_GPU_EXCLUSIVITY` (0x0903) on the `LEASE_ALLOC` request.
/// Unsupported classes fail closed — the node rejects the lease.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GpuExclusivityClass {
    /// Device may multiplex other tenants.
    Shared,
    /// Exclusive residency for the session lifetime.
    SessionExclusive,
    /// Whole device for the lease lifetime.
    DeviceExclusive,
    /// Reserved for future MIG/partition isolation.
    PartitionExclusive,
}

impl GpuBuilder {
    /// Create a new builder with no VRAM constraint and no exclusivity
    /// preference (daemon default applies).
    pub fn new() -> Self {
        GpuBuilder {
            _min_vram: 0,
            _lease_secs: 300,
            _exclusivity: None,
            _affinities: Vec::new(),
        }
    }

    /// Set the minimum VRAM required in bytes.
    pub fn min_vram(mut self, n: u64) -> Self {
        self._min_vram = n;
        self
    }

    /// Request a specific GPU exclusivity class for this lease.
    ///
    /// When set, the daemon emits `TLV_LEASE_GPU_EXCLUSIVITY` (0x0903) on the
    /// `LEASE_ALLOC` request.  When omitted, the daemon's `--gpu-share-mode`
    /// default applies.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_std::gpu::{GpuBuilder, GpuExclusivityClass};
    ///
    /// # grafos_std::host::reset_mock();
    /// let lease = GpuBuilder::new()
    ///     .min_vram(1024)
    ///     .exclusivity(GpuExclusivityClass::DeviceExclusive)
    ///     .acquire()?;
    /// # Ok::<(), grafos_std::FabricError>(())
    /// ```
    pub fn exclusivity(mut self, class: GpuExclusivityClass) -> Self {
        self._exclusivity = Some(class);
        self
    }

    /// Add an affinity constraint (toward the target).
    pub fn affinity(mut self, a: crate::affinity::Affinity) -> Self {
        self._affinities.push(a);
        self
    }

    /// Add an anti-affinity constraint (away from the target).
    pub fn anti_affinity(
        mut self,
        strength: crate::affinity::Strength,
        target: crate::affinity::Target,
    ) -> Self {
        self._affinities
            .push(crate::affinity::Affinity::anti(strength, target));
        self
    }

    /// Set the lease TTL in seconds.
    pub fn lease_secs(mut self, secs: u32) -> Self {
        self._lease_secs = secs.max(1);
        self
    }

    /// Acquire a GPU lease.
    ///
    /// On wasm32, allocates via the `gpu_lease_alloc` host import.
    /// On native, uses the local mock lease allocator.
    ///
    /// # Errors
    ///
    /// Returns [`crate::error::FabricError::CapacityExceeded`] or
    /// [`crate::error::FabricError::Disconnected`] if the host cannot satisfy the request.
    pub fn acquire(self) -> Result<GpuLease> {
        // resource_id 0 = unspecified / any available GPU
        let host_lease_id = host::gpu_lease_alloc(0, self._lease_secs)?;
        let created_at = host::unix_time_secs();
        let expires_at = created_at.saturating_add(self._lease_secs as u64);
        let state = lease::new_shared_lease_from_parts(
            host_lease_id,
            created_at,
            expires_at,
            LeaseStatus::Active,
        );
        let lease = GpuLease { state };
        #[cfg(feature = "observe")]
        crate::observe_hooks::on_lease_acquired(LEASE_TAG_GPU, lease.lease_id(), "host", 0);
        Ok(lease)
    }
}

impl Default for GpuBuilder {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// GPU Session — persistent context tied to lease lifetime
// ---------------------------------------------------------------------------

/// RAII handle to a device memory allocation within a GPU session.
///
/// # Ownership model (SDK polish wave, post-Phase 48.15)
///
/// `GpuMemHandle` is a RAII handle. It is [`Clone`] but **not** `Copy`.
/// When the last clone is dropped it automatically calls
/// `gpu_session_mem_free` for the underlying device allocation, unless
/// you explicitly freed it first via [`GpuSession::mem_free`].
///
/// The intended pattern is:
///
/// ```text
/// let h = sess.mem_alloc(1024)?;
/// sess.mem_write(&h, 0, &data)?;
/// // either explicitly free:
/// sess.mem_free(h)?;
/// // ...or just let `h` go out of scope; Drop frees it.
/// ```
///
/// **Migration note:** prior to the SDK polish wave this type was
/// `Copy`, which meant a `?` between `mem_alloc` and `mem_free` could
/// silently leak the device allocation until lease expiry. Removing
/// `Copy` is what makes the RAII guarantee work — if you previously
/// relied on copying handles around, clone them explicitly. Cloned
/// handles share the same underlying allocation and the same `freed`
/// state; only the *last* clone to drop will issue the `mem_free`.
///
/// Drop is best-effort: it silently swallows hostcall errors and never
/// panics. If the lease has already expired, the daemon's
/// W3.5-validated lease-expiry teardown chain has already freed the
/// device memory, so Drop suppresses the call entirely in that case.
///
/// The raw u64 device pointer is exposed via [`GpuMemHandle::raw`] for
/// callers who need to pack it into kernel argument buffers manually.
/// Prefer [`KernelArgs::push_buffer`] which handles the marshalling.
#[derive(Debug, Clone)]
pub struct GpuMemHandle {
    inner: alloc::rc::Rc<core::cell::Cell<GpuHandleInner>>,
    lease_state: SharedLeaseState,
}

#[derive(Debug, Clone, Copy)]
struct GpuHandleInner {
    handle: u64,
    freed: bool,
}

impl GpuMemHandle {
    /// Raw device pointer value as returned by the bridge
    /// (`cuMemAlloc` on CUDA, `hipMalloc` on HIP).
    pub fn raw(&self) -> u64 {
        self.inner.get().handle
    }

    fn mark_freed(&self) {
        let mut v = self.inner.get();
        v.freed = true;
        self.inner.set(v);
    }
}

impl PartialEq for GpuMemHandle {
    fn eq(&self, other: &Self) -> bool {
        self.raw() == other.raw()
    }
}
impl Eq for GpuMemHandle {}

impl Drop for GpuMemHandle {
    fn drop(&mut self) {
        // Only the *last* clone of the inner Rc should issue the free.
        if alloc::rc::Rc::strong_count(&self.inner) > 1 {
            return;
        }
        let inner = self.inner.get();
        if inner.freed {
            return;
        }
        // If the lease is no longer active the daemon's lease-expiry
        // teardown has already freed device memory; suppress the call
        // entirely. Best-effort, never panic.
        if lease::status(&self.lease_state) != LeaseStatus::Active {
            return;
        }
        let lease_id = lease::lease_id(&self.lease_state);
        let _ = host::gpu_session_mem_free(lease_id, inner.handle);
    }
}

/// RAII handle to a loaded GPU module within a GPU session.
///
/// # Ownership model (Phase 48.16 SDK polish)
///
/// `GpuModule` is [`Clone`] but **not** `Copy`. When the last clone is
/// dropped it automatically calls `gpu_session_module_unload` for the
/// underlying CUDA module, unless you explicitly unloaded it first via
/// [`GpuSession::module_unload`].
///
/// The intended pattern is symmetric with [`GpuMemHandle`]:
///
/// ```text
/// let m = sess.module_load(&ptx)?;
/// sess.launch(&m, "k", [1,1,1], [1,1,1], &args, &sizes)?;
/// // either explicitly unload:
/// sess.module_unload(m)?;
/// // ...or just let `m` go out of scope; Drop unloads it.
/// ```
///
/// Drop is best-effort: it silently swallows hostcall errors and never
/// panics. If the lease has already expired, the daemon's
/// W3.5-validated lease-expiry teardown chain has already unloaded the
/// module, so Drop suppresses the call entirely in that case.
///
/// **Migration note:** prior to the SDK polish wave this type was
/// `Copy`. Cloned handles share the same underlying CUmodule and the
/// same `freed` state; only the *last* clone to drop will issue the
/// `module_unload`.
#[derive(Debug, Clone)]
pub struct GpuModule {
    inner: alloc::rc::Rc<core::cell::Cell<GpuHandleInner>>,
    lease_state: SharedLeaseState,
}

impl GpuModule {
    /// Raw module id as returned by the bridge.
    pub fn raw(&self) -> u64 {
        self.inner.get().handle
    }

    fn mark_freed(&self) {
        let mut v = self.inner.get();
        v.freed = true;
        self.inner.set(v);
    }
}

impl PartialEq for GpuModule {
    fn eq(&self, other: &Self) -> bool {
        self.raw() == other.raw()
    }
}
impl Eq for GpuModule {}

impl Drop for GpuModule {
    fn drop(&mut self) {
        // Only the *last* clone of the inner Rc should issue the unload.
        if alloc::rc::Rc::strong_count(&self.inner) > 1 {
            return;
        }
        let inner = self.inner.get();
        if inner.freed {
            return;
        }
        // If the lease is no longer active the daemon's lease-expiry
        // teardown has already unloaded the module; suppress the call
        // entirely. Best-effort, never panic.
        if lease::status(&self.lease_state) != LeaseStatus::Active {
            return;
        }
        let lease_id = lease::lease_id(&self.lease_state);
        let _ = host::gpu_session_module_unload(lease_id, inner.handle);
    }
}

/// Persistent GPU session backed by a lease's CUDA context.
///
/// `GpuSession` provides persistent device memory and loaded modules for
/// multi-kernel workloads. The session is tied to the [`GpuLease`] and
/// all resources are freed when the lease expires.
pub struct GpuSession {
    lease_state: SharedLeaseState,
}

// Phase 48.15 W1 — `GpuSession` is source-portable across host and wasm32.
//
// On native host targets the `host::gpu_session_*` family is mock-backed
// (see `crates/grafos-std/src/host.rs`). On wasm32 targets the same Rust
// signatures lower to the `fabricbios_gpu_v1` extern block and translate
// the i32 wire status into `FabricError` / `FabricError::GpuSessionFailed`.
//
// Because both targets expose the identical Rust API, every method below
// is on the Phase 48.14 **safe-list** and needs no per-target rustdoc tag.
// The P3.1 (`c078c502`) wasm32 cfg-gate has been removed.
impl GpuSession {
    /// Create a session handle from a GPU lease.
    ///
    /// # Canonical lifecycle example
    ///
    /// This doctest exercises the full session lifecycle against the
    /// host mock so it runs on every `cargo test` build, with no GPU
    /// hardware required.
    ///
    /// ```rust
    /// use grafos_std::gpu::{GpuBuilder, GpuSession, KernelArgs};
    ///
    /// # grafos_std::host::reset_mock();
    /// // 1. Acquire a GPU lease.
    /// let lease = GpuBuilder::new().min_vram(1024).acquire()?;
    ///
    /// // 2. Open a persistent session on the lease.
    /// let mut sess = GpuSession::new(&lease);
    ///
    /// // 3. Allocate device memory and write input data.
    /// let buf = sess.mem_alloc(1024)?;
    /// sess.mem_write(&buf, 0, &[1u8, 2, 3, 4])?;
    ///
    /// // 4. Load a module (PTX/cubin bytes — placeholder under the mock).
    /// let module = sess.module_load(&[0x7f, 0x45, 0x4c, 0x46])?;
    ///
    /// // 5. Launch a kernel using the typed argument builder.
    /// let args = KernelArgs::new()
    ///     .push_u32(42)
    ///     .push_buffer(&buf);
    /// sess.launch_with_args(&module, "vector_add", [256, 1, 1], [64, 1, 1], args)?;
    ///
    /// // 6. Synchronize and read results back.
    /// sess.sync()?;
    /// let _out = sess.mem_read(&buf, 0, 4)?;
    ///
    /// // 7. `buf` and `module` drop here. Both are RAII-freed (the
    /// //    daemon calls `cuMemFree` / `cuModuleUnload` on real
    /// //    hardware). The lease itself drops at end of scope.
    /// # Ok::<(), grafos_std::FabricError>(())
    /// ```
    pub fn new(lease: &GpuLease) -> Self {
        GpuSession {
            lease_state: lease.state.clone(),
        }
    }

    /// Lease ID for this session (session_id == lease_id).
    pub fn lease_id(&self) -> u128 {
        lease::lease_id(&self.lease_state)
    }

    fn ensure_active(&self) -> Result<()> {
        lease::ensure_active(&self.lease_state)
    }

    /// Allocate device memory.
    pub fn mem_alloc(&mut self, size: u64) -> Result<GpuMemHandle> {
        self.ensure_active()?;
        let handle = host::gpu_session_mem_alloc(self.lease_id(), size)?;
        Ok(GpuMemHandle {
            inner: alloc::rc::Rc::new(core::cell::Cell::new(GpuHandleInner {
                handle,
                freed: false,
            })),
            lease_state: self.lease_state.clone(),
        })
    }

    /// Write data to device memory.
    pub fn mem_write(&mut self, handle: &GpuMemHandle, offset: u64, data: &[u8]) -> Result<()> {
        self.ensure_active()?;
        host::gpu_session_mem_write(self.lease_id(), handle.raw(), offset, data)
    }

    /// Read data from device memory.
    pub fn mem_read(&mut self, handle: &GpuMemHandle, offset: u64, size: u32) -> Result<Vec<u8>> {
        self.ensure_active()?;
        host::gpu_session_mem_read(self.lease_id(), handle.raw(), offset, size)
    }

    /// Free a device memory allocation explicitly.
    ///
    /// Consumes the handle. After this returns successfully (or even
    /// on error), the handle's `Drop` will be a no-op — explicit free
    /// and Drop are both safe; the latter never double-frees.
    pub fn mem_free(&mut self, handle: GpuMemHandle) -> Result<()> {
        self.ensure_active()?;
        let raw = handle.raw();
        // Mark freed first so Drop (running at end of this fn) is a no-op,
        // even if the hostcall errors.
        handle.mark_freed();
        host::gpu_session_mem_free(self.lease_id(), raw)
    }

    /// Load a GPU module (PTX/cubin).
    pub fn module_load(&mut self, binary: &[u8]) -> Result<GpuModule> {
        self.ensure_active()?;
        let module_id = host::gpu_session_module_load(self.lease_id(), binary)?;
        Ok(GpuModule {
            inner: alloc::rc::Rc::new(core::cell::Cell::new(GpuHandleInner {
                handle: module_id,
                freed: false,
            })),
            lease_state: self.lease_state.clone(),
        })
    }

    /// Unload a GPU module explicitly (Phase 48.16 SDK polish).
    ///
    /// Consumes the handle. After this returns successfully (or even
    /// on error), the handle's `Drop` will be a no-op — explicit
    /// unload and Drop are both safe; the latter never double-unloads.
    pub fn module_unload(&mut self, module: GpuModule) -> Result<()> {
        self.ensure_active()?;
        let raw = module.raw();
        // Mark freed first so Drop (running at end of this fn) is a no-op,
        // even if the hostcall errors.
        module.mark_freed();
        host::gpu_session_module_unload(self.lease_id(), raw)
    }

    /// Launch a kernel from a loaded module (async — does not wait).
    pub fn launch(
        &mut self,
        module: &GpuModule,
        kernel: &str,
        grid: [u32; 3],
        block: [u32; 3],
        args: &[u8],
        arg_sizes: &[u32],
    ) -> Result<()> {
        self.ensure_active()?;
        host::gpu_session_launch(
            self.lease_id(),
            module.raw(),
            kernel,
            grid,
            block,
            args,
            arg_sizes,
        )
    }

    /// Launch a kernel using a [`KernelArgs`] builder.
    ///
    /// Convenience over [`GpuSession::launch`]: builds the
    /// `(args, arg_sizes)` byte/size pair from the typed builder and
    /// forwards to the same underlying hostcall.
    pub fn launch_with_args(
        &mut self,
        module: &GpuModule,
        kernel: &str,
        grid: [u32; 3],
        block: [u32; 3],
        args: KernelArgs,
    ) -> Result<()> {
        let (bytes, sizes) = args.build();
        self.launch(module, kernel, grid, block, &bytes, &sizes)
    }

    /// Synchronize: wait for all outstanding launches to complete.
    pub fn sync(&mut self) -> Result<()> {
        self.ensure_active()?;
        host::gpu_session_sync(self.lease_id())
    }
}

// ---------------------------------------------------------------------------
// KernelArgs — typed launch argument builder
// ---------------------------------------------------------------------------

/// Typed builder for GPU kernel launch arguments.
///
/// `KernelArgs` packs scalar values and device-buffer pointers into the
/// `(args, arg_sizes)` byte/size pair that the
/// [`fabricbios_gpu_v1`] bridge expects. Each `push_*` records the
/// raw bytes and their per-arg size so the daemon can build a CUDA
/// `kernelParams` array with one pointer per argument.
///
/// # Example
///
/// ```rust
/// use grafos_std::gpu::{GpuBuilder, GpuSession, KernelArgs};
///
/// # grafos_std::host::reset_mock();
/// let lease = GpuBuilder::new().acquire()?;
/// let mut sess = GpuSession::new(&lease);
/// let buf = sess.mem_alloc(64)?;
/// let module = sess.module_load(&[0u8; 8])?;
/// let args = KernelArgs::new()
///     .push_u32(7)
///     .push_f32(2.5)
///     .push_buffer(&buf);
/// sess.launch_with_args(&module, "k", [1, 1, 1], [1, 1, 1], args)?;
/// # Ok::<(), grafos_std::FabricError>(())
/// ```
///
/// # Wire format
///
/// - Scalars are written little-endian.
/// - Buffer args are written as 8-byte little-endian device pointers
///   (the raw `cuMemAlloc` return value carried by [`GpuMemHandle`]).
/// - The daemon's launch path packs these into a contiguous byte
///   buffer and walks `arg_sizes` to compute per-arg pointer offsets
///   for `kernelParams`. See
///   `crates/fabricbios-platform-linux/src/gpu.rs::launch`.
#[derive(Debug, Default, Clone)]
pub struct KernelArgs {
    bytes: Vec<u8>,
    sizes: Vec<u32>,
}

impl KernelArgs {
    /// Create an empty argument list.
    pub fn new() -> Self {
        Self::default()
    }

    /// Push a `u32` argument.
    pub fn push_u32(mut self, v: u32) -> Self {
        self.bytes.extend_from_slice(&v.to_le_bytes());
        self.sizes.push(4);
        self
    }

    /// Push a `u64` argument.
    pub fn push_u64(mut self, v: u64) -> Self {
        self.bytes.extend_from_slice(&v.to_le_bytes());
        self.sizes.push(8);
        self
    }

    /// Push an `i32` argument.
    pub fn push_i32(mut self, v: i32) -> Self {
        self.bytes.extend_from_slice(&v.to_le_bytes());
        self.sizes.push(4);
        self
    }

    /// Push an `i64` argument.
    pub fn push_i64(mut self, v: i64) -> Self {
        self.bytes.extend_from_slice(&v.to_le_bytes());
        self.sizes.push(8);
        self
    }

    /// Push an `f32` argument.
    pub fn push_f32(mut self, v: f32) -> Self {
        self.bytes.extend_from_slice(&v.to_le_bytes());
        self.sizes.push(4);
        self
    }

    /// Push an `f64` argument.
    pub fn push_f64(mut self, v: f64) -> Self {
        self.bytes.extend_from_slice(&v.to_le_bytes());
        self.sizes.push(8);
        self
    }

    /// Push a device buffer argument (the raw 8-byte device pointer).
    ///
    /// This is the typed equivalent of manually appending
    /// `handle.raw().to_le_bytes()` with size 8. The daemon will pass
    /// the underlying device pointer to the kernel as the
    /// corresponding `kernelParams` entry.
    pub fn push_buffer(mut self, handle: &GpuMemHandle) -> Self {
        self.bytes.extend_from_slice(&handle.raw().to_le_bytes());
        self.sizes.push(8);
        self
    }

    /// Push a value of arbitrary type by reinterpreting its bytes.
    ///
    /// # Safety
    ///
    /// `T` must satisfy all of the following:
    /// - `Copy` and trivially-destructible.
    /// - No interior padding bytes whose value matters to the kernel.
    ///   For C-ABI structs, this generally means `#[repr(C)]` plus
    ///   manual layout review.
    /// - No interior pointers, references, or non-`'static` lifetimes.
    /// - Endianness matches what the kernel expects (host order is
    ///   little-endian on every platform fabricBIOS targets).
    ///
    /// In short: the value must be a "POD" (plain-old-data) blob that
    /// can be byte-copied into a CUDA `kernelParams` slot. Prefer the
    /// typed `push_*` helpers for primitive scalars; reach for
    /// `push_raw` only when packing a `#[repr(C)]` struct argument.
    pub unsafe fn push_raw<T: Copy>(mut self, value: &T) -> Self {
        let size = core::mem::size_of::<T>();
        let ptr = value as *const T as *const u8;
        // SAFETY: caller upholds the contract above; the slice
        // is read-only and lives for the duration of this call.
        let slice = core::slice::from_raw_parts(ptr, size);
        self.bytes.extend_from_slice(slice);
        self.sizes.push(size as u32);
        self
    }

    /// Consume the builder and return the raw `(bytes, sizes)` pair
    /// suitable for [`GpuSession::launch`].
    pub fn build(self) -> (Vec<u8>, Vec<u32>) {
        (self.bytes, self.sizes)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::host;

    #[test]
    fn gpu_builder_acquires() {
        host::reset_mock();
        let lease = GpuBuilder::new().min_vram(1024).acquire();
        assert!(lease.is_ok());
    }

    #[test]
    fn gpu_lease_lifecycle_alloc_status_renew_free() {
        host::reset_mock();
        let lease = GpuBuilder::new().lease_secs(60).acquire().expect("acquire");
        let id = lease.lease_id();
        assert_ne!(id, 0);

        // Status should be Active right after acquisition.
        assert_eq!(lease.status(), LeaseStatus::Active);

        // Renew should succeed for an active lease.
        assert!(lease.renew(30).is_ok());

        // Free and verify host-side status becomes expired.
        lease.free();
        let status = host::gpu_lease_query(id).expect("query");
        assert_eq!(status, 1); // expired (removed from active set)
    }

    #[test]
    fn gpu_lease_drop_frees_host_lease() {
        host::reset_mock();
        let id;
        {
            let lease = GpuBuilder::new().acquire().expect("acquire");
            id = lease.lease_id();
            assert_eq!(host::gpu_lease_query(id).unwrap(), 0); // active
        }
        // After drop, lease should be freed on host side.
        assert_eq!(host::gpu_lease_query(id).unwrap(), 1); // expired
    }

    // ── Phase 48.16 SDK polish: GpuMemHandle RAII tests ──────────────

    #[test]
    fn mem_handle_explicit_free_works() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let h = sess.mem_alloc(1024).expect("alloc");
        let before = host::test_mock::_gpu_session_mem_free_count();
        sess.mem_free(h).expect("explicit free");
        let after = host::test_mock::_gpu_session_mem_free_count();
        assert_eq!(
            after,
            before + 1,
            "explicit mem_free issues exactly one host call"
        );
    }

    #[test]
    fn mem_handle_drop_frees_on_early_return() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let before = host::test_mock::_gpu_session_mem_free_count();
        // Simulate the early-return-after-alloc path: alloc then drop.
        {
            let _h = sess.mem_alloc(1024).expect("alloc");
            // _h drops here.
        }
        let after = host::test_mock::_gpu_session_mem_free_count();
        assert_eq!(after, before + 1, "Drop on GpuMemHandle issues mem_free");
    }

    #[test]
    fn mem_handle_double_free_is_noop() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let h = sess.mem_alloc(1024).expect("alloc");
        let before = host::test_mock::_gpu_session_mem_free_count();
        sess.mem_free(h).expect("explicit free");
        // The handle was consumed by mem_free; nothing else to drop.
        // Verify exactly one free call was issued — i.e. Drop did not
        // issue a second free during the in-scope drop of `h` inside
        // `mem_free`.
        let after = host::test_mock::_gpu_session_mem_free_count();
        assert_eq!(
            after,
            before + 1,
            "explicit free + Drop = exactly one host free"
        );
    }

    #[test]
    fn mem_handle_drop_after_lease_expiry_is_silent() {
        host::reset_mock();
        host::mock_set_unix_time_secs(5_000);
        let lease = GpuBuilder::new().lease_secs(5).acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let h = sess.mem_alloc(1024).expect("alloc");
        // Advance past lease expiry — Drop should observe the expired
        // lease and skip the hostcall entirely (the daemon's
        // lease-expiry teardown owns cleanup at that point).
        host::mock_advance_time_secs(10);
        let before = host::test_mock::_gpu_session_mem_free_count();
        drop(h); // must not panic, must not log, must not call free.
        let after = host::test_mock::_gpu_session_mem_free_count();
        assert_eq!(after, before, "Drop after lease expiry suppresses mem_free");
    }

    #[test]
    fn mem_handle_clone_only_last_drops() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let h = sess.mem_alloc(1024).expect("alloc");
        let h2 = h.clone();
        assert_eq!(h, h2);
        let before = host::test_mock::_gpu_session_mem_free_count();
        drop(h);
        let mid = host::test_mock::_gpu_session_mem_free_count();
        assert_eq!(mid, before, "non-last clone does not free");
        drop(h2);
        let after = host::test_mock::_gpu_session_mem_free_count();
        assert_eq!(after, before + 1, "last clone frees exactly once");
    }

    // ── Phase 48.16 SDK polish: GpuModule RAII tests ─────────────────

    #[test]
    fn module_explicit_unload_works() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let m = sess.module_load(&[0u8; 8]).expect("module_load");
        let before = host::test_mock::_gpu_session_module_unload_count();
        sess.module_unload(m).expect("explicit unload");
        let after = host::test_mock::_gpu_session_module_unload_count();
        assert_eq!(
            after,
            before + 1,
            "explicit module_unload issues exactly one host call"
        );
    }

    #[test]
    fn module_drop_unloads_on_early_return() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let before = host::test_mock::_gpu_session_module_unload_count();
        {
            let _m = sess.module_load(&[0u8; 8]).expect("module_load");
            // _m drops here.
        }
        let after = host::test_mock::_gpu_session_module_unload_count();
        assert_eq!(after, before + 1, "Drop on GpuModule issues module_unload");
    }

    #[test]
    fn module_double_unload_is_noop() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let m = sess.module_load(&[0u8; 8]).expect("module_load");
        let before = host::test_mock::_gpu_session_module_unload_count();
        sess.module_unload(m).expect("explicit unload");
        let after = host::test_mock::_gpu_session_module_unload_count();
        assert_eq!(
            after,
            before + 1,
            "explicit unload + Drop = exactly one host unload"
        );
    }

    #[test]
    fn module_drop_after_lease_expiry_is_silent() {
        host::reset_mock();
        host::mock_set_unix_time_secs(5_000);
        let lease = GpuBuilder::new().lease_secs(5).acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let m = sess.module_load(&[0u8; 8]).expect("module_load");
        host::mock_advance_time_secs(10);
        let before = host::test_mock::_gpu_session_module_unload_count();
        drop(m); // must not panic, must not call unload.
        let after = host::test_mock::_gpu_session_module_unload_count();
        assert_eq!(
            after, before,
            "Drop after lease expiry suppresses module_unload"
        );
    }

    #[test]
    fn module_clone_only_last_drops() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let m = sess.module_load(&[0u8; 8]).expect("module_load");
        let m2 = m.clone();
        assert_eq!(m, m2);
        let before = host::test_mock::_gpu_session_module_unload_count();
        drop(m);
        let mid = host::test_mock::_gpu_session_module_unload_count();
        assert_eq!(mid, before, "non-last clone does not unload");
        drop(m2);
        let after = host::test_mock::_gpu_session_module_unload_count();
        assert_eq!(after, before + 1, "last clone unloads exactly once");
    }

    // ── KernelArgs tests ─────────────────────────────────────────────

    #[test]
    fn kernel_args_typed_push_layout() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let buf = sess.mem_alloc(1024).expect("alloc");
        let raw = buf.raw();

        let (bytes, sizes) = KernelArgs::new()
            .push_u32(0x1122_3344)
            .push_f32(2.5)
            .push_buffer(&buf)
            .build();

        // Sizes: u32=4, f32=4, buffer=8.
        assert_eq!(sizes, vec![4, 4, 8]);
        // Bytes: little-endian concatenation.
        let mut expected = Vec::new();
        expected.extend_from_slice(&0x1122_3344u32.to_le_bytes());
        expected.extend_from_slice(&2.5f32.to_le_bytes());
        expected.extend_from_slice(&raw.to_le_bytes());
        assert_eq!(bytes, expected);

        // explicit free so the test doesn't lean on Drop semantics.
        sess.mem_free(buf).expect("free");
    }

    #[test]
    fn kernel_args_launch_with_args_round_trip() {
        host::reset_mock();
        let lease = GpuBuilder::new().acquire().expect("acquire");
        let mut sess = GpuSession::new(&lease);
        let module = sess.module_load(&[0u8; 8]).expect("module_load");
        let buf = sess.mem_alloc(64).expect("alloc");
        let args = KernelArgs::new().push_u64(99).push_buffer(&buf);
        sess.launch_with_args(&module, "k", [1, 1, 1], [1, 1, 1], args)
            .expect("launch");
        sess.mem_free(buf).expect("free");
    }
}