engine/device.js

// `immediate_address_space` is a WGSL language extension rather than a
// requestDevice feature. Browsers exposing it also expose
// GPUComputePassEncoder.setImmediates(); record() still checks the method at
// use time so an inconsistent implementation fails loudly.
export function supportsImmediates(gpu = globalThis.navigator?.gpu) {
  return !!gpu?.wgslLanguageFeatures?.has?.('immediate_address_space');
}

const immediateDevices = new WeakSet();

export function deviceSupportsImmediates(device) {
  return !!device && immediateDevices.has(device);
}

export async function initDevice({ requireF16 = false } = {}) {
  if (!navigator.gpu) throw new Error('WebGPU unavailable');
  const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' });
  if (!adapter) throw new Error('no adapter');
  const want = ['shader-f16', 'timestamp-query', 'subgroups'];
  const features = want.filter((f) => adapter.features.has(f));
  if (requireF16 && !features.includes('shader-f16')) throw new Error('shader-f16 unsupported');
  const device = await adapter.requestDevice({
    requiredFeatures: features,
    requiredLimits: {
      // Ask for generous limits (clamped to what the adapter offers) so large
      // decode batches fit in single bindings; granted values are surfaced in
      // ctx.limits for the maxBatchForLimits guard.
      maxStorageBufferBindingSize: Math.min(536870912, adapter.limits.maxStorageBufferBindingSize),
      maxBufferSize: Math.min(1073741824, adapter.limits.maxBufferSize),
    },
  });
  const hasImmediates = supportsImmediates(navigator.gpu);
  if (hasImmediates) immediateDevices.add(device);
  return {
    device,
    adapterInfo: { vendor: adapter.info?.vendor ?? '', architecture: adapter.info?.architecture ?? '', device: adapter.info?.device ?? '', description: adapter.info?.description ?? '' },
    hasF16: features.includes('shader-f16'),
    hasTimestamps: features.includes('timestamp-query'),
    hasSubgroups: features.includes('subgroups'),
    hasImmediates,
    // Kernel SG variants assume a TK-slice never straddles a subgroup:
    // routing must require TK ≤ subgroupMinSize (power-of-two sizes ⇒ that
    // also gives sgSize % TK == 0). 0 when the adapter doesn't report it.
    subgroupMinSize: adapter.info?.subgroupMinSize ?? 0,
    subgroupMaxSize: adapter.info?.subgroupMaxSize ?? 0,
    limits: {
      maxStorageBufferBindingSize: device.limits.maxStorageBufferBindingSize,
      maxBufferSize: device.limits.maxBufferSize,
      // Read-only extras for the benchmark payload: workgroup-storage size is
      // exactly the ceiling suspected in the Adreno mega+sg miscompile (#3).
      maxComputeWorkgroupStorageSize: device.limits.maxComputeWorkgroupStorageSize,
      maxComputeInvocationsPerWorkgroup: device.limits.maxComputeInvocationsPerWorkgroup,
    },
    // Raw adapter capability (NOT the granted device limits above): limits
    // we never request come back as spec defaults on the device, which is
    // useless for driver-quirk analysis. Enumerate the prototype getters —
    // the key set follows the browser's spec version automatically.
    adapterLimits: (() => {
      const out = {};
      try {
        for (const k of Object.getOwnPropertyNames(Object.getPrototypeOf(adapter.limits))) {
          const v = adapter.limits[k];
          if (typeof v === 'number' && Number.isFinite(v)) out[k] = v;
        }
      } catch { /* keep whatever was collected */ }
      return out;
    })(),
    features: (() => { try { return [...adapter.features]; } catch { return []; } })(),
  };
}