Skip to content

Tracking: Apple zero-copy CPU/Metal shared memory and Metal FFT #1394

Description

@shinaoka

Closure scope amendment (2026-08-08)

This tracking issue is narrowed to the completed Apple zero-copy CPU/Metal
shared-allocation and CPU/Metal FFT work.

CUDA/cuFFT is explicitly split out to #967 and is not a closure condition
for this issue. Any CUDA/cuFFT requirements, Phase 5 CUDA items, or cross-CUDA
acceptance criteria in the historical plan below are superseded by this
amendment.

Completed scope

  • Apple host-visible allocation domain with explicit paired CPU and
    Metal-adapter WebGpuBackend executors.
  • Stable allocation identity, guarded CPU mapping, CPU/GPU exclusion, and
    zero post-creation tensor-byte transfers across CPU -> Metal -> CPU.
  • Runtime-owned FFT capability/cache contract and RustFFT CPU execution for
    F32/F64/C32/C64.
  • CubeK Metal CFFT/RFFT/IRFFT execution for the supported F32/C32,
    power-of-two surface, with typed unsupported errors and no implicit fallback.
  • Same-domain output preservation, mapped CPU Cholesky, integration tests,
    runnable examples, and user/design documentation.
  • Reviewed CubeCL/CubeK revisions pinned as Git dependencies; publication
    remains separate release work and is not a closure condition.

Completion evidence

Deferred scope


Historical original plan (CUDA clauses superseded by the amendment above)

Goal

Add explicit CPU and Metal execution over one shared Apple allocation domain, then extend the existing tenferro-fft extension from its CPU host implementation to explicit native CPU, CUDA, and Metal/WebGPU execution.

The Apple memory contract is:

  • One initial copy when a unified tensor is created is allowed.
  • After creation, switching that tensor between CPU and Metal must not copy tensor bytes.
  • Synchronization and map/unmap transitions are allowed and expected.
  • Backend selection is always explicit. No dtype-, size-, or capability-based CPU/Metal fallback is allowed.
  • Outputs allocated by either executor remain in the same Apple shared allocation domain.

Current baseline

tenferro-fft already provides concrete and traced fft, ifft, rfft, and irfft APIs backed by RustFFT for host tensors. It supports F32/F64/C32/C64 and already defines normalization, shape inference, and AD behavior.

The current Metal-capable tenferro type is WebGpuBackend; there is no public MetalBackend type. The first Apple API should therefore construct a Metal-adapter WebGpuBackend rather than duplicate the backend surface.

The current RustFFT plan cache is not a valid foundation for additional backends. Open bug #1406 documents that it is a hidden, unbounded process-global OnceLock<Mutex<HashMap<...>>> with no owner, capacity, clear path, or stats. The earlier lock-poison bug #1350 is closed, but the ownership and lifetime bug remains open.

Confirmed scope

Explicit Apple execution

Do not implement an auto-dispatching AppleBackend: TensorBackend.

Introduce an Apple context/shared-allocation owner that constructs and configures:

  • a CpuBackend;
  • a WebGpuBackend pinned to a Metal adapter;
  • one shared Apple allocation domain used by both.

The caller selects the executor for every operation. A CPU operation maps the shared allocation and waits for outstanding Metal work. A Metal operation unmaps CPU access and waits for outstanding CPU access. Concurrent CPU mapping and GPU use of the same resource is rejected by guards.

Use whole-resource pooling initially. Do not suballocate unrelated tensors from one mappable WGPU resource because mapping excludes GPU use of the whole resource.

CPU FFT

Use the existing RustFFT implementation as the only initial CPU FFT provider.

  • Keep F32/F64/C32/C64 behavior.
  • Keep the existing CFFT/RFFT/IRFFT semantics and public normalization modes.
  • Do not add Accelerate/vDSP.
  • Do not add process-global OpenBLAS/Accelerate provider selection.
  • Do not add a realfft dependency unless a separate measured issue demonstrates a need; the existing RustFFT implementation already covers the required transforms.

Native GPU FFT

Define FFT as an explicit backend extension capability, following the linalg extension/runtime pattern without adding FFT to the core tensor operation trait surface.

Initial implementations:

  • CPU: existing RustFFT execution;
  • CUDA: cuFFT;
  • Metal/WebGPU: published t4a-cubek-fft.

