vyre 0.4.0

GPU compute intermediate representation with a standard operation library
Documentation 
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

@group(0) @binding(0) var<storage, read> input_words: array<u32>;
@group(0) @binding(1) var<storage, read_write> output_words: array<u32>;

const B3_IV: array<u32, 8> = array<u32, 8>(0x6a09e667u,0xbb67ae85u,0x3c6ef372u,0xa54ff53au,0x510e527fu,0x9b05688cu,0x1f83d9abu,0x5be0cd19u);
const B3_MSG_PERM: array<u32, 16> = array<u32, 16>(2u,6u,3u,10u,7u,0u,4u,13u,1u,11u,12u,5u,9u,14u,15u,8u);

fn b3_rotr(x: u32, n: u32) -> u32 { return (x >> n) | (x << ((32u - n) & 31u)); }
fn b3_g(v_in: array<u32, 16>, a: u32, b: u32, c: u32, d: u32, mx: u32, my: u32) -> array<u32, 16> {
  var v = v_in;
  v[a] = v[a] + v[b] + mx; v[d] = b3_rotr(v[d] ^ v[a], 16u);
  v[c] = v[c] + v[d];      v[b] = b3_rotr(v[b] ^ v[c], 12u);
  v[a] = v[a] + v[b] + my; v[d] = b3_rotr(v[d] ^ v[a], 8u);
  v[c] = v[c] + v[d];      v[b] = b3_rotr(v[b] ^ v[c], 7u);
  return v;
}
fn b3_round(v_in: array<u32, 16>, m: array<u32, 16>) -> array<u32, 16> {
  var v = v_in;
  v = b3_g(v, 0u,4u,8u,12u, m[0u], m[1u]);
  v = b3_g(v, 1u,5u,9u,13u, m[2u], m[3u]);
  v = b3_g(v, 2u,6u,10u,14u, m[4u], m[5u]);
  v = b3_g(v, 3u,7u,11u,15u, m[6u], m[7u]);
  v = b3_g(v, 0u,5u,10u,15u, m[8u], m[9u]);
  v = b3_g(v, 1u,6u,11u,12u, m[10u], m[11u]);
  v = b3_g(v, 2u,7u,8u,13u, m[12u], m[13u]);
  v = b3_g(v, 3u,4u,9u,14u, m[14u], m[15u]);
  return v;
}
fn b3_permute(m_in: array<u32, 16>) -> array<u32, 16> {
  var out: array<u32, 16>;
  for (var i = 0u; i < 16u; i = i + 1u) { out[i] = m_in[B3_MSG_PERM[i]]; }
  return out;
}

@compute @workgroup_size(1, 1, 1)
fn hash_blake3(@builtin(global_invocation_id) id: vec3<u32>) {
  if (id.x != 0u) { return; }
  var m: array<u32, 16>;
  for (var i = 0u; i < 16u; i = i + 1u) { m[i] = 0u; }
  let words = min(arrayLength(&input_words), 16u);
  for (var i = 0u; i < words; i = i + 1u) { m[i] = input_words[i]; }
  var cv: array<u32, 8>;
  for (var i = 0u; i < 8u; i = i + 1u) { cv[i] = B3_IV[i]; }
  var v: array<u32, 16>;
  for (var i = 0u; i < 8u; i = i + 1u) { v[i] = cv[i]; v[i + 8u] = B3_IV[i]; }
  v[12u] = 0u; v[13u] = 0u; v[14u] = words * 4u; v[15u] = 11u;
  for (var r = 0u; r < 7u; r = r + 1u) {
    v = b3_round(v, m);
    if (r != 6u) { m = b3_permute(m); }
  }
  for (var i = 0u; i < 8u; i = i + 1u) { output_words[i] = v[i] ^ v[i + 8u]; }
}