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
52
53
54
55
56
57
58
59
60

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

const KECCAK_RC_256: array<u64, 24> = array<u64, 24>(
  0x0000000000000001u, 0x0000000000008082u, 0x800000000000808au, 0x8000000080008000u,
  0x000000000000808bu, 0x0000000080000001u, 0x8000000080008081u, 0x8000000000008009u,
  0x000000000000008au, 0x0000000000000088u, 0x0000000080008009u, 0x000000008000000au,
  0x000000008000808bu, 0x800000000000008bu, 0x8000000000008089u, 0x8000000000008003u,
  0x8000000000008002u, 0x8000000000000080u, 0x000000000000800au, 0x800000008000000au,
  0x8000000080008081u, 0x8000000000008080u, 0x0000000080000001u, 0x8000000080008008u);

const KECCAK_R_256: array<u32, 25> = array<u32, 25>(
  0u, 1u, 62u, 28u, 27u, 36u, 44u, 6u, 55u, 20u, 3u, 10u, 43u, 25u, 39u, 41u, 45u, 15u, 21u, 8u, 18u, 2u, 61u, 56u, 14u);

fn k256_rot(x: u64, n: u32) -> u64 {
  if (n == 0u) { return x; }
  return (x << n) | (x >> (64u - n));
}

fn k256_round(a_in: array<u64, 25>, rc: u64) -> array<u64, 25> {
  var a = a_in;
  var c: array<u64, 5>;
  var d: array<u64, 5>;
  for (var x = 0u; x < 5u; x = x + 1u) { c[x] = a[x] ^ a[x + 5u] ^ a[x + 10u] ^ a[x + 15u] ^ a[x + 20u]; }
  for (var x = 0u; x < 5u; x = x + 1u) { d[x] = c[(x + 4u) % 5u] ^ k256_rot(c[(x + 1u) % 5u], 1u); }
  for (var x = 0u; x < 5u; x = x + 1u) { for (var y = 0u; y < 5u; y = y + 1u) { a[x + 5u * y] = a[x + 5u * y] ^ d[x]; } }
  var b: array<u64, 25>;
  for (var x = 0u; x < 5u; x = x + 1u) {
    for (var y = 0u; y < 5u; y = y + 1u) {
      let src = x + 5u * y;
      let dst = y + 5u * ((2u * x + 3u * y) % 5u);
      b[dst] = k256_rot(a[src], KECCAK_R_256[src]);
    }
  }
  for (var x = 0u; x < 5u; x = x + 1u) {
    for (var y = 0u; y < 5u; y = y + 1u) {
      let lane = x + 5u * y;
      a[lane] = b[lane] ^ ((~b[((x + 1u) % 5u) + 5u * y]) & b[((x + 2u) % 5u) + 5u * y]);
    }
  }
  a[0u] = a[0u] ^ rc;
  return a;
}

@compute @workgroup_size(1, 1, 1)
fn hash_sha3_256(@builtin(global_invocation_id) id: vec3<u32>) {
  if (id.x != 0u) { return; }
  var s: array<u64, 25>;
  for (var i = 0u; i < 25u; i = i + 1u) { s[i] = 0u; }
  let words = min(arrayLength(&input_lanes), 16u);
  for (var i = 0u; i < words; i = i + 1u) { s[i] = s[i] ^ input_lanes[i]; }
  let pad_lane = words;
  if (pad_lane < 17u) { s[pad_lane] = s[pad_lane] ^ 0x0000000000000006u; }
  s[16u] = s[16u] ^ 0x8000000000000000u;
  for (var r = 0u; r < 24u; r = r + 1u) { s = k256_round(s, KECCAK_RC_256[r]); }
  for (var i = 0u; i < 4u; i = i + 1u) {
    output_words[i * 2u] = u32(s[i] & 0xffffffffu);
    output_words[i * 2u + 1u] = u32(s[i] >> 32u);
  }
}