↩ 018. Mean-shift clustering ↩

Tue, 29 Apr 2025 16:32:56 +0200

In Mean Shift Clustering, Gneiss Name, 2025 ▶, he describes a model for parametric palette reduction that appears, to me, to be best-in-class. It's a good paper video, and this is not a replacement therefor — this instead serves

1. to rigorously describe the core algorithm in text,
2. to provide a general implementation — download, run, instruxions in README,
3. to describe the optimisations that can be applied to this problem
4. (also those that shouldn't because they lose),
5. as a commentary on the state of Vulkan development "on-boarding" (I hadn't written a GPU program before this).

Compare the measured performance on my AMD Radeon RX 5700 XT THICC II/i7-2600/DDR3-1600 vs an AMD Radeon RX 6500 XT (RADV NAVI24)/5800X3D/DDR4-3600 (weighted average 14% faster) and a laptop with an iGPU (Intel(R) Iris(R) Xe Graphics (ADL GT2)/i5-1235U/SDDDR4 3200 MHz (0.3 ns)) (average 3× slower):

(colour count is (clearly) not the only factor affecting speed; measurements from the test dataset). (All time figures are mean of at least 5 measurements, usually more.)

A CPU implementation processes that 16'755-colour image in 5.5s (on the i7-2600 which has 8 CPUs). The others are a "don't even dream about it" scale.

A naive GPU implementation processes it in 160ms, times out after 8s after processing ~50/242'931 colours, and times out after 8s after processing ~150/1'103'347 colours (measured from f39afc3). This can be analysed as a 2×/1080×/2450× improvement and that is basically real.

I would like to thank Cicada, rightfold, ratchetfreak and to acknowledge the prior art of shaderposting by Sebastian Lague, Acerola, and their ilk. This drew heavily on my largely-faded memory of techniques described thereby. I would like to thank Dr. Hoel Kervadec, Orson R. L. Peters and Jade Master for their notational insights.

# Problem statement

-3. With operator [​[​\circ​]​] as an element-wise application: [​[​\forall{i} : (f \circ A)_{i} = f(A_{i})​]​],
-2. with operator [​[​\odot​]​] as the Hadamard product: [​[​\forall{i} : (A \odot B)_{i} = A_{i}B_{i}​]​],
-1. with [​[​[]​]​] as the Iverson bracket: [​[​[x] = (1\text{ if }x\text{ otherwise }0)​]​],
0. with [​[​\bar{\ }​]​] as the set of preceding naturals: [​[​\bar{n} = [0, n ) \cap \mathbb{Z}​]​],
1. let colour [​[​C = ([0, 255] \cap \mathbb{Z})^3​]​],
2. given input [​[​n​]​]-dimensional image [​[​I \in (C_i)_{i \in \mathbb{N}_+^n}​]​],
3. construct a map [​[​f: C \to C​]​] such that
4. the output image [​[​O​]​] is given by [​[​O = f \circ I​]​].

Map [​[​f​]​] depends on deciding on a convex coordinate space for processing [​[​P \subset \mathbb{R}^3​]​] with mapping [​[​p : C \to P \text{ & } p^{-1} : P \to C​]​] such that [​[​p^{-1}(p(x)) = x​]​]; the original paper uses Oklab for this, described in A perceptual color space for image processing, Björn Ottosson, 2020, which appears to be best-in-class due to its favourable properties for image processing. This gives [​[​P = [0, 1] \times [\frac{-1}{2}, \frac{1}{2}] \times [\frac{-1}{2}, \frac{1}{2}]​]​] and

Then [​[​f​]​] depends on radius [​[​r​]​] (with a functional range of [​[​[0, \sqrt{3}]​]​] in Oklab), and deciding on variation [​[​palette​]​] or [​[​*freq​]​]: [​[​*freq​]​] ends up on the center of mass by weighting by how many times it occurs in [​[​I​]​], [​[​palette​]​] ends up directly at the mean (as-if all colours had the same weight). With this given:

This can be generalised to transparent images ([​[​C \times {\mkern 1.5mu\overline{\mkern-1.5mu{256}\mkern-1.5mu}\mkern 1.5mu}​]​]) by replicating the alpha channel directly: [​[​I_i = (\color{var(--r)}r\color{initial}, \color{var(--g)}g\color{initial}, \color{var(--b)}b\color{initial}) \to O_i = (\color{var(--r)}r\color{initial}', \color{var(--g)}g\color{initial}', \color{var(--b)}b\color{initial}', a)​]​] where [​[​(\color{var(--r)}r\color{initial}', \color{var(--g)}g\color{initial}', \color{var(--b)}b\color{initial}') = f((\color{var(--r)}r\color{initial}, \color{var(--g)}g\color{initial}, \color{var(--b)}b\color{initial}))​]​]. Fully-transparent pixels should be discarded altogether.

Seeing as [​[​D​]​] is constant for either variation, it follows trivially that [​[​q_t = q_{t-1} \Leftrightarrow M_t = M_{t-1}​]​] ("the point didn't move between iterations" is equivalent to "the same points are within [​[​r​]​] as last iteration").

It may be more convenient computationally to consider [​[​D = p \circ Multiset(I)​]​] with [​[​palette​]​] reducing the multiplicity of each element to 1. God only knows how to notate this.

It is illustrative at this point to draw out a demonstration. To this end, the following results can be observed (enlarged to show texture):

For demonstrative simplicity, the source colours are all sampled from the Oklab plane at l=0.797447133. Due to quantisation unavoidable since [​[​|P| \gg |C|​]​] and them not naturally falling on that plane out of [​[​p​]​], the intermediate [​[​q_t​]​]s aren't on that plane (but remain relatively close), and those have been projected onto it (but this is why the points don't seem to follow a perfect geometric mean in 2D).

# [​[​palette, r=0.1: D=p \circ ((\mathrm{CC}_{16}, \mathrm{4E}_{16}, \mathrm{FF}_{16}), (\mathrm{FF}_{16}, \mathrm{2D}_{16}, \mathrm{FF}_{16}), (\mathrm{FF}_{16}, \mathrm{2C}_{16}, \mathrm{A6}_{16}))​]​]

# [​[​*freq, r=0.1: D=p \circ \begin{pmatrix}(\mathrm{CC}_{16}, \mathrm{4E}_{16}, \mathrm{FF}_{16}) & (\mathrm{FF}_{16}, \mathrm{2D}_{16}, \mathrm{FF}_{16})\\(\mathrm{FF}_{16}, \mathrm{2C}_{16}, \mathrm{A6}_{16}) & (\mathrm{FF}_{16}, \mathrm{2C}_{16}, \mathrm{A6}_{16})\end{pmatrix}​]​]

In both cases [​[​t​]​] ends up being 1 for [​[​f((\mathrm{FF}_{16}, \mathrm{2D}_{16}, \mathrm{FF}_{16}))​]​] (top-right) and 2 for the others. A constant [​[​t​]​], smaller than would converge, may be used to achieve more-visually-pleasing output, but does not accomplish the functionality described by the original paper video.

# Base case

A scaffolding for the implementation is obtained trivially by copy-pasting the sdl2 documentation examples, mashing them together hap-hazardly until a basic GUI appears, and then never touching them again unless they're broken. The SDL2 experience is not a good one on a good day, and the QoI that falls out of this is even worse, but whatever.

More importantly, the pixel layout format is somehow even worse and unpredictable and you will observe some rows' colour data being wrong on odd-sized images because row-pitch-in-bytes is not always pixel-stride-in-bytes · image-width, which appears undocumented in the bindings, barely-documented upstream, and the pixel format is a &[u8] and a "go to hell". Additionally, it's either per-channel bytes or int with channels allocated from LSB, and the int is in native byte order which makes RGB24 [​[​[\color{var(--r)}r\color{initial}, \color{var(--g)}g\color{initial}, \color{var(--b)}b\color{initial}]​]​] but RGBA32 either [​[​[a, \color{var(--b)}b\color{initial}, \color{var(--g)}g\color{initial}, \color{var(--r)}r\color{initial}]​]​] or [​[​[\color{var(--r)}r\color{initial}, \color{var(--g)}g\color{initial}, \color{var(--b)}b\color{initial}, a]​]​] depending on host endianness, so you need to explicitly pick the opposite: ABGR32 on little-endian to get [​[​[\color{var(--r)}r\color{initial}, \color{var(--g)}g\color{initial}, \color{var(--b)}b\color{initial}, a]​]​] and vice versa. That's enough on SDL2.

The core insights informing the implementation is that [​[​f​]​] does not depend on the pixel position — just its colour — and that [​[​D​]​] is constant for a given [​[​I​]​]. This means that computing [​[​O​]​] can take place in three distinct steps: analysis by computing [​[​D​]​], pre-computing [​[​f​]​] for each distinct element in [​[​D​]​], mapping the pixels directly into [​[​O​]​]. The analysis step can happen just once per [​[​I​]​], regardless of variant and [​[​r​]​] selection.

Since [​[​|C| = 2^{24}​]​] (or, equivalently, [​[​C \equiv {\color{var(--r)}r\color{initial}}{\color{var(--r)}r\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--b)}b\color{initial}}{\color{var(--b)}b\color{initial}}_{16}​]​]), and computer memory is many orders of magnitude larger than 16Mi · any-reasonable-primitive-size, a map from any colour can easily take the form of a look-up table of size 2²⁴. For example, the multiplicity map of [​[​D​]​] (frequencies) and the final map [​[​f​]​] (outputs) are of types [u32; 0xFFFFFF + 1] and [[u8; 3]; 0xFFFFFF + 1] totalling a mere 112MiB. Extending this methodology to higher colour bit depths may require slower-but-more-memory-conscious collection types (10-bit colour ([​[​C = {\mkern 1.5mu\overline{\mkern-1.5mu{1024}\mkern-1.5mu}\mkern 1.5mu}^3​]​]) would grow them to to 7GiB). A 32-bit integer for the multiplicity of each colour means that an entire 4096² image would need to be the same colour to overflow, which isn't a case we care about.

Thus, the analysis step (with [​[​p​]​] as linear_srgb_to_oklab) [CW gore @ link, initially building in "debug" (underoptimised 🙄) mode made the counting appear to be faster in parallel with atomics — this is not the case — this sample has been reduced to the final single-threaded += model, the rest are temporally accurate]:

let frequencies = [0u32; 0xFFFFFF + 1];
for rgb in image.pixel_data.chunks_exact(pixel_stride) {
	let (r, g, b) = (rgb[0], rgb[1], rgb[2]);
	frequencies[(r << 16) | (g << 8) | (b << 0)] += 1;
}

#[repr(C)]
struct Lab { l: f32, a: f32, b: f32 }
let known_rgbs: Vec<(u32, Lab, u32)> = frequencies.iter()
	.enumerate()
	.filter(|(_, freq)| freq != 0)
	.map(|(rgb, freq)| {
		let (r, g, b): (f32, f32, f32) =
			((rgb >> 16) & 0xFF, ((rgb >> 8) & 0xFF), ((rgb >> 0) & 0xFF));
		let lab = linear_srgb_to_oklab((r / 255, g / 255, b / 255));
		(rgb, lab, freq)
	})
	.collect();

yields in known_rgbs a sequence of work items [​[​({\color{var(--r)}r\color{initial}}{\color{var(--r)}r\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--b)}b\color{initial}}{\color{var(--b)}b\color{initial}}_{16}, p(I_{\color{var(--r)}r\color{initial},\color{var(--g)}g\color{initial},\color{var(--b)}b\color{initial}}), \text{multiplicity of}\ \color{var(--r)}r\color{initial},\color{var(--g)}g\color{initial},\color{var(--b)}b\color{initial})​]​].

The pre-computation step:

let radiussy = 0.02;
let mut outputs = [Lab { .. }; 0xFFFFFF + 1];

let radiussy_squared = radiussy * radiussy;
for (rgb, lab, _) in known_rgbs {
	let mut center = lab;

	loop {
		let (newcenter_count, mut newcenter) = known_rgbs.iter()
			.filter(|(_, lab, _)| {
				let distance_squared = (lab.l - center.l).sqr() + (lab.a - center.a).sqr() + (lab.b - center.b).sqr();
				distance_squared <= radiussy_squared
			})
			.fold((0, Lab { l: 0, a: 0, b: 0 }), |(cnt, acc), (_, lab, freq)| {
				#[cfg(palette)]
				(cnt + 1,
				 Lab {
				 	l: acc.l + lab.l,
				 	a: acc.a + lab.a,
				 	b: acc.b + lab.b,
				 })
				#[cfg(mul_freq)]
				(cnt + freq,
				 Lab {
				 	l: acc.l + lab.l * freq,
				 	a: acc.a + lab.a * freq,
				 	b: acc.b + lab.b * freq,
				 })
			});

		newcenter.l /= newcenter_count;
		newcenter.a /= newcenter_count;
		newcenter.b /= newcenter_count;
		if newcenter == center {
			break;
		}
		center = newcenter;
	}

	outputs[rgb] = center;
}

is a very direct implementation of the [​[​q_t​]​] search, with [​[​t​]​] being the iteration count of the loop (the only and obvious optimisation is not computing the square root). All elements of [​[​D​]​] are evaluated every time. The variant is selected at compile-time but could easily be an if, and the outer-most loop parallelises trivially; for the first test image:

the speed-up is less than half of expected, I'd blame this on low memory bandwidth and overall decrodedness of the i7-2600. I'd thought that [​[​*freq​]​] converged ~40% faster due to the specific data contained in the test image, which makes direct comparison of the two difficult; henceforth, images annotated _d are the "degenerate" versions of their base case, with every colour appearing once, making [​[​*freq​]​] equivalent to [​[​palette​]​]:

One would expect (more floating-point operations executed, more data loaded) [​[​*freq​]​] to be, at best, the same speed as [​[​palette​]​], but the optimiser actually makes the [​[​palette​]​] implementation 30% worse. Science cannot explain this.

And the output mapping step (with [​[​p^{-1}​]​] as oklab_to_linear_srgb):

for (to_rgb, from_rgb) in image_processed.pixel_data.chunks_exact_mut(pixel_stride)
                     .zip(image          .pixel_data.chunks_exact    (pixel_stride)) {
	let (r, g, b) = (from_rgb[0], from_rgb[1], from_rgb[2]);
	let rgb = oklab_to_linear_srgb(outputs[(r << 16) | (g << 8) | (b << 0)]);
	to_rgb[0] = (rgb.r * 255).round();
	to_rgb[1] = (rgb.g * 255).round();
	to_rgb[2] = (rgb.b * 255).round();
}

This, well, works. It's not "fast" but since practically every Minecraft texture has <100 distinct colours I expect the original author to have finished here since at that scale each iteration is basically instant. But since each [​[​f(c)​]​] inspects each distinct element of [​[​D​]​] [​[​t​]​] times, the total complexity is [​[​\propto (\overline{Set(D)})^2 t_{avg}​]​] which for more colourful images this simply won't do.

To wit: the unified Minecraft Alpha v1.2.2 block texture atlas has 400 distinct colours and converges in [​[​\Sigma t = 423378​]​]/[​[​\Sigma t = 664403​]​], respectively with the same settings, in <1ms.

Output mapping is a little sub-optimal for large images, since it re-computes [​[​p'​]​] for every pixel every time (a large (2160²) image takes 176ms to map). By in situ applying [​[​p'​]​] to turn [Lab; 0xFFFFFF + 1] into [[u8; 3 + padding]; 0xFFFFFF + 1]:

type RgbStuffedInLab = [u8; mem::size_of::<Lab>()];
for (rgb, _, _) in known_rgbs {
	let result = oklab_to_linear_srgb(outputs[rgb]);

	let to_rgb: &mut RgbStuffedInLab = unsafe { mem::transmute(&mut outputs[rgb]) };
	to_rgb[0] = (result.r * 255).round();
	to_rgb[1] = (result.g * 255).round();
	to_rgb[2] = (result.b * 255).round();
}
let outputs_rgb: &[RgbStuffedInLab; 0xFFFFFF + 1] = unsafe { mem::transmute(&outputs) };

for (to_rgb, from_rgb) in image_processed.pixel_data.chunks_exact_mut(pixel_stride)
                     .zip(image          .pixel_data.chunks_exact(pixel_stride)) {
	let (r, g, b) = (from_rgb[0], from_rgb[1], from_rgb[2]);
	to_rgb.copy_from_slice(outputs_rgb[(r << 16) | (g << 8) | (b << 0)]);
});

This drops to 16ms, and correspondingly less on smaller images — immaterial on the scale of seconds.

The data flow, at this point, can be analysed trivially as

This allows one, like goodyear, to put the

# rubber to the road

When looking for bindings, I was looking for something strongly-typed and with reasonable documentation. To that end, I evaluated wgpu — first by fixing the broken links to the examples, then by getting whiplash from the "Learn" document which simultanously covers "what is Rust", "this implements the WebGPU API to look like the Vulkan API and translates Some how to Some back-end", "how to create a Rust crate", "you must have a window handle and use our logger", and "use this specific version of an intermediate dependency otherwise it all explodes". No googleable compute shaders seem to be extant for it, either.

After a few less remarkable ones, I'd settled on vulkanalia because, in reverse order of import:

1. he he kanalia,
2. it appeared to competently implement both a direct Vulkan binding and a 10% more structured and 40% less verbose "builder" scheme that supposedly avoided dangling references (in retrospect, it's done that quite well), thus allowing easy adaptation of other Vulkan posts,
3. it has a direct port — Vulkan Tutorial (Rust), Kyle Mayes, trans., Alexander Overvoorde, &a. — of the upstream Vulkan tutorial, which is basically best-in-class.

Admittedly, the tutorial also wants to make a window and render but this can be strategically ignored, and which bits are strictly for windowing is relatively obvious. The port doesn't include the Compute Shader sexion, but, counter to the rest of the text, I've found it to be quite poor. To this end, A Simple Vulkan Compute Example in C++, Baked Bits, 2021 performs quite well, working out from "you have a functioning vulkan context" holistically through to "you are executing the shader".

As part of the Example, a shader compilation method is specified with dxc, which doesn't appear to be distributed by any-one (defined here-in as "Debian or MSYS2"). The standard and universal solution appears to be glslc, which, contranomenclaturally, actually compiles GLSL and HLSL, and also has a C-style preprocessor. Not a great one, but, better yet, a good-enough one. This allows a simple build.rs:

fn main() {
	let processes = vec![];
	for f in fs::read_dir("src").map(|r| r.path()).filter(|p| p.extension() == "hlsl") {
		println!("cargo::rerun-if-changed={}", f.display());
		proceses.push(std::process::Command::new("glslc")
			.args(&["-fshader-stage=compute", "-g", "-O"])
			.args(&["-o", Path::new(&env::var_os("OUT_DIR")).join(f.file_name()).with_extension("spv")])
			.args(&[f])
			.spawn());
	}
	for proc in processes
		proc.wait().exit_ok().unwrap();
	}
}

which is equivalent to a glslc -fshader-stage=compute -g -O -o $(OUT_DIR)/$(subst .hlsl,.spv,$@) $< rule for every file matching src/*.hlsl. -fshader-stage could've been substituted by replacing .hlsl with .comp, and thankfully, in contrast to the devilish DXC cmdline, -O and -g are like in normal compilers (though, with the "targeting IR" compilation model, -g has an apt side-effect of including the full source and per-instruction line number annotations in the shader binary! this is convenient when you need to figure out if the optimiser is fucking you, but irrelevant most of the time). Overall, the tooling is pretty mature, just the subject matter largely hateful.

What is relevant but not immediately obvious is that DXC requires the entry point to be specified explicitly (as Main in the Example), but glslc defaults to main, and this needs to agree with both the source and vk::PipelineShaderStageCreateInfoBuilder::name. HLSL in general has, to me, a similar stench to DXC, what with the Microsoft. I haven't yet been able to produce GLSL by any means other than decompilation, and, while uniquely Microsofty, documentation for HLSL "exists" and "is useable". Ah well.

More confusingly, vk::ShaderModuleCreateInfo::code and ::code_size takes, separately, an array of u32s and a length of that array in u32s (not bytes!). This array must be aligned, even on amd64 where misalignment doesn't usually trap. As an added bonus from a testing perspective, you do just get a SIGBUS, but not on Win32, because the linker always aligns the shader (well, did for me). Nothing an include_typed!(u32, ".../shader.spv") macro can't fix, but the Cargo compilation model and the way the implementation works being "magic" did not help. Some would say that Rust is not any better than C++ at this (but, well, I would be able to slap a [[alignas(std::uint32_t)]] on it there, so this is actually kind of an L).

The terminology is a little sneaky as well: what Vulkan calls a "pipeline" is actually just a shader in human language, and a "command buffer" is a pipeline. One may also that think "command buffers"… buffer commands. They don't, you can only put one pipeline (though an arbitrarily complex one!) in there at a time. Most-everything else scans relatively sanely.

I had been a little slow on the up-take on validation layers (largely owing to the fact that the MSYS2 package didn't make them work and I didn't really want to fuck my whole shit up by installing something random; this is of course not an issue for users of normal distributions). My concern was misplaced: the upstream LunarG Vulkan SDK Just Worked and didn't interfere with anything else. The validation layers are, frankly, a marvel of a dynamic analysis suite (and they let you have printf-from-shader).

What's a bit less of a wonder is the non-printf debugging sitch. If you're rendering you use RenderDoc, but that doesn't capture compute pipelines. Well it claims it does, experimentally, but I haven't managed to have that happen. Vendor-specific tooling — in my case the Radeon Developer Panel — is presumably the only option, but I haven't gotten it to work on my computer and with my GPU and with Vulkan. Some part of the fire triangle's always failed to materialise for me.

The Example's seemingly-only fault is that it does not use memory barriers, and it's true that it seems to work fine without them, but more complex programs (however they're defined, this one definitely fits!) do actually need them, even in coherent memory and if no concurrent access occurs and everything is fence-synchronised. These have been elided from the samples here because they're pretty verbose.

With this pre-amble/summary in mind and no fan-fare, Example's shader works. Since known_rgbs could have anywhere [​[​\in [1, 2^{24}]​]​] elements (but is very unlikely to have actually that many), a way to bump the buffers is nice.

In a few simple steps it can be expanded and adapted to implement the pre-computation step:

struct params_t {
	uint  known_rgbs_len;
	float radiussy;
	uint  iter_limit;
};

struct known_rgbs_pack {
	float3 lab;
	uint   rgb;
};

[[vk::binding(0)]] RWStructuredBuffer<float3>          outputs;
[[vk::binding(1)]]   StructuredBuffer<params_t>        params;
[[vk::binding(2)]]   StructuredBuffer<known_rgbs_pack> known_rgbs_bundle;

[numthreads(64, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
	const uint  known_rgbs_len = params[0].known_rgbs_len;
	const float radiussy       = params[0].radiussy;
	const uint  iter_limit     = params[0].iter_limit;

	const uint id = DTid.x;

	known_rgbs_pack this = known_rgbs_bundle[id];
	float3 center = this.lab;
	for(uint iter = 0; iter < iter_limit; ++iter) {
		uint newcenter_count = 0;
		float3 newcenter     = float3(0, 0, 0);
		for(uint i = 0; i < known_rgbs_len; ++i) {
			float3 lab = known_rgbs_bundle[i].lab;
			if(length(lab - center) <= radiussy) {
				newcenter_count += 1;
				newcenter       += lab;
			}
		}

		newcenter /= newcenter_count;
		bool same = newcenter == center;
		if(same)
			break;
		center = newcenter;
	}
	outputs[this.rgb] = center;
}

As driven by (initially the nomenclature was derived from posterisation instead of mean-shift clustering; this is be corrected later):

static SHADER_DATA: &[u32] = include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise.spv"));

#[repr(C)]
struct ParamsT {
	known_rgbs_len: u32,
	radiussy: f32,
	iter_limit: u32,
}
#[repr(C)]
struct KnownRgbsPack {
	lab: Lab,
	rgb: u32,
}

pub struct PosteriseGpu {
	// ...
	device: vulkanalia::Device,
	compute_pipeline: vulkanalia::Pipeline,

	memory_type_index: u32,                 // HOST_VISIBLE & HOST_COHERENT
	buffer_outputs: BufferBundle,           // [[vk::binding(0)]] SB<float3>          outputs;
	buffer_params: BufferBundle,            // [[vk::binding(1)]] SB<params_t>        params;
	buffer_known_rgbs_bundle: BufferBundle, // [[vk::binding(2)]] SB<known_rgbs_pack> known_rgbs_bundle;

	buffer_outputs_map: &'static mut [[f32; 4]; 0xFFFFFF + 1],
	buffer_params_map: &'static mut ParamsT,
}

fn PosteriseGpu::init() -> PosteriseGpu {
	// ...

	let vk_shader_module = vk_device.create_shader_module(&vk::ShaderModuleCreateInfo::builder()
		.code_size(SHADER_DATA.len() * 4).code(SHADER_DATA), None);
	dev.compute_pipeline = {
		let pipeline_shader_create_info = vk::PipelineShaderStageCreateInfo::builder()
			.stage(vk::ShaderStageFlags::COMPUTE)
			.module(vk_shader_module)
			.name(b"main\0");
		let compute_pipeline_create_info = vk::ComputePipelineCreateInfo::builder()
			.stage(pipeline_shader_create_info) /*...*/;
		vk_device.create_compute_pipelines(vk_pipeline_cache, &[compute_pipeline_create_info], None).0.swap_remove(0)
	};

	ret.buffer_outputs = ret.ensure_buffer(BufferBundle{}, mem::size_of::<[[f32; 4]; 0xFFFFFF + 1]>());
	ret.buffer_outputs_map = &mut *ret.device.map_memory(ret.buffer_outputs);

	ret.buffer_params = ret.ensure_buffer(BufferBundle{}, mem::size_of::<ParamsT>());
	ret.buffer_params_map = &mut *ret.device.map_memory(ret.buffer_params);
	ret
}

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	self.buffer_known_rgbs_bundle = self.ensure_buffer(self.buffer_known_rgbs_bundle, known_rgbs.len() * mem::size_of::<KnownRgbsPack>());
	assert!(mem::size_of::<KnownRgbsPack>() == 4 * 4);

	*self.buffer_params_map = ParamsT {
		known_rgbs_len: known_rgbs.len(),
		radiussy: radiussy,
		iter_limit: 0xFFFFFFFF,
	};

	let buffer_known_rgbs_bundle_map = self.device.map_memory(self.buffer_known_rgbs_bundle);
	let known_rgbs_bundle: &mut [KnownRgbsPack] = slice::from_raw_parts_mut(buffer_known_rgbs_bundle_map, known_rgbs.len());
	for ((rgb, lab, _), dest_lab) in known_rgbs.iter().zip(known_rgbs_bundle.iter_mut()) {
		dest_lab = KnownRgbsPack {
			rgb: rgb,
			lab: lab,
		};
	}
	self.device.unmap_memory(self.buffer_known_rgbs_bundle.buffer_memory);

	self.descriptor_set = self.device.update_descriptor_sets(vk::WriteDescriptorSet{[self.buffer_outputs, self.buffer_params, self.buffer_known_rgbs_bundle]});

	self.device.begin_command_buffer(self.cmd_buffer);
	self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE, self.compute_pipeline);
	self.device.cmd_bind_descriptor_sets(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE, self.pipeline_layout, 0, &[self.descriptor_set], &[]);
	// cmd_pipeline_barrier(HOST -> COMPUTE_SHADER: (self.buffer_params, self.buffer_known_rgbs_bundle): HOST_WRITE -> SHADER_READ)
	self.device.cmd_dispatch(self.cmd_buffer, known_rgbs.len() /* x */, 1 /* y */, 1 /* z */);
	// cmd_pipeline_barrier(COMPUTE_SHADER -> HOST: self.buffer_outputs: SHADER_WRITE -> HOST_READ)
	self.device.end_command_buffer(self.cmd_buffer);

	self.device.reset_fences(&[self.fence]).unwrap();
	self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::builder().command_buffers(&[self.cmd_buffer])], self.fence);
	self.device.wait_for_fences(&[self.fence], true, u64::MAX).unwrap();

	for (rgb, _, _) in known_rgbs {
		outputs[rgb].l = self.buffer_outputs_map[rgb][0];
		outputs[rgb].a = self.buffer_outputs_map[rgb][1];
		outputs[rgb].b = self.buffer_outputs_map[rgb][2];
	}
}

The Example mapped and unmapped the buffers every time at point of use, which seems to have no up-side: params and outputs are always mapped.

It uses [​[​\text{numthreads} = (1, 1, 1), \text{dispatch}(k, 1, 1)​]​], which is fine for its application. On expert advice for "why is this so damn slow still", I'd naively upgraded to [​[​\text{numthreads} = (64, 1, 1), \text{dispatch}(k, 1, 1)​]​] (where [​[​k ={\ }​]​]known_rgbs.len()), which exposes a glaring bug: the "threads" dispatched in each dimension are [​[​\prod_{x \in \text{numthreads} \odot \text{dispatch}} [0, x )\ (\text{all in}\ \mathbb{Z})​]​], so, in this case, [​[​[0, 64n ) \times [0, 1 ) \times [0, 1 )​]​] instead of the desired [​[​[0, \mathbin{\approx}k ) \times [0, 1 )^2​]​] (which took even longer, hence why iter_limit to avoid spinning forever in uninitialised memory out of bounds).

I've only just now realised Vulkan Tutorial, Compute Shader, Alexander Overvoorde &a. describes this (hint: it's not Running compute commands, Dispatch! it's actually Compute Space!) and I don't think it does a great job of it: the description is incredibly wordy and disconnected from where the other "256" is, defines only the edge-case when [​[​k \bmod 64 = 0​]​], and the shader is not prepared for the ignore-if-past-the-end situation. Also, the way this maps to the variables/globals is not referenced at all.

I've found Direct3D / HLSL / numthreads, Microsoft Learn to surprisingly be the best here, and the explanatory figure, while more schizo, actually more informative. The correct implementation is [​[​\text{numthreads} = (64, 1, 1), \text{dispatch}(\lceil\frac{k}{64}\rceil, 1, 1)​]​] which yields [​[​[1, k+\varepsilon ) \times [0, 1 )^2​]​] where [​[​\varepsilon \in [0, 64 )​]​]. The former can be conveniently expressed builtinly as known_rgbs.len().div_ceil(64) and the latter rejected with

	const uint id = DTid.x;
	if(id >= current_rgbs_len)
		return;

this appears canonical, but honestly I haven't found a model would satisfactorily define the appx. 10 different views into the job dispatch matrix. (Personal failing.)

Why 64 specifically? I measured and everything else is worse. Something to do with compute unit sizes or whatever? Unclear.

That bit about the Example having one fault? Not true actually. It doesn't mention fence resetting. If you're like me, god forbid, and don't intently read through 60% of the Vulkan Tutorial (since you're actually reading Vulkan Tutorial (Rust) which only mentions it in Rendering and presentation), and be directed to consume Vulkan Guide / Using Vulkan / Synchronization / Synchronization Examples in its entirety; it's instructive and honestly the best piece of documentation mentioned so far, but still doesn't mention that you indeed have missed the need for the reset_fences() call and this was worth two hours.

The fence overall is unnecessary if not scheduling multiple shaders at once, and can be replaced with a queue_wait_idle() call, but if you're waiting on it and it's already completed once then it'll just not block again.

One may note the mismatch of outputs being SB<float3> but then also [[f32; 4]; …], as well as params also being in a SB even though there's only one of them, ever. This is due to a confluence of two HLSL oddities:

• arrays have either a 4-byte or a 16-byte alignment
• globals not part of a SB are collected and encapsulated into a single global you can't control

thus, SB<float3> is actually SB<(float3, 4 unused bytes), which also means that known_rgbs_bundle is packed optimally. Arrays whose elements are 1/2 scalars (4/8 bytes) are 4-byte-aligned. Larger arrays are 16-byte-aligned. Structures are, functionally, arrays of their elements as well: struct { float3, float3 } has a padding scalar between them, which you can use if you want (this explains the goofy lay-out of params_t later). This, AIUI, corresponds to GLSL std430/std140 layouts, except determined automagically and you have to validate the layout yourself sometimes. Fun!

This is a direct port of the CPU implementation above (sans square distance optimisation), and only of [​[​palette​]​]. The hawk-tuah'deyed reader has spotted that center is redundant, and so is bool same, except the assignment from bool3 acts like an all() (every operator is piece-wise, and == is no exception; some-how for other types this is obvious to me, but for booln it's consistently failed to internalise. Personal failing 2). With the dispatching fixed (all times measured hence include copying buffer_outputs_map to outputs):

The memory lay-out is largely replicated directly as well (this is not fast or good, but it is drop-in): outputs is fixed-size since it's a scatter output, but known_rgbs_bundle is only as big as known_rgbs (distinct colours in [​[​I​]​]), so it can be bumped if the new size required is larger than the last size. Since reallocating means "free, allocate", it can't be mapped for any longer. At least we get the real allocated size as SOP, so if there's any rounding up (there is), we are aware and can skip extraneous allocations.

This could, in principle, function as Just A Faster Implementation, but my Win32 amdgpu-equivalent driver kills the shader and returns TIMEOUT if it takes longer than ~8s (which is honestly better than any alternative and we see how many values were fully-computed since outputs was updated by the threads that did manage to finish).

By observing that length(v) [​[​= \sqrt{v_x^2 + v_y^2 + v_z^2} = \sqrt{v_xv_x + v_yv_y + v_zv_z} = \sqrt{v \cdot v} =​]​] sqrt(dot(v, v)) (though it actually lowers directly into OpExtInst %float (OpExtInstImport "GLSL.std.450") Length %v instead of OpExtInst %float %… Sqrt %(OpDot %float %v %v)), the obvious optimisation can be restored (with dispatching fixed likewise):

struct params_t {
	uint  known_rgbs_len;
	float radiussy_squared;
	uint  iter_limit;
};
// ...
			float3 tmp = lab - center;
			if(dot(tmp, tmp) <= radiussy_squared) {

Now all it takes for the

# base-line

port is implementing [​[​*freq​]​]. glslc's preprocessor implementation allows a de-duplicated implementation, much like the #[cfg]s prior; posterise_kernel.hlsl.inc:

struct params_t {
	uint  known_rgbs_len;
	float radiussy_squared;
	int   iter_limit;
};

struct known_rgbs_pack {
	float3 lab;
	uint   rgb;
};

[[vk::binding(0)]] RWStructuredBuffer<float3>          outputs;
[[vk::binding(1)]]   StructuredBuffer<params_t>        params;
[[vk::binding(2)]]   StructuredBuffer<known_rgbs_pack> known_rgbs_bundle;
KNOWN_RGBS_FREQS

[numthreads(64, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
	const uint  known_rgbs_len   = params[0].known_rgbs_len;
	const float radiussy_squared = params[0].radiussy_squared;
	const uint  iter_limit       = params[0].iter_limit;

	const uint id = DTid.x;
	if(id >= known_rgbs_len)
		return;

	known_rgbs_pack this = known_rgbs_bundle[id];
	float3 center = this.lab;
	for(int iter = 0; iter != iter_limit; ++iter) {
		uint newcenter_count = 0;
		float3 newcenter     = float3(0, 0, 0);
		for(uint i = 0; i < known_rgbs_len; ++i) {
			float3 lab = known_rgbs_bundle[i].lab;

			float3 tmp = lab - center;
			if(dot(tmp, tmp) <= radiussy_squared) {
				FREQ_LOAD
				newcenter_count += FREQ;
				newcenter       += lab TIMES_FREQ;
			}
		}

		newcenter /= newcenter_count;
		bool same = newcenter == center;
		if(same)
			break;
		center = newcenter;
	}
	outputs[this.rgb] = center;
}

posterise.hlsl:

#define KNOWN_RGBS_FREQS
#define FREQ_LOAD
#define FREQ             1
#define TIMES_FREQ
#include "posterise_kernel.hlsl.inc"

posterise_mul.hlsl:

#define KNOWN_RGBS_FREQS [[vk::binding(3)]] StructuredBuffer<uint4> known_rgbs_freqs;
#define FREQ_LOAD        uint freq = known_rgbs_freqs[i / 4][i % 4];
#define FREQ             freq
#define TIMES_FREQ       * freq
#include "posterise_kernel.hlsl.inc"

These expand to

for(uint i = 0; i < known_rgbs_len; ++i) {
	float3 lab = known_rgbs_bundle[i].lab;

	float3 tmp = lab - center;
	if(dot(tmp, tmp) <= radiussy_squared) {
		newcenter_count += 1;
		newcenter       += lab;
	}
}

and

for(uint i = 0; i < known_rgbs_len; ++i) {
	float3 lab = known_rgbs_bundle[i].lab;

	float3 tmp = lab - center;
	if(dot(tmp, tmp) <= radiussy_squared) {
		uint freq = known_rgbs_freqs[i / 4][i % 4];
		newcenter_count += freq;
		newcenter       += lab * freq;
	}
}

respectively.

known_rgbs_freqs being an array of uint4s is unfortunate and unnecessary, but the minutiae of automagic layout bull shit were not obvious to me at the time. This is functionally-equivalent to an array of uints being indexed by i directly.

By storing the population data separately with the same indexing, the optimal zero-padding layout is preserved (adding it into a known_rgbs_pack would mean { (float3, uint), (uint, 3 4-byte padding primitives) }), and the population can be only loaded when necessary, which minimises performance impact.

Since the uniforms ([[vk::binding]]-tagged StructuredBuffers; it's unclear why some literature calls them "uniforms", Vulkan calls them bindings or bound descriptor sets and the like universally AFAICT) use a global indexing system defined by the host, any shader can use any subset of them. This will be (ab)used extensively, and every shader will be given all of them, but will pick out and RW-mark only the ones it cares about. Nominally, Vulkan only has 4 bindings in baseline, but you just ask for however-many you need and abort; it may be nice to check before-hand, but the value is marginal since all GPUs mentioned heretofore support 32.

The driver requires minimal modification, functionally duplicating most of it (this can be instrumented well with [Т; N].map(); suffers a bit with Rust's lack of C-style auto-incrementing self-referential enums. ah well):

pub struct PosteriseGpu {
	// ...
	device: vulkanalia::Device,
	compute_pipeline: vulkanalia::Pipeline,     // posterise.hlsl
	compute_pipeline_mul: vulkanalia::Pipeline, // posterise_mul.hlsl

	memory_type_index: u32,                 // HOST_VISIBLE & HOST_COHERENT
	buffer_outputs: BufferBundle,           // [[vk::binding(0)]] SB<float3>          outputs;
	buffer_params: BufferBundle,            // [[vk::binding(1)]] SB<params_t>        params;
	buffer_known_rgbs_bundle: BufferBundle, // [[vk::binding(2)]] SB<known_rgbs_pack> known_rgbs_bundle;
	buffer_known_rgbs_freqs: BufferBundle,  // [[vk::binding(3)]] SB<uint4>           known_rgbs_freqs;

	buffer_outputs_map: &'static mut [[f32; 4]; 0xFFFFFF + 1],
	buffer_params_map: &'static mut ParamsT,
}

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, freq_weighting: bool, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	self.buffer_known_rgbs_bundle = self.ensure_buffer(self.buffer_known_rgbs_bundle, known_rgbs.len() * mem::size_of::<KnownRgbsPack>());
	assert!(mem::size_of::<KnownRgbsPack>() == 4 * 4);
	if freq_weighting {
		self.buffer_known_rgbs_freqs = self.ensure_buffer(self.buffer_known_rgbs_freqs, known_rgbs.len().div_ceil(4) * mem::size_of::<[f32; 4]>());
	}

	self.buffer_params_map = ParamsT {
		known_rgbs_len: known_rgbs.len(),
		radiussy: radiussy,
		iter_limit: -1,
	};

	let buffer_known_rgbs_bundle_map = self.device.map_memory(self.buffer_known_rgbs_bundle);
	let known_rgbs_bundle: &mut [KnownRgbsPack] = slice::from_raw_parts_mut(buffer_known_rgbs_bundle_map, known_rgbs.len());
	for ((rgb, lab, _), dest_lab) in known_rgbs.iter().zip(known_rgbs_bundle.iter_mut()) {
		dest_lab = KnownRgbsPack {
			rgb: rgb,
			lab: lab,
		};
	}
	self.device.unmap_memory(self.buffer_known_rgbs_bundle.buffer_memory);

	if freq_weighting {
		let buffer_known_rgbs_freqs_map = self.device.map_memory(self.buffer_known_rgbs_freqs);
		let known_rgbs_freqs: &mut [[u32; 4]] = slice::from_raw_parts_mut(buffer_known_rgbs_freqs_map, known_rgbs.len().div_ceil(4));
		for ((_, _, freq), dest_freq) in known_rgbs.iter().zip(known_rgbs_freqs.iter_mut().flatten()) {
			dest_freq = freq;
		}
		self.device.unmap_memory(self.buffer_known_rgbs_freqs.buffer_memory);
	}

	self.descriptor_set = self.device.update_descriptor_sets(vk::WriteDescriptorSet{[/* … */, self.buffer_known_rgbs_freqs]});

	self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE,
		if freq_weighting { self.compute_pipeline_mul } else { self.compute_pipeline });
	// cmd_pipeline_barrier(HOST -> COMPUTE_SHADER: also self.buffer_known_rgbs_freqs: HOST_WRITE -> SHADER_READ)
	self.device.cmd_dispatch(self.cmd_buffer, known_rgbs.len().div_ceil(64) /* x */, 1 /* y */, 1 /* z */);

	self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::builder().command_buffers(&[self.cmd_buffer])], self.fence);
	self.device.wait_for_fences(&[self.fence], true, u64::MAX).unwrap();

	for (rgb, _, _) in known_rgbs {
		outputs[rgb].l = self.buffer_outputs_map[rgb][0];
		outputs[rgb].a = self.buffer_outputs_map[rgb][1];
		outputs[rgb].b = self.buffer_outputs_map[rgb][2];
	}
}

[​[​*freq​]​] being 16% slower than [​[​palette​]​] finally matches the expectation based on complexity; [​[​palette​]​] hasn't changed, and it's dead-on the same.

So far, all buffers have been allocated in HOST_VISIBLE & HOST_COHERENT memory like in the Example, which was basically fine, but proper usage of

# memory types

is required to achieve best performance with the widest device support matrix. vk_instance.get_physical_device_memory_properties(physical_device) yields a list of memory types, already sorted by preference. Thus, "a [​[​Q​]​] memory" is commonly taken to mean "the first memory type that satisfies predicate [​[​Q​]​]", and since memory features are typified by bitmask, the predicate is mostly expressed as a conjunction of desired bits. Listing the memory types on the GPUs in question yields:

For the purposes of writing the shaders, I was targeting the 5700 XT in my desktop; conveniently, this is also the most restrictive device. It appears that GPUs can only see up to 32GiB of host memory (or, indeed, memory at all), regardless of how much is actually attached (I can come up with excuses for the discrete 6500 XT, but the Intel iGPU? no clue). It's likewise unclear to me why the 6500 XT doubles up on some memory types (it's not the heap flags, those are also the same (DEVICE_LOCAL and empty, resp.)). Doesn't really matter or affect funxionality in any way, it's just a little goofy.

All three input buffers are accessed constantly (and written directly by the host), so they need to be in DEVICE_LOCAL & HOST_VISIBLE & HOST_COHERENT memory (taking up [​[​\operatorname{sizeof}(\texttt{params}) + k (\operatorname{sizeof}(\texttt{}\text{(}\texttt{float3}\text{, }\texttt{uint}\text{)}) + \operatorname{sizeof}(\texttt{uint}))​]​], so 20B per unique colour). This is a precious resource, so outputs can't fit here ([​[​2^{24} \operatorname{sizeof}(\texttt{float4})​]​] is the whole [​[​256\text{MiB}​]​]!), but from the perspective of the GPU, outputs is write-once write-only, and the last iteration of [​[​f​]​], functionally, uploads the result, so it doesn't need to be DEVICE_LOCAL anyway, since the latency or speed don't matter.

Fall-backs are implemented from more specific to more generic memory types (DEVICE_LOCAL & HOST_VISIBLE & HOST_COHERENT → HOST_VISIBLE & HOST_COHERENT), but I haven't encountered a device where this fall-back is taken. Using these improved memory types yields

# who up branchin they execution

If there's half of one potential thing I know about the GPU execution model, it's that branches are either expensive or not real or some mixture of the two. [​[​palette​]​]'s branch can be flattened out; so can [​[​*freq​]​]'s (but shouldn't, it's worse), with a trivial change to the common kernel:

#define KNOWN_RGBS_FREQS
#define BRANCHLESS_ADD   1
#define FREQ_LOAD
#define FREQ             matched
#include "posterise_kernel.hlsl.inc"

#define KNOWN_RGBS_FREQS [[vk::binding(3)]] StructuredBuffer<uint4> known_rgbs_freqs; // DEVICE_LOCAL & HOST_VISIBLE & HOST_COHERENT
#define BRANCHLESS_ADD   0
#define FREQ_LOAD        uint freq = known_rgbs_freqs[i / 4][i % 4];
#define FREQ             freq
#include "posterise_kernel.hlsl.inc"

