@jax-js/jax 0.1.2 → 0.1.4

package/dist/{backend-BqymqzuU.js → backend-tngXtWe4.js}

@@ -68,6 +68,9 @@ function zipn(...arrays) {
 	const minLength = Math.min(...arrays.map((x) => x.length));
 	return Array.from({ length: minLength }, (_, i) => arrays.map((arr) => arr[i]));
 }
+function sorted(arr) {
+	return [...arr].sort((a, b) => a - b);
+}
 function rep(length, value) {
 	if (value instanceof Function) return new Array(length).fill(0).map((_, i) => value(i));
 	return new Array(length).fill(value);
@@ -144,7 +147,7 @@ function normalizeAxis(axis, ndim) {
 			if (seen.has(ca)) throw new Error(`Duplicate axis ${ca} passed to function`);
 			seen.add(ca);
 		}
-		return [...seen].sort();
+		return sorted(seen);
 	}
 }
 function range(start, stop, step = 1) {
@@ -557,16 +560,16 @@ var AluExp = class AluExp {
 		});
 	}
 	/** Reindex gid values in this expression as needed. */
-	reindexGids(gidMap) {
+	reindexGids(newGids) {
 		return this.rewrite((exp) => {
 			if (exp.op === AluOp.GlobalIndex) {
 				const [gid, len] = exp.arg;
-				const newGid = gidMap.get(gid);
-				if (newGid !== void 0 && newGid !== gid) return AluExp.globalIndex(exp.dtype, newGid, len, exp.src[0]);
+				const newGid = newGids[gid];
+				if (newGid !== gid) return AluExp.globalIndex(exp.dtype, newGid, len, exp.src[0]);
 			} else if (exp.op === AluOp.GlobalView) {
 				const gid = exp.arg[0];
-				const newGid = gidMap.get(gid);
-				if (newGid !== void 0 && newGid !== gid) return AluExp.globalView(exp.dtype, newGid, exp.arg[1], exp.src);
+				const newGid = newGids[gid];
+				if (newGid !== gid) return AluExp.globalView(exp.dtype, newGid, exp.arg[1], exp.src);
 			}
 		});
 	}
@@ -780,7 +783,7 @@ var AluExp = class AluExp {
 			if (op === AluOp.Sub && i === 1 && x === 0) return src[1 - i];
 			if (op === AluOp.Mul && x === 1) return src[1 - i];
 			if (op === AluOp.Mul && x === 0) return AluExp.const(this.dtype, 0);
-			if (op === AluOp.Idiv && i === 1 && x === 1) return src[1 - i];
+			if (op === AluOp.Idiv && i === 1 && x === 1 && !isFloatDtype(this.dtype)) return src[1 - i];
 			if (op === AluOp.Cmpne && src[i].dtype === DType.Bool && x === 0) return src[1 - i];
 		}
 		if ((op === AluOp.Add || op === AluOp.Sub) && src[1].op === AluOp.Mul) {
@@ -1327,7 +1330,7 @@ var Reduction = class {
 	/** Evaluate this operation on CPU. */
 	evaluate(...values) {
 		if (this.dtype === DType.Bool) {
-			if (this.op === AluOp.Add || this.op === AluOp.Max) return values.reduce((a, b) => a || b, true);
+			if (this.op === AluOp.Add || this.op === AluOp.Max) return values.reduce((a, b) => a || b, false);
 			else if (this.op === AluOp.Mul || this.op === AluOp.Min) return values.reduce((a, b) => a && b, true);
 		} else if (this.dtype === DType.Int32) {
 			if (this.op === AluOp.Add) return values.reduce((a, b) => a + b | 0, 0);
@@ -1437,6 +1440,112 @@ function erfc(x) {
 	else return 2 - _erfapprox$1(-x);
 }
 
+//#endregion
+//#region src/routine.ts
+/**
+* Advanced operations that don't fit into the `AluExp` compiler representation.
+*
+* Some routines like iterative matrix algorithms, FFTs, or sorting may not be
+* easy to express efficiently as a `Kernel` object. These also tend to be
+* somewhat expensive, so the benefit of kernel fusion and inlining is less
+* relevant.
+*
+* For these operations, we dispatch them as a custom operation on the backend,
+* which each backend implements in a specific way. These are listed in the
+* `Routines` enum below.
+*
+* Routines cannot be fused into other kernels and always operate on contiguous
+* arrays (default `ShapeTracker`).
+*/
+var Routine = class {
+	constructor(name, type, params) {
+		this.name = name;
+		this.type = type;
+		this.params = params;
+	}
+};
+/** One of the valid `Routine` that can be dispatched to backend. */
+let Routines = /* @__PURE__ */ function(Routines$1) {
+	/** Stable sorting algorithm along the last axis. */
+	Routines$1["Sort"] = "Sort";
+	/** Returns `int32` indices of the stably sorted array. */
+	Routines$1["Argsort"] = "Argsort";
+	/** Solve a triangular system of questions. */
+	Routines$1["TriangularSolve"] = "TriangularSolve";
+	/** Cholesky decomposition of 2D positive semi-definite matrices. */
+	Routines$1["Cholesky"] = "Cholesky";
+	return Routines$1;
+}({});
+function runCpuRoutine(routine, inputs, outputs) {
+	const { name, type } = routine;
+	const inputAr = inputs.map((buf, i) => dtypedArray(type.inputDtypes[i], buf));
+	const outputAr = outputs.map((buf, i) => dtypedArray(type.outputDtypes[i], buf));
+	switch (name) {
+		case Routines.Sort: return runSort(type, inputAr, outputAr);
+		case Routines.Argsort: return runArgsort(type, inputAr, outputAr);
+		case Routines.TriangularSolve: return runTriangularSolve(type, inputAr, outputAr, routine.params);
+		case Routines.Cholesky: return runCholesky(type, inputAr, outputAr);
+		default:
+	}
+}
+function runSort(type, [x], [y]) {
+	const xs = type.inputShapes[0];
+	if (xs.length === 0) throw new Error("sort: cannot sort a scalar");
+	const n = xs[xs.length - 1];
+	y.set(x);
+	for (let i = 0; i < y.length; i += n) y.subarray(i, i + n).sort();
+}
+function runArgsort(type, [x], [y, yi]) {
+	const xs = type.inputShapes[0];
+	if (xs.length === 0) throw new Error("argsort: cannot sort a scalar");
+	const n = xs[xs.length - 1];
+	for (let offset = 0; offset < y.length; offset += n) {
+		const ar = x.subarray(offset, offset + n);
+		const out = y.subarray(offset, offset + n);
+		const outi = yi.subarray(offset, offset + n);
+		for (let i = 0; i < n; i++) outi[i] = i;
+		outi.sort((a, b) => ar[a] - ar[b]);
+		for (let i = 0; i < n; i++) out[i] = ar[outi[i]];
+	}
+}
+function runTriangularSolve(type, [a, b], [x], { unitDiagonal }) {
+	const as = type.inputShapes[0];
+	const bs = type.inputShapes[1];
+	if (as.length < 2) throw new Error(`triangular_solve: a must be at least 2D, got ${as}`);
+	if (bs.length < 2) throw new Error(`triangular_solve: b must be at least 2D, got ${bs}`);
+	const n = as[as.length - 2];
+	if (n !== as[as.length - 1] || n !== bs[bs.length - 1]) throw new Error(`triangular_solve: incompatible shapes a=${as}, b=${bs}`);
+	const batch = bs[bs.length - 2];
+	for (let counter = 0; counter < a.length / (n * n); counter++) {
+		const a1 = a.subarray(counter * n * n, (counter + 1) * n * n);
+		for (let t = 0; t < batch; t++) {
+			const b1 = b.subarray((counter * batch + t) * n, (counter * batch + t + 1) * n);
+			const x1 = x.subarray((counter * batch + t) * n, (counter * batch + t + 1) * n);
+			for (let i = n - 1; i >= 0; i--) {
+				let sum = b1[i];
+				for (let j = i + 1; j < n; j++) sum -= a1[i * n + j] * x1[j];
+				x1[i] = unitDiagonal ? sum : sum / a1[i * n + i];
+			}
+		}
+	}
+}
+function runCholesky(type, [x], [y]) {
+	const xs = type.inputShapes[0];
+	if (xs.length < 2) throw new Error("cholesky: input must be at least 2D");
+	const n = xs[xs.length - 2];
+	const m = xs[xs.length - 1];
+	if (n !== m) throw new Error(`cholesky: input must be square, got [${n}, ${m}]`);
+	for (let offset = 0; offset < y.length; offset += n * n) {
+		const ar = x.subarray(offset, offset + n * n);
+		const out = y.subarray(offset, offset + n * n);
+		for (let i = 0; i < n; i++) for (let j = 0; j <= i; j++) {
+			let sum = ar[i * n + j];
+			for (let k = 0; k < j; k++) sum -= out[i * n + k] * out[j * n + k];
+			out[i * n + j] = i === j ? Math.sqrt(sum) : sum / out[j * n + j];
+		}
+	}
+}
+
 //#endregion
 //#region src/shape.ts
 const jstr = JSON.stringify;
@@ -1907,7 +2016,7 @@ var ShapeTracker = class ShapeTracker {
 		let st = this;
 		if (axis.length > 0) {
 			const unsqueezed = [...st.shape];
-			for (const i of axis.toSorted()) unsqueezed.splice(i, 0, 1);
+			for (const i of sorted(axis)) unsqueezed.splice(i, 0, 1);
 			st = st.reshape(unsqueezed);
 		}
 		return st.expand(newShape);
@@ -2066,7 +2175,8 @@ function tuneNullopt(kernel) {
 	if (kernel.reduction) vars.ridx = AluExp.special(DType.Int32, "ridx", kernel.reduction.size);
 	return {
 		exp: kernel.exp.substitute(vars).rewriteGlobalViews().simplify(),
-		outputIdxExp: AluExp.special(DType.Int32, "gidx", kernel.size),
+		epilogue: kernel.reduction?.epilogue.substitute({ gidx: vars.gidx }).rewriteGlobalViews().simplify(),
+		outputIdxExp: vars.gidx,
 		threadCount: kernel.size,
 		size: { reduce: kernel.reduction ? kernel.reduction.size : 0 }
 	};
@@ -2099,7 +2209,11 @@ function tuneWebgpu(kernel) {
 	while (prod(dim.st.shape.slice(0, dim.groups)) >= 1024) {
 		const choices = [];
 		const composedSts = sts.map((st) => st.compose(dim.st));
-		for (let axis = 0; axis < dim.groups; axis++) for (const amount of [3, 4]) if (!upcastedAxis.has(axis) && dim.st.shape[axis] % amount === 0 && composedSts.some((st) => st.lastStrides[axis] === 0 && st.lastStrides.slice(dim.unroll).every((stride) => stride > 0))) {
+		for (let axis = 0; axis < dim.groups; axis++) for (const amount of [
+			3,
+			4,
+			5
+		]) if (!upcastedAxis.has(axis) && dim.st.shape[axis] % amount === 0 && composedSts.some((st) => st.lastStrides[axis] === 0 && st.lastStrides.slice(dim.unroll).every((stride) => stride > 0))) {
 			let nonzeroStrides = 0;
 			let totalStrides = 0;
 			for (const st of composedSts) {
@@ -2127,7 +2241,7 @@ function tuneWebgpu(kernel) {
 			break;
 		}
 	}
-	for (const ax of Array.from(upcastedAxis).sort()) {
+	for (const ax of sorted(upcastedAxis)) {
 		const s = dim.st.shape[ax];
 		for (const amount of [8, 4]) if (s % amount === 0) {
 			dim.applyLocal(ax, amount);
@@ -2175,7 +2289,15 @@ function tuneWebgpu(kernel) {
 	});
 	const outputGidx = dim.outputSt.shape.slice(0, dim.groups);
 	const outputUpcast = dim.outputSt.shape.slice(dim.groups);
-	const [outputIdxExp, _] = dim.outputSt.toAluExp([...unravelAlu(outputGidx, AluExp.special(DType.Int32, "gidx", prod(outputGidx))), ...unravelAlu(outputUpcast, AluVar.upcast)]);
+	const outputIndices = [...unravelAlu(outputGidx, AluExp.special(DType.Int32, "gidx", prod(outputGidx))), ...unravelAlu(outputUpcast, AluVar.upcast)];
+	const [outputIdxExp, _] = dim.outputSt.toAluExp(outputIndices);
+	const newEpilogue = reduction.epilogue.rewrite((exp$1) => {
+		if (exp$1.op === AluOp.GlobalView) {
+			const gid = exp$1.arg[0];
+			const st = exp$1.arg[1];
+			return accessorGlobal(exp$1.dtype, gid, st.compose(dim.outputSt), outputIndices);
+		}
+	});
 	if (prod(dim.st.shape.slice(dim.groups, dim.upcast)) !== reduction.size) throw new Error(`Invariant violation: reduction size ${reduction.size} does not match tuned dims ${JSON.stringify(dim.st.shape.slice(dim.groups, dim.upcast))}`);
 	const size = {
 		groups: prod(dim.st.shape.slice(dim.groups, dim.reduce)),
@@ -2185,6 +2307,7 @@ function tuneWebgpu(kernel) {
 	};
 	return {
 		exp: newExp.simplify(),
+		epilogue: newEpilogue.simplify(),
 		outputIdxExp: outputIdxExp.simplify(),
 		threadCount: kernel.size / size.upcast * size.groups,
 		size
@@ -2236,17 +2359,25 @@ var CpuBackend = class {
 		if (count === void 0) count = buffer.byteLength - start;
 		return buffer.slice(start, start + count);
 	}
-	async prepare(kernel) {
-		return this.prepareSync(kernel);
+	async prepareKernel(kernel) {
+		return this.prepareKernelSync(kernel);
 	}
-	prepareSync(kernel) {
+	prepareKernelSync(kernel) {
 		return new Executable(kernel, void 0);
 	}
-	dispatch({ kernel }, inputs, outputs) {
-		const { exp } = tuneNullopt(kernel);
+	async prepareRoutine(routine) {
+		return this.prepareRoutineSync(routine);
+	}
+	prepareRoutineSync(routine) {
+		return new Executable(routine, void 0);
+	}
+	dispatch(exe, inputs, outputs) {
+		if (exe.source instanceof Routine) return runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
+		const kernel = exe.source;
+		const { exp, epilogue } = tuneNullopt(kernel);
 		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
 		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
-		const usedArgs = new Map(exp.collect((exp$1) => exp$1.op === AluOp.GlobalIndex).map((exp$1) => [exp$1.arg[0], exp$1.dtype]));
+		const usedArgs = new Map([...exp.collect((exp$1) => exp$1.op === AluOp.GlobalIndex), ...epilogue ? epilogue.collect((exp$1) => exp$1.op === AluOp.GlobalIndex) : []].map((exp$1) => [exp$1.arg[0], exp$1.dtype]));
 		const inputArrays = inputBuffers.map((buf, i) => {
 			const dtype = usedArgs.get(i);
 			if (!dtype) return null;
@@ -2268,7 +2399,10 @@ var CpuBackend = class {
 				}, globals);
 				acc = kernel.reduction.evaluate(acc, item);
 			}
-			outputArray[i] = kernel.reduction.epilogue.evaluate({ acc });
+			outputArray[i] = epilogue.evaluate({
+				acc,
+				gidx: i
+			}, globals);
 		}
 	}
 	#getBuffer(slot) {
@@ -2297,8 +2431,10 @@ var WasmAllocator = class {
 		const sizeClass = this.#findSizeClass(size);
 		const freeList = this.#freeLists.get(sizeClass);
 		let ptr;
-		if (freeList && freeList.length > 0) ptr = freeList.pop();
-		else ptr = this.#bumpAlloc(sizeClass);
+		if (freeList && freeList.length > 0) {
+			ptr = freeList.pop();
+			new Uint8Array(this.#memory.buffer, ptr, sizeClass).fill(0);
+		} else ptr = this.#bumpAlloc(sizeClass);
 		this.#allocatedBuffers.set(ptr, sizeClass);
 		return ptr;
 	}
@@ -2431,7 +2567,7 @@ function wasm_log(cg) {
 		const t2 = cg.local.declare(cg.f32);
 		cg.local.get(0);
 		cg.f32.const(0);
-		cg.f32.le();
+		cg.f32.lt();
 		cg.if(cg.void);
 		cg.f32.const(NaN);
 		cg.return();
@@ -2446,6 +2582,20 @@ function wasm_log(cg) {
 		cg.i32.const(127);
 		cg.i32.sub();
 		cg.local.set(e);
+		cg.local.get(e);
+		cg.i32.const(-127);
+		cg.i32.eq();
+		cg.if(cg.void);
+		cg.f32.const(-Infinity);
+		cg.return();
+		cg.end();
+		cg.local.get(e);
+		cg.i32.const(128);
+		cg.i32.eq();
+		cg.if(cg.void);
+		cg.local.get(0);
+		cg.return();
+		cg.end();
 		cg.local.get(bits);
 		cg.i32.const(8388607);
 		cg.i32.and();
@@ -2511,7 +2661,7 @@ function _sincos(cg) {
 	cg.f32.mul();
 	cg.f32.nearest();
 	cg.local.tee(qf);
-	cg.i32.trunc_f32_s();
+	cg.i32.trunc_sat_f32_s();
 	cg.local.set(q);
 	cg.local.get(y);
 	cg.local.get(qf);
@@ -3598,6 +3748,7 @@ var F32x4 = class extends V128 {
 
 //#endregion
 //#region src/backend/wasm.ts
+const moduleCache = /* @__PURE__ */ new Map();
 /** Backend that compiles into WebAssembly bytecode for immediate execution. */
 var WasmBackend = class {
 	type = "wasm";
@@ -3649,15 +3800,25 @@ var WasmBackend = class {
 		if (count === void 0) count = buffer.byteLength - start;
 		return buffer.slice(start, start + count);
 	}
-	async prepare(kernel) {
-		return this.prepareSync(kernel);
+	async prepareKernel(kernel) {
+		return this.prepareKernelSync(kernel);
 	}
-	prepareSync(kernel) {
-		const bytes = codegenWasm(kernel);
-		const module = new WebAssembly.Module(bytes);
+	prepareKernelSync(kernel) {
+		const kernelHash = FpHash.hash(kernel);
+		const module = runWithCache(moduleCache, kernelHash.toString(), () => {
+			const bytes = codegenWasm(kernel);
+			return new WebAssembly.Module(bytes);
+		});
 		return new Executable(kernel, { module });
 	}
+	async prepareRoutine(routine) {
+		return this.prepareRoutineSync(routine);
+	}
+	prepareRoutineSync(routine) {
+		return new Executable(routine, void 0);
+	}
 	dispatch(exe, inputs, outputs) {
+		if (exe.source instanceof Routine) return runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
 		const instance = new WebAssembly.Instance(exe.data.module, { env: { memory: this.#memory } });
 		const func = instance.exports.kernel;
 		const ptrs = [...inputs, ...outputs].map((slot) => this.#buffers.get(slot).ptr);
@@ -3675,7 +3836,7 @@ function codegenWasm(kernel) {
 	if (DEBUG >= 3) console.info(`kernel.exp: ${kernel.exp}\ntune.exp: ${tune.exp}`);
 	const cg = new CodeGenerator();
 	cg.memory.import("env", "memory");
-	const distinctOps = mapSetUnion(tune.exp.distinctOps(), re?.epilogue.distinctOps());
+	const distinctOps = mapSetUnion(tune.exp.distinctOps(), tune.epilogue?.distinctOps());
 	const funcs = {};
 	if (distinctOps.has(AluOp.Sin)) funcs.sin = wasm_sin(cg);
 	if (distinctOps.has(AluOp.Cos)) funcs.cos = wasm_cos(cg);
@@ -3753,7 +3914,10 @@ function codegenWasm(kernel) {
 			cg.br(1);
 			cg.end();
 			cg.end();
-			translateExp(cg, funcs, kernel.reduction.epilogue, { acc });
+			translateExp(cg, funcs, tune.epilogue, {
+				acc,
+				gidx
+			});
 		} else translateExp(cg, funcs, tune.exp, { gidx });
 		dty(cg, null, kernel.dtype).store(Math.log2(byteWidth(kernel.dtype)));
 		cg.local.get(gidx);
@@ -4002,7 +4166,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await import("./webgpu-BGuG58KZ.js");
+		const { WebGPUBackend } = await import("./webgpu-ChVgx3b6.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4036,8 +4200,8 @@ function getBackend(device) {
 	return backend;
 }
 var Executable = class {
-	constructor(kernel, data) {
-		this.kernel = kernel;
+	constructor(source, data) {
+		this.source = source;
 		this.data = data;
 	}
 };
@@ -4053,7 +4217,11 @@ var UnsupportedOpError = class extends Error {
 		super(msg);
 	}
 };
+var UnsupportedRoutineError = class extends Error {
+	constructor(name, device) {
+		super(`routine '${name}' is not supported in ${device} backend`);
+	}
+};
 
 //#endregion
-export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, ShapeTracker, SlotError, UnsupportedOpError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneWebgpu, unravelAlu, unzip2, zip, zipn };
-//# sourceMappingURL=backend-BqymqzuU.js.map
+export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneWebgpu, unravelAlu, unzip2, zip, zipn };