GPU-powered ORAS TinyMT seed finding via soaring fidgets

main.py

from math import log2
from time import perf_counter_ns
import numpy as np
import pyopencl as cl

BASE_ADVANCE = 45

class TinyMT:
    def __init__(self, seed):
        self.state = [seed, 0x8F7011EE, 0xFC78FF1F, 0x3793FDFF]
        for i in range(1, 8):
            self.state[i & 3] ^= (
                0x6C078965 * (self.state[(i - 1) & 3] ^ (self.state[(i - 1) & 3] >> 30))
                + i
            )
            self.state[i & 3] &= 0xFFFFFFFF
        self.advance(8)

    def next_state(self):
        y = self.state[3]
        x = (self.state[0] & 0x7FFFFFFF) ^ self.state[1] ^ self.state[2]

        x ^= x << 1
        x &= 0xFFFFFFFF
        y ^= (y >> 1) ^ x

        self.state = [
            self.state[1],
            self.state[2] ^ ((y & 1) * 0x8F7011EE),
            x ^ ((y << 10) & 0xFFFFFFFF) ^ ((y & 1) * 0xFC78FF1F),
            y,
        ]

    def advance(self, n):
        for _ in range(n):
            self.next_state()

    def next_uint(self):
        self.next_state()
        t0 = self.state[3]
        t1 = self.state[0] + (self.state[2] >> 8)
        t1 &= 0xFFFFFFFF
        t0 ^= t1
        if t1 & 1:
            t0 ^= 0x3793FDFF
        return t0


with open("shader.cl", "r", encoding="utf-8") as f:
    shader_code = f.read()

context = cl.create_some_context()
queue = cl.CommandQueue(context)

input("Enter to start")

data = []
old = perf_counter_ns() - 3e9
denominator = 1

while log2(denominator) < 40:
    input("Enter when fidget")
    new = perf_counter_ns()
    seconds = (new - old) / 1e9
    jump = round((new - old) / 3e9) - 2
    data.append(jump)
    denominator *= 3 ** (jump + 1) / (2**jump)
    print(f"{seconds=}")
    print(f"{jump=}")
    print(f"{log2(denominator):.02f}/40")
    print(data)
    old = new

program = cl.Program(context, shader_code).build(
    [
        f"-D JUMP_COUNT={len(data)}",
        f"-D JUMP_DATA={','.join(map(str, data))}",
        f"-D BASE_ADVANCE={BASE_ADVANCE}",
    ]
)

host_results = np.zeros(150, np.uint64)
host_count = np.zeros(1, np.int32)

device_results = cl.Buffer(context, cl.mem_flags.READ_WRITE, host_results.nbytes)
device_count = cl.Buffer(context, cl.mem_flags.READ_WRITE, host_count.nbytes)

cl.enqueue_copy(queue, device_results, host_results)
cl.enqueue_copy(queue, device_count, host_count)

kernel = program.find_initial_seeds
kernel(
    queue,
    (2**16, 2**16),
    None,
    device_count,
    device_results,
)

host_results = np.empty_like(host_results)
host_count = np.empty_like(host_count)

print("Processing ....")

cl.enqueue_copy(queue, host_results, device_results)
cl.enqueue_copy(queue, host_count, device_count)

host_results = host_results[: host_count[0]]
print(host_results)

initial_seed = int(host_results[0])

rng = TinyMT(initial_seed)
rng.advance(3)
pkrd_state = " ".join(f"[{i}]{s:08X}" for i, s in reversed(list(enumerate(rng.state))))

print(f"{initial_seed=:08X}")
print(f"Pokereader state: {pkrd_state}")

shader.cl

#ifndef JUMP_COUNT
#define JUMP_COUNT 0
#endif
#ifndef JUMP_DATA
#define JUMP_DATA
#endif
#ifndef BASE_ADVANCE
#define BASE_ADVANCE 0
#endif
__constant unsigned char JUMPS[JUMP_COUNT] = { JUMP_DATA };

struct tinymt {
  uint state[4];
};

inline void advance(struct tinymt *rng) {
  uint y = rng->state[3];
  uint x = (rng->state[0] & 0x7FFFFFFF) ^ rng->state[1] ^ rng->state[2];
  x ^= x << 1;
  y ^= (y >> 1) ^ x;

  rng->state[0] = rng->state[1];
  rng->state[1] = rng->state[2] ^ ((y & 1) * 0x8F7011EE);
  rng->state[2] = x ^ ((y << 10) & 0xFFFFFFFF) ^ ((y & 1) * 0xFC78FF1F);
  rng->state[3] = y;
}

inline uint next_uint(struct tinymt *rng) {
  advance(rng);
  uint t0 = rng->state[3];
  uint t1 = rng->state[0] + (rng->state[2] >> 8);
  t0 ^= t1;
  if (t1 & 1) {
    t0 ^= 0x3793FDFF;
  }
  return t0;
}

inline void init(struct tinymt *rng, unsigned int seed) {
  rng->state[0] = seed;
  rng->state[1] = 0x8F7011EE;
  rng->state[2] = 0xFC78FF1F;
  rng->state[3] = 0x3793FDFF;

  for (int i = 1; i < 8; i++) {
    rng->state[i & 3] ^= (0x6C078965 * (rng->state[(i - 1) & 3] ^ (rng->state[(i - 1) & 3] >> 30)) + i);
  }

  for (int i = 0; i < 8; i++) {
    advance(rng);
  }
}

__kernel void find_initial_seeds(__global uint *cnt, __global ulong *res_g) {
  unsigned int seed = get_global_id(0) | (get_global_id(1) << 16);
  struct tinymt rng;
  init(&rng, seed);
  for (int i = 0; i < BASE_ADVANCE; i++) {
    advance(&rng);
  }

  bool valid = true;
  for (int i = 0; i < JUMP_COUNT; i++) {
    for (unsigned char j = 0; j < JUMPS[i]; j++) {
      if ((next_uint(&rng) % 3) == 0) {
        valid = false;
        break;
      }
    }
    valid &= (next_uint(&rng) % 3) == 0;
    if (!valid) {
      break;
    }
  }

  if (valid) {
    res_g[atomic_inc(&cnt[0])] = seed;
  }
}