for(uint i = 0; i < known_rgbs_len; ++i) {
	float3 lab = known_rgbs_bundle[i].lab;

	float3 tmp = lab - center;
	const bool matched = dot(tmp, tmp) <= radiussy_squared;
#if !BRANCHLESS_ADD
	if(matched)
#endif
	{
		FREQ_LOAD
		newcenter_count += FREQ;
		newcenter       += FREQ * lab;
	}
}

which expand to, effectively,

float3 tmp = lab - center;
const bool matched = dot(tmp, tmp) <= radiussy_squared;
newcenter_count += matched;
newcenter       += matched * lab;

float3 tmp = lab - center;
const bool matched = dot(tmp, tmp) <= radiussy_squared;
if(matched) {
	uint freq = known_rgbs_freqs[i / 4][i % 4];
	newcenter_count += freq;
	newcenter       += freq * lab;
}

yielding

For reference, "worse" means, giving [​[​*freq​]​] BRANCHLESS_ADD = 1 and FREQ = matched * freq:

# Hitting the griddy? In this economy?

Any particle system, which this is, is improved by spatially partitioning constituent particles, and iterating only those whose partitions fall within the radius of influence. This applies basically-straight-forwardly to this problem, except that since the underlying set of [​[​Multiset(D)​]​] is constant, the elements of [​[​D​]​] can be assigned to constant cells (or, equivalently but more importantly, cells can be said to contain a constant subset of [​[​D​]​], where no two cells share any elements and the union of all cells is [​[​D​]​]) a priori, but each iteration of [​[​q_t​]​] moves around, so the cell it's found in will differ.

For optimal (well, "in any way meaningful") occupancy, we must additionally observe a nigh-universal factoid about the shape of the data:

As is hopefully obvious from the figures, the range of [​[​P​]​] actually used by any given input tends to be quite narrow, so subdividing [​[​P​]​] directly will yield poor occupancy. Indeed, for the datasets visualised above:

so, with Oklab [​[​P​]​]'s volume being 1, the average bounds are just 5.33% of that!

For a cartesian grid of dimension [​[​G \in \mathbb{N}_+​]​] we can define a bijective map [​[​g : P \to [0, G )^3​]​] such that

a bare definition of [​[​g​]​] (without [​[​\min​]​]) would also include [​[​{G}^3​]​] in the codomain for points on the far extrema of the bounding box; we abuse the fact that we don't actually use [​[​g^{-1}(g(p))​]​] as data, just as optimisation metadata, to make it much more useful (the error induced is imperceptible after quantisation later): the grid cell each [​[​p​]​] falls in is given by [​[​\operatorname{trunc}(g(p)) \in \bar{G}^3​]​]. The points within each cell [​[​h : \bar{G}^3 \to \mathcal{P}(D)​]​] can then be obtained by computing [​[​h(x) = \{p | p \in D \wedge \operatorname{trunc}(g(p)) = x\}​]​]; the desired properties [​[​h(x) \cap h(y) \ne \varnothing \Leftrightarrow x = y \text{ & } \bigcup (h \circ {\operatorname{trunc}} \circ g \circ D) = D​]​] are preserved under these mappings.

# box bondage 🥺

Computing [​[​(B_m, B_M)​]​] is more-usefully analysed as a [​[​(\min, \max)​]​] fold over [​[​D​]​] starting at [​[​(+\infty^3, -\infty^3)​]​]. To this end, [​[​\min​]​] and [​[​\max​]​] are available as atomics — HLSL calls them InterlockedMin/-Max even though no locking, inter- or otherwise, takes place; GLSL nomenclature is reasonable, so this is an unforced error on Microsoft's part — but only on integers, which coordinates in [​[​D​]​] aren't.

But they are IEEE754 floats, which have a well-behaved structure when interpreted as-if they were integers. Indeed, if [​[​P​]​] were in the form [​[​[0, x]^3​]​], no further processing would be required, since positive floats compare the same as uints (other relations can be drawn, cf. No. Two floating point values (IEEE 754 binary64) cannot compare, chux, 2015). But it isn't, so we desire a bijection [​[​u : float \to uint \text{ or } int​]​] monotonic in [​[​(-\infty, +\infty)​]​], ideally something I don't have to derive because that sucks battle-tested. The literature provides Bullet 2.81, GPU Soft Body Solvers, DX11, Erwin Coumans &a., 2012 (and [​[​u^{-1}​]​]), for 32-bit floats and uints (notationally, this is where real math (Wikipedia lists operators in [​[​\LaTeX​]​]) ends and "computer science" (Wikipedia l. op. in Courier) begins):

u(f)   = asuint(f) ^ ((~(asuint(f) >> 31) + 1) | 0x8000`0000) 
u-1(u) = asfloat(u ^ (( (       u  >> 31) - 1) | 0x8000`0000))

Within the [​[​(-\infty, +\infty)​]​] domain, IEEE754 floats consist of [​[​(1\ \text{sign bit})(31\ \text{bits of distance from}\ 0\ \text{in natural }\texttt{uint}\text{ sort order})​]​]. There's more stuff in there but it's not relevant here; Google float \infty to learn more. The [​[​l​]​] axis from example β can thus be transcoded as

This reveals some core insights:

• [​[​\begin{align} \text{If}\ sgn(f) \text{:}\\ u(f) &= f \oplus \mathrm{FFFFFFFF}_{16} = \neg f\\ u^{-1}(u) &= u \oplus \mathrm{FFFFFFFF}_{16} = \neg u\\ \text{Otherwise:}\\ u(f) &= f \oplus \mathrm{80000000}_{16}\\ u^{-1}(u) &= u \oplus \mathrm{80000000}_{16} \end{align}​]​]
• the sign bit is always flipped, yielding the top bit set for positive fs and clear for negative fs, inducing a clear partition:
  u( 0.0) = 0x8000`0000 = 0b1000`0000`0000`0000`0000`0000`0000`0000
  u(−0.0) = 0x7FFF`FFFF = 0b0111`1111`1111`1111`1111`1111`1111`1111
  but this means that the condition is mirrored between u and u^-1
