qsm_core/inversion/
tile_budget.rs

1//! How much heap a tiled network's tiles take, and how many a browser keeps in flight.
2//!
3//! The arithmetic the tiled drivers (`inversion::tiled`, `onnx` feature) budget with, kept out of
4//! the `onnx` gate so a host that only estimates memory (and does not run networks) computes the
5//! same figures. A network's cost per tile comes from the registry
6//! ([`ModelSpec::tile_cost`](crate::models::ModelSpec::tile_cost)).
7
8use crate::models::TileCost;
9
10/// Tiling parameters for the tiled drivers (`tiled_field_inversion`, `tiled_volume_algorithm`).
11#[derive(Clone, Copy, Debug)]
12pub struct TileConfig {
13    /// Output core size per axis — the region each patch contributes to the result.
14    pub core: usize,
15    /// Context margin (voxels) included on every side of each patch. Larger halos reduce
16    /// tile-boundary artifacts at the cost of more work per patch.
17    pub halo: usize,
18}
19
20impl Default for TileConfig {
21    /// 128³ cores with an 8-voxel halo → 144³ patches. Empirically the halo barely affects
22    /// accuracy here — the tiling error is dominated by *global* low-frequency drift (a patch
23    /// can't see distant susceptibility sources), not tile-boundary seams — so this favours the
24    /// large core (few patches, ~4× less overlap compute than 64³/32) over a big halo. On real
25    /// 3 T data this matched whole-volume xQSM at r≈0.94 in ~1/4 the time.
26    ///
27    /// Sized for **native** memory: a 144³ patch of xQSM holds about 3.4 GB while it runs (see
28    /// [`ModelSpec::tile_bytes_per_voxel`](crate::models::ModelSpec::tile_bytes_per_voxel)), which a 32-bit WASM host cannot fit at all. Browser
29    /// hosts pass a far smaller config — QSMbly uses `core: 56, halo: 4` (64³ patches).
30    fn default() -> Self {
31        Self { core: 128, halo: 8 }
32    }
33}
34
35/// Heap budget for all concurrent tiles on WASM — see `tile_concurrency` (`onnx` feature).
36///
37/// A threaded WASM host has one shared 32-bit heap with a hard 4 GB ceiling and an allocator
38/// that cannot compact, and that heap also holds the model bytes, the full-volume inputs and
39/// outputs and the host's own data. 1.5 GB of tiles leaves room for all of it plus fragmentation.
40pub const WASM_TILE_ACTIVATION_BUDGET: u64 = 1_500_000_000;
41
42/// Largest footprint a single patch may have on WASM, running alone.
43///
44/// `tile_concurrency` (`onnx` feature) can bring the patches in flight down to one but not below, so a patch
45/// bigger than this cannot run in a browser at all. The rest of the 4 GB heap is not free: a 104³
46/// xQSM patch running alone peaked at 1.74 GB, so about 0.5 GB goes to the model, the volumes, the
47/// host and fragmentation. 2.5 GB keeps one patch well clear of the 4.24 GB at which a run stalled
48/// (#89).
49pub const WASM_MAX_PATCH_ACTIVATION_BYTES: u64 = 2_500_000_000;
50
51/// Heap one tile of this config holds while it runs, for a network of this cost
52/// ([`ModelSpec::tile_cost`](crate::models::ModelSpec::tile_cost)): its patch, rounded up to the
53/// network's divisor as [`tile_patch_size`] does, times the network's bytes per patch voxel. `u64`,
54/// because `p³ · bytes` overflows a 32-bit `usize` for large patches.
55pub fn tile_bytes(cfg: &TileConfig, cost: &TileCost) -> u64 {
56    let d = cost.divisor.max(1);
57    let edge = cfg.core.max(1).saturating_add(cfg.halo.saturating_mul(2)).div_ceil(d).saturating_mul(d);
58    patch_bytes(edge, cost.bytes_per_voxel)
59}
60
61/// Heap one cubic patch of this edge holds.
62fn patch_bytes(patch_edge: usize, bytes_per_voxel: u64) -> u64 {
63    (patch_edge as u64).saturating_pow(3).saturating_mul(bytes_per_voxel)
64}
65
66/// Largest patch edge that fits [`WASM_MAX_PATCH_ACTIVATION_BYTES`] for a network costing
67/// `bytes_per_voxel` per patch voxel.
68pub fn max_wasm_patch_edge(bytes_per_voxel: u64) -> usize {
69    let mut p = 1usize;
70    while patch_bytes(p + 1, bytes_per_voxel) <= WASM_MAX_PATCH_ACTIVATION_BYTES {
71        p += 1;
72    }
73    p
74}
75
76/// Tiles a WASM host with `threads` workers keeps in flight: the threads, capped by
77/// [`tiles_within_wasm_budget`]. This is what `tile_concurrency` (`onnx` feature) returns on WASM; it is callable
78/// anywhere, so a host estimating a browser run's memory can use the same count.
79pub fn wasm_tiles_in_flight(cfg: &TileConfig, cost: &TileCost, threads: usize) -> usize {
80    threads.max(1).min(tiles_within_wasm_budget(cfg, cost))
81}
82
83/// How many tiles of this config fit a browser's share of heap — at least one, however big the
84/// patch — for a network of this cost ([`ModelSpec::tile_cost`](crate::models::ModelSpec::tile_cost)).
85pub fn tiles_within_wasm_budget(cfg: &TileConfig, cost: &TileCost) -> usize {
86    (WASM_TILE_ACTIVATION_BUDGET / tile_bytes(cfg, cost).max(1)).max(1) as usize
87}
88
89/// Padded input-patch size per axis for a config and the net's `size_divisor`: `core + 2·halo`
90/// rounded up to a multiple of `divisor`. Isotropic, so one value serves all axes. Callers use
91/// this to build a reusable `OnnxModel::plan_for` plan of shape `[1, 1, p, p, p]`.
92pub fn tile_patch_size(cfg: &TileConfig, divisor: usize) -> usize {
93    (cfg.core.max(1) + 2 * cfg.halo).div_ceil(divisor.max(1)) * divisor.max(1)
94}
qsm_core/inversion/tile_budget.rs

qsm_core/inversion/
tile_budget.rs
