@jax-js/jax 0.1.10 → 0.1.12

package/dist/{webgpu-DMSx7a6M.cjs → webgpu-C4S8Uq9e.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-DMauYnfl.cjs');
+const require_backend = require('./backend-x-6vqzIM.cjs');
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -72,6 +72,45 @@ const headerWgsl = String.raw`
 fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }
 fn inf() -> f32 { let bits = 0x7f800000u; return bitcast<f32>(bits); }
 `.trim();
+/** Builder class for simple program generation with WGSL. */
+var WgslBuilder = class {
+	pushIndent = Symbol("pushIndent");
+	popIndent = Symbol("popIndent");
+	lines = [];
+	#indent = "";
+	emit(...lines) {
+		for (const line of lines) if (line === this.pushIndent) this.#indent += "  ";
+		else if (line === this.popIndent) this.#indent = this.#indent.slice(0, -2);
+		else this.lines.push(line ? this.#indent + line : "");
+	}
+	emitPreamble(device, exps) {
+		let hasFloat16 = false;
+		let distinctOps = /* @__PURE__ */ new Map();
+		for (const exp of exps) {
+			if (exp == null) continue;
+			hasFloat16 ||= exp.some((e) => e.dtype === require_backend.DType.Float16);
+			distinctOps = require_backend.mapSetUnion(distinctOps, exp.distinctOps());
+		}
+		if (hasFloat16) {
+			if (!device.features.has("shader-f16")) throw new Error("WebGPU device does not support shader-f16 feature");
+			this.emit("enable f16;");
+		}
+		this.emit(headerWgsl);
+		if (distinctOps.has(require_backend.AluOp.Threefry2x32)) this.emit(threefrySrc);
+		if (distinctOps.has(require_backend.AluOp.Erf) || distinctOps.has(require_backend.AluOp.Erfc)) this.emit(erfSrc);
+		this.emit("");
+	}
+	/**
+	* Insert phony assignments, in case some inputs are not in use.
+	* <https://github.com/gpuweb/gpuweb/discussions/4582#discussioncomment-9146686>
+	*/
+	emitPhonyAssignments(args) {
+		if (args.length > 0) this.emit(args.map((arg) => `_ = &${arg};`).join(" "));
+	}
+	toString() {
+		return this.lines.join("\n");
+	}
+};
 function dtypeToWgsl(dtype, storage = false) {
 	switch (dtype) {
 		case require_backend.DType.Bool: return storage ? "i32" : "bool";
@@ -108,6 +147,136 @@ function constToWgsl(dtype, value) {
 	}
 	throw new Error(`Unsupported const dtype: ${dtype}`);
 }
+/** Codegen for WebGPU expressions, linearizing AluOp into a kernel. */
+var WgslExpCodegen = class {
+	#gensymCount = 0;
+	#references = /* @__PURE__ */ new Map();
+	#seen = /* @__PURE__ */ new Set();
+	#context = /* @__PURE__ */ new Map();
+	constructor(wb, args) {
+		this.wb = wb;
+		this.args = args;
+	}
+	#gensym() {
+		return `alu${this.#gensymCount++}`;
+	}
+	#isGensym(text) {
+		return text.match(/^alu[0-9]+$/);
+	}
+	/**
+	* Count references for an expression.
+	*
+	* Used to get an ahead-of-time reference count for each node in the AluExp.
+	* Expressions with reference count greater than 1 are stored in temporary
+	* variables to avoid recomputation.
+	*/
+	countReferences(exp) {
+		this.#references.set(exp, (this.#references.get(exp) ?? 0) + 1);
+		if (!this.#seen.has(exp)) {
+			this.#seen.add(exp);
+			for (const src of exp.src) this.countReferences(src);
+		}
+	}
+	reset() {
+		this.#references.clear();
+		this.#seen.clear();
+		this.#context.clear();
+	}
+	/**
+	* Generate code for an expression.
+	*
+	* Calls itself recursively and eliminates common subexpressions by storing
+	* them in temporary variables, emitted to the current builder scope. This is
+	* a side-effect that leads to multiline code generation.
+	*/
+	run(exp) {
+		if (this.#context.has(exp)) return this.#context.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 = this.run(src[0]);
+			const b = this.run(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) if (dtype === require_backend.DType.Bool) source = `(${a} && ${b})`;
+			else source = `min(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
+			else if (op === require_backend.AluOp.Max) if (dtype === require_backend.DType.Bool) source = `(${a} || ${b})`;
+			else source = `max(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
+			else if (op === require_backend.AluOp.BitCombine) if (arg === "and") source = `(${a} & ${b})`;
+			else if (arg === "or") source = `(${a} | ${b})`;
+			else source = dtype === require_backend.DType.Bool ? `(${a} != ${b})` : `(${a} ^ ${b})`;
+			else if (op === require_backend.AluOp.BitShift) if (arg === "shl") source = `(${a} << ${b})`;
+			else source = `(${a} >> ${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 = this.#isGensym(a) ? a : this.#gensym();
+				if (x !== a) this.wb.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 = this.run(src[0].src[0]);
+			source = `inverseSqrt(${a})`;
+		} else {
+			const a = this.run(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) {
+				const srcTy = dtypeToWgsl(src[0].dtype);
+				const dstTy = dtypeToWgsl(dtype);
+				if (require_backend.isFloatDtype(src[0].dtype) && !(require_backend.isFloatDtype(dtype) || dtype === require_backend.DType.Bool)) {
+					const maxVal = maxValueWgsl(dtype);
+					const x = this.#isGensym(a) ? a : this.#gensym();
+					if (x !== a) this.wb.emit(`let ${x}: ${srcTy} = ${require_backend.strip1(a)};`);
+					source = `select(${dstTy}(${x}), ${maxVal}, ${x} >= ${srcTy}(${maxVal}))`;
+				} else source = `${dstTy}(${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(this.run(src[2]))}, ${require_backend.strip1(this.run(src[1]))}, ${require_backend.strip1(this.run(src[0]))})`;
+		else if (op === require_backend.AluOp.Threefry2x32) {
+			const x = this.#gensym();
+			const [k0, k1, c0, c1] = src.map((x$1) => require_backend.strip1(this.run(x$1)));
+			this.wb.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 = `${this.args[arg[0]]}[${require_backend.strip1(this.run(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 ((this.#references.get(exp) ?? 0) > 1) {
+			const name = this.#gensym();
+			this.#context.set(exp, name);
+			this.wb.emit(`let ${name}: ${typeName} = ${require_backend.strip1(source)};`);
+			return name;
+		} else {
+			this.#context.set(exp, source);
+			return source;
+		}
+	}
+};
 const gridOffsetY = 16384;
 function calculateGrid(gridSize) {
 	let gridX = gridSize;
@@ -119,6 +288,91 @@ function calculateGrid(gridSize) {
 	return [gridX, gridY];
 }
 
+//#endregion
+//#region src/backend/webgpu/nullaryKernel.ts
+function uniformDtype(dtype) {
+	if (dtype === require_backend.DType.Float16) return require_backend.DType.Float32;
+	if (dtype === require_backend.DType.Bool) return require_backend.DType.Int32;
+	return dtype;
+}
+function replacementFor({ name, dtype, uniformDtype: uniformDtype$1 }) {
+	const value = require_backend.AluExp.variable(uniformDtype$1, `uniforms.${name}`);
+	if (dtype === require_backend.DType.Float16) return require_backend.AluExp.cast(require_backend.DType.Float16, value);
+	if (dtype === require_backend.DType.Bool) return require_backend.AluExp.cmpne(value, require_backend.AluExp.i32(0));
+	return value;
+}
+function writeUniform(view, offset, dtype, value) {
+	switch (dtype) {
+		case require_backend.DType.Float32:
+			view.setFloat32(offset, value, true);
+			break;
+		case require_backend.DType.Int32:
+			view.setInt32(offset, value, true);
+			break;
+		case require_backend.DType.Uint32:
+			view.setUint32(offset, value, true);
+			break;
+		default: throw new Error(`Unsupported dtype for constant uniform: ${dtype}`);
+	}
+}
+function liftConstants(exp) {
+	const uniforms = [];
+	const lifted = exp.rewrite((node) => {
+		if (node.op !== require_backend.AluOp.Const || node.arg === 0) return;
+		const uniform = {
+			name: `c${uniforms.length}`,
+			dtype: node.dtype,
+			uniformDtype: uniformDtype(node.dtype),
+			value: node.arg
+		};
+		uniforms.push(uniform);
+		return replacementFor(uniform);
+	});
+	return [lifted, uniforms];
+}
+function uniformsData(uniforms) {
+	const data = new Uint8Array(uniforms.length * 4);
+	const view = new DataView(data.buffer);
+	uniforms.forEach((u, i) => writeUniform(view, i * 4, u.uniformDtype, u.value));
+	return data;
+}
+function nullaryKernelSource(device, kernel) {
+	if (kernel.nargs !== 0 || kernel.reduction) return null;
+	let exp = kernel.exp.substitute({ gidx: require_backend.AluExp.special(require_backend.DType.Int32, "gidx", kernel.size) }).simplify();
+	let uniforms = [];
+	[exp, uniforms] = liftConstants(exp);
+	const wb = new WgslBuilder();
+	wb.emitPreamble(device, [exp]);
+	if (uniforms.length > 0) wb.emit("struct Uniforms {", wb.pushIndent, ...uniforms.map((u) => `${u.name}: ${dtypeToWgsl(u.uniformDtype)},`), wb.popIndent, "}\n");
+	const resultTy = dtypeToWgsl(kernel.dtype, true);
+	wb.emit(`@group(0) @binding(0) var<storage, read_write> result : array<${resultTy}>;`);
+	if (uniforms.length > 0) wb.emit(`@group(1) @binding(0) var<uniform> uniforms: Uniforms;`);
+	const workgroupSize = require_backend.findPow2(kernel.size, 256);
+	const gridSize = Math.ceil(kernel.size / workgroupSize);
+	const [gridX, gridY] = calculateGrid(gridSize);
+	wb.emit("", `@compute @workgroup_size(${workgroupSize})`, "fn main(@builtin(global_invocation_id) id : vec3<u32>) {", wb.pushIndent);
+	if (gridY === 1) wb.emit(`if (id.x >= ${kernel.size}) { return; }`, "let gidx: i32 = i32(id.x);");
+	else {
+		const sizeX = gridX * workgroupSize;
+		wb.emit(`if (${sizeX} * id.y + id.x >= ${kernel.size}) { return; }`, `let gidx: i32 = i32(${sizeX} * id.y + id.x);`);
+	}
+	const gen = new WgslExpCodegen(wb, []);
+	gen.countReferences(exp);
+	let rhs = require_backend.strip1(gen.run(exp));
+	if (resultTy !== dtypeToWgsl(exp.dtype)) rhs = `${resultTy}(${rhs})`;
+	wb.emit(`result[gidx] = ${rhs};`, wb.popIndent, "}");
+	return {
+		code: wb.toString(),
+		numInputs: 0,
+		numOutputs: 1,
+		hasUniform: uniforms.length > 0,
+		passes: [{
+			grid: [gridX, gridY],
+			uniform: uniforms.length > 0 ? uniformsData(uniforms) : void 0
+		}]
+	};
+}
+
 //#endregion
 //#region src/backend/webgpu/reader.ts
 /**
@@ -1024,28 +1278,14 @@ var WebGPUBackend = class {
 * and y axes, to run the kernel.
 */
 function pipelineSource(device, kernel) {
+	const nullaryKernel = nullaryKernelSource(device, kernel);
+	if (nullaryKernel) return nullaryKernel;
 	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(headerWgsl);
-	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 wb = new WgslBuilder();
+	wb.emitPreamble(device, [tune.exp, tune.epilogue]);
 	const usedArgs = Array.from({ length: nargs }, () => null);
 	tune.exp.fold((exp) => {
 		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
@@ -1055,127 +1295,33 @@ function pipelineSource(device, kernel) {
 	});
 	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}>;`);
+		wb.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}>;`);
+	wb.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);
 	const [gridX, gridY] = calculateGrid(gridSize);
-	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);");
+	wb.emit("", `@compute @workgroup_size(${workgroupSize})`, "fn main(@builtin(global_invocation_id) id : vec3<u32>) {", wb.pushIndent);
+	if (gridY === 1) wb.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);`);
+		wb.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) if (dtype === require_backend.DType.Bool) source = `(${a} && ${b})`;
-			else source = `min(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
-			else if (op === require_backend.AluOp.Max) if (dtype === require_backend.DType.Bool) source = `(${a} || ${b})`;
-			else 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) {
-				const srcTy = dtypeToWgsl(src[0].dtype);
-				const dstTy = dtypeToWgsl(dtype);
-				if (require_backend.isFloatDtype(src[0].dtype) && !(require_backend.isFloatDtype(dtype) || dtype === require_backend.DType.Bool)) {
-					const maxVal = maxValueWgsl(dtype);
-					const x = isGensym(a) ? a : gensym();
-					if (x !== a) emit(`let ${x}: ${srcTy} = ${require_backend.strip1(a)};`);
-					source = `select(${dstTy}(${x}), ${maxVal}, ${x} >= ${srcTy}(${maxVal}))`;
-				} else source = `${dstTy}(${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;
-		}
-	};
+	wb.emitPhonyAssignments(args);
+	const gen = new WgslExpCodegen(wb, args);
 	if (!re) {
-		countReferences(tune.exp);
-		let rhs = require_backend.strip1(gen(tune.exp));
+		gen.countReferences(tune.exp);
+		let rhs = require_backend.strip1(gen.run(tune.exp));
 		if (resultTy !== dtypeToWgsl(tune.exp.dtype)) rhs = `${resultTy}(${rhs})`;
-		emit(`result[gidx] = ${rhs};`);
+		wb.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);
+		for (let i = 0; i < upcast; i++) wb.emit(`var ${acc[i]}: ${dtypeToWgsl(re.dtype)} = ${constToWgsl(re.dtype, re.identity)};`);
+		wb.emit(`for (var ridx: i32 = 0; ridx < ${tune.size.reduce}; ridx++) {`, wb.pushIndent);
 		const exps = [];
 		const cache = /* @__PURE__ */ new Map();
 		for (let up = 0; up < upcast; up++) {
@@ -1186,10 +1332,10 @@ function pipelineSource(device, kernel) {
 					unroll: require_backend.AluExp.i32(un)
 				});
 				exps[up].push(exp.simplify(cache));
-				countReferences(exps[up][un]);
+				gen.countReferences(exps[up][un]);
 			}
 		}
-		const items = exps.map((ar) => ar.map(gen).map(require_backend.strip1));
+		const items = exps.map((ar) => ar.map((x) => gen.run(x)).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]}`;
@@ -1197,40 +1343,38 @@ function pipelineSource(device, kernel) {
 			else if (re.op === require_backend.AluOp.Min) rhs = re.dtype === require_backend.DType.Bool ? `(${rhs} && ${items[i][j]})` : `min(${rhs}, ${items[i][j]})`;
 			else if (re.op === require_backend.AluOp.Max) rhs = re.dtype === require_backend.DType.Bool ? `(${rhs} || ${items[i][j]})` : `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) if (re.dtype === require_backend.DType.Bool) emit(`${acc[i]} = ${acc[i]} && ${rhs};`);
-			else emit(`${acc[i]} = min(${acc[i]}, ${rhs});`);
-			else if (re.op === require_backend.AluOp.Max) if (re.dtype === require_backend.DType.Bool) emit(`${acc[i]} = ${acc[i]} || ${rhs};`);
-			else emit(`${acc[i]} = max(${acc[i]}, ${rhs});`);
+			if (re.op === require_backend.AluOp.Add) wb.emit(`${acc[i]} += ${rhs};`);
+			else if (re.op === require_backend.AluOp.Mul) wb.emit(`${acc[i]} *= ${rhs};`);
+			else if (re.op === require_backend.AluOp.Min) if (re.dtype === require_backend.DType.Bool) wb.emit(`${acc[i]} = ${acc[i]} && ${rhs};`);
+			else wb.emit(`${acc[i]} = min(${acc[i]}, ${rhs});`);
+			else if (re.op === require_backend.AluOp.Max) if (re.dtype === require_backend.DType.Bool) wb.emit(`${acc[i]} = ${acc[i]} || ${rhs};`);
+			else wb.emit(`${acc[i]} = max(${acc[i]}, ${rhs});`);
 			else throw new Error(`Unsupported reduction op: ${re.op}`);
 		}
-		emit(popIndent, "}");
-		expContext.clear();
-		references.clear();
-		seen.clear();
+		wb.emit(wb.popIndent, "}");
+		gen.reset();
 		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]);
+			gen.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]);
+			gen.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]));
+			const index = require_backend.strip1(gen.run(outputIdxExps[i]));
+			let rhs = require_backend.strip1(gen.run(fusionExps[i]));
 			if (resultTy !== dtypeToWgsl(fusionExps[i].dtype)) rhs = `${resultTy}(${rhs})`;
-			emit(`result[${index}] = ${rhs};`);
+			wb.emit(`result[${index}] = ${rhs};`);
 		}
 	}
-	emit(popIndent, "}");
+	wb.emit(wb.popIndent, "}");
 	return {
-		code: shader.join("\n"),
+		code: wb.toString(),
 		numInputs: nargs,
 		numOutputs: 1,
 		hasUniform: false,
@@ -1274,11 +1418,17 @@ function pipelineSubmit(device, exe, inputs, outputs) {
 		}
 		for (let i = 0; i < filteredPasses.length; i++) {
 			const { grid } = filteredPasses[i];
-			const passEncoder = commandEncoder.beginComputePass({ timestampWrites: slot ? {
-				querySet: slot.batch.querySet,
-				beginningOfPassWriteIndex: i === 0 ? slot.beginIndex : void 0,
-				endOfPassWriteIndex: i === filteredPasses.length - 1 ? slot.endIndex : void 0
-			} : void 0 });
+			let timestampWrites;
+			if (slot) {
+				const isFirst = i === 0;
+				const isLast = i === filteredPasses.length - 1;
+				if (isFirst || isLast) timestampWrites = {
+					querySet: slot.batch.querySet,
+					...isFirst ? { beginningOfPassWriteIndex: slot.beginIndex } : {},
+					...isLast ? { endOfPassWriteIndex: slot.endIndex } : {}
+				};
+			}
+			const passEncoder = commandEncoder.beginComputePass({ timestampWrites });
 			passEncoder.setPipeline(pipeline);
 			passEncoder.setBindGroup(0, bindGroup);
 			if (uniformBindGroup) passEncoder.setBindGroup(1, uniformBindGroup, [i * uniformAlignment]);

package/package.json

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