@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/patches/webgpu_program.js

@@ -0,0 +1,293 @@
+import { Ln as e, Ms as t, Ns as n, Ps as r } from "../dist-BewPQWjc.js";
+import { a as i, c as a, f as o, i as s, o as c, r as l, u } from "../webgpu_program-WOyIVMlZ.js";
+//#region lib/patches/webgpu_program.ts
+var d = /* @__PURE__ */ function(e) {
+	return e[e.FROM_PIXELS = 0] = "FROM_PIXELS", e[e.DRAW = 1] = "DRAW", e;
+}({}), f = (e, n, r, i, a) => {
+	let o = p(r, {
+		dtype: i.dtype,
+		shape: i.shape
+	}, n), s = e.createShaderModule({
+		code: o,
+		label: n.constructor.name
+	}), c = t().get("WEBGPU_PRINT_SHADER");
+	if (c !== "") {
+		c = c.toLowerCase();
+		let e = c.split(",");
+		(c === "all" || e.some((e) => n.shaderKey.toLowerCase().includes(e))) && (console.group(n.shaderKey), console.debug(o), console.groupEnd());
+	}
+	return a ? e.createComputePipelineAsync({
+		compute: {
+			module: s,
+			entryPoint: "_start"
+		},
+		label: n.constructor.name,
+		layout: "auto"
+	}) : e.createComputePipeline({
+		compute: {
+			module: s,
+			entryPoint: "_start"
+		},
+		label: n.constructor.name,
+		layout: "auto"
+	});
+};
+function p(e, t, n) {
+	let r = [], a = n.workgroupSize[0] * n.workgroupSize[1] * n.workgroupSize[2];
+	if (n.outputComponent = n.outputComponent ? n.outputComponent : 1, r.push(`
+
+      var<private> localId: vec3<u32>;
+      var<private> localIndex: u32;
+      var<private> globalId: vec3<u32>;
+      var<private> numWorkgroups: vec3<u32>;
+      var<private> workgroupId: vec3<u32>;
+        ${n.subgroupBuiltins ? "var<private> subgroupInvocationId: u32;" : ""}
+        ${n.subgroupBuiltins ? "var<private> subgroupSize: u32;" : ""}
+
+      // Only used when the y/z dimension of workgroup size is 1.
+      fn getGlobalIndex() -> i32 {
+        ${x(n) ? "  return i32(globalId.x);" : `  return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +
+                workgroupId.y * numWorkgroups.x + workgroupId.x) * ${a}u +
+                localIndex);
+        `}
+      }
+    `), n.pixelsOpType != null) {
+		let a = n.pixelsOpType === 0 ? `@group(0) @binding(0) var<storage, read_write> result: array<${l(t.dtype, n.outputComponent)}>;` : `@group(0) @binding(1) var<storage, read> inBuf : array<${l(e[0].dtype, n.outputComponent)}>;`, o = t.shape.length === 3 ? "vec2<i32>" : "i32";
+		r.push(`
+        struct Uniform {
+          outShapeStrides : ${o},
+          size            : i32,
+          numChannels     : i32,
+          alpha           : f32,
+        };
+
+        ${a}
+        @group(0) @binding(2) var<uniform> uniforms: Uniform;
+      `);
+		let s = w(n);
+		return [
+			m,
+			r.join("\n"),
+			i(t.shape),
+			n.getUserCode(),
+			T(s, n)
+		].join("\n");
+	}
+	let o, c, u = "struct Uniforms { NAN : f32, INFINITY : f32, ";
+	n.variableNames.forEach((t, n) => {
+		let r = s(e[n].shape.length);
+		u += `${t.charAt(0).toLowerCase() + t.slice(1)}Shape : ${r}, `, o = e[n].shape.length - 1, c = s(o), u += `${t.charAt(0).toLowerCase() + t.slice(1)}ShapeStrides: ${c}, `;
+	});
+	let d = s(t.shape.length);
+	u += `outShape : ${d}, `, o = t.shape.length - 1, c = s(o), u += `
+         outShapeStrides: ${c}, `, n.size && (u += "size : i32, "), n.uniforms && (u += n.uniforms), u += "};", u = C(u), r.push(u), n.atomic ? r.push("\n      @group(0) @binding(0) var<storage, read_write> result: array<atomic<i32>>;\n    ") : r.push(`
+      @group(0) @binding(0) var<storage, read_write> result: array<${l(t.dtype, n.outputComponent)}>;
+    `), n.variableNames.forEach((t, i) => {
+		r.push(`
+      @group(0) @binding(${1 + i}) var<storage, read> ${t}: array<${n.variableComponents ? l(e[i].dtype, n.variableComponents[i]) : l(e[i].dtype, n.outputComponent)}>;
+        `);
+	}), u !== "" && r.push(`
+      @group(0) @binding(${1 + n.variableNames.length}) var<uniform> uniforms: Uniforms;
+      `);
+	let f = y(t.shape, n.dispatchLayout), p = [
+		n.subgroups ? "enable subgroups;" : "",
+		m,
+		r.join("\n") + h,
+		i(t.shape),
+		f,
+		b(t.shape.length)
+	];
+	n.atomic || p.push(S(t.shape, t.dtype, n.outputComponent)), n.variableNames.forEach((t, n) => {
+		p.push(`${i(e[n].shape, t)}`);
+	});
+	let g = e.map((e, r) => v(e, t.shape, n.variableComponents ? n.variableComponents[r] : n.outputComponent, n.dispatchLayout.x.length === t.shape.length)).join("\n");
+	p.push(g), p.push(n.getUserCode());
+	let _ = w(n);
+	return p.push(T(_, n)), p.join("\n");
+}
+var m = "\n  struct vec5 {x: i32, y: i32, z: i32, w: i32, u: i32};\n  struct vec6 {x: i32, y: i32, z: i32, w: i32, u: i32, v: i32};\n\n  // Checks whether coordinates lie within the bounds of the shape.\n  fn coordsInBounds2D(coord : vec2<i32>, shape : vec2<i32>) -> bool {\n    return all(coord >= vec2<i32>(0)) && all(coord < shape);\n  }\n  fn coordsInBounds3D(coord : vec3<i32>, shape : vec3<i32>) -> bool {\n    return all(coord >= vec3<i32>(0)) && all(coord < shape);\n  }\n  fn coordsInBounds4D(coord : vec4<i32>, shape : vec4<i32>) -> bool {\n    return all(coord >= vec4<i32>(0)) && all(coord < shape);\n  }\n\n  fn getIndexFromCoords1D(coord : i32, shape : i32) -> i32 {\n    return coord;\n  }\n  fn getIndexFromCoords2D(coords : vec2<i32>, shape : vec2<i32>) -> i32 {\n    return dot(coords, vec2<i32>(shape.y, 1));\n  }\n  fn getIndexFromCoords3D(coords : vec3<i32>, shape : vec3<i32>) -> i32 {\n    return dot(coords, vec3<i32>(shape.y * shape.z, shape.z, 1));\n  }\n  fn getIndexFromCoords4D(coords : vec4<i32>, shape : vec4<i32>) -> i32 {\n    return dot(coords, vec4<i32>(\n        shape.y * shape.z * shape.w, shape.z * shape.w, shape.w, 1));\n  }\n  fn getIndexFromCoords5D(coords : vec5, shape : vec5) -> i32 {\n    let shapeStrides: vec5 = vec5(shape.y * shape.z * shape.w * shape.u, shape.z * shape.w * shape.u, shape.w * shape.u, shape.u, 1);\n    return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u;\n  }\n  fn getIndexFromCoords6D(coords : vec6, shape : vec6) -> i32 {\n    let shapeStrides: vec6 = vec6(shape.y * shape.z * shape.w * shape.u * shape.v, shape.z * shape.w * shape.u * shape.v, shape.w * shape.u * shape.v, shape.u * shape.v, shape.v, 1);\n    return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u + coords.v*shapeStrides.v;\n  }\n\n  // NaN defination in IEEE 754-1985 is :\n  //   - sign = either 0 or 1.\n  //   - biased exponent = all 1 bits.\n  //   - fraction = anything except all 0 bits (since all 0 bits represents infinity).\n  // https://en.wikipedia.org/wiki/IEEE_754-1985#Representation_of_non-numbers\n  fn isnan(val: f32) -> bool {\n    let floatToUint: u32 = bitcast<u32>(val);\n    return (floatToUint & 0x7fffffffu) > 0x7f800000u;\n  }\n  fn isnanVec4(val : vec4<f32>) -> vec4<bool> {\n    let floatToUint: vec4<u32> = bitcast<vec4<u32>>(val);\n    return (floatToUint & vec4<u32>(0x7fffffffu)) > vec4<u32>(0x7f800000u);\n  }\n", h = "\n  fn isinf(val: f32) -> bool {\n    return abs(val) == uniforms.INFINITY;\n  }\n";
+function g(e, t) {
+	let n = e.name, r = e.shape.length, i = s(r), a = "get" + n.charAt(0).toUpperCase() + n.slice(1), o = [
+		"d0",
+		"d1",
+		"d2",
+		"d3",
+		"d4",
+		"d5"
+	].slice(0, r), c = o.map((e) => `${e} : i32`).join(", ");
+	if (r < 1) return `
+      fn ${a}() -> ${u(t)} {
+        return ${u(t)}(${n}[0]);
+      }
+    `;
+	let l = `uniforms.${n.charAt(0).toLowerCase() + n.slice(1)}Shape`, d = `${r}D`;
+	return r === 0 && (d = "1D"), `
+    fn ${a}(${c}) -> ${u(t)} {
+      return ${u(t)}(${n}[getIndexFromCoords${d}(${i}(${o.join(",")}),
+        ${l})${t === 1 ? "" : ` / ${t}`}]);
+    }
+   `;
+}
+function _(t, r, i, a) {
+	let o = t.name, l = o.charAt(0).toUpperCase() + o.slice(1), d = "get" + l + "ByOutput", f = t.shape.length, p = r.length, m = s(p);
+	if (n(t.shape, r) && a) return `
+    fn ${d}Index(globalIndex : i32) -> ${u(i)} {
+      return ${u(i)}(${o}[globalIndex]);
+    }
+
+    fn ${d}Coords(coords : ${m}) -> ${u(i)} {
+      return ${u(i)}(${o}[${p > 1 ? "getOutputIndexFromCoords(coords)" : "coords"}${i === 1 ? "" : ` / ${i}`}]);
+    }
+    `;
+	let h = e(t.shape, r), g = p - f, _ = "";
+	if (f === 0) return `
+    fn ${d}Index(globalIndex : i32) -> ${u(i)}{
+      return get${l}();
+    }
+
+    fn ${d}Coords(coords : ${m}) -> ${u(i)}{
+      return get${l}();
+    }
+  `;
+	_ = p < 2 && h.length >= 1 ? "coords = 0;" : h.map((e) => `coords.${c(e + g)} = 0;`).join("\n");
+	let v = "";
+	v = p < 2 && f > 0 ? "coords" : p > 1 ? `${s(f)}(${t.shape.map((e, t) => `coords.${c(t + g)}`).join(", ")})` : "coords";
+	let y = `uniforms.${o.charAt(0).toLowerCase() + o.slice(1)}Shape`, b = `${f}D`;
+	return `
+  fn ${d}Index(globalIndex : i32) -> ${u(i)} {
+    var coords = getCoordsFromIndex(globalIndex);
+    ${_}
+    return ${u(i)}(${o}[getIndexFromCoords${b}(${v}, ${y})${i === 1 ? "" : ` / ${i}`}]);
+  }
+
+  fn ${d}Coords(coordsIn : ${m}) -> ${u(i)} {
+    var coords = coordsIn;
+    ${_}
+    return ${u(i)}(${o}[getIndexFromCoords${b}(${v}, ${y})${i === 1 ? "" : ` / ${i}`}]);
+  }
+`;
+}
+function v(e, t, n, r) {
+	let i = g(e, n);
+	return e.shape.length <= t.length && (i += _(e, t, n, r)), i;
+}
+function y(e, t) {
+	let { x: n, y: r = [], z: i = [] } = t, a = e.length, c = n.length + r.length + i.length;
+	if (c !== a) return "";
+	if (n.length === a) return `fn getOutputCoords() -> ${s(a)}{
+    let globalIndex = getGlobalIndex();
+    return getCoordsFromIndex(globalIndex);
+  }
+  `;
+	let l = "", u = [
+		n,
+		r,
+		i
+	];
+	for (let e = 0; e < u.length; e++) {
+		let t = u[e];
+		if (t.length !== 0) if (t.length === 1) l += `let d${t[0]} = i32(globalId[${e}]);`;
+		else {
+			let n = o(t, "uniforms.outShape");
+			l += `var index${e} = i32(globalId[${e}]);`;
+			for (let r = 0; r < n.length; r++) l += `let d${t[r]} = index${e} / ${n[r]};`, r === n.length - 1 ? l += `let d${t[r + 1]} = index${e} - d${t[r]} * ${n[r]};` : l += `index${e} = index${e} - d${t[r]} * ${n[r]};`;
+		}
+	}
+	let d = [];
+	for (let e = 0; e < c; e++) d.push(`d${e}`);
+	let f = s(c), p = `fn getOutputCoords() -> ${f} {
+  ${l}
+`;
+	return d.length === 0 ? p += `return ${f}(0); }` : p += `return ${f}(${d.join(",")}); }`, p;
+}
+function b(e) {
+	let t = "";
+	switch (e) {
+		case 0:
+		case 1:
+			t += "\n        fn getOutputIndexFromCoords(coords : i32) -> i32 {\n          return coords;\n        }\n        ";
+			break;
+		case 2:
+			t += "\n        fn getOutputIndexFromCoords(coords : vec2<i32>) -> i32 {\n          return dot(coords, vec2<i32>(uniforms.outShapeStrides, 1));\n        }\n        ";
+			break;
+		case 3:
+			t += "\n        fn getOutputIndexFromCoords(coords : vec3<i32>) -> i32 {\n          return dot(coords, vec3<i32>(uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, 1));\n        }\n        ";
+			break;
+		case 4:
+			t += "\n        fn getOutputIndexFromCoords(coords : vec4<i32>) -> i32 {\n          return dot(coords, vec4<i32>(\n            uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, uniforms.outShapeStrides.z, 1));\n        }\n        ";
+			break;
+		case 5:
+			t += "\n        fn getOutputIndexFromCoords(coords : vec5) -> i32 {\n          return coords.x * uniforms.outShapeStrides.x +\n              coords.y * uniforms.outShapeStrides.y +\n              coords.z * uniforms.outShapeStrides.z +\n              coords.w * uniforms.outShapeStrides.w +\n              coords.u;\n        }\n        ";
+			break;
+		case 6:
+			t += "\n        fn getOutputIndexFromCoords(coords : vec6) -> i32 {\n          return coords.x * uniforms.outShapeStrides.x +\n              coords.y * uniforms.outShapeStrides.y +\n              coords.z * uniforms.outShapeStrides.z +\n              coords.w * uniforms.outShapeStrides.w +\n              coords.u * uniforms.outShapeStrides.u +\n              coords.v;\n        }\n        ";
+			break;
+		default:
+			r(!1, () => `Unsupported ${e}D shape`);
+			break;
+	}
+	return t;
+}
+function x(e) {
+	return e.dispatch[1] === 1 && e.dispatch[2] === 1;
+}
+function S(e, t, n) {
+	let r = e.length, i = l(t, n), a = `fn setOutputAtIndex(flatIndex : i32, value : ${u(n)}) {
+      result[flatIndex] = ${i}(value);
+    }
+
+    fn setOutputAtIndexI32(flatIndex : i32, value : ${u(n, "i32")}) {
+      result[flatIndex] = ${i}(value);
+    }
+    `;
+	if (r >= 2) {
+		let e = [
+			"d0",
+			"d1",
+			"d2",
+			"d3",
+			"d4",
+			"d5"
+		].slice(0, r), t = s(r);
+		a += `
+      fn setOutputAtCoords(${e.map((e) => `${e} : i32`).join(", ")}, value : ${u(n)}) {
+        let flatIndex = getOutputIndexFromCoords(${t}(${e.join(", ")}));
+        setOutputAtIndex(flatIndex${n === 1 ? "" : ` / ${n}`}, value);
+      }
+      fn setOutputAtCoordsI32(${e.map((e) => `${e} : i32`).join(", ")}, value : ${u(n, "i32")}) {
+        let flatIndex = getOutputIndexFromCoords(${t}(${e.join(", ")}));
+        setOutputAtIndexI32(flatIndex${n === 1 ? "" : ` / ${n}`}, value);
+      }
+    `;
+	}
+	return a;
+}
+function C(e) {
+	return e = e.replace(/(\w+)\s*:\s*vec(5|6)/g, (e) => "@align(16) " + e), e = e.replace(/vec(5|6)\s*,\s*(\w+)/g, (e, t, n) => `vec${t}, @align(16) ${n}`), e;
+}
+function w(e) {
+	return !(e.dispatchLayout.hasOwnProperty("y") && e.dispatchLayout.y?.length !== 0 || e.dispatchLayout.hasOwnProperty("z") && e.dispatchLayout.z?.length !== 0);
+}
+function T(e, t) {
+	return `
+     ${a(t)}
+      fn _start(@builtin(local_invocation_id) LocalId : vec3<u32>,
+                @builtin(global_invocation_id) GlobalId : vec3<u32>,
+                @builtin(local_invocation_index) LocalIndex: u32,
+                @builtin(workgroup_id) WorkgroupId : vec3<u32>,
+                ${t.subgroupBuiltins ? "@builtin(subgroup_invocation_id) SubgroupInvocationId : u32," : ""}
+                ${t.subgroupBuiltins ? "@builtin(subgroup_size) SubgroupSize : u32," : ""}
+                @builtin(num_workgroups) NumWorkgroups : vec3<u32>) {
+        localId = LocalId;
+        localIndex = LocalIndex;
+        globalId = GlobalId;
+        numWorkgroups = NumWorkgroups;
+        workgroupId = WorkgroupId;
+        ${t.subgroupBuiltins ? "subgroupInvocationId = SubgroupInvocationId" : ""};
+        ${t.subgroupBuiltins ? "subgroupSize = SubgroupSize" : ""};
+        ${e ? "main(getGlobalIndex());" : "main();"};
+      }
+    `;
+}
+//#endregion
+export { d as PixelsOpType, f as compileProgram, T as getStartHeaderString };