(Py)OpenCL implementation of searching for a TID/SID in Pokemon DPPT/HGSS (find an initial seed that generates a given 32bit rand as the second mersenne twister call)

gen4ids.py

import pyopencl as cl
import numpy as np
import time

TID = int(input("TID: "))
SID = int(input("SID: "))
MAX_HOUR_VALUE = input("Max Hour Value (empty for full search): ")
MAX_HOUR_VALUE = int(MAX_HOUR_VALUE or 0xFF)

def untemper(y):
    y ^= (y >> 18)
    y ^= (y << 15) & 0xEFC60000
    y ^= (
        ((y << 7) & 0x9D2C5680)
        ^ ((y << 14) & 0x94284000)
        ^ ((y << 21) & 0x14200000)
        ^ ((y << 28) & 0x10000000)
    )
    y ^= (y >> 11) ^ (y >> 22)
    return y

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

program = cl.Program(
    context,
    f"__constant uint TARGET = {untemper((SID << 16) | TID)};"
    + """
    __kernel void find_seeds(__global uint *cnt, __global uint *res_g) {
        unsigned int day_month = get_global_id(0) << 24;
        unsigned int hour_value = get_global_id(1) << 16;
        unsigned int delay = get_global_id(2);
        unsigned int seed = day_month | hour_value | delay;
        unsigned int s1 = 0x6C078965 * (seed ^ (seed >> 30)) + 1;
        unsigned int s2 = 0x6C078965 * (s1 ^ (s1 >> 30)) + 2;
        unsigned int s398 = s2;
        for (int i = 3; i < 399; i++) {
            s398 = 0x6C078965 * (s398 ^ (s398 >> 30)) + i;
        }
        unsigned int y = (s1 & 0x80000000) | (s2 & 0x7FFFFFFF);
        y = s398 ^ (y >> 1) ^ ((y & 1) * 0x9908B0DF);
        if (y == TARGET) {
            res_g[atomic_inc(&cnt[0])] = seed;
        }
    }
    """
).build()

# heavy over estimate
host_results = np.zeros(128, np.uint32)
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)

kernel = program.find_seeds
kernel(
    queue,
    (0x100, MAX_HOUR_VALUE + 1, 0x10000),
    None,
    device_count,
    device_results,
)

print("Processing ....")
start = time.perf_counter()
cl.enqueue_copy(queue, host_results, device_results)
cl.enqueue_copy(queue, host_count, device_count)
end = time.perf_counter()

host_results = host_results[: host_count[0]]

print(" | ".join(f"{result:08X}" for result in host_results))
print(f"{end - start:.02f}s search time")
print(f"{(end - start) * 1000000000 / ((MAX_HOUR_VALUE + 1) << 24):.02f}ns per seed")