Backend layout and parity contract

Backend layout and parity contract#

quantem.gpu is organized by scientific domain first and accelerator second. This is deliberate: IO, detector reductions, DPC, display, and SSB each expose one scientific contract, while CUDA, Python MPS, native Swift/Metal, and WebGPU provide implementations of that contract. A backend must not invent a second public workflow or result type.

Ownership rules#

  • Public Python callers enter through quantem.gpu.io, quantem.gpu.detector, quantem.gpu.dpc, quantem.gpu.display, and quantem.gpu.direct_ptycho.

  • Backend modules implement those contracts and are not consumer APIs.

  • Native clients consume the repository-root Swift package. Swift and Swift-only Metal sources live in native/swift, outside the Python wheel; Metal sources and data Python also reads live once in the Python package, and Package.swift copies them into the resource bundles.

  • Browser consumers import the domain-owned TypeScript sources listed in src/quantem/gpu/webgpu/sources.json; there is no separate browser entry module.

  • Browser kernels remain beside their scientific domain. WebGPU is a browser runtime, not a Python device name and not a synonym for Metal.

  • CPU/NumPy is an explicit reference backend for tests. Production scientific work must never silently fall back to it.

  • UI, view state, cache scheduling, and resource-policy choices remain in the consuming application. Reusable math, kernels, resource estimation, and scientific provenance remain in quantem.gpu.

IO representation layout#

The migration separates representation-specific implementation names without changing the public Python verbs, load defaults, metadata, or kernels. The Python files below were moved, not copied.

Backend

HDF5 decode implementation

Python CUDA

io/hdf5/cuda/decode.py kernels, consumed by io/encoded.py, io/paired.py, and io/precision.py

Python MPS

io/hdf5/mps/decode.py, consumed by io/encoded.py and io/precision.py

Explicit CPU reference

io/hdf5/cpu.py (dense)

WebGPU

io/hdf5/webgpu/local-h5.ts (dense)

Python callers continue to use io.load(..., representation=...). The representation contract remains authoritative for supported formats and operations. Representation names do not imply that every backend accepts every source or offers every operation.

io/hdf5/webgpu/local-h5.ts exports setShow4DSTEMLocalFiles, loadShow4DSTEMLocalH5Master, and loadShow4DSTEMLocalH5MaskedSum; detector/webgpu/backend.ts exports DetectorCompute; browser display is quantem.widget’s. Register the selected HDF5 master/chunk files with setShow4DSTEMLocalFiles before using the dense loader. Each registration replaces the previous file registry; it does not release GPU buffers.

Python backend imports use the canonical modules above. The old cpu.reference and mps.decoder compatibility modules are removed, and the lossless-pack readers (cuda.packed, mps.packed) were deleted with their container format. Mutable backend state has one owner; tests that inject allocation, cache, or device state target that canonical module. Browser build scripts must export the complete dependency graph.

The loader separates result/ownership types (io/dataset.py), bounded region reads (io/read.py), header readers (formats/hdf5/master.py) and pinned staging ownership (device/cuda_runtime.py).

load.py keeps the public route choice. Bounded HDF5-to-ANS streaming lives in io/encoded.py, the CUDA paired-count layout in io/paired.py, scaled precision storage in io/precision.py, and the dense CPU reference in io/hdf5/cpu.py; GPU dense loading is not a route. Further decomposition must preserve source-range, cancellation, and failure behavior.

Current source tree#

Only implemented directories are shown. A directory does not imply every operation is supported by every backend:

src/quantem/gpu/                  # layers import downward only:
                                  # device < formats < resident < io < detector, geometry
                                  # < dpc, parallax < screening < ssb < remote < cli
  device/                         # device detection, CUDA and Metal runtime plumbing
  formats/                        # files on disk, GPU-free: hdf5/, qem/, emd, empad, ...
  resident/                       # acquisitions held on the GPU and their exact reductions
    cuda/, mps/                   # backend residents; kernels/ beside the Python that reads them
  io/
    __init__.py                   # public API only
    load.py                       # public load and route choice
    dataset.py                    # loaded data and ownership contract
    representation.py             # representation is independent of dtype
    hdf5/{cpu.py,cuda,mps,webgpu} # bitshuffle+LZ4 codecs and browser sources
  detector/
    __init__.py                   # public BF/ADF/DF/mean-DP API
    cuda/, mps/, webgpu/
  dpc/
    __init__.py                   # public CoM, rotation, and iDPC API
    results.py
    webgpu/
  display/
    __init__.py                   # shared display-math contract
    colormaps.json                # colormap control points for every backend
    goldens/parity.json           # shared display goldens
    metal/display.metal           # native Metal display shader
  ssb/
    __init__.py                   # one SSB workflow and result contract
    contract.py, cuda/, mps/, webgpu/
  geometry/cuda.py
  parallax/cuda/
  optics/physics.py
  movie/{export,cuda,mps}.py
  remote/                         # transport of exact scientific arrays
  webgpu/                         # the sources.json build manifest

native/                           # not part of the Python wheel
  swift/
    Sources/                      # SwiftPM products mirroring the domains
    Tests/
    Benchmarks/
    Vendor/                       # CHDF5.xcframework

tests/
  contracts/                      # public API, provenance, and failure rules
  parity/
    backend_matrix.json           # machine-readable required coverage
    fixtures/                     # retained frozen scientific evidence
    webgpu/                       # executable TypeScript checks
  hardware/{cuda,mps,metal}/       # device-dependent gates
  e2e/                            # consumer-local override and real-data gates
  infrastructure/                 # docs, packaging, benchmark checks

mps remains the Python backend selector for compatibility. Its accelerated implementation may use MLX or Metal. Native Swift products use Metal in their names because they expose Metal buffers and command encoding directly.

Canonical owners#

Implementations have one canonical owner. Backend code lives in a cuda/, mps/ or webgpu/ directory (or one cuda.py, mps.py or cpu.py module) directly under the package that owns the science; there is no backends/ or compute/ level, and no compatibility shims.

Keep one support matrix, tests/parity/backend_matrix.json, and keep historical result paths and trial IDs intact. Moving source code does not make pending hardware cells pass.

Browser build integration#

Do not maintain a consumer-side list of individual kernel files. The package ships webgpu/sources.json, the list of package-owned browser sources as paths relative to quantem.gpu. Build scripts read that manifest (for example through importlib.resources.files("quantem.gpu")), copy the listed files into a generated directory, and import the modules they need. Qualify the complete consumer bundle against an exact package revision before adoption. Manifest completeness alone is not application compatibility.

Reproduce the layout checks#

Run from the repository root. These checks validate imports, numerical regressions, package resources, native contracts, and documentation. They do not replace real-device timing or app acceptance.

# Python: imports, lazy accelerators, dense and encoded behavior.
python scripts/run_tests.py --list
python scripts/run_tests.py contracts parity infrastructure -q
python scripts/run_tests.py hardware/cuda -q
python scripts/run_tests.py hardware/mps -q
python scripts/run_tests.py e2e -q

# Suite names map to test directories; pytest paths and node IDs pass through.
python scripts/run_tests.py tests/contracts/io/test_load.py -q

# Browser: install the pinned development tools, then check the source manifest.
npm ci
python scripts/run_tests.py tests/contracts/test_webgpu_export.py -q
# Every source in webgpu/sources.json, with quantem.widget's compiler options.
npm run typecheck

# Native Apple: unchanged SwiftPM entry.
swift test -j 4

jupyter-book build docs --warningiserror --nitpick
python scripts/check_docs_links.py --html-root docs/_build/html

The browser contract tests read their tools from node_modules/ in the repository root. They are skipped when Node.js is not installed (the logical-pixel hash contract also needs Node.js 22 or newer), and a skip is not a passing WebGPU contract; a missing tool fails with the npm ci instruction.

One cross-language scientific contract#

Every backend result bundle must record:

  • source identity or fixture hash;

  • source scan and detector shape and dtype;

  • half-open scan and detector regions;

  • scan bin, detector bin, output shape, output dtype, and accumulation dtype;

  • bad-pixel policy and detector-mask definition;

  • (row, column) ≡ (r, c) coordinate convention;

  • backend, device, source revision, and kernel revision;

  • whether the result is native resolution, explicitly binned, or explicitly cropped; and

  • output array hashes plus the parity metric used.

Real-space crop is never an implicit memory or speed policy. Detector binning must be explicit and count-preserving, including incomplete edge bins. A binned array must never be described as native-resolution evidence.

Parity layers#

Parity is cumulative. A source-presence or compile test does not replace a hardware or real-data gate.

  1. Contract: imports, signatures, shapes, dtypes, coordinate order, provenance, and honest unsupported/failure behavior.

  2. Synthetic numerical: small odd and rectangular arrays; partial edge bins; masks; nonfinite display inputs; deterministic seeds.

  3. Frozen cross-backend: every backend reads the same input fixture and writes the same versioned result bundle. Integer decode, bin, masks, sums, and RGBA/histogram outputs are byte-exact.

  4. Real-data: full source shape and dtype, no unreported crop/bin/precision change, exact source hashes, peak memory, and output hashes.

  5. Physical end-to-end: the real consumer uses a local package override or exact revision pin; cold first-source, warm source, and saved-result reopen are reported separately.

Floating-point operations use a frozen, operation-specific metric. CoM, DPC, FFT, and SSB must report maximum and high-percentile error as applicable; a tolerance may not be widened to make a new backend pass. Goldens are generated only by an explicit recapture command and never by the backend being adjudicated.

Migration gate#

Move one domain at a time in this order: contract/fixtures, IO, detector/DPC, display, then SSB. For each move:

  1. freeze the old output bundles;

  2. move the module and update every caller in the same coordinated change: this repository, quantem.widget, quantem.live, and denoise. No import-only compatibility shim or alias keeps the old path;

  3. run CPU-reference, CUDA, MPS/Metal, Swift, and WebGPU gates listed in tests/parity/backend_matrix.json;

  4. test supported native and browser clients through a local package override; and

  5. commit the move and pin consumers to that exact revision.

No folder cleanup is complete merely because unit tests pass on one host.