The initial Metal/WebGPU scope is F32/C32. C64 remains available through explicitly selected CPU or CUDA execution.

The current CubeK FFT implementation uses split real/imaginary buffers and power-of-two transform lengths. Publish an interleaved C32 ABI matching num_complex::Complex32:

  • cfft: C32 -> C32;
  • rfft: F32 -> C32;
  • irfft: C32 -> F32.

Internal scratch may remain split. Unsupported dtype, layout, or transform-length combinations return capability errors; they never trigger a host download or backend switch.

Metal linalg

The Metal-adapter WebGpuBackend does not implement LinalgBackend.

  • Statically typed execution expresses this through extension bounds or registration.
  • Dynamic eager execution returns an explicit unsupported-extension error.
  • The user may explicitly run linalg with the shared CpuBackend without transferring tensor bytes.

FFT plan and cache contract

Resolving #1406 is a prerequisite for adding CUDA or Metal FFT plans.

  • Remove the process-global RustFFT plan caches.
  • Store graph/eager FFT plans in the owning ExtensionCacheStore.
  • Use backend-defined exact plan keys containing only the fields relevant to plan compatibility. RustFFT needs dtype, length, and direction; vendor plans additionally include backend/device and required batch/layout fields.
  • Verify exact key equality before reuse; do not trust a bare hash collision.
  • Use the existing bounded extension-cache limits and aggregate clear/stats paths.
  • Account for retained plan/workspace bytes.
  • Return typed errors for poisoned runtime state.
  • Keep the host reference implementation stateless.
  • For repeated direct concrete execution outside GraphExecutor or EagerRuntime, add an explicitly owned prepared-plan API. Dropping that object releases its plan. The one-shot convenience API must not recreate hidden global state.

No new backend may introduce a static/thread-local plan or workspace cache.

Delivery decision (2026-07-19)

  • This implementation may pin reviewed CubeCL and CubeK fork commits as Git dependencies.
  • Publishing new t4a-cubecl-*, t4a-cubek-*, or tenferro crate releases is deferred to a separate release task.
  • The implementation PRs stop at tested code, documentation, and runnable examples.

Out of scope

  • Web CPU/WASM FFT; track it separately if a concrete web consumer needs it.
  • Accelerate/vDSP FFT.
  • Process-global OpenBLAS/Accelerate switching.
  • General WGSL/MSL complex lowering.
  • Automatic CPU/Metal fallback.
  • C64 Metal/WebGPU FFT in the first implementation.
  • Non-power-of-two CubeK FFT algorithms in the first implementation.

Dependency-ordered implementation phases

Phase 0: repair FFT plan ownership (#1406)

  • Replace the static RustFFT caches with runtime-owned extension-cache entries.
  • Add exact plan-key verification, retained-byte accounting, bounded eviction, clear, and stats tests.
  • Add an explicitly owned prepared plan for repeated direct concrete FFT calls.
  • Preserve current FFT values, shapes, normalization, AD behavior, and lock-poison errors.

This phase must merge before backend-specific FFT plans are added.

Phase 1: CubeCL host-visible primary allocations

  • Add a generic host-visible/mappable primary allocation class to the tensor4all CubeCL fork.
  • Implement it for Metal-backed WGPU resources.
  • Add whole-resource pooling and stable allocation identity.
  • Add read/write mapping guards, synchronization, and CPU/GPU mutual exclusion.
  • Test map -> GPU -> map transitions and rejection of overlapping access.
  • Pin the reviewed CubeCL fork commit as a Git dependency; publishing is deferred.

Phase 2: tenferro Apple shared allocation domain

  • Extend backend-buffer/storage contracts with guarded host mapping without exposing WGPU or Metal types to tenferro-tensor or tenferro-cpu.
  • Represent Apple shared storage as managed/shared placement with a stable allocation-domain identity.
  • Allow CpuBackend to read, write, and allocate through an optional shared allocation domain.
  • Allow a Metal-adapter WebGpuBackend to use the same domain.
  • Add AppleContext construction returning the configured CPU and WebGPU executors.
  • Preserve the allocation domain on every output.
  • Instrument allocation identity and upload/download byte counts for zero-copy tests.

Phase 3: FFT backend capability and CPU migration

  • Extract the current monolithic tenferro-fft/src/lib.rs into operation/spec, CPU execution, cache/prepared-plan, AD, and runtime-registration modules.
  • Introduce a one-dimensional FftPlanSpec matching the existing API: operation kind, axis, optional length, direction, normalization, dtype, shape, and layout requirements.
  • Introduce the backend FFT capability used by concrete and extension runtime execution.
  • Move the existing RustFFT path behind the CPU capability without changing numerical semantics.
  • Ensure RustFFT can operate on a guarded mapping of an Apple shared buffer and allocate its output in the same domain.
  • Return capability errors for incompatible placements instead of requesting a transfer.

Phase 4: CubeK interleaved C32 FFT

  • Make CFFT a public CubeK FFT operation.
  • Add interleaved C32 input/output adapters for CFFT/RFFT/IRFFT.
  • Keep split internal scratch where useful.
  • Test C32 ABI layout, forward/inverse round trips, batches, axes, zero-padding, normalization hooks, and supported power-of-two lengths.
  • Pin the reviewed CubeK fork commit as a Git dependency; publishing is deferred.

Phase 5: Metal/WebGPU and CUDA adapters

  • Register the CubeK FFT runtime for Metal/WebGPU F32/C32 execution.
  • Register the cuFFT runtime for CUDA F32/F64/C32/C64 execution.
  • Keep plans/workspaces in the runtime-owned cache and key them by backend/device.
  • Add explicit errors for unsupported dtypes, lengths, layouts, missing runtime registration, and Metal linalg.
  • Verify that neither adapter invokes download/upload as a fallback.

Phase 6: integration and documentation

  • Add CPU -> Metal -> CPU tests that retain one allocation identity.
  • Add instrumented assertions that backend switches copy zero tensor bytes after unified creation.
  • Add cross-backend FFT conformance tests for shapes, values, normalization, zero-padding/truncation, and round trips.
  • Add compile/run examples for explicit Apple executor selection and explicit CPU linalg fallback.
  • Update GPU/FFT design docs and the device guide.
  • Pin reviewed CubeCL/CubeK fork commits for this integration; perform crate releases separately.

PR boundaries

Keep review and rollback boundaries explicit:

  1. tenferro-rs: fix [Bug]: Unbounded process-global FFT plan cache (<code>static FftPlanCache</code>) with no owner, bound, clear, or stats — violates REPOSITORY_RULES.md cache-ownership contract #1406 and add prepared direct FFT plans.
  2. CubeCL fork: host-visible allocation API and Metal mapping/pool implementation.
  3. CubeK fork: public interleaved C32 FFT API.
  4. tenferro-rs: Apple shared allocation domain and AppleContext.
  5. tenferro-rs: FFT capability refactor plus RustFFT migration.
  6. tenferro-rs: Metal/WebGPU and CUDA FFT runtimes.
  7. tenferro-rs: integration tests, docs, examples, and pinned Git-dependency verification.

Phases 0, 1, and 4 may proceed independently. Phase 3 begins after Phase 0; Phase 2 depends on Phase 1; the Apple-mapped RustFFT integration in Phase 3 depends on Phase 2; Phase 5 depends on Phases 2, 3, and 4.

Acceptance criteria

  • A unified tensor performs CPU -> Metal -> CPU operations while retaining the same allocation identity.
  • Instrumented tests observe zero tensor-byte upload/download copies after unified tensor creation.
  • CPU/Metal transitions wait for outstanding work and enforce exclusive mapped/GPU access.
  • Outputs from both executors preserve the Apple shared allocation domain.
  • No API performs a hidden transfer or CPU fallback.
  • Direct Metal linalg returns an explicit unsupported-extension error.
  • The RustFFT plan path has no process-global/thread-local cache and satisfies owner, bound, clear, stats, retained-byte, and poison-error contracts.
  • Existing CPU FFT public behavior and AD tests remain unchanged.
  • FFT shape and normalization semantics agree across CPU, CUDA, and supported Metal/WebGPU cases.
  • C32 CFFT/RFFT/IRFFT round trips pass on Metal/WebGPU for supported power-of-two lengths.
  • Unsupported C64 Metal/WebGPU or unsupported-length requests fail cleanly and remain usable through an explicitly selected supported backend.
  • The implementation PRs resolve reviewed, pinned CubeCL/CubeK Git revisions; crate publication is handled separately.

Related issues

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions