@jax-js/jax 0.0.1 → 0.0.2

package/dist/webgpu-QNXDOQZP.js

@@ -1,559 +0,0 @@
-import {
-  AluExp,
-  AluGroup,
-  DEBUG,
-  Executable,
-  FpHash,
-  SlotError,
-  findPow2,
-  strip1,
-  tuneWebgpu
-} from "./chunk-B2GFURUN.js";
-
-// src/backend/webgpu.ts
-var WebGPUBackend = class {
-  constructor(device) {
-    this.device = device;
-    if (DEBUG >= 3 && device.adapterInfo) {
-      console.info(
-        "webgpu adapter:",
-        device.adapterInfo.vendor,
-        device.adapterInfo.architecture
-      );
-    }
-    this.maxArgs = this.device.limits.maxStorageBuffersPerShaderStage - 1;
-    this.pipelines = new ShaderPipelineCache(device);
-    this.syncReader = new SyncReader(device);
-    this.buffers = /* @__PURE__ */ new Map();
-    this.nextSlot = 1;
-  }
-  type = "webgpu";
-  maxArgs;
-  pipelines;
-  syncReader;
-  buffers;
-  nextSlot;
-  #cachedShaderMap = /* @__PURE__ */ new Map();
-  malloc(size, initialData) {
-    let buffer;
-    if (initialData) {
-      if (initialData.byteLength !== size) {
-        throw new Error("initialData size does not match buffer size");
-      }
-      if (initialData.byteLength < 4096) {
-        buffer = this.#createBuffer(size, { mapped: true });
-        new Uint8Array(buffer.getMappedRange()).set(
-          new Uint8Array(initialData)
-        );
-        buffer.unmap();
-      } else {
-        buffer = this.#createBuffer(size);
-        this.device.queue.writeBuffer(buffer, 0, initialData);
-      }
-    } else {
-      buffer = this.#createBuffer(size);
-    }
-    const slot = this.nextSlot++;
-    this.buffers.set(slot, { buffer, ref: 1 });
-    return slot;
-  }
-  incRef(slot) {
-    const buffer = this.buffers.get(slot);
-    if (!buffer) throw new SlotError(slot);
-    buffer.ref++;
-  }
-  decRef(slot) {
-    const buffer = this.buffers.get(slot);
-    if (!buffer) throw new SlotError(slot);
-    buffer.ref--;
-    if (buffer.ref === 0) {
-      this.buffers.delete(slot);
-      buffer.buffer.destroy();
-    }
-  }
-  async read(slot, start, count) {
-    const buffer = this.#getBuffer(slot);
-    if (start === void 0) start = 0;
-    if (count === void 0) count = buffer.size - start;
-    const staging = this.#createBuffer(count, { read: true });
-    try {
-      const commandEncoder = this.device.createCommandEncoder();
-      commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, count);
-      this.device.queue.submit([commandEncoder.finish()]);
-      await staging.mapAsync(GPUMapMode.READ);
-      const arrayBuffer = staging.getMappedRange();
-      return arrayBuffer.slice();
-    } finally {
-      staging.destroy();
-    }
-  }
-  readSync(slot, start, count) {
-    const buffer = this.#getBuffer(slot);
-    if (start === void 0) start = 0;
-    if (count === void 0) count = buffer.size - start;
-    return this.syncReader.read(buffer, start, count);
-  }
-  #cachedShader(kernel) {
-    const cacheKey = FpHash.hash(kernel);
-    let result = this.#cachedShaderMap.get(cacheKey);
-    if (!result) {
-      result = pipelineSource(this.device, kernel);
-      this.#cachedShaderMap.set(cacheKey, result);
-    }
-    return result;
-  }
-  async prepare(kernel) {
-    const { shader, grid } = this.#cachedShader(kernel);
-    const pipeline = await this.pipelines.prepare(shader);
-    return new Executable(kernel, { shader, grid, pipeline });
-  }
-  prepareSync(kernel) {
-    const { shader, grid } = this.#cachedShader(kernel);
-    const pipeline = this.pipelines.prepareSync(shader);
-    return new Executable(kernel, { shader, grid, pipeline });
-  }
-  dispatch(exe, inputs, outputs) {
-    const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
-    const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
-    pipelineSubmit(this.device, exe.data, inputBuffers, outputBuffers);
-  }
-  #getBuffer(slot) {
-    const buffer = this.buffers.get(slot);
-    if (!buffer) throw new SlotError(slot);
-    return buffer.buffer;
-  }
-  /**
-   * Create a GPU buffer.
-   *
-   * By default, this creates a general-purpose buffer with the given size.
-   *
-   * - If `mapped` is true, initialize the buffer in mapped mode so that it can
-   *   be populated with data from the CPU. (Call `.unmap()` later.)
-   * - If `read` is true, create a staging buffer for returning data to CPU.
-   *   (Call `.mapAsync()` later.)
-   */
-  #createBuffer(size, { mapped = false, read = false } = {}) {
-    if (read && mapped) {
-      throw new Error("mapped and read cannot both be true");
-    }
-    const buffer = this.device.createBuffer({
-      size,
-      usage: read ? GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST : GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
-      mappedAtCreation: mapped
-    });
-    return buffer;
-  }
-};
-function dtypeToWgsl(dtype, storage = false) {
-  switch (dtype) {
-    case "bool" /* Bool */:
-      return storage ? "i32" : "bool";
-    // WebGPU does not support bools in buffers.
-    case "int32" /* Int32 */:
-      return "i32";
-    case "float32" /* Float32 */:
-      return "f32";
-    default:
-      throw new Error(`Unsupported dtype: ${dtype}`);
-  }
-}
-function constToWgsl(dtype, value) {
-  if (dtype === "bool" /* Bool */) return value ? "true" : "false";
-  if (dtype === "int32" /* Int32 */) return value.toString();
-  if (dtype === "float32" /* Float32 */) {
-    let s = value.toString();
-    if (!s.includes(".")) s += ".0";
-    return s;
-  }
-  throw new Error(`Unsupported const dtype: ${dtype}`);
-}
-function pipelineSource(device, kernel) {
-  const tune = tuneWebgpu(kernel);
-  if (DEBUG >= 3) {
-    console.info(`kernel.exp: ${kernel.exp}
-tune.exp: ${tune.exp}`);
-  }
-  const { nargs, reduction: re } = kernel;
-  const args = Array.from({ length: nargs }, (_, i) => `in${i}`);
-  const shader = [];
-  let indent = "";
-  const pushIndent = Symbol("pushIndent");
-  const popIndent = Symbol("popIndent");
-  const emit = (...lines) => {
-    for (const line of lines) {
-      if (line === pushIndent) indent += "  ";
-      else if (line === popIndent) indent = indent.slice(0, -2);
-      else shader.push(line ? indent + line : line);
-    }
-  };
-  const usedArgs = Array.from({ length: nargs }, () => null);
-  tune.exp.fold((exp) => {
-    if (exp.op === "GlobalIndex" /* GlobalIndex */) usedArgs[exp.arg] = exp.dtype;
-  });
-  for (let i = 0; i < nargs; i++) {
-    const ty = dtypeToWgsl(usedArgs[i] ?? "float32" /* Float32 */, true);
-    emit(
-      `@group(0) @binding(${i}) var<storage, read> ${args[i]} : array<${ty}>;`
-    );
-  }
-  const resultTy = dtypeToWgsl(re?.dtype ?? tune.exp.dtype, true);
-  emit(
-    `@group(0) @binding(${nargs}) var<storage, read_write> result : array<${resultTy}>;`
-  );
-  const workgroupSize = findPow2(tune.threadCount, 256);
-  emit(
-    "",
-    `@compute @workgroup_size(${workgroupSize})`,
-    "fn main(@builtin(global_invocation_id) id : vec3<u32>) {",
-    pushIndent,
-    `if (id.x >= ${tune.threadCount}) { return; }`,
-    "let gidx: i32 = i32(id.x);"
-  );
-  let gensymCount = 0;
-  const gensym = () => `alu${gensymCount++}`;
-  for (let i = 0; i < args.length; i++) {
-    if (!usedArgs[i]) emit(`_ = &${args[i]};`);
-  }
-  const references = /* @__PURE__ */ new Map();
-  const seen = /* @__PURE__ */ new Set();
-  const countReferences = (exp) => {
-    references.set(exp, (references.get(exp) ?? 0) + 1);
-    if (!seen.has(exp)) {
-      seen.add(exp);
-      for (const src of exp.src) countReferences(src);
-    }
-  };
-  const expContext = /* @__PURE__ */ new Map();
-  const gen = (exp) => {
-    if (expContext.has(exp)) return expContext.get(exp);
-    const { op, src, dtype, arg } = exp;
-    let source = "";
-    if (AluGroup.Binary.has(op) || AluGroup.Compare.has(op)) {
-      const a = gen(src[0]);
-      const b = gen(src[1]);
-      if (op === "Add" /* Add */) {
-        if (dtype === "bool" /* Bool */) source = `(${a} || ${b})`;
-        else source = `(${a} + ${b})`;
-      } else if (op === "Sub" /* Sub */) source = `(${a} - ${b})`;
-      else if (op === "Mul" /* Mul */) {
-        if (dtype === "bool" /* Bool */) source = `(${a} && ${b})`;
-        else source = `(${a} * ${b})`;
-      } else if (op === "Idiv" /* Idiv */)
-        source = dtype === "int32" /* Int32 */ ? `(${a} / ${b})` : `floor(${a} / ${b})`;
-      else if (op === "Mod" /* Mod */) source = `(${a} % ${b})`;
-      else if (op === "Min" /* Min */) source = `min(${strip1(a)}, ${strip1(b)})`;
-      else if (op === "Max" /* Max */) source = `max(${strip1(a)}, ${strip1(b)})`;
-      else if (op === "Cmplt" /* Cmplt */) source = `(${a} < ${b})`;
-      else if (op === "Cmpne" /* Cmpne */) source = `(${a} != ${b})`;
-    } else if (AluGroup.Unary.has(op)) {
-      const a = gen(src[0]);
-      if (op === "Sin" /* Sin */) source = `sin(${a})`;
-      else if (op === "Cos" /* Cos */) source = `cos(${a})`;
-      else if (op === "Exp" /* Exp */) source = `exp(${a})`;
-      else if (op === "Log" /* Log */) source = `log(${a})`;
-      else if (op === "Reciprocal" /* Reciprocal */) source = `(1.0 / ${a})`;
-      else if (op === "Cast" /* Cast */)
-        source = `${dtypeToWgsl(dtype)}(${strip1(a)})`;
-    } else if (op === "Where" /* Where */) {
-      source = `select(${strip1(gen(src[2]))}, ${strip1(gen(src[1]))}, ${strip1(gen(src[0]))})`;
-    } else if (op === "Const" /* Const */) {
-      return constToWgsl(dtype, arg);
-    } else if (op === "Special" /* Special */) {
-      return arg[0];
-    } else if (op === "Variable" /* Variable */) {
-      return arg;
-    } else if (op === "GlobalIndex" /* GlobalIndex */) {
-      source = `${args[arg]}[${strip1(gen(src[0]))}]`;
-      if (dtype === "bool" /* Bool */) source = `(${source} != 0)`;
-    }
-    if (!source) throw new Error(`Missing impl for op: ${op}`);
-    const typeName = dtypeToWgsl(dtype);
-    if ((references.get(exp) ?? 0) > 1) {
-      const name = gensym();
-      expContext.set(exp, name);
-      emit(`let ${name}: ${typeName} = ${strip1(source)};`);
-      return name;
-    } else {
-      expContext.set(exp, source);
-      return source;
-    }
-  };
-  if (!kernel.reduction) {
-    countReferences(tune.exp);
-    let rhs = strip1(gen(tune.exp));
-    if (resultTy !== dtypeToWgsl(tune.exp.dtype)) rhs = `${resultTy}(${rhs})`;
-    emit(`result[gidx] = ${rhs};`);
-  } else {
-    const re2 = kernel.reduction;
-    if ((tune.size.groups ?? 1) > 1) {
-      throw new Error("WebGPU backend does not support group optimization yet");
-    }
-    const unroll = tune.size.unroll ?? 1;
-    const upcast = tune.size.upcast ?? 1;
-    const acc = [...Array(upcast)].map((_, i) => `acc${i}`);
-    for (let i = 0; i < upcast; i++) {
-      emit(
-        `var ${acc[i]}: ${dtypeToWgsl(tune.exp.dtype)} = ${constToWgsl(re2.dtype, re2.identity)};`
-      );
-    }
-    emit(
-      `for (var ridx: i32 = 0; ridx < ${tune.size.reduce}; ridx++) {`,
-      pushIndent
-    );
-    const exps = [];
-    const cache = /* @__PURE__ */ new Map();
-    for (let up = 0; up < upcast; up++) {
-      exps.push([]);
-      for (let un = 0; un < unroll; un++) {
-        const exp = tune.exp.substitute({
-          upcast: AluExp.i32(up),
-          unroll: AluExp.i32(un)
-        });
-        exps[up].push(exp.simplify(cache));
-        countReferences(exps[up][un]);
-      }
-    }
-    const items = exps.map((ar) => ar.map(gen).map(strip1));
-    for (let i = 0; i < upcast; i++) {
-      let rhs = items[i][0];
-      for (let j = 1; j < unroll; j++) {
-        if (re2.op === "Add" /* Add */) rhs = `${rhs} + ${items[i][j]}`;
-        else if (re2.op === "Mul" /* Mul */) rhs = `${rhs} * ${items[i][j]}`;
-        else if (re2.op === "Min" /* Min */) rhs = `min(${rhs}, ${items[i][j]})`;
-        else if (re2.op === "Max" /* Max */) rhs = `max(${rhs}, ${items[i][j]})`;
-        else throw new Error(`Unsupported reduction op: ${re2.op}`);
-      }
-      if (re2.op === "Add" /* Add */) emit(`${acc[i]} += ${rhs};`);
-      else if (re2.op === "Mul" /* Mul */) emit(`${acc[i]} *= ${rhs};`);
-      else if (re2.op === "Min" /* Min */) emit(`${acc[i]} = min(${acc[i]}, ${rhs});`);
-      else if (re2.op === "Max" /* Max */) emit(`${acc[i]} = max(${acc[i]}, ${rhs});`);
-      else throw new Error(`Unsupported reduction op: ${re2.op}`);
-    }
-    emit(popIndent, "}");
-    expContext.clear();
-    references.clear();
-    seen.clear();
-    const outputIdxExps = [];
-    const fusionExps = [];
-    for (let i = 0; i < upcast; i++) {
-      const exp = tune.outputIdxExp.substitute({ upcast: AluExp.i32(i) });
-      outputIdxExps.push(exp.simplify(cache));
-      countReferences(outputIdxExps[i]);
-      fusionExps.push(
-        re2.fusion.substitute({ acc: AluExp.variable(re2.dtype, acc[i]) }).simplify(cache)
-      );
-      countReferences(fusionExps[i]);
-    }
-    for (let i = 0; i < upcast; i++) {
-      const index = strip1(gen(outputIdxExps[i]));
-      let rhs = strip1(gen(fusionExps[i]));
-      if (resultTy !== dtypeToWgsl(fusionExps[i].dtype))
-        rhs = `${resultTy}(${rhs})`;
-      emit(`result[${index}] = ${rhs};`);
-    }
-  }
-  emit(popIndent, "}");
-  return {
-    shader: shader.join("\n"),
-    grid: [Math.ceil(tune.threadCount / workgroupSize), 1]
-  };
-}
-function pipelineSubmit(device, { pipeline, grid }, inputs, outputs) {
-  if (inputs.length + outputs.length > device.limits.maxStorageBuffersPerShaderStage) {
-    const actual = inputs.length + outputs.length;
-    const max = device.limits.maxStorageBuffersPerShaderStage;
-    throw new Error(
-      `Too many buffers (${actual}) for WebGPU pipeline (max: ${max})`
-    );
-  }
-  const bindGroup = device.createBindGroup({
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      ...inputs.map((buffer, i) => {
-        return { binding: i, resource: { buffer } };
-      }),
-      { binding: inputs.length, resource: { buffer: outputs[0] } }
-    ]
-  });
-  const commandEncoder = device.createCommandEncoder();
-  const passEncoder = commandEncoder.beginComputePass();
-  passEncoder.setPipeline(pipeline);
-  passEncoder.setBindGroup(0, bindGroup);
-  passEncoder.dispatchWorkgroups(grid[0], grid[1]);
-  passEncoder.end();
-  device.queue.submit([commandEncoder.finish()]);
-}
-var ShaderPipelineCache = class {
-  constructor(device) {
-    this.device = device;
-    this.cache = /* @__PURE__ */ new Map();
-    this.inProgress = /* @__PURE__ */ new Map();
-  }
-  cache;
-  inProgress;
-  async prepare(code) {
-    const existingPipeline = this.cache.get(code);
-    if (existingPipeline) return existingPipeline;
-    const existingPromise = this.inProgress.get(code);
-    if (existingPromise) return await existingPromise;
-    if (DEBUG >= 2) {
-      console.info("=========== WebGPU shader ===========\n" + code);
-    }
-    const shaderModule = this.device.createShaderModule({ code });
-    const promise = (async () => {
-      this.device.pushErrorScope("validation");
-      try {
-        const pipeline2 = await this.device.createComputePipelineAsync({
-          layout: "auto",
-          compute: {
-            module: shaderModule,
-            entryPoint: "main"
-          }
-        });
-        await this.device.popErrorScope();
-        return pipeline2;
-      } catch (_error) {
-        const scope = await this.device.popErrorScope();
-        const emsg = await compileError(shaderModule, scope, code);
-        throw new Error(emsg);
-      }
-    })();
-    this.inProgress.set(code, promise);
-    const pipeline = await promise;
-    this.cache.set(code, pipeline);
-    return pipeline;
-  }
-  prepareSync(code) {
-    const existingPipeline = this.cache.get(code);
-    if (existingPipeline) return existingPipeline;
-    if (DEBUG >= 2) {
-      console.info("=========== WebGPU shader ===========\n" + code);
-    }
-    const shaderModule = this.device.createShaderModule({ code });
-    this.device.pushErrorScope("validation");
-    const pipeline = this.device.createComputePipeline({
-      layout: "auto",
-      compute: {
-        module: shaderModule,
-        entryPoint: "main"
-      }
-    });
-    this.device.popErrorScope().then(async (scope) => {
-      if (scope !== null) {
-        const emsg = await compileError(shaderModule, scope, code);
-        console.error(emsg);
-      }
-    });
-    this.cache.set(code, pipeline);
-    return pipeline;
-  }
-};
-async function compileError(shaderModule, scope, code) {
-  let message = `Failed to compile shader: ${scope ? scope.message : "(no error scope)"}`;
-  const info = await shaderModule.getCompilationInfo();
-  for (const msg of info.messages) {
-    message += `
-  [${msg.type} at ${msg.lineNum}:${msg.linePos}] ${msg.message}`;
-  }
-  if (code) {
-    message += `
-
-${code}`;
-  }
-  return message;
-}
-var SyncReader = class _SyncReader {
-  constructor(device) {
-    this.device = device;
-  }
-  static alphaModes = [
-    "opaque",
-    "premultiplied"
-  ];
-  static width = 256;
-  static height = 256;
-  initialized = false;
-  deviceStorage;
-  deviceContexts;
-  hostStorage;
-  hostContext;
-  #init() {
-    const makeCanvas = () => new OffscreenCanvas(_SyncReader.width, _SyncReader.height);
-    this.deviceStorage = _SyncReader.alphaModes.map(makeCanvas);
-    this.deviceContexts = this.deviceStorage.map((canvas, i) => {
-      const context = canvas.getContext("webgpu");
-      context.configure({
-        device: this.device,
-        // rgba8unorm is not supported on Chrome for macOS.
-        // https://bugs.chromium.org/p/chromium/issues/detail?id=1298618
-        format: "bgra8unorm",
-        usage: GPUTextureUsage.COPY_DST,
-        alphaMode: _SyncReader.alphaModes[i]
-      });
-      return context;
-    });
-    this.hostStorage = makeCanvas();
-    this.hostContext = this.hostStorage.getContext("2d", {
-      willReadFrequently: true
-    });
-    this.initialized = true;
-  }
-  read(buffer, start, count) {
-    if (!this.initialized) this.#init();
-    if (count % 4 !== 0) {
-      throw new Error("Read size must be a multiple of 4 bytes");
-    }
-    const deviceStorage = this.deviceStorage;
-    const deviceContexts = this.deviceContexts;
-    const hostContext = this.hostContext;
-    const pixelsSize = count / 4;
-    const bytesPerRow = _SyncReader.width * 4;
-    const valsGPU = new ArrayBuffer(count);
-    for (let i = 0; i < deviceContexts.length; i++) {
-      const texture = deviceContexts[i].getCurrentTexture();
-      const readData = (width, height, offset2) => {
-        const encoder = this.device.createCommandEncoder();
-        encoder.copyBufferToTexture(
-          { buffer, bytesPerRow, offset: offset2 + start },
-          { texture },
-          { width, height, depthOrArrayLayers: 1 }
-        );
-        const commandBuffer = encoder.finish();
-        this.device.queue.submit([commandBuffer]);
-        hostContext.clearRect(0, 0, width, height);
-        hostContext.drawImage(deviceStorage[i], 0, 0);
-        const values = hostContext.getImageData(0, 0, width, height).data;
-        const span = new Uint8ClampedArray(valsGPU, offset2, 4 * width * height);
-        const alphaMode = _SyncReader.alphaModes[i];
-        for (let k = 0; k < span.length; k += 4) {
-          if (alphaMode === "premultiplied") {
-            span[k + 3] = values[k + 3];
-          } else {
-            span[k] = values[k + 2];
-            span[k + 1] = values[k + 1];
-            span[k + 2] = values[k];
-          }
-        }
-      };
-      const pixelsPerCanvas = _SyncReader.width * _SyncReader.height;
-      const wholeChunks = Math.floor(pixelsSize / pixelsPerCanvas);
-      let remainder = pixelsSize % pixelsPerCanvas;
-      const remainderRows = Math.floor(remainder / _SyncReader.width);
-      remainder = remainder % _SyncReader.width;
-      let offset = 0;
-      for (let j = 0; j < wholeChunks; j++) {
-        readData(_SyncReader.width, _SyncReader.height, offset);
-        offset += pixelsPerCanvas * 4;
-      }
-      if (remainderRows > 0) {
-        readData(_SyncReader.width, remainderRows, offset);
-        offset += remainderRows * _SyncReader.width * 4;
-      }
-      if (remainder > 0) readData(remainder, 1, offset);
-    }
-    return valsGPU;
-  }
-};
-export {
-  WebGPUBackend
-};