Skip to content

Metal interop: retained native handles and external buffer-write initialization #9928

Description

@jcwal1516

Problem

A native Metal producer can write directly into a wgpu::Buffer, but wgpu 29 does not expose the two ownership and tracking contracts needed to do that without a staging copy.

The concrete use case is a GPU image decoder filling a Burn tensor backed by wgpu. Decode submission is asynchronous and must remain a true direct destination: no decoded-pixel host round trip and no relabeled CPU decode.

  1. as_hal only lends the HAL device, queue, or buffer for the callback. An asynchronous external submission needs an ownership-safe retained Metal handle that can outlive that callback. The Metal buffer raw object is also private today.
  2. After the external producer writes the bytes, the wgpu memory initialization tracker still considers the buffer range uninitialized. A later wgpu use can therefore clear bytes that were initialized outside wgpu.

This is a focused case under #4067. It does not ask wgpu to own codec submission or synchronization policy.

Minimum contracts

For the Metal backend, an equivalent of an ownership-safe retained handle for device, queue, and buffer:

  • acquiring the handle performs the required Objective-C retain
  • the caller has an explicit, documented release obligation, or receives an RAII retained wrapper
  • the handle remains valid for the documented retained lifetime, independent of the as_hal callback borrow
  • no handle is available when the backend does not match

At the wgpu::Buffer level, an unsafe range-scoped operation equivalent to:

unsafe fn mark_external_write_initialized(
    &self,
    range: impl RangeBounds<BufferAddress>,
) -> Result<(), BufferAsyncError>;

The exact name and error type are flexible. The important behavior is:

  • validate the buffer and normalize the range with overflow and bounds checks
  • mark only that range initialized in the core memory-init tracker
  • perform no submission or synchronization itself
  • require the caller to guarantee the external write and queue ordering before subsequent wgpu use

A local proof implements the retained handles by cloning the Metal retained objects and transferring the +1 reference, and routes the buffer method through wgpu-core to MemoryInitTracker::drain for the validated range.

Tests needed

  • retained device, queue, and buffer handles stay valid after the as_hal callback and release exactly once
  • range syntax, empty ranges, overflow, and out-of-bounds ranges are handled explicitly
  • marking one range prevents a later lazy clear of that range while adjacent unwritten bytes remain tracked as uninitialized
  • drop safety and queue ordering are covered for an asynchronous external producer

Questions

Would these narrow APIs be acceptable for a PR, or would maintainers prefer a scoped/RAII interop guard that combines handle lifetime and initialization registration? Is a release-compatible form suitable for the v29 line?

The matching CubeCL CUDA-stream contract is being proposed separately in tracel-ai/cubecl#1443.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    Status
    Todo

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions