typegpu 0.10.0 → 0.10.2

package/core/root/init.js

@@ -0,0 +1,457 @@
+import { $getNameForward, $internal } from "../../shared/symbols.js";
+import { setName } from "../../shared/meta.js";
+import { invariant } from "../../errors.js";
+import { isAccessor, isMutableAccessor } from "../slot/slotTypes.js";
+import { u32 } from "../../data/numeric.js";
+import { vec3f, vec3u } from "../../data/vector.js";
+import { fn } from "../function/tgpuFn.js";
+import { ConfigurableImpl } from "./configurableImpl.js";
+import { INTERNAL_createBuffer, isBuffer } from "../buffer/buffer.js";
+import { INTERNAL_createComparisonSampler, INTERNAL_createSampler, isComparisonSampler, isSampler } from "../sampler/sampler.js";
+import { clearTextureUtilsCache } from "../texture/textureUtils.js";
+import { INTERNAL_createTexture, isTexture, isTextureView } from "../texture/texture.js";
+import { TgpuBindGroupImpl, isBindGroup, isBindGroupLayout } from "../../tgpuBindGroupLayout.js";
+import { ceil } from "../../std/numeric.js";
+import { isComputePipeline, isRenderPipeline } from "../pipeline/typeGuards.js";
+import { builtin } from "../../builtin.js";
+import { INTERNAL_createQuerySet, isQuerySet } from "../querySet/querySet.js";
+import { WeakMemo } from "../../memo.js";
+import { TgpuBufferShorthandImpl } from "../buffer/bufferShorthand.js";
+import { computeFn } from "../function/tgpuComputeFn.js";
+import { applyBindGroups, applyVertexBuffers } from "../pipeline/applyPipelineState.js";
+import { INTERNAL_createComputePipeline } from "../pipeline/computePipeline.js";
+import { isVertexLayout } from "../vertexLayout/vertexLayout.js";
+import { INTERNAL_createRenderPipeline } from "../pipeline/renderPipeline.js";
+import { allEq } from "../../std/boolean.js";
+
+//#region src/core/root/init.ts
+/**
+* Changes the given array to a vec of 3 numbers, filling missing values with 1.
+*/
+function toVec3(arr) {
+	if (arr.includes(0)) throw new Error("Size and workgroupSize cannot contain zeroes.");
+	return vec3u(arr[0] ?? 1, arr[1] ?? 1, arr[2] ?? 1);
+}
+const workgroupSizeConfigs = [
+	vec3u(1, 1, 1),
+	vec3u(256, 1, 1),
+	vec3u(16, 16, 1),
+	vec3u(8, 8, 4)
+];
+var TgpuGuardedComputePipelineImpl = class TgpuGuardedComputePipelineImpl {
+	#root;
+	#pipeline;
+	#sizeUniform;
+	#workgroupSize;
+	#lastSize;
+	constructor(root, pipeline, sizeUniform, workgroupSize) {
+		this.#root = root;
+		this.#pipeline = pipeline;
+		this.#sizeUniform = sizeUniform;
+		this.#workgroupSize = workgroupSize;
+		this.#lastSize = vec3u();
+	}
+	with(bindGroup) {
+		return new TgpuGuardedComputePipelineImpl(this.#root, this.#pipeline.with(bindGroup), this.#sizeUniform, this.#workgroupSize);
+	}
+	dispatchThreads(...threads) {
+		const sanitizedSize = toVec3(threads);
+		const workgroupCount = ceil(vec3f(sanitizedSize).div(vec3f(this.#workgroupSize)));
+		if (!allEq(sanitizedSize, this.#lastSize)) {
+			this.#lastSize = sanitizedSize;
+			this.#sizeUniform.write(sanitizedSize);
+		}
+		this.#pipeline.dispatchWorkgroups(workgroupCount.x, workgroupCount.y, workgroupCount.z);
+	}
+	get pipeline() {
+		return this.#pipeline;
+	}
+	get sizeUniform() {
+		return this.#sizeUniform;
+	}
+	[$internal] = true;
+	get [$getNameForward]() {
+		return this.#pipeline;
+	}
+	$name(label) {
+		setName(this, label);
+		return this;
+	}
+};
+var WithBindingImpl = class WithBindingImpl {
+	constructor(_getRoot, _slotBindings) {
+		this._getRoot = _getRoot;
+		this._slotBindings = _slotBindings;
+	}
+	with(slot, value) {
+		return new WithBindingImpl(this._getRoot, [...this._slotBindings, [isAccessor(slot) || isMutableAccessor(slot) ? slot.slot : slot, value]]);
+	}
+	withCompute(entryFn) {
+		return new WithComputeImpl(this._getRoot(), this._slotBindings, entryFn);
+	}
+	createComputePipeline(descriptor) {
+		return INTERNAL_createComputePipeline(this._getRoot(), this._slotBindings, descriptor);
+	}
+	createRenderPipeline(descriptor) {
+		return INTERNAL_createRenderPipeline({
+			root: this._getRoot(),
+			slotBindings: this._slotBindings,
+			descriptor
+		});
+	}
+	createGuardedComputePipeline(callback) {
+		const root = this._getRoot();
+		if (callback.length >= 4) throw new Error("Guarded compute callback only supports up to three dimensions.");
+		const workgroupSize = workgroupSizeConfigs[callback.length];
+		const wrappedCallback = fn([
+			u32,
+			u32,
+			u32
+		])(callback);
+		const sizeUniform = root.createUniform(vec3u);
+		const mainCompute = computeFn({
+			workgroupSize,
+			in: { id: builtin.globalInvocationId }
+		})`{
+  if (any(in.id >= sizeUniform)) {
+    return;
+  }
+  wrappedCallback(in.id.x, in.id.y, in.id.z);
+}`.$uses({
+			sizeUniform,
+			wrappedCallback
+		});
+		return new TgpuGuardedComputePipelineImpl(root, this.createComputePipeline({ compute: mainCompute }), sizeUniform, workgroupSize);
+	}
+	withVertex(entryFn, attribs) {
+		return new WithVertexImpl(this._getRoot(), this._slotBindings, {
+			vertex: entryFn,
+			attribs
+		});
+	}
+	pipe(transform) {
+		const newCfg = transform(new ConfigurableImpl([]));
+		return new WithBindingImpl(this._getRoot, [...this._slotBindings, ...newCfg.bindings]);
+	}
+};
+var WithComputeImpl = class {
+	constructor(_root, _slotBindings, _entryFn) {
+		this._root = _root;
+		this._slotBindings = _slotBindings;
+		this._entryFn = _entryFn;
+	}
+	createPipeline() {
+		return INTERNAL_createComputePipeline(this._root, this._slotBindings, { compute: this._entryFn });
+	}
+};
+var WithVertexImpl = class WithVertexImpl {
+	#root;
+	#slotBindings;
+	#partialDescriptor;
+	constructor(root, slotBindings, partialDescriptor) {
+		this.#root = root;
+		this.#slotBindings = slotBindings;
+		this.#partialDescriptor = partialDescriptor;
+	}
+	withFragment(fragmentFn, targets, _mismatch) {
+		invariant(typeof fragmentFn !== "string", "Just type mismatch validation");
+		invariant(targets === void 0 || typeof targets !== "string", "Just type mismatch validation");
+		return new WithFragmentImpl(this.#root, this.#slotBindings, {
+			...this.#partialDescriptor,
+			fragment: fragmentFn,
+			targets
+		});
+	}
+	withPrimitive(primitive) {
+		return new WithVertexImpl(this.#root, this.#slotBindings, {
+			...this.#partialDescriptor,
+			primitive
+		});
+	}
+	withDepthStencil(depthStencil) {
+		return new WithVertexImpl(this.#root, this.#slotBindings, {
+			...this.#partialDescriptor,
+			depthStencil
+		});
+	}
+	withMultisample(multisample) {
+		return new WithVertexImpl(this.#root, this.#slotBindings, {
+			...this.#partialDescriptor,
+			multisample
+		});
+	}
+	createPipeline() {
+		return INTERNAL_createRenderPipeline({
+			root: this.#root,
+			slotBindings: this.#slotBindings,
+			descriptor: this.#partialDescriptor
+		});
+	}
+};
+var WithFragmentImpl = class WithFragmentImpl {
+	#root;
+	#slotBindings;
+	#descriptor;
+	constructor(root, slotBindings, descriptor) {
+		this.#root = root;
+		this.#slotBindings = slotBindings;
+		this.#descriptor = descriptor;
+	}
+	withPrimitive(primitive) {
+		return new WithFragmentImpl(this.#root, this.#slotBindings, {
+			...this.#descriptor,
+			primitive
+		});
+	}
+	withDepthStencil(depthStencil) {
+		return new WithFragmentImpl(this.#root, this.#slotBindings, {
+			...this.#descriptor,
+			depthStencil
+		});
+	}
+	withMultisample(multisample) {
+		return new WithFragmentImpl(this.#root, this.#slotBindings, {
+			...this.#descriptor,
+			multisample
+		});
+	}
+	createPipeline() {
+		return INTERNAL_createRenderPipeline({
+			root: this.#root,
+			slotBindings: this.#slotBindings,
+			descriptor: this.#descriptor
+		});
+	}
+};
+/**
+* Holds all data that is necessary to facilitate CPU and GPU communication.
+* Programs that share a root can interact via GPU buffers.
+*/
+var TgpuRootImpl = class extends WithBindingImpl {
+	"~unstable";
+	_unwrappedBindGroupLayouts = new WeakMemo((key) => key.unwrap(this));
+	_unwrappedBindGroups = new WeakMemo((key) => key.unwrap(this));
+	[$internal];
+	constructor(device, nameRegistrySetting, _ownDevice, logOptions, shaderGenerator) {
+		super(() => this, []);
+		this.device = device;
+		this.nameRegistrySetting = nameRegistrySetting;
+		this._ownDevice = _ownDevice;
+		this.shaderGenerator = shaderGenerator;
+		this["~unstable"] = this;
+		this[$internal] = { logOptions };
+	}
+	configureContext(options) {
+		const context = options.canvas.getContext("webgpu");
+		if (!context) throw new Error("Unable to initialize 'webgpu' context on the provided canvas.");
+		context.configure({
+			...options,
+			device: this.device,
+			format: options.format ?? navigator.gpu.getPreferredCanvasFormat()
+		});
+		return context;
+	}
+	get enabledFeatures() {
+		return new Set(this.device.features);
+	}
+	createBuffer(typeSchema, initialOrBuffer) {
+		return INTERNAL_createBuffer(this, typeSchema, initialOrBuffer);
+	}
+	createUniform(typeSchema, initialOrBuffer) {
+		return new TgpuBufferShorthandImpl("uniform", INTERNAL_createBuffer(this, typeSchema, initialOrBuffer).$usage("uniform"));
+	}
+	createMutable(typeSchema, initialOrBuffer) {
+		return new TgpuBufferShorthandImpl("mutable", INTERNAL_createBuffer(this, typeSchema, initialOrBuffer).$usage("storage"));
+	}
+	createReadonly(typeSchema, initialOrBuffer) {
+		return new TgpuBufferShorthandImpl("readonly", INTERNAL_createBuffer(this, typeSchema, initialOrBuffer).$usage("storage"));
+	}
+	createQuerySet(type, count, rawQuerySet) {
+		return INTERNAL_createQuerySet(this, type, count, rawQuerySet);
+	}
+	createBindGroup(layout, entries) {
+		return new TgpuBindGroupImpl(layout, entries);
+	}
+	destroy() {
+		clearTextureUtilsCache(this.device);
+		if (this._ownDevice) this.device.destroy();
+	}
+	createTexture(props) {
+		return INTERNAL_createTexture(props, this);
+	}
+	createSampler(props) {
+		return INTERNAL_createSampler(props, this);
+	}
+	createComparisonSampler(props) {
+		return INTERNAL_createComparisonSampler(props, this);
+	}
+	unwrap(resource) {
+		if (isComputePipeline(resource)) return resource[$internal].rawPipeline;
+		if (isRenderPipeline(resource)) return resource[$internal].core.unwrap().pipeline;
+		if (isBindGroupLayout(resource)) return this._unwrappedBindGroupLayouts.getOrMake(resource);
+		if (isBindGroup(resource)) return this._unwrappedBindGroups.getOrMake(resource);
+		if (isBuffer(resource)) return resource.buffer;
+		if (isTexture(resource)) return resource[$internal].unwrap();
+		if (isTextureView(resource)) {
+			if (!resource[$internal].unwrap) throw new Error("Cannot unwrap laid-out texture view as it has no underlying resource.");
+			return resource[$internal].unwrap();
+		}
+		if (isVertexLayout(resource)) return resource.vertexLayout;
+		if (isSampler(resource) || isComparisonSampler(resource)) {
+			if (resource[$internal].unwrap) return resource[$internal].unwrap();
+			throw new Error("Cannot unwrap laid-out sampler.");
+		}
+		if (isQuerySet(resource)) return resource.querySet;
+		throw new Error(`Unknown resource type: ${resource}`);
+	}
+	createDrawablePassProxy(encoder) {
+		const bindGroups = /* @__PURE__ */ new Map();
+		const vertexBuffers = /* @__PURE__ */ new Map();
+		let currentPipeline;
+		let dirty = true;
+		const applyPipelineState = () => {
+			if (!currentPipeline) throw new Error("Cannot draw without a call to pass.setPipeline");
+			if (!dirty) return;
+			dirty = false;
+			const { core, priors } = currentPipeline[$internal];
+			const memo = core.unwrap();
+			encoder.setPipeline(memo.pipeline);
+			applyBindGroups(encoder, this, memo.usedBindGroupLayouts, memo.catchall, (layout) => priors.bindGroupLayoutMap?.get(layout) ?? bindGroups.get(layout));
+			applyVertexBuffers(encoder, this, memo.usedVertexLayouts, (vertexLayout) => {
+				const priorBuffer = priors.vertexLayoutMap?.get(vertexLayout);
+				return priorBuffer ? {
+					buffer: priorBuffer,
+					offset: void 0,
+					size: void 0
+				} : vertexBuffers.get(vertexLayout);
+			});
+		};
+		return {
+			setPipeline(pipeline) {
+				currentPipeline = pipeline;
+				dirty = true;
+			},
+			setIndexBuffer: (buffer, indexFormat, offset, size) => {
+				if (isBuffer(buffer)) encoder.setIndexBuffer(this.unwrap(buffer), indexFormat, offset, size);
+				else encoder.setIndexBuffer(buffer, indexFormat, offset, size);
+			},
+			setVertexBuffer(vertexLayout, buffer, offset, size) {
+				vertexBuffers.set(vertexLayout, {
+					buffer,
+					offset,
+					size
+				});
+				dirty = true;
+			},
+			setBindGroup(bindGroupLayout, bindGroup) {
+				bindGroups.set(bindGroupLayout, bindGroup);
+				dirty = true;
+			},
+			draw(vertexCount, instanceCount, firstVertex, firstInstance) {
+				applyPipelineState();
+				encoder.draw(vertexCount, instanceCount, firstVertex, firstInstance);
+			},
+			drawIndexed(...args) {
+				applyPipelineState();
+				encoder.drawIndexed(...args);
+			},
+			drawIndirect(...args) {
+				applyPipelineState();
+				encoder.drawIndirect(...args);
+			},
+			drawIndexedIndirect(...args) {
+				applyPipelineState();
+				encoder.drawIndexedIndirect(...args);
+			}
+		};
+	}
+	beginRenderPass(descriptor, callback) {
+		const commandEncoder = this.device.createCommandEncoder();
+		const pass = commandEncoder.beginRenderPass(descriptor);
+		callback({
+			setViewport(...args) {
+				pass.setViewport(...args);
+			},
+			setScissorRect(...args) {
+				pass.setScissorRect(...args);
+			},
+			setBlendConstant(...args) {
+				pass.setBlendConstant(...args);
+			},
+			setStencilReference(...args) {
+				pass.setStencilReference(...args);
+			},
+			beginOcclusionQuery(...args) {
+				pass.beginOcclusionQuery(...args);
+			},
+			endOcclusionQuery(...args) {
+				pass.endOcclusionQuery(...args);
+			},
+			executeBundles(...args) {
+				pass.executeBundles(...args);
+			},
+			...this.createDrawablePassProxy(pass)
+		});
+		pass.end();
+		this.device.queue.submit([commandEncoder.finish()]);
+	}
+	beginRenderBundleEncoder(descriptor, callback) {
+		const bundleEncoder = this.device.createRenderBundleEncoder(descriptor);
+		callback(this.createDrawablePassProxy(bundleEncoder));
+		return bundleEncoder.finish();
+	}
+	flush() {
+		console.warn("flush() has been deprecated, and has no effect.");
+	}
+};
+/**
+* Requests a new GPU device and creates a root around it.
+* If a specific device should be used instead, use @see initFromDevice.
+*
+* @example
+* When given no options, the function will ask the browser for a suitable GPU device.
+* ```ts
+* const root = await tgpu.init();
+* ```
+*
+* @example
+* If there are specific options that should be used when requesting a device, you can pass those in.
+* ```ts
+* const adapterOptions: GPURequestAdapterOptions = ...;
+* const deviceDescriptor: GPUDeviceDescriptor = ...;
+* const root = await tgpu.init({ adapter: adapterOptions, device: deviceDescriptor });
+* ```
+*/
+async function init(options) {
+	const { adapter: adapterOpt, device: deviceOpt, unstable_names: names = "strict", unstable_logOptions } = options ?? {};
+	if (!navigator.gpu) throw new Error("WebGPU is not supported by this browser.");
+	const adapter = await navigator.gpu.requestAdapter(adapterOpt);
+	if (!adapter) throw new Error("Could not find a compatible GPU");
+	const availableFeatures = [];
+	for (const feature of deviceOpt?.requiredFeatures ?? []) {
+		if (!adapter.features.has(feature)) throw new Error(`Requested feature "${feature}" is not supported by the adapter.`);
+		availableFeatures.push(feature);
+	}
+	for (const feature of deviceOpt?.optionalFeatures ?? []) if (adapter.features.has(feature)) availableFeatures.push(feature);
+	else console.warn(`Optional feature "${feature}" is not supported by the adapter.`);
+	return new TgpuRootImpl(await adapter.requestDevice({
+		...deviceOpt,
+		requiredFeatures: availableFeatures
+	}), names, true, unstable_logOptions ?? {}, options?.shaderGenerator);
+}
+/**
+* Creates a root from the given device, instead of requesting it like @see init.
+*
+* @example
+* ```ts
+* const device: GPUDevice = ...;
+* const root = tgpu.initFromDevice({ device });
+* ```
+*/
+function initFromDevice(options) {
+	const { device, unstable_names: names = "strict", unstable_logOptions } = options ?? {};
+	return new TgpuRootImpl(device, names, false, unstable_logOptions ?? {}, options?.shaderGenerator);
+}
+
+//#endregion
+export { init, initFromDevice };