• since the remaining bits are "distance from zero", so more negative numbers are larger, they are flipped so more negative numbers are smaller
• functionally, this is basically most of a (sign + magnitude) ↔ (two's complement) implementation (just missing the + 1 at the end for negative numbers)
• both major optimisers (and, accd'g to my measurements, all three GPU driver optimisers) turn ~(f >> 31) + 1 into f >>_arithmetic 31 (since the desired state equvalent to the top bit being copied to every other bit, which a sign-extension all the way down will do); this reduces u from (lsr31,neg,add,or,xor) to (asr31,or,xor) which is unbeatable
• this could be more clearly (but less goodly since it induces a branch) be expressed ideologically as

u(f)   = asuint(f) ^ (0xFFFF`FFFF if (asuint(f) & 0x8000`0000) else 0x8000`0000)
u-1(u) = asfloat(u ^ (0xFFFF`FFFF if not    (u  & 0x8000`0000) else 0x8000`0000)

[​[​B_m​]​] and [​[​B_M​]​]/[​[​B_r​]​] can be allocated directly, minding the free scalars in each 4-primitive block. Since the reduxion wants to do atomic operations on integers (and no aliasing is possible), these must be declared as uint3s for that shader, and may be float3 in others; theoretically it's possible to keep known_rgbs_pack::lab as a float3 and bit-cast with asuint() but there's no reason to do this. So, lifting the declaration to a common shader.h:

struct params_t {
	uint    known_rgbs_len;
	float   radiussy_squared;
	int     iter_limit;
	uint    _padding1;

	point_t bbox_min;
	uint    _padding2;

	point_t bbox_range;
	uint    _padding3;
#define bbox_max bbox_range  // find_bbox.hlsl only
};

struct known_rgbs_pack {
	point_t lab;
	uint    rgb;
};

thus, find_bbox.hlsl can typedef point_t one way, and everything else — the other (also note how Bullet 2.81 does it a little goofy and applies u per-component; there is no reason to do this):

typedef uint3 point_t;
#include "shader.h"

[[vk::binding(1)]] RWStructuredBuffer<params_t>        params;             // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(2)]]   StructuredBuffer<known_rgbs_pack> known_rgbs_bundle;  // DEVICE_LOCAL & HOST_VISIBLE

[numthreads(64, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
	const uint known_rgbs_len = params[0].known_rgbs_len;

	const uint id = DTid.x;
	if(id >= known_rgbs_len)
		return;

	uint3 tmp = known_rgbs_bundle[id].lab;
	tmp ^= (~(tmp >> 31) + 1) | 0x80000000;

	InterlockedMin(params[0].bbox_min.x, tmp.x);
	InterlockedMin(params[0].bbox_min.y, tmp.y);
	InterlockedMin(params[0].bbox_min.z, tmp.z);

	InterlockedMax(params[0].bbox_max.x, tmp.x);
	InterlockedMax(params[0].bbox_max.y, tmp.y);
	InterlockedMax(params[0].bbox_max.z, tmp.z);
}

and the driver just needs to supply the integral equivalent of [​[​(+\infty^3, -\infty^3)​]​] and post-process with u^-1 (this is where a real enum would've been nice):

const FIND_BBOX: usize = 2;
const SHADER_COUNT: usize = FIND_BBOX + 1;
static SHADER_DATA: [&[u32]; SHADER_COUNT] = [include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise_mul.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/find_bbox.spv"))];
pub struct PosteriseGpu {
	// ...
	compute_pipeline: Vec<vk::Pipeline>, // [SHADER_COUNT]
}

#[repr(C)]
struct ParamsT {
	known_rgbs_len: u32,
	radiussy_squared: f32,
	iter_limit: i32,
	_padding1: u32,

	bbox_min: [f32; 3],
	_padding2: u32,

	bbox_range: [f32; 3],
	_padding3: u32,
}

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, freq_weighting: bool, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	self.buffer_params_map = ParamsT {
		known_rgbs_len: known_rgbs.len(),
		radiussy: radiussy,
		iter_limit: -1,
		bbox_min:                     [f32::from_bits(0xFFFFFFFF); 3];
		bbox_range /* as bbox_max */: [f32::from_bits(0x00000000); 3];
		_padding1: 0, _padding2: 0, _padding3: 0,
	};
	// ...

	// find_bbox.hlsl
	self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE, self.compute_pipeline[FIND_BBOX]);
	self.device.cmd_dispatch(self.cmd_buffer, known_rgbs.len().div_ceil(64), 1, 1);
	self.submit_buffer(self.cmd_buffer);

	fn become_float3(bbox: &[f32; 3]) -> [f32; 3] {
		bbox.map(f32::to_bits).map(|bu| f32::from_bits(bu ^ ((bu >> 31).wrapping_sub(1) | 0x80000000)))
	}
	{
		let bbox_min = become_float3(&self.buffer_params_map.bbox_min);
		let bbox_range = {
			let bbox_max = become_float3(&self.buffer_params_map.bbox_range);
			[bbox_max[0] - bbox_min[0], bbox_max[1] - bbox_min[1], bbox_max[2] - bbox_min[2]]
		};
		self.buffer_params_map.bbox_min   = bbox_min;
		self.buffer_params_map.bbox_range = bbox_range;
	}

	// ... posterise.hlsl/posterise_mul.hlsl &c.
}

This takes anywhere between 120-400µs on the images in the test dataset on the 5700 XT, but variance is so high at this scale that this is barely correlated with the input size (amount of distinct colours). Functionally this is as good as instant. Bench-marking on orders-of-magnitude more data shows throughput of around 414'419'783 uint3/s. The next generation (NVIDIA RTX 4070 SUPER/AMD Radeon RX 6500 XT (RADV NAVI24)) gets 8'221'135'253/8'383'953'229!

# he do b hittin it doe

With all the parameters to grid [​[​g​]​] figured out, we can compute "which points live in each grid cell" [​[​h​]​]. This is relatively trivial, except the GPU has no structured programming capabilities, so the natural uint3 → Vec<indices of known_rgbs/[​[​D​]​]> mapping is inexpressible. The obvious-to-me solution is partitioning all indices by the grid cell they belong to, then the map becomes uint3 → [begin, end) iteratorsindices; this seems canonical. The original example (if we force [​[​B_m = (0,\frac{-1}{2},\frac{-1}{2}); B_r = (1,1,1)​]​] and specify [​[​G=10​]​])

Thus, the grid cells occupied are, resp. [​[​(7, 6, 3)​]​], [​[​(7, 6, 3)​]​], and [​[​(7, 6, 4)​]​]. The backing storage can then be [​[​(1, 2, 3)​]​], and the map can take the form

since all other cells are empty. (The backing storage could also be [​[​(3, 2, 1)​]​], and the map return [​[​[2, 4)​]​]/[​[​[1, 1)​]​].)

Realistically, both the map and backing storage are "a big array" — of uint2s with dimension [​[​G^3​]​] and uints with dimension [​[​\overline{Set(D)}​]​] respectively. It's theoretically possible to implement [​[​h​]​] by collecting the points for each grid cell coordinate separately, like the definition, but this would be [​[​O(n^2)​]​], so pre-calculating the grid cell coordinates for each point, then collecting them, is [​[​O(n)​]​]. To that end, with [​[​g​]​] as lab_to_gridpos and [​[​G​]​] as GRID_SIZE, shader.h:

#define GRID_SIZE 100
struct groupdesc_t {
	uint start_idx, end_idx;
};

struct params_t {
	// ...
	groupdesc_t group_ranges[GRID_SIZE][GRID_SIZE][GRID_SIZE];
};
// ...
float3 lab_to_gridpos(float3 pos, float3 bbox_min, float3 bbox_range) {
	pos = (pos - bbox_min) / bbox_range;
	return min(pos * GRID_SIZE, GRID_SIZE - 0.001);
}

groupdesc_t has the same lay-out as uint2, and putting it in an array like this doesn't induce additional padding. The total size of params_t, and, thus, usage of precious DEVICE_LOCAL & HOST_VISIBLE & HOST_COHERENT memory, grows by [​[​G^3 \operatorname{sizeof}(\texttt{uint2}) = 8\text{MB}​]​], which is good to know but otherwise negligible.

[​[​g​]​]'s [​[​\varepsilon = 0.001​]​] is just as good as any other value (so long as it actually exists around [​[​G​]​]); the optimiser lowers this to a constant. The pre-computation can thus be [​[​\operatorname{trunc} \circ g​]​] directly in assign_cell.hlsl:

typedef float3 point_t;
#include "shader.h"

[[vk::binding(0)]] RWStructuredBuffer<uint>            outputs;            // HOST_VISIBLE
[[vk::binding(1)]]   StructuredBuffer<params_t>        params;             // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(2)]]   StructuredBuffer<known_rgbs_pack> known_rgbs_bundle;  // DEVICE_LOCAL & HOST_VISIBLE

[numthreads(64, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
	const uint   known_rgbs_len = params[0].known_rgbs_len;
	const float3 bbox_min       = params[0].bbox_min;
	const float3 bbox_range     = params[0].bbox_range;

	const uint id = DTid.x;
	if(id >= known_rgbs_len)
	    return;

	uint3 gridpos = trunc(lab_to_gridpos(known_rgbs_bundle[id].lab, bbox_min, bbox_range));
	outputs[id] = (gridpos.x << (8 * 2)) |
	              (gridpos.y << (8 * 1)) |
	              (gridpos.z << (8 * 0));
}

this abuses outputs as an output buffer — this is fine, since the size used is [​[​\texttt{known_rgbs.len()} \cdot \operatorname{sizeof}(\texttt{uint})​]​], and the allocation is [​[​2^{24} \operatorname{sizeof}(\texttt{float4})​]​], which is guaranteed to be at least 4× larger. Obviously the packing would need to be changed for [​[​G > \mathrm{FF}_{16}​]​]. This will not pose an issue.

Shaders can include shader.h to have one source of truth for [​[​G​]​], but the Rust driver can't. This is emulated in build.rs:

fn main() {
	let grid_size = fs::read_to_string("src/shader.h").lines()
		.find(|l| l.contains("#define GRID_SIZE")).split_whitespace().last();
	fs::write(Path::new(&env::var_os("OUT_DIR")).join("grid_size.rs"), &grid_size);
	// ... glslc
}

So, with, [​[​h​]​] as groups_data:

const ASSIGN_CELL: usize = FIND_BBOX + 1; // &c.

const GRID_SIZE: usize = include!(concat!(env!("OUT_DIR"), "/grid_size.rs"));
#[repr(C)]
struct GroupDesc {
	start_idx: u32,
	end_idx: u32,
}
#[repr(C)]
struct ParamsT {
	// ...
	group_ranges: [[[GroupDesc; GRID_SIZE]; GRID_SIZE]; GRID_SIZE],
}

const SHADER_MAX_BINDING: u32 = 4;
pub struct PosteriseGpu {
	// ...
	buffer_outputs: BufferBundle,           // [[vk::binding(0)]] SB<float3>          outputs;           [HOST_VISIBLE]
	buffer_params: BufferBundle,            // [[vk::binding(1)]] SB<params_t>        params;            [DEVICE_LOCAL & HOST_VISIBLE]
	buffer_known_rgbs_bundle: BufferBundle, // [[vk::binding(2)]] SB<known_rgbs_pack> known_rgbs_bundle; [DEVICE_LOCAL & HOST_VISIBLE]
	buffer_known_rgbs_freqs: BufferBundle,  // [[vk::binding(3)]] SB<uint>            known_rgbs_freqs;  [DEVICE_LOCAL & HOST_VISIBLE]
	buffer_groups_data: BufferBundle,       // [[vk::binding(4)]] SB<uint>            groups_data;       [DEVICE_LOCAL]

	groups_accumulator: Box<[[[Vec<u32>; GRID_SIZE]; GRID_SIZE]; GRID_SIZE]>,
}

fn PosteriseGpu::init() -> PosteriseGpu {
	// ...
	let physical_device_properties = vk_instance.get_physical_device_properties(physical_device);
	assert!(SHADER_MAX_BINDING < physical_device_properties.limits.max_bound_descriptor_sets);

	let output_memory_type_index = memory_types.iter()
		.position(|memory_type| {
			memory_type.property_flags.contains(vk::MemoryPropertyFlags::HOST_VISIBLE) &&
			memory_type.property_flags.contains(vk::MemoryPropertyFlags::HOST_COHERENT)
		})
		.expect("No output memory type");
	let input_memory_type_index = memory_types.iter()
		.position(|memory_type| {
			memory_type.property_flags.contains(vk::MemoryPropertyFlags::DEVICE_LOCAL) &&
			memory_type.property_flags.contains(vk::MemoryPropertyFlags::HOST_VISIBLE) &&
			memory_type.property_flags.contains(vk::MemoryPropertyFlags::HOST_COHERENT)
		})
		.unwrap_or(output_memory_type_index);
	let input_memory_type_index = memory_types.iter()
		.position(|memory_type| memory_type.property_flags.contains(vk::MemoryPropertyFlags::DEVICE_LOCAL))
		.unwrap_or(output_memory_type_index);

	// vk::BufferUsageFlags::STORAGE_BUFFER              | ::TRANSFER_SRC ⬎       ⬐ ::TRANSFER_DST
	ret.buffer_outputs           = BufferBundle(output_memory_type_index,  true,  false);
	ret.buffer_params            = BufferBundle(input_memory_type_index,   false, true);
	ret.buffer_known_rgbs_bundle = BufferBundle(input_memory_type_index,   false, false);
	ret.buffer_known_rgbs_freqs  = BufferBundle(input_memory_type_index,   false, false);
	ret.buffer_groups_data       = BufferBundle(private_memory_type_index, false, true);

	unsafe {
		ret.groups_accumulator = Box::<mem::MaybeUninit<_>>::assume_init(Box::new_uninit()),
		for ga in ret.groups_accumulator.iter_mut().flatten().flatten() {
			ptr::write(ga as *mut _, Vec::new());
		}
	}
}

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, freq_weighting: bool, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	self.buffer_groups_data.ensure_buffer(known_rgbs.len() * mem::size_of::<u32>());
	if freq_weighting {
		self.buffer_known_rgbs_freqs.ensure_buffer(known_rgbs.len() * mem::size_of::<u32>());
	}
	// ...

	// memset(self.buffer_params.buffer, 0, *); ⇔ self.buffer_params_map = mem::zeroed();
	self.device.cmd_fill_buffer(self.cmd_buffer, self.buffer_params.buffer, 0, vk::WHOLE_SIZE, 0);
	self.submit_buffer(self.cmd_buffer);
	{
		self.buffer_params_map.known_rgbs_len   = known_rgbs.len();
		self.buffer_params_map.bbox_min         = [f32::from_bits(0xFFFFFFFF); 3];
		// bbox_range/bbox_max          already = 0x00000000 from memset()
		self.buffer_params_map.radiussy_squared = radiussy * radiussy;
		self.buffer_params_map.iter_limit       = -1;
	}
	// ... through find_bbox.hlsl &c.

	// assign_cell.hlsl
	self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE, self.compute_pipeline[ASSIGN_CELL]);
	self.device.cmd_dispatch(self.cmd_buffer, known_rgbs.len().div_ceil(64), 1, 1);
	self.submit_buffer(self.cmd_buffer);

	// grouping
	let known_rgbs_groupcoords: &mut [u32] = slice::from_raw_parts_mut(self.buffer_outputs_map.as_ptr(), known_rgbs.len());

	for group in self.groups_accumulator.iter_mut().flatten().flatten() {
		group.clear();
	}
	for (id, groupcoord) in known_rgbs_groupcoords.iter().enumerate() {
		let (x, y, z) = ((groupcoord >> (8 * 2)),
		                 (groupcoord >> (8 * 1)) & 0xFF,
		                 (groupcoord >> (8 * 0)) & 0xFF);
		self.groups_accumulator[x][y][z].push(id);
	}

	// copying
	let mut known_rgbs_groups: *mut u32 = self.buffer_outputs_map.as_ptr() as _;
	let mut cur = 0;
	for (gr_z, ga_z) in self.buffer_params_map.group_ranges.iter_mut().flatten().flatten()
	               .zip(self.groups_accumulator            .iter()    .flatten().flatten()) {
		if ga_z.len() != 0 {
			let start = cur;

			// memcpy(known_rgbs_groups,               ga_z.as_ptr(), ga_z.len())
			known_rgbs_groups.copy_from_nonoverlapping(ga_z.as_ptr(), ga_z.len());
			cur += ga_z.len();
			known_rgbs_groups = known_rgbs_groups.add(ga_z.len());

			*gr_z = GroupDesc {
				start_idx: start,
				end_idx: cur,
			};
		}
	}
	assert!(cur == known_rgbs.len());

	// memcpy(self.buffer_groups_data, self.buffer_outputs, known_rgbs * mem::size_of::<u32>());
	self.device.cmd_copy_buffer(self.cmd_buffer, self.buffer_outputs.buffer, self.buffer_groups_data.buffer,
	                            &[vk::BufferCopy {
	                              	src_offset: 0,
	                              	dst_offset: 0,
	                              	size: known_rgbs * mem::size_of::<u32>(),
	                              }]);
	self.submit_buffer(self.cmd_buffer);

	// ... posterise.hlsl/posterise_mul.hlsl &c.
}

This is the first time more than the aforementioned 4 bindings are used; checking this a priori potentially avoids a more confusing error later, but this will never be triggered on real implementations.

groups_data is, from the perspective of the host, write-only/write-once for each image, so it can live in DEVICE_LOCAL memory and be uploaded via outputs (AIUI this technique is called the "staging buffer"; bounds used same as find_bbox.hlsl's). This also means that they need to be allocated with the right TRANSFER_direction flags.

groups_accumulator is a similar size to the [​[​8\text{MB}​]​] — [​[​G^3 \cdot \operatorname{sizeof}(\texttt{Vec<_>}) = G^3 \cdot 3 \operatorname{sizeof}(\texttt{usize}) = 24\text{MB}_\text{LP64} \vee 12\text{MB}_\text{ILP32}​]​] would blow any stack and it's impossible to use objects too large to spill in Rust (I tried!), so it needs to go in an additional allocation and then be initialised piece-meal. Vecs can't be all-0 so is it in principle UB to assume_init(new_uninit()), then force-initialise them? Maybe. Is this a real concern? no.

Like-wise due to its size, params can no longer be cleared by re-assignment (it'd be a memcpy(mapping ← stack)), and it makes no sense to clear it with an actual memset(), effectively uploading [​[​8\text{MB}​]​] of zeroes to DEVICE_LOCAL memory. Thus, a device-side memset — vkCmdFillBuffer() — is dispatched, working exactly as you'd expect (except it sets whole uints instead of bytes), and completing too quickly to measure above the fixed submission overhead of a couple dozen μs. From the point of buffer usage, this is a transfer from a null page, hence TRANSFER_DST.

Collecting the indices in each cell (grouping) is expensive but can't be meaningfully sped up. Neither can copying the data, beyond only copying the index ranges for non-empty cells (after the memset, the default value is [​[​[0, 0 ) = \varnothing​]​] which is what we'd be writing anyway), but that is relatively quick:

So grouping, on the i7-2600, quickly arrives at a consistent ~11Mp/s throughput. Bounding box scaling clearly works, increasing the occupancy, in the case of the 1'103'347 image, from ~5% (average population of ~20 in non-empty cells, median 9) to ~17% (~7, 3).

# The grid-based mean-shift clustering shader implementation,

thus, can take advantage of the grid by, instead of checking if all elements of [​[​D​]​] are closer to [​[​q_t​]​] than [​[​r​]​], considering only the ones that could feasibly fulfil that predicate by virtue of being contained within near-by grid cells. The actual range is complicated slightly by the bounding box scale being different along each axis, so the sphere of influence (with a cubic bounding box) turns into an ellipsoid of influence with "radiuses" [​[​R​]​] (with a prismatic bounding box from [​[​g_{tm}​]​] to [​[​g_{tM}​]​]) rounded up to encompas any cell it intersects after the [​[​B_r​]​] scaling:

With [​[​R​]​] as grid_bothdirs, [​[​g_t​]​] as gridpos, [​[​g_{tm}​]​] as gridpos_from, and [​[​g_{tM}​]​] as gridpos_to, posterise_grid_kernel.hlsl.inc can implement the [​[​\bigcup​]​] literally:

typedef float3 point_t;
#include "shader.h"

[[vk::binding(0)]] RWStructuredBuffer<float3>          outputs;            // HOST_VISIBLE
[[vk::binding(1)]]   StructuredBuffer<params_t>        params;             // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(2)]]   StructuredBuffer<known_rgbs_pack> known_rgbs_bundle;  // DEVICE_LOCAL & HOST_VISIBLE
/*            3 */   KNOWN_RGBS_FREQS
[[vk::binding(4)]]   StructuredBuffer<uint>            groups_data;        // DEVICE_LOCAL

[numthreads(64, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
	const uint   known_rgbs_len   = params[0].known_rgbs_len;
	const float  radiussy_squared = params[0].radiussy_squared;
	const uint   iter_limit       = params[0].iter_limit;
	const float3 bbox_min         = params[0].bbox_min;
	const float3 bbox_range       = params[0].bbox_range;
	const uint3  grid_bothdir     = params[0].grid_bothdirs;

	const uint id = DTid.x;
	if(id >= known_rgbs_len)
		return;

	known_rgbs_pack this = known_rgbs_bundle[id];
	float3 center = this.lab;
	for(int iter = 0; iter != iter_limit; ++iter) {
		const uint3 gridpos      = trunc(lab_to_gridpos(center, bbox_min, bbox_range));
		const uint3 gridpos_from = gridpos - min(grid_bothdirs, gridpos);
		const uint3 gridpos_to   = min(gridpos + grid_bothdirs, GRID_SIZE - 1);

		uint newcenter_count = 0;
		float3 newcenter     = float3(0, 0, 0);
		for(uint gx = gridpos_from.x; gx <= gridpos_to.x; ++gx)
			for(uint gy = gridpos_from.y; gy <= gridpos_to.y; ++gy)
				for(uint gz = gridpos_from.z; gz <= gridpos_to.z; ++gz) {
					groupdesc_t group = params[0].group_ranges[gx][gy][gz];
					for(; group.start_idx != group.end_idx; ++group.start_idx) {
						const uint i = groups_data[group.start_idx];
						float3 lab = known_rgbs_bundle[i].lab;

						float3 tmp = lab - center;
						const bool matched = dot(tmp, tmp) <= radiussy_squared;
#if !BRANCHLESS_ADD
						if(matched)
#endif
						{
							FREQ_LOAD
							newcenter_count += FREQ;
							newcenter       += FREQ_MUL lab;
						}
					}
				}

		newcenter /= newcenter_count;
		bool same = newcenter == center;
		if(same)
			break;
		center = newcenter;
	}
	outputs[this.rgb] = center;
}

posterise_grid.hlsl and posterise_grid_mul.hlsl are unchanged, and so is both the outer-most [​[​q_t​]​] loop and inner-most [​[​\Sigma​]​]. To avoid overflow if [​[​R > g_t​]​] we observe [​[​\max(g_t - R, (0, 0, 0)) = g_t - \min(R, g_t)​]​].

Reordering the scalar parameters to forced-padding locations allows retaining the same size while fitting grid_bothdirs/[​[​R​]​]:

struct params_t {
	uint3   grid_bothdirs;
	uint    known_rgbs_len;

	point_t bbox_min;
	float   radiussy_squared;

	point_t bbox_range;
	int     iter_limit;
#define bbox_max bbox_range

	groupdesc_t group_ranges[GRID_SIZE][GRID_SIZE][GRID_SIZE];
};

[​[​R​]​] computation is can likewise be implemented directly:

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, freq_weighting: bool, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	// ...

	// find_bbox.hlsl ...
	{
		let bbox_min = become_float3(&self.buffer_params_map.bbox_min);
		let bbox_range = {
			let bbox_max = become_float3(&self.buffer_params_map.bbox_range);
			[bbox_max[0] - bbox_min[0], bbox_max[1] - bbox_min[1], bbox_max[2] - bbox_min[2]]
		};
		self.buffer_params_map.bbox_min   = bbox_min;
		self.buffer_params_map.bbox_range = bbox_range;

		let scaled_radiussies = bbox_range.map(|scale| radiussy / scale);
		let direction_cells   = scaled_radiussies.map(|sr| (sr * GRID_SIZE).ceil());
		self.buffer_params_map.grid_bothdirs = direction_cells;
	}

	// ...
}

By mildly instrumenting (dbg!(direction_cells)) or recalling the bounding boxes from before, we can compute the following and use it as the best-case prediction (additionally assuming every cell inspected will be non-empty, which it won't, but non-empty cells are where the points are so they're roughly where [​[​q_t​]​]s are and they're pretty dense, so.):

To contrast with real performance data:

[​[​*freq​]​] becomes more slower than [​[​palette​]​] in the degenerate case and less slower than [​[​palette​]​] on the real image. What this means is unclear to me.

The predictions are entirely uncorrelated with reality. I love economics.

# Failure: GPU-side grid-cell-to-index-list mapping

While the collection (grouping) can't be sped up on the host, it could, theoretically, be executed on the GPU by collecting each value of [​[​h​]​] directly from the definition by looking at each element of [​[​D​]​] for each cell, after pre-computing the sizes to fix the allocation question. To that end, the latter can be summed up easily in assign_cell.hlsl:

// ...
[[vk::binding(1)]]   StructuredBuffer<params_t>        params;             // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(2)]]   StructuredBuffer<known_rgbs_pack> known_rgbs_bundle;  // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(3)]] RWStructuredBuffer<uint>            known_rgbs_freqs;   // DEVICE_LOCAL & HOST_VISIBLE

[numthreads(64, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
  // ...
	outputs[id] = (gridpos.x << (8 * 2)) |
	              (gridpos.y << (8 * 1)) |
	              (gridpos.z << (8 * 0));

	InterlockedAdd(params[0].group_ranges[gridpos.x][gridpos.y][gridpos.z].group_len, 1);
}

(noting that the intermediate map now goes to a non-HOST_VISIBLE memory since it doesn't need to be), after easily restructuring the group index range to [begin, begin+size):

struct groupdesc_t {
	uint start_idx, group_len;
};

Thus, the allocator can dump the indices into groups_data in the same manner as the host-side implementation, but instead of doing it in order, it can use Some Global to get the next unused chunk. This is much more like a classical bump allocator. Conveniently, the Some Global can just be the first integer in groups_data. So, the driver is greatly simplified:

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, freq_weighting: bool, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	self.buffer_groups_data.ensure_buffer((known_rgbs.len() + 1) * mem::size_of::<u32>());

	// ... through assign_cell.hlsl &c.

	// self.buffer_groups_data[0] = 1;
	self.device.cmd_fill_buffer(self.cmd_buffer, self.buffer_params.buffer, 0, mem::size_of::<u32>(), 1);
	// cicada_grid.hlsl
	self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE, self.compute_pipeline[CICADA_GRID]);
	self.device.cmd_dispatch(self.cmd_buffer, 1, GRID_SIZE, GRID_SIZE);
	self.submit_buffer(self.cmd_buffer);

	// ... posterise.hlsl/posterise_mul.hlsl &c.
}

The initial offset must be initialised (to 1 instead of 0 to not overwrite itself), but this can be done in the same command buffer as dispatching the shader and the "couple dozen μs" (like 79±32 in this case) become free. cicada_grid.hlsl is dispatched at [​[​\text{numthreads} = (G, 1, 1), \text{dispatch}(1, G, G)​]​] to evaluate every integer in [​[​h​]​]'s domain of [​[​[0, G )^3​]​]:

typedef float3 point_t;
#include "shader.h"

[[vk::binding(1)]] RWStructuredBuffer<params_t>        params;             // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(3)]]   StructuredBuffer<uint>            known_rgbs_freqs;   // DEVICE_LOCAL & HOST_VISIBLE
[[vk::binding(4)]] RWStructuredBuffer<uint>            groups_data;        // DEVICE_LOCAL

#define next_start groups_data[0]
[numthreads(GRID_SIZE, 1, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
	const uint3 gridpos        = DTid;
	const uint  known_rgbs_len = params[0].known_rgbs_len;
	const uint  my_rgbs        = params[0].group_ranges[gridpos.x][gridpos.y][gridpos.z].group_len;

	if(my_rgbs == 0)
		return;

	uint start_idx;
	InterlockedAdd(next_start, my_rgbs, start_idx);
	params[0].group_ranges[gridpos.x][gridpos.y][gridpos.z].start_idx = start_idx;

	const uint gridpos_saved = (gridpos.x << (8 * 2)) |
	                           (gridpos.y << (8 * 1)) |
	                           (gridpos.z << (8 * 0));
	for(uint rgb = 0; rgb < known_rgbs_len; ++rgb)
		if(known_rgbs_freqs[rgb] == gridpos_saved) {
			groups_data[start_idx] = rgb;
			++start_idx;
		}
}

Visiting every point in [​[​Set(D)​]​] for every occupied cell (up to [​[​G^3​]​] in theory but cf. occupancy rates above; there are [​[​G^3​]​] (so, 1'000'000), threads dispatched though), is, expectedly, ruinous:

Interestingly enough, the cost of uploading to system RAM actually makes up for the additional atomic in all but the smallest cases. But everything else is eye-wateringly worser.

Also, this makes the allocation order basically random. Functionally this is immaterial, but may be tangentially relevant later.

Even another 40× speed-up wouldn't guarantee we can compute every [​[​f​]​] within the ~8s driver-imposed limit. But, each shader call could iterate [​[​t​]​], like, 10 times. The points that converge can be committed as the final [​[​f​]​] mapping as previous, but the ones that don't can be saved and

# processed piece-meal.

The first step here is getting rid of the mis-leading iter_limit in the current shader:

// ...
	known_rgbs_pack this = known_rgbs_bundle[id];
#define center this.lab
	for(;;) {
// ...

Since the amount of unconverged points decreases each iteration, the [​[​q_t​]​] points-to-process array must be compacted between iterations. Well, it doesn't need to be and a thread that just exits instantly could be spawned for it, but this sucks, actively makes occupancy (how many GPU threads are working vs waiting) worse, and basically defeats the purpose; this is already a bit of an issue in the current driver, where it may take 64 pregnant threads where 63 have few fast-to-solve points and 1 that takes long to converge, as long as the latter to bear an output (this is hopefully attenuated slightly by the points being sorted by [​[​{\color{var(--r)}r\color{initial}}{\color{var(--r)}r\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--b)}b\color{initial}}{\color{var(--b)}b\color{initial}}_{16}​]​] so similar colours in similar natal situations are carried together, but.). We can store [​[​q_t​]​]s and [​[​q_{t+10}​]​]s in alternating ranges of a big array and page-flip the input/output iterators (this makes initialisation a memcpy from known_rgbs_bundle as [​[​q_0​]​], which is quite nice from a point-and-woo-at-definition perspective):

In the current driver, known_rgbs_freqs is indexed by the same offset as known_rgbs_bundle. This is no longer possible once original-index data is lost when the state compacts, so in this case we want to index known_rgbs_freqs by [​[​{\color{var(--r)}r\color{initial}}{\color{var(--r)}r\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--g)}g\color{initial}}{\color{var(--b)}b\color{initial}}{\color{var(--b)}b\color{initial}}_{16}​]​] (like outputs), which is kept with the state.

const USEONESHOT: usize = 2;
const FIND_BBOX: usize = USEONESHOT + 2;
const ASSIGN_CELL: usize = FIND_BBOX + 1;
const SHADER_COUNT: usize = ASSIGN_CELL + 1;
static SHADER_DATA: [&[u32]; SHADER_COUNT] = [include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise_grid.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise_grid.mul.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise_grid_oneshot.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/posterise_grid_oneshot.mul.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/find_bbox.spv")),
                                              include_typed!(u32, concat!(env!("OUT_DIR"), "/assign_cell.spv"))];
const SHADER_MAX_BINDING: u32 = 5;

#[repr(C)]
struct ParamsT {
	grid_bothdirs: [u32; 3],
	known_rgbs_len: u32,

	bbox_min: [f32; 3],
	radiussy_squared: f32,

	bbox_range: [f32; 3],
	_padding1: i32,

	current_rgbs_start: u32, // [​[​q_t \in \texttt{current_rgbs_bundle} \circ [\texttt{current_rgbs_start}, \texttt{current_rgbs_start}+\texttt{current_rgbs_len} )​]​]
	current_rgbs_len: u32,
	current_rgbs_next: u32,  // [​[​q_{t+\texttt{iter_limit}} \to \texttt{current_rgbs_bundle} \circ [\texttt{current_rgbs_next}, \ldots )​]​]
	iter_limit: i32,

	group_ranges: [[[GroupDesc; GRID_SIZE]; GRID_SIZE]; GRID_SIZE],
}

pub struct PosteriseGpu {
	// ...
	buffer_outputs: BufferBundle,             // [[vk::binding(0)]] SB<float3>          outputs;             [HOST_VISIBLE]
	buffer_params: BufferBundle,              // [[vk::binding(1)]] SB<params_t>        params;              [DEVICE_LOCAL & HOST_VISIBLE]
	buffer_known_rgbs_bundle: BufferBundle,   // [[vk::binding(2)]] SB<known_rgbs_pack> known_rgbs_bundle;   [DEVICE_LOCAL & HOST_VISIBLE]
	buffer_known_rgbs_freqs: BufferBundle,    // [[vk::binding(3)]] SB<uint>            known_rgbs_freqs;    [DEVICE_LOCAL & HOST_VISIBLE]
	buffer_groups_data: BufferBundle,         // [[vk::binding(4)]] SB<uint>            groups_data;         [DEVICE_LOCAL]
	buffer_current_rgbs_bundle: BufferBundle, // [[vk::binding(5)]] SB<float4>          current_rgbs_bundle; [DEVICE_LOCAL]
	// ...
}

fn PosteriseGpu::init() -> PosteriseGpu {
	// ...
	// vk::BufferUsageFlags::STORAGE_BUFFER                | ::TRANSFER_SRC ⬎       ⬐ ::TRANSFER_DST
	ret.buffer_known_rgbs_bundle   = BufferBundle(input_memory_type_index,   true,  false);  // was false, false
	ret.buffer_current_rgbs_bundle = BufferBundle(private_memory_type_index, false, true);
	// ...
}

fn PosteriseGpu::susmit(&mut self, known_rgbs: &[(u32, Lab, u32)], radiussy: f32, freq_weighting: bool, outputs: &mut [Lab; 0xFFFFFF + 1]) {
	let use_oneshot_driver = known_rgbs.len() > 50_000;

	self.buffer_known_rgbs_bundle.ensure_buffer(&self.device, known_rgbs.len() * mem::size_of::<KnownRgbsPack>());
	assert!(mem::size_of::<KnownRgbsPack>() == 4 * 4);
	self.buffer_groups_data.ensure_buffer(&self.device, known_rgbs.len() * mem::size_of::<u32>());
	if use_oneshot_driver {
		self.buffer_current_rgbs_bundle.ensure_buffer(&self.device, (known_rgbs.len() * 2) * mem::size_of::<KnownRgbsPack>());
	}
	let buffer_known_rgbs_freqs_len = if use_oneshot_driver { 0xFFFFFF + 1 } else { known_rgbs.len() } * mem::size_of::<u32>());
	if freq_weighting {
		self.buffer_known_rgbs_freqs.ensure_buffer(&self.device, buffer_known_rgbs_freqs_len);
	}

	let buffer_known_rgbs_bundle_map = self.device.map_memory(self.buffer_known_rgbs_bundle);
	let known_rgbs_bundle: &mut [KnownRgbsPack] = slice::from_raw_parts_mut(buffer_known_rgbs_bundle_map, known_rgbs.len());
	for ((rgb, lab, _), dest_lab) in known_rgbs.iter().zip(known_rgbs_bundle.iter_mut()) {
		dest_lab = KnownRgbsPack {
			rgb: rgb,
			lab: lab,
		};
	}
	self.buffer_known_rgbs_bundle.unmap(&self.device);

	if freq_weighting {
		let buffer_known_rgbs_freqs_map = self.device.map_memory(self.buffer_known_rgbs_freqs);
		let known_rgbs_freqs: &mut [u32] = slice::from_raw_parts_mut(buffer_known_rgbs_freqs_map, buffer_known_rgbs_freqs_len / mem::size_of::<u32>());
		if use_oneshot_driver {
			for (rgb, _, freq) in known_rgbs {
				known_rgbs_freqs[rgb] = freq;
			}
		} else {
			for ((_, _, freq), dest_freq) in known_rgbs.iter().zip(known_rgbs_freqs.iter_mut()) {
				dest_freq = freq;
			}
		}
		self.buffer_known_rgbs_freqs.unmap(&self.device);
	}

	// ...

	// posterise_grid_oneshot.kernel.hlsl.inc or posterise_grid.kernel.hlsl.inc
	if use_oneshot_driver {
		self.buffer_params_map.current_rgbs_len = known_rgbs.len();

		let il = |len| if len <=  10_000 { -1 }
		          else if len <= 100_000 { 20 }
		          else                   { 10 };
		let mut first = true;
		for (orig_start, orig_next) in [(0,                known_rgbs.len()),
		                                (known_rgbs.len(), 0               )].iter().cycle() {
			self.buffer_params_map.current_rgbs_start = orig_start;
			self.buffer_params_map.current_rgbs_next  = orig_next;
			self.buffer_params_map.iter_limit = il(self.buffer_params_map.current_rgbs_len);

			if first {
				// memcpy(self.buffer_current_rgbs, self.buffer_known_rgbs, known_rgbs.len() * mem::size_of::<KnownRgbsPack>());
				self.device.cmd_copy_buffer(self.cmd_buffer,
					self.buffer_known_rgbs_bundle.buffer, self.buffer_current_rgbs_bundle.buffer,
					&[vk::BufferCopy {
					  	src_offset: 0,
					  	dst_offset: 0,
					  	size: known_rgbs.len() * mem::size_of::<KnownRgbsPack>(),
					  }]);
			}
			self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE,
				self.compute_pipeline[USEONESHOT + freq_weighting]);
			self.device.cmd_dispatch(self.cmd_buffer, self.buffer_params_map.current_rgbs_len.div_ceil(64), 1, 1);
			self.submit_buffer(self.cmd_buffer);

			let still_to_do = self.buffer_params_map.current_rgbs_next - orig_next;
			if still_to_do == 0 {
				break;
			}

			println!("did {}/{} ({} steps); {} left",
			         self.buffer_params_map.current_rgbs_len - still_to_do,
			         self.buffer_params_map.current_rgbs_len,
			         self.buffer_params_map.iter_limit,
			         still_to_do);
			self.buffer_params_map.current_rgbs_len = still_to_do;
			first = false;
		}
	} else {
		self.device.cmd_bind_pipeline(self.cmd_buffer, vk::PipelineBindPoint::COMPUTE,
			self.compute_pipeline[freq_weighting]);
		self.device.cmd_dispatch(self.cmd_buffer, known_rgbs.len().div_ceil(64), 1, 1);
		self.submit_buffer(self.cmd_buffer);
	}

	// ... output copying
}