matrix/js/webgpu/utils.js

const loadTexture = async (device, cache, url) => {
	const key = url;
	if (cache.has(key)) {
		return cache.get(key);
	}

	let texture;

	if (url == null) {
		texture = device.createTexture({
			size: [1, 1, 1],
			format: "rgba8unorm",
			usage:
				GPUTextureUsage.TEXTURE_BINDING |
				GPUTextureUsage.COPY_DST |
				GPUTextureUsage.RENDER_ATTACHMENT,
		});
	} else {
		let imageURL;
		if (typeof cache.get(`import::${url}`) === "function") {
			imageURL = (await cache.get(`import::${url}`)()).default;
		} else {
			imageURL = url;
		}

		const response = await fetch(imageURL);
		const data = await response.blob();
		const source = await createImageBitmap(data);
		const size = [source.width, source.height, 1];

		texture = device.createTexture({
			size,
			format: "rgba8unorm",
			usage:
				GPUTextureUsage.TEXTURE_BINDING |
				GPUTextureUsage.COPY_DST |
				GPUTextureUsage.RENDER_ATTACHMENT,
		});

		device.queue.copyExternalImageToTexture({ source, flipY: true }, { texture }, size);
	}

	cache.set(key, texture);

	return texture;
};

const makeRenderTarget = (device, size, format, mipLevelCount = 1) =>
	device.createTexture({
		size: [...size, 1],
		mipLevelCount,
		format,
		usage:
			GPUTextureUsage.TEXTURE_BINDING |
			GPUTextureUsage.COPY_SRC |
			GPUTextureUsage.COPY_DST |
			GPUTextureUsage.RENDER_ATTACHMENT,
	});

const makeComputeTarget = (device, size, mipLevelCount = 1) =>
	device.createTexture({
		size: [...size, 1],
		mipLevelCount,
		format: "rgba8unorm",
		usage:
			GPUTextureUsage.TEXTURE_BINDING |
			GPUTextureUsage.COPY_SRC |
			GPUTextureUsage.COPY_DST |
			GPUTextureUsage.STORAGE_BINDING,
	});

const loadShader = async (device, cache, url) => {
	const key = url;
	if (cache.has(key)) {
		return cache.get(key);
	}
	let textURL;
	if (typeof cache.get(`import::${url}`) === "function") {
		textURL = (await cache.get(`import::${url}`)()).default;
	} else {
		textURL = url;
	}
	const response = await fetch(textURL);
	const code = await response.text();
	return {
		code,
		module: device.createShaderModule({ code }),
	};
};

const makeUniformBuffer = (device, uniforms, data = null) => {
	const buffer = device.createBuffer({
		size: uniforms.minSize,
		usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
		mappedAtCreation: data != null,
	});
	if (data != null) {
		uniforms.toBuffer(data, buffer.getMappedRange());
		buffer.unmap();
	}
	return buffer;
};

const make1DTexture = (device, rgbas) => {
	const size = [rgbas.length];
	const texture = device.createTexture({
		size,
		// dimension: "1d",
		format: "rgba8unorm",
		usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
	});
	const data = new Uint8ClampedArray(rgbas.map((color) => color.map((f) => f * 0xff)).flat());
	device.queue.writeTexture({ texture }, data, {}, size);
	return texture;
};

const makeBindGroup = (device, pipeline, index, entries) =>
	device.createBindGroup({
		layout: pipeline.getBindGroupLayout(index),
		entries: entries
			.map((resource) => (resource instanceof GPUBuffer ? { buffer: resource } : resource))
			.map((resource, binding) => ({
				binding,
				resource,
			})),
	});

const makePass = (name, loaded, build, run) => ({
	loaded: loaded ?? Promise.resolve(),
	build: build ?? ((size, inputs) => inputs),
	run: (encoder, shouldRender) => {
		encoder.pushDebugGroup(`Pass "${name}"`);
		run?.(encoder, shouldRender);
		encoder.popDebugGroup();
	},
});

const makePipeline = async (context, steps) => {
	steps = steps.filter((f) => f != null).map((f) => f(context));
	await Promise.all(steps.map((step) => step.loaded));
	return {
		steps,
		build: (canvasSize) => steps.reduce((outputs, step) => step.build(canvasSize, outputs), null),
		run: (encoder, shouldRender) => steps.forEach((step) => step.run(encoder, shouldRender)),
	};
};

export {
	makeRenderTarget,
	makeComputeTarget,
	make1DTexture,
	loadTexture,
	loadShader,
	makeUniformBuffer,
	makePass,
	makePipeline,
	makeBindGroup,
};