@jax-js/jax 0.1.3 → 0.1.5

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jax-js/jax",
-  "version": "0.1.3",
+  "version": "0.1.5",
   "description": "Numerical computing and ML in the browser",
   "keywords": [
     "machine learning",

package/dist/webgpu-BVns4DbI.cjs

@@ -1,663 +0,0 @@
-const require_backend = require('./backend-CmaidnkQ.cjs');
-
-//#region src/backend/webgpu/builtins.ts
-const threefrySrc = `
-fn threefry2x32(key: vec2<u32>, ctr: vec2<u32>) -> vec2<u32> {
-  let ks0: u32 = key.x;
-  let ks1: u32 = key.y;
-  let ks2: u32 = ks0 ^ ks1 ^ 0x1BD11BDAu;
-
-  var x0: u32 = ctr.x + ks0;
-  var x1: u32 = ctr.y + ks1;
-
-  x0 += x1; x1 = (x1 << 13u) | (x1 >> 19u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 15u) | (x1 >> 17u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 26u) | (x1 >> 6u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 6u) | (x1 >> 26u); x1 ^= x0;
-  x0 += ks1;
-  x1 += ks2 + 1u;
-
-  x0 += x1; x1 = (x1 << 17u) | (x1 >> 15u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 29u) | (x1 >> 3u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 16u) | (x1 >> 16u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 24u) | (x1 >> 8u); x1 ^= x0;
-  x0 += ks2;
-  x1 += ks0 + 2u;
-
-  x0 += x1; x1 = (x1 << 13u) | (x1 >> 19u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 15u) | (x1 >> 17u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 26u) | (x1 >> 6u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 6u) | (x1 >> 26u); x1 ^= x0;
-  x0 += ks0;
-  x1 += ks1 + 3u;
-
-  x0 += x1; x1 = (x1 << 17u) | (x1 >> 15u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 29u) | (x1 >> 3u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 16u) | (x1 >> 16u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 24u) | (x1 >> 8u); x1 ^= x0;
-  x0 += ks1;
-  x1 += ks2 + 4u;
-
-  x0 += x1; x1 = (x1 << 13u) | (x1 >> 19u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 15u) | (x1 >> 17u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 26u) | (x1 >> 6u); x1 ^= x0;
-  x0 += x1; x1 = (x1 << 6u) | (x1 >> 26u); x1 ^= x0;
-  x0 += ks2;
-  x1 += ks0 + 5u;
-
-  return vec2<u32>(x0, x1);
-}`;
-const erfSrc = `
-const _erf_p: f32 = 0.3275911;
-const _erf_a1: f32 = 0.254829592;
-const _erf_a2: f32 = -0.284496736;
-const _erf_a3: f32 = 1.421413741;
-const _erf_a4: f32 = -1.453152027;
-const _erf_a5: f32 = 1.061405429;
-fn erf(x: f32) -> f32 {
-  let t = 1.0 / (1.0 + _erf_p * abs(x));
-  let P_t = fma(fma(fma(fma(_erf_a5, t, _erf_a4), t, _erf_a3), t, _erf_a2), t, _erf_a1) * t;
-  return sign(x) * (1.0 - P_t * exp(-x * x));
-}
-fn erfc(x: f32) -> f32 {
-  let t = 1.0 / (1.0 + _erf_p * abs(x));
-  let P_t = fma(fma(fma(fma(_erf_a5, t, _erf_a4), t, _erf_a3), t, _erf_a2), t, _erf_a1) * t;
-  let E = P_t * exp(-x * x);
-  return select(2.0 - E, E, x >= 0.0);
-}`;
-
-//#endregion
-//#region src/backend/webgpu/reader.ts
-/**
-* Graphics state used to synchronously read data from WebGPU buffers.
-*
-* This trick is borrowed from TensorFlow.js. Basically, the idea is to create
-* an offscreen canvas with one pixel for every 4 bytes ("device storage"), then
-* configure it with a WebGPU context. Copy the buffer to a texture, then draw
-* the canvas onto another offscreen canvas with '2d' context ("host storage").
-*
-* Once it's on host storage, we can use `getImageData()` to read the pixels
-* from the image directly.
-*
-* We use 256x256 canvases here (256 KiB). The performance of this is bad
-* because it involves multiple data copies, but it still works. We also
-* actually need to copy the image twice: once in "opaque" mode for the RGB
-* values, and once in "premultiplied" mode for the alpha channel.
-*
-* https://github.com/tensorflow/tfjs/blob/tfjs-v4.22.0/tfjs-backend-webgpu/src/backend_webgpu.ts#L379
-*/
-var SyncReader = class SyncReader {
-	static alphaModes = ["opaque", "premultiplied"];
-	static width = 256;
-	static height = 256;
-	initialized = false;
-	deviceStorage;
-	deviceContexts;
-	hostStorage;
-	hostContext;
-	constructor(device) {
-		this.device = device;
-	}
-	#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,
-				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();
-		const deviceStorage = this.deviceStorage;
-		const deviceContexts = this.deviceContexts;
-		const hostContext = this.hostContext;
-		const pixelsSize = Math.ceil(count / 4);
-		const bytesPerRow = SyncReader.width * 4;
-		const valsGPU = /* @__PURE__ */ new ArrayBuffer(pixelsSize * 4);
-		for (let i = 0; i < deviceContexts.length; i++) {
-			const texture = deviceContexts[i].getCurrentTexture();
-			const readData = (width, height, offset$1) => {
-				const encoder = this.device.createCommandEncoder();
-				encoder.copyBufferToTexture({
-					buffer,
-					bytesPerRow,
-					offset: offset$1 + 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, offset$1, 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 new Uint8Array(valsGPU, 0, count);
-	}
-};
-
-//#endregion
-//#region src/backend/webgpu.ts
-/** Implementation of `Backend` that uses WebGPU in browsers. */
-var WebGPUBackend = class {
-	type = "webgpu";
-	maxArgs;
-	pipelines;
-	syncReader;
-	buffers;
-	nextSlot;
-	#cachedShaderMap = /* @__PURE__ */ new Map();
-	constructor(device) {
-		this.device = device;
-		if (require_backend.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;
-	}
-	malloc(size, initialData) {
-		let buffer;
-		const paddedSize = Math.ceil(size / 4) * 4;
-		if (initialData) {
-			if (initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
-			if (initialData.byteLength < 4096) {
-				buffer = this.#createBuffer(paddedSize, { mapped: true });
-				new Uint8Array(buffer.getMappedRange(), 0, size).set(initialData);
-				buffer.unmap();
-			} else {
-				buffer = this.#createBuffer(paddedSize);
-				if (initialData.byteLength % 4 === 0) this.device.queue.writeBuffer(buffer, 0, initialData);
-				else {
-					const aligned = initialData.byteLength - initialData.byteLength % 4;
-					this.device.queue.writeBuffer(buffer, 0, initialData, 0, aligned);
-					const remainder = new Uint8Array(4);
-					remainder.set(initialData.subarray(aligned));
-					this.device.queue.writeBuffer(buffer, aligned, remainder);
-				}
-			}
-		} else buffer = this.#createBuffer(paddedSize);
-		const slot = this.nextSlot++;
-		this.buffers.set(slot, {
-			buffer,
-			size,
-			ref: 1
-		});
-		return slot;
-	}
-	incRef(slot) {
-		const buffer = this.buffers.get(slot);
-		if (!buffer) throw new require_backend.SlotError(slot);
-		buffer.ref++;
-	}
-	decRef(slot) {
-		const buffer = this.buffers.get(slot);
-		if (!buffer) throw new require_backend.SlotError(slot);
-		buffer.ref--;
-		if (buffer.ref === 0) {
-			this.buffers.delete(slot);
-			buffer.buffer.destroy();
-		}
-	}
-	async read(slot, start, count) {
-		const { buffer, size } = this.#getBuffer(slot);
-		if (start === void 0) start = 0;
-		if (count === void 0) count = size - start;
-		const paddedSize = Math.ceil(count / 4) * 4;
-		const staging = this.#createBuffer(paddedSize, { read: true });
-		try {
-			const commandEncoder = this.device.createCommandEncoder();
-			commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, paddedSize);
-			this.device.queue.submit([commandEncoder.finish()]);
-			await staging.mapAsync(GPUMapMode.READ);
-			const arrayBuffer = staging.getMappedRange();
-			return new Uint8Array(arrayBuffer.slice(), 0, count);
-		} finally {
-			staging.destroy();
-		}
-	}
-	readSync(slot, start, count) {
-		const { buffer, size } = this.#getBuffer(slot);
-		if (start === void 0) start = 0;
-		if (count === void 0) count = size - start;
-		return this.syncReader.read(buffer, start, count);
-	}
-	#cachedShader(kernel) {
-		const cacheKey = require_backend.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 require_backend.Executable(kernel, {
-			shader,
-			grid,
-			pipeline
-		});
-	}
-	prepareSync(kernel) {
-		const { shader, grid } = this.#cachedShader(kernel);
-		const pipeline = this.pipelines.prepareSync(shader);
-		return new require_backend.Executable(kernel, {
-			shader,
-			grid,
-			pipeline
-		});
-	}
-	dispatch(exe, inputs, outputs) {
-		if (inputs.length !== exe.kernel.nargs) throw new Error(`webgpu: dispatch with ${inputs.length} inputs, expected ${exe.kernel.nargs}`);
-		if (exe.kernel.size === 0) return;
-		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot).buffer);
-		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot).buffer);
-		pipelineSubmit(this.device, exe.data, inputBuffers, outputBuffers);
-	}
-	#getBuffer(slot) {
-		const buffer = this.buffers.get(slot);
-		if (!buffer) throw new require_backend.SlotError(slot);
-		return {
-			buffer: buffer.buffer,
-			size: buffer.size
-		};
-	}
-	/**
-	* 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 require_backend.DType.Bool: return storage ? "i32" : "bool";
-		case require_backend.DType.Int32: return "i32";
-		case require_backend.DType.Uint32: return "u32";
-		case require_backend.DType.Float32: return "f32";
-		case require_backend.DType.Float16: return "f16";
-		default: throw new Error(`Unsupported dtype for WebGPU: ${dtype}`);
-	}
-}
-function constToWgsl(dtype, value) {
-	if (dtype === require_backend.DType.Bool) return value ? "true" : "false";
-	if (dtype === require_backend.DType.Int32) return value.toString();
-	if (dtype === require_backend.DType.Uint32) return value.toString() + "u";
-	if (dtype === require_backend.DType.Float32) {
-		if (Number.isNaN(value)) return "nan()";
-		if (!Number.isFinite(value)) return value > 0 ? "inf()" : "-inf()";
-		return "f32(" + value.toString() + ")";
-	}
-	if (dtype === require_backend.DType.Float16) {
-		if (Number.isNaN(value)) return "f16(nan())";
-		if (!Number.isFinite(value)) return value > 0 ? "f16(inf())" : "f16(-inf())";
-		return "f16(" + value.toString() + ")";
-	}
-	throw new Error(`Unsupported const dtype: ${dtype}`);
-}
-/**
-* Compiles an expression into WebGPU shader source code.
-*
-* Returns the shader source and the number of workgroups to dispatch along x
-* and y axes, to run the kernel.
-*/
-function pipelineSource(device, kernel) {
-	const tune = require_backend.tuneWebgpu(kernel);
-	if (require_backend.DEBUG >= 3) console.info(`kernel.exp: ${kernel.exp}\ntune.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);
-	};
-	if (tune.exp.some((exp) => exp.dtype === require_backend.DType.Float16) || tune.epilogue?.some((exp) => exp.dtype === require_backend.DType.Float16)) {
-		if (!device.features.has("shader-f16")) throw new Error("WebGPU device does not support shader-f16 feature");
-		emit("enable f16;");
-	}
-	emit("fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }", "fn inf() -> f32 { let bits = 0x7f800000u; return bitcast<f32>(bits); }");
-	const distinctOps = require_backend.mapSetUnion(tune.exp.distinctOps(), tune.epilogue?.distinctOps());
-	if (distinctOps.has(require_backend.AluOp.Threefry2x32)) emit(threefrySrc);
-	if (distinctOps.has(require_backend.AluOp.Erf) || distinctOps.has(require_backend.AluOp.Erfc)) emit(erfSrc);
-	emit("");
-	const usedArgs = Array.from({ length: nargs }, () => null);
-	tune.exp.fold((exp) => {
-		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
-	});
-	tune.epilogue?.fold((exp) => {
-		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
-	});
-	for (let i = 0; i < nargs; i++) {
-		const ty = dtypeToWgsl(usedArgs[i] ?? require_backend.DType.Float32, true);
-		emit(`@group(0) @binding(${i}) var<storage, read> ${args[i]} : array<${ty}>;`);
-	}
-	const resultTy = dtypeToWgsl(kernel.dtype, true);
-	emit(`@group(0) @binding(${nargs}) var<storage, read_write> result : array<${resultTy}>;`);
-	const workgroupSize = require_backend.findPow2(tune.threadCount, 256);
-	const gridSize = Math.ceil(tune.threadCount / workgroupSize);
-	let gridX = gridSize;
-	let gridY = 1;
-	if (gridSize > device.limits.maxComputeWorkgroupsPerDimension) {
-		gridX = 16384;
-		gridY = Math.ceil(gridSize / gridX);
-	}
-	emit("", `@compute @workgroup_size(${workgroupSize})`, "fn main(@builtin(global_invocation_id) id : vec3<u32>) {", pushIndent);
-	if (gridY === 1) emit(`if (id.x >= ${tune.threadCount}) { return; }`, "let gidx: i32 = i32(id.x);");
-	else {
-		const sizeX = gridX * workgroupSize;
-		emit(`if (${sizeX} * id.y + id.x >= ${tune.threadCount}) { return; }`, `let gidx: i32 = i32(${sizeX} * id.y + id.x);`);
-	}
-	let gensymCount = 0;
-	const gensym = () => `alu${gensymCount++}`;
-	const isGensym = (text) => text.match(/^alu[0-9]+$/);
-	if (args.length > 0) emit(args.map((arg) => `_ = &${arg};`).join(" "));
-	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 (require_backend.AluGroup.Binary.has(op) || require_backend.AluGroup.Compare.has(op)) {
-			const a = gen(src[0]);
-			const b = gen(src[1]);
-			if (op === require_backend.AluOp.Add) if (dtype === require_backend.DType.Bool) source = `(${a} || ${b})`;
-			else source = `(${a} + ${b})`;
-			else if (op === require_backend.AluOp.Sub) source = `(${a} - ${b})`;
-			else if (op === require_backend.AluOp.Mul) if (dtype === require_backend.DType.Bool) source = `(${a} && ${b})`;
-			else source = `(${a} * ${b})`;
-			else if (op === require_backend.AluOp.Idiv) source = require_backend.isFloatDtype(dtype) ? `trunc(${a} / ${b})` : `(${a} / ${b})`;
-			else if (op === require_backend.AluOp.Mod) source = `(${a} % ${b})`;
-			else if (op === require_backend.AluOp.Min) source = `min(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
-			else if (op === require_backend.AluOp.Max) source = `max(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
-			else if (op === require_backend.AluOp.Cmplt) source = `(${a} < ${b})`;
-			else if (op === require_backend.AluOp.Cmpne) if (require_backend.isFloatDtype(src[0].dtype)) {
-				const x = isGensym(a) ? a : gensym();
-				if (x !== a) emit(`let ${x} = ${a};`);
-				source = `(${x} != ${b} || min(${x}, ${dtypeToWgsl(src[0].dtype)}(inf())) != ${x})`;
-			} else source = `(${a} != ${b})`;
-		} else if (require_backend.AluGroup.Unary.has(op)) if (op === require_backend.AluOp.Reciprocal && src[0].op === require_backend.AluOp.Sqrt) {
-			const a = gen(src[0].src[0]);
-			source = `inverseSqrt(${a})`;
-		} else {
-			const a = gen(src[0]);
-			if (op === require_backend.AluOp.Sin) source = `sin(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Cos) source = `cos(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Asin) source = `asin(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Atan) source = `atan(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Exp) source = `exp(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Log) source = `log(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Erf || op === require_backend.AluOp.Erfc) {
-				const funcName = op === require_backend.AluOp.Erf ? "erf" : "erfc";
-				if (dtype !== require_backend.DType.Float32) source = `${dtypeToWgsl(dtype)}(${funcName}(f32(${require_backend.strip1(a)})))`;
-				else source = `${funcName}(${require_backend.strip1(a)})`;
-			} else if (op === require_backend.AluOp.Sqrt) source = `sqrt(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Reciprocal) source = `(1.0 / ${a})`;
-			else if (op === require_backend.AluOp.Floor) source = `floor(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Ceil) source = `ceil(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Cast) source = `${dtypeToWgsl(dtype)}(${require_backend.strip1(a)})`;
-			else if (op === require_backend.AluOp.Bitcast) source = `bitcast<${dtypeToWgsl(dtype)}>(${require_backend.strip1(a)})`;
-		}
-		else if (op === require_backend.AluOp.Where) source = `select(${require_backend.strip1(gen(src[2]))}, ${require_backend.strip1(gen(src[1]))}, ${require_backend.strip1(gen(src[0]))})`;
-		else if (op === require_backend.AluOp.Threefry2x32) {
-			const x = gensym();
-			const [k0, k1, c0, c1] = src.map((x$1) => require_backend.strip1(gen(x$1)));
-			emit(`let ${x} = threefry2x32(vec2(${k0}, ${k1}), vec2(${c0}, ${c1}));`);
-			if (arg === "xor") source = `(${x}.x ^ ${x}.y)`;
-			else if (arg === 0) source = `${x}.x`;
-			else if (arg === 1) source = `${x}.y`;
-			else throw new require_backend.UnsupportedOpError(op, dtype, "webgpu", arg);
-		} else if (op === require_backend.AluOp.Const) return constToWgsl(dtype, arg);
-		else if (op === require_backend.AluOp.Special) return arg[0];
-		else if (op === require_backend.AluOp.Variable) return arg;
-		else if (op === require_backend.AluOp.GlobalIndex) {
-			source = `${args[arg[0]]}[${require_backend.strip1(gen(src[0]))}]`;
-			if (dtype === require_backend.DType.Bool) source = `(${source} != 0)`;
-		}
-		if (!source) throw new require_backend.UnsupportedOpError(op, dtype, "webgpu", arg);
-		const typeName = dtypeToWgsl(dtype);
-		if ((references.get(exp) ?? 0) > 1) {
-			const name = gensym();
-			expContext.set(exp, name);
-			emit(`let ${name}: ${typeName} = ${require_backend.strip1(source)};`);
-			return name;
-		} else {
-			expContext.set(exp, source);
-			return source;
-		}
-	};
-	if (!re) {
-		countReferences(tune.exp);
-		let rhs = require_backend.strip1(gen(tune.exp));
-		if (resultTy !== dtypeToWgsl(tune.exp.dtype)) rhs = `${resultTy}(${rhs})`;
-		emit(`result[gidx] = ${rhs};`);
-	} else {
-		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(re.dtype)} = ${constToWgsl(re.dtype, re.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: require_backend.AluExp.i32(up),
-					unroll: require_backend.AluExp.i32(un)
-				});
-				exps[up].push(exp.simplify(cache));
-				countReferences(exps[up][un]);
-			}
-		}
-		const items = exps.map((ar) => ar.map(gen).map(require_backend.strip1));
-		for (let i = 0; i < upcast; i++) {
-			let rhs = items[i][0];
-			for (let j = 1; j < unroll; j++) if (re.op === require_backend.AluOp.Add) rhs = `${rhs} + ${items[i][j]}`;
-			else if (re.op === require_backend.AluOp.Mul) rhs = `${rhs} * ${items[i][j]}`;
-			else if (re.op === require_backend.AluOp.Min) rhs = `min(${rhs}, ${items[i][j]})`;
-			else if (re.op === require_backend.AluOp.Max) rhs = `max(${rhs}, ${items[i][j]})`;
-			else throw new Error(`Unsupported reduction op: ${re.op}`);
-			if (re.op === require_backend.AluOp.Add) emit(`${acc[i]} += ${rhs};`);
-			else if (re.op === require_backend.AluOp.Mul) emit(`${acc[i]} *= ${rhs};`);
-			else if (re.op === require_backend.AluOp.Min) emit(`${acc[i]} = min(${acc[i]}, ${rhs});`);
-			else if (re.op === require_backend.AluOp.Max) emit(`${acc[i]} = max(${acc[i]}, ${rhs});`);
-			else throw new Error(`Unsupported reduction op: ${re.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: require_backend.AluExp.i32(i) });
-			outputIdxExps.push(exp.simplify(cache));
-			countReferences(outputIdxExps[i]);
-			fusionExps.push(tune.epilogue.substitute({
-				acc: require_backend.AluExp.variable(re.dtype, acc[i]),
-				upcast: require_backend.AluExp.i32(i)
-			}).simplify(cache));
-			countReferences(fusionExps[i]);
-		}
-		for (let i = 0; i < upcast; i++) {
-			const index = require_backend.strip1(gen(outputIdxExps[i]));
-			let rhs = require_backend.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: [gridX, gridY]
-	};
-}
-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()]);
-}
-/**
-* A cache for compiled GPU compute pipelines, keyed by the shader source.
-*
-* This supports both async compilation (recommended) and a synchronous variant.
-* If the pipeline is not in the cache, it will be compiled and added. For async
-* compilation, only one compilation will be in progress at a time for a given
-* shader source.
-*/
-var ShaderPipelineCache = class {
-	cache;
-	inProgress;
-	constructor(device) {
-		this.device = device;
-		this.cache = /* @__PURE__ */ new Map();
-		this.inProgress = /* @__PURE__ */ new Map();
-	}
-	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 (require_backend.DEBUG >= 2) console.info("=========== WebGPU shader ===========\n" + code);
-		const shaderModule = this.device.createShaderModule({ code });
-		const promise = (async () => {
-			this.device.pushErrorScope("validation");
-			try {
-				const pipeline$1 = await this.device.createComputePipelineAsync({
-					layout: "auto",
-					compute: {
-						module: shaderModule,
-						entryPoint: "main"
-					}
-				});
-				await this.device.popErrorScope();
-				return pipeline$1;
-			} 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 (require_backend.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;
-	}
-};
-/** Gather information about a compilation error and format it. */
-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 += `\n  [${msg.type} at ${msg.lineNum}:${msg.linePos}] ${msg.message}`;
-	if (code) message += `\n\n${code}`;
-	return message;
-}
-
-//#endregion
-exports.WebGPUBackend = WebGPUBackend;
-//# sourceMappingURL=webgpu-BVns4DbI.cjs.map