oarust 0.8.6

GPU-first semantic computing engine written in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

import attributes;
import storage;

struct PushConstants {
	uint latent_index;
	uint indices_index;
	uint embed_sum_input_index;
	uint cluster_size_input_index;
	uint codebook_input_index;
	uint embed_sum_output_index;
	uint cluster_size_output_index;
	uint codebook_output_index;
	uint rows;
	uint code_dim;
	uint num_codes;
	float decay;
	float epsilon;
	float dead_threshold;
	uint seed;
	uint normalize;
};

[[vk::push_constant]] PushConstants push;
[[vk::binding(0, 0)]] RWByteAddressBuffer storage_buffers[];

[kernel_name("vq_ema_update")]
[domain("ml")]
[variant("functional_state")]
[dtype("f32")]
[status("experimental")]
[shader("compute")]
[numthreads(256, 1, 1)]
void main(uint3 dispatch_thread_id : SV_DispatchThreadID) {
	uint code = dispatch_thread_id.x;
	if (code >= push.num_codes) {
		return;
	}
	uint count = 0;
	for (uint row = 0; row < push.rows; ++row) {
		if (load_i32(storage_buffers[push.indices_index], row) == int(code)) {
			++count;
		}
	}
	float cluster_size = push.decay
		* load_f32(storage_buffers[push.cluster_size_input_index], code)
		+ (1.0 - push.decay) * float(count);
	uint code_base = code * push.code_dim;
	float sum_of_squares = 0.0;
	for (uint column = 0; column < push.code_dim; ++column) {
		float sum = 0.0;
		for (uint row = 0; row < push.rows; ++row) {
			if (load_i32(storage_buffers[push.indices_index], row) == int(code)) {
				sum += load_f32(storage_buffers[push.latent_index], row * push.code_dim + column);
			}
		}
		float embed_sum = push.decay
			* load_f32(storage_buffers[push.embed_sum_input_index], code_base + column)
			+ (1.0 - push.decay) * sum;
		store_f32(storage_buffers[push.embed_sum_output_index], code_base + column, embed_sum);
		float value = embed_sum / max(cluster_size, push.epsilon);
		store_f32(storage_buffers[push.codebook_output_index], code_base + column, value);
		sum_of_squares += value * value;
	}
	store_f32(storage_buffers[push.cluster_size_output_index], code, cluster_size);
	if (push.normalize != 0) {
		float rms = sqrt(sum_of_squares / float(push.code_dim));
		if (rms > push.epsilon) {
			for (uint column = 0; column < push.code_dim; ++column) {
				uint index = code_base + column;
				store_f32(
					storage_buffers[push.codebook_output_index], index,
					load_f32(storage_buffers[push.codebook_output_index], index) / rms);
			}
		}
	}
	if (cluster_size < push.dead_threshold) {
		uint hash = ((code + 1u) * 2654435761u) ^ (push.seed * 2246822519u) ^ (push.seed >> 15u);
		hash ^= hash >> 13u;
		hash *= 3266489917u;
		hash ^= hash >> 16;
		uint revival_row = hash % push.rows;
		for (uint column = 0; column < push.code_dim; ++column) {
			float value = load_f32(
				storage_buffers[push.latent_index], revival_row * push.code_dim + column);
			store_f32(storage_buffers[push.codebook_output_index], code_base + column, value);
			store_f32(storage_buffers[push.embed_sum_output_index], code_base + column, value);
		}
		store_f32(storage_buffers[push.cluster_size_output_index], code, 1.0);
	}
}