diff --git a/ocean/affine_lock/affine_lock.cu b/ocean/affine_lock/affine_lock.cu new file mode 100644 index 0000000000..adb550eb5f --- /dev/null +++ b/ocean/affine_lock/affine_lock.cu @@ -0,0 +1,808 @@ +// Vibe coded by OpenAI Codex. +// GPU Affine Lock environment. This is intentionally standalone from +// affine_lock.h: --gpu builds include this file instead of the CPU source. +#ifndef PUFFER_AFFINE_LOCK_GPU_CU +#define PUFFER_AFFINE_LOCK_GPU_CU + +#define PUF_BACKEND PUF_GPU + +#include +#include + +#include +#include +#include +#include +#include + +// Environment observations are fixed bf16. All bit observations are +/-1 +// (exact in bf16); only the timer is rounded. Keeping them bf16 halves the +// rollout bandwidth compared with the CPU float representation. +typedef __nv_bfloat16 obs_t; +#include "pufferenv.h" +#include "affine_lock_visible_targets.h" + +#define BITS 16 +#define TIMER_INDEX (2 * BITS) +#define OBS_SIZE (TIMER_INDEX + 1) +#define NUM_ATNS 1 +#define NUM_ACTIONS 8 +#define MAX_SOLUTION_DEPTH 16 +#define CURRICULUM_DEPTH_COUNT 6 +#define STEP_REWARD (-0.01f) +#ifndef VISIBLE_TARGET_TABLE_PATH +#define VISIBLE_TARGET_TABLE_PATH "ocean/affine_lock/generated/affine_lock_8action_visible_targets.bin" +#endif +#define ACT_SIZES {NUM_ACTIONS} +#define PUF_STEPS_PER_SEC 2 + +#define PERF_WEIGHTING_LINEAR 0 +#define PERF_WEIGHTING_QUADRATIC 1 + +#ifndef AFFINE_LOCK_GPU_SHARED_OBS +#define AFFINE_LOCK_GPU_SHARED_OBS 1 +#endif +#ifndef AFFINE_LOCK_GPU_SHARED_BLOCK +#define AFFINE_LOCK_GPU_SHARED_BLOCK 128 +#endif +#define AFFINE_LOCK_GPU_DEPTH_LUT_SIZE (MAX_SOLUTION_DEPTH + 1) + +static_assert(AFFINE_LOCK_GPU_SHARED_BLOCK >= 32 && + AFFINE_LOCK_GPU_SHARED_BLOCK <= 256 && + AFFINE_LOCK_GPU_SHARED_BLOCK % 32 == 0, + "AFFINE_LOCK_GPU_SHARED_BLOCK must contain whole warps"); + +#if !AFFINE_LOCK_GPU_SHARED_OBS +#ifndef AFFINE_LOCK_GPU_LANES +#define AFFINE_LOCK_GPU_LANES 4 +#endif +#define AFFINE_LOCK_GPU_BLOCK 256 +static_assert(AFFINE_LOCK_GPU_LANES == 4 || AFFINE_LOCK_GPU_LANES == 8 || + AFFINE_LOCK_GPU_LANES == 16 || AFFINE_LOCK_GPU_LANES == 32, + "AFFINE_LOCK_GPU_LANES must be a power-of-two subwarp"); +static_assert(AFFINE_LOCK_GPU_BLOCK % AFFINE_LOCK_GPU_LANES == 0, + "block size must contain whole environments"); +#endif + +struct Log { + float perf; + float score; + float solve_rate; + float max_depth_solve; + float episode_return; + float episode_length; + float solve_steps; + float timeout_rate; + float solve_efficiency; + float target_distance; + float solved_target_distance; + float d6_rate; + float d6_solve_rate; + float d8_rate; + float d8_solve_rate; + float d16_rate; + float d16_solve_rate; + float n; +}; + +static_assert(sizeof(Log) == 18 * sizeof(float), + "trainer log reduction requires a packed float-only Log"); + +// The trainer only reads Env::log for a GPU backend. Runtime state is kept in +// a separate compact array so log scans do not pull state into cache and state +// updates do not stride over the relatively large log payload. +struct Env { + Log log; + Agent agents[1]; + int num_agents; + int tag; + int boundary_reached; + unsigned int rng; +}; + +// Exactly 32 bytes: four adjacent environment records fit in one 128-byte +// transaction. The default shared-observation kernel reads one record per env. +typedef struct GpuAffineLockState { + uint32_t rng; + uint16_t state; + uint16_t target; + int step_count; + int max_steps; + int scramble_depth; + int curriculum_depth; + int target_distance; + float episode_return; +} GpuAffineLockState; + +static_assert(sizeof(GpuAffineLockState) == 32, + "GpuAffineLockState layout is performance-sensitive"); + +typedef struct GpuAffineLockConfig { + int start_depth; + int max_depth; + int step_grace; + int perf_weighting; + uint32_t depth_first[AFFINE_LOCK_GPU_DEPTH_LUT_SIZE]; + uint32_t depth_counts[AFFINE_LOCK_GPU_DEPTH_LUT_SIZE]; +} GpuAffineLockConfig; + +__constant__ GpuAffineLockConfig d_affine_lock_config; + +static struct { + Env* envs; + GpuAffineLockState* states; + uint32_t* target_pairs; + int n; + obs_t* observations; + float* actions; + float* rewards; + float* terminals; + cudaStream_t stream; + GpuAffineLockConfig config; +} g_gpu; + +static void gpu_affine_lock_check(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + std::fprintf(stderr, "Affine Lock CUDA: %s failed: %s\n", + operation, cudaGetErrorString(status)); + std::exit(1); + } +} + +#if !AFFINE_LOCK_GPU_SHARED_OBS +static int gpu_affine_lock_grid(int threads) { + return (threads + AFFINE_LOCK_GPU_BLOCK - 1) / AFFINE_LOCK_GPU_BLOCK; +} +#endif + +__device__ __forceinline__ uint32_t gpu_affine_lock_random_mixed_u32( + GpuAffineLockState* env) { + env->rng = env->rng * 1664525u + 1013904223u; + uint32_t x = env->rng; + x ^= x >> 16; + x *= 0x7feb352du; + x ^= x >> 15; + x *= 0x846ca68bu; + x ^= x >> 16; + return x; +} + +__device__ __forceinline__ int gpu_affine_lock_random_bounded( + GpuAffineLockState* env, int bound) { + uint32_t ubound = (uint32_t)bound; + uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; + uint32_t value = gpu_affine_lock_random_mixed_u32(env); + while (value >= limit) { + value = gpu_affine_lock_random_mixed_u32(env); + } + return (int)(value % ubound); +} + +__device__ __forceinline__ void gpu_affine_lock_reset_state( + GpuAffineLockState* env, const uint32_t* target_pairs) { + env->scramble_depth = env->curriculum_depth; + env->step_count = 0; + env->episode_return = 0.0f; + int depth = env->scramble_depth; + uint32_t count = d_affine_lock_config.depth_counts[depth]; + int choice = gpu_affine_lock_random_bounded(env, (int)count); + uint32_t record_index = d_affine_lock_config.depth_first[depth] + + (uint32_t)choice; + uint32_t pair = target_pairs[record_index]; + env->state = (uint16_t)(pair & 0xffffu); + env->target = (uint16_t)(pair >> 16); + env->target_distance = env->scramble_depth; + env->max_steps = env->target_distance + d_affine_lock_config.step_grace; +} + +__device__ __forceinline__ uint16_t gpu_affine_lock_apply_action( + uint16_t state, int action) { + uint32_t next = state; + switch (action) { + case 0: + next = (state >> 1) | ((state & 1u) << 15); + break; + case 1: + next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); + break; + case 2: + next = state ^ 0xfe00u; + break; + case 3: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + break; + case 4: + next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); + break; + case 5: + next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); + break; + case 6: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + break; + case 7: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); + break; + } + return (uint16_t)(next & 0xffffu); +} + +__device__ __forceinline__ int gpu_affine_lock_next_curriculum_depth( + int current_depth) { + constexpr int curriculum_depths[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; +#pragma unroll + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = curriculum_depths[i]; + if (depth > current_depth) { + return depth < d_affine_lock_config.max_depth + ? depth : d_affine_lock_config.max_depth; + } + } + return d_affine_lock_config.max_depth; +} + +__device__ __forceinline__ void gpu_affine_lock_add_log( + Env* trainer_env, const GpuAffineLockState* env, int solved) { + int log_depth = env->target_distance; + int at_max_depth = log_depth == d_affine_lock_config.max_depth; + float ratio = log_depth / (float)d_affine_lock_config.max_depth; + float solve_credit = 0.0f; + if (solved) { + solve_credit = d_affine_lock_config.perf_weighting == PERF_WEIGHTING_QUADRATIC + ? ratio * ratio : ratio; + } + Log* log = &trainer_env->log; + log->perf += solve_credit; + log->score += solve_credit; + log->solve_rate += solved; + log->max_depth_solve += solved && at_max_depth; + log->episode_return += env->episode_return; + log->episode_length += env->step_count; + log->solve_steps += solved ? env->step_count : 0; + log->timeout_rate += !solved; + log->solve_efficiency += solved + ? env->step_count / (float)log_depth : 0.0f; + log->target_distance += env->target_distance; + log->solved_target_distance += solved ? env->target_distance : 0; + log->d6_rate += log_depth == 6; + log->d6_solve_rate += solved && log_depth == 6; + log->d8_rate += log_depth == 8; + log->d8_solve_rate += solved && log_depth == 8; + log->d16_rate += log_depth == 16; + log->d16_solve_rate += solved && log_depth == 16; + log->n += 1; +} + +__device__ __forceinline__ uint32_t gpu_affine_lock_step_one( + Env* trainer_env, GpuAffineLockState* env, + const uint32_t* target_pairs, float action, + float* reward_out, float* terminal_out, float* timer_out) { + float reward = STEP_REWARD; + float terminal = 0.0f; + int solved = 0; + env->step_count += 1; + int invalid = !isfinite(action) || action < 0.0f || action > NUM_ACTIONS - 1; + if (invalid) { + reward = -1.0f; + terminal = 1.0f; + } else { + env->state = gpu_affine_lock_apply_action(env->state, (int)action); + if (env->state == env->target) { + reward = 1.0f; + terminal = 1.0f; + solved = 1; + } else if (env->step_count >= env->max_steps) { + reward = -1.0f; + terminal = 1.0f; + } + } + env->episode_return += reward; + if (terminal != 0.0f) { + gpu_affine_lock_add_log(trainer_env, env, solved); + env->curriculum_depth = solved + ? gpu_affine_lock_next_curriculum_depth(env->scramble_depth) + : d_affine_lock_config.start_depth; + gpu_affine_lock_reset_state(env, target_pairs); + } + *reward_out = reward; + *terminal_out = terminal; + *timer_out = env->step_count / (float)env->max_steps; + return (uint32_t)env->state | ((uint32_t)env->target << 16); +} + +#if !AFFINE_LOCK_GPU_SHARED_OBS +__device__ __forceinline__ void gpu_affine_lock_write_observations( + obs_t* observations, uint32_t packed_bits, float timer, int lane) { +#pragma unroll + for (int bit = lane; bit < 2 * BITS; bit += AFFINE_LOCK_GPU_LANES) { + observations[bit] = __float2bfloat16( + (packed_bits & (1u << bit)) ? 1.0f : -1.0f); + } + if (lane == 0) { + observations[TIMER_INDEX] = __float2bfloat16(timer); + } +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_BLOCK) +void gpu_affine_lock_reset_kernel(Env* envs, GpuAffineLockState* states, + const uint32_t* target_pairs, obs_t* observations, + float* rewards, float* terminals, int num_envs) { + int thread = blockIdx.x * blockDim.x + threadIdx.x; + int relative_env = thread / AFFINE_LOCK_GPU_LANES; + int lane = thread & (AFFINE_LOCK_GPU_LANES - 1); + int active = relative_env < num_envs; + uint32_t packed_bits = 0; + float timer = 0.0f; + if (active && lane == 0) { + GpuAffineLockState* env = &states[relative_env]; + gpu_affine_lock_reset_state(env, target_pairs); + rewards[relative_env] = 0.0f; + terminals[relative_env] = 0.0f; + packed_bits = (uint32_t)env->state | ((uint32_t)env->target << 16); + } + int leader = (threadIdx.x & 31) & ~(AFFINE_LOCK_GPU_LANES - 1); + packed_bits = __shfl_sync(0xffffffffu, packed_bits, leader); + timer = __shfl_sync(0xffffffffu, timer, leader); + if (active) { + gpu_affine_lock_write_observations( + observations + (size_t)relative_env * OBS_SIZE, + packed_bits, timer, lane); + } + (void)envs; +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_BLOCK) +void gpu_affine_lock_step_kernel(Env* envs, GpuAffineLockState* states, + const uint32_t* target_pairs, const float* actions, + obs_t* observations, float* rewards, float* terminals, + int num_envs) { + int thread = blockIdx.x * blockDim.x + threadIdx.x; + int relative_env = thread / AFFINE_LOCK_GPU_LANES; + int lane = thread & (AFFINE_LOCK_GPU_LANES - 1); + int active = relative_env < num_envs; + uint32_t packed_bits = 0; + float timer = 0.0f; + if (active && lane == 0) { + packed_bits = gpu_affine_lock_step_one( + &envs[relative_env], &states[relative_env], target_pairs, + actions[(size_t)relative_env * NUM_ATNS], + &rewards[relative_env], &terminals[relative_env], &timer); + } + int leader = (threadIdx.x & 31) & ~(AFFINE_LOCK_GPU_LANES - 1); + packed_bits = __shfl_sync(0xffffffffu, packed_bits, leader); + timer = __shfl_sync(0xffffffffu, timer, leader); + if (active) { + gpu_affine_lock_write_observations( + observations + (size_t)relative_env * OBS_SIZE, + packed_bits, timer, lane); + } +} +#endif + +#if AFFINE_LOCK_GPU_SHARED_OBS +// One simulation thread per environment writes +// a conflict-free 33-float shared-memory row, then the whole block converts and +// stores a linear bf16 tile with fully coalesced global writes. +__global__ __launch_bounds__(AFFINE_LOCK_GPU_SHARED_BLOCK) +void gpu_affine_lock_shared_reset_kernel(Env* envs, + GpuAffineLockState* states, const uint32_t* target_pairs, + obs_t* observations, float* rewards, float* terminals, int num_envs) { + __shared__ float observation_tile[AFFINE_LOCK_GPU_SHARED_BLOCK * OBS_SIZE]; + int block_start = blockIdx.x * AFFINE_LOCK_GPU_SHARED_BLOCK; + int relative_env = block_start + threadIdx.x; + int active_count = num_envs - block_start; + if (active_count > AFFINE_LOCK_GPU_SHARED_BLOCK) { + active_count = AFFINE_LOCK_GPU_SHARED_BLOCK; + } + if (active_count < 0) { + active_count = 0; + } + if (threadIdx.x < active_count) { + GpuAffineLockState* env = &states[relative_env]; + gpu_affine_lock_reset_state(env, target_pairs); + rewards[relative_env] = 0.0f; + terminals[relative_env] = 0.0f; + uint32_t bits = (uint32_t)env->state | ((uint32_t)env->target << 16); + float* row = observation_tile + threadIdx.x * OBS_SIZE; +#pragma unroll + for (int bit = 0; bit < 2 * BITS; bit++) { + row[bit] = (bits & (1u << bit)) ? 1.0f : -1.0f; + } + row[TIMER_INDEX] = 0.0f; + } + __syncthreads(); + int tile_values = active_count * OBS_SIZE; + for (int value = threadIdx.x; value < tile_values; value += blockDim.x) { + observations[(size_t)block_start * OBS_SIZE + value] = + __float2bfloat16(observation_tile[value]); + } + (void)envs; +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_SHARED_BLOCK) +void gpu_affine_lock_shared_step_kernel(Env* envs, + GpuAffineLockState* states, const uint32_t* target_pairs, + const float* actions, obs_t* observations, + float* rewards, float* terminals, int num_envs) { + __shared__ float observation_tile[AFFINE_LOCK_GPU_SHARED_BLOCK * OBS_SIZE]; + int block_start = blockIdx.x * AFFINE_LOCK_GPU_SHARED_BLOCK; + int relative_env = block_start + threadIdx.x; + int active_count = num_envs - block_start; + if (active_count > AFFINE_LOCK_GPU_SHARED_BLOCK) { + active_count = AFFINE_LOCK_GPU_SHARED_BLOCK; + } + if (active_count < 0) { + active_count = 0; + } + if (threadIdx.x < active_count) { + float timer = 0.0f; + uint32_t bits = gpu_affine_lock_step_one( + &envs[relative_env], &states[relative_env], target_pairs, + actions[(size_t)relative_env * NUM_ATNS], + &rewards[relative_env], &terminals[relative_env], &timer); + float* row = observation_tile + threadIdx.x * OBS_SIZE; +#pragma unroll + for (int bit = 0; bit < 2 * BITS; bit++) { + row[bit] = (bits & (1u << bit)) ? 1.0f : -1.0f; + } + row[TIMER_INDEX] = timer; + } + __syncthreads(); + int tile_values = active_count * OBS_SIZE; + for (int value = threadIdx.x; value < tile_values; value += blockDim.x) { + observations[(size_t)block_start * OBS_SIZE + value] = + __float2bfloat16(observation_tile[value]); + } +} +#endif + +void puf_log(Log* log, Dict* out) { + float nsolve = log->solve_rate; + float solved_min_win_moves = nsolve + ? log->solved_target_distance / nsolve : 0; + float conditional_solve_steps = nsolve ? log->solve_steps / nsolve : 0; + float conditional_solve_efficiency = nsolve + ? log->solve_efficiency / nsolve : 0; + + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "solve_rate", log->solve_rate); + dict_set(out, "max_depth_solve", log->max_depth_solve); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "timeout_rate", log->timeout_rate); + dict_set(out, "min_win_moves", log->target_distance); + dict_set(out, "solved_min_win_moves", solved_min_win_moves); + dict_set(out, "conditional_solve_steps", conditional_solve_steps); + dict_set(out, "conditional_solve_efficiency", conditional_solve_efficiency); + dict_set(out, "d6_solve_rate", log->d6_rate + ? log->d6_solve_rate / log->d6_rate : 0); + dict_set(out, "d8_solve_rate", log->d8_rate + ? log->d8_solve_rate / log->d8_rate : 0); + dict_set(out, "d16_solve_rate", log->d16_rate + ? log->d16_solve_rate / log->d16_rate : 0); + dict_set(out, "n", log->n); +} + +static int gpu_affine_lock_host_has_depth( + const GpuAffineLockConfig* config, int depth) { + return depth >= 0 && depth < AFFINE_LOCK_GPU_DEPTH_LUT_SIZE + && config->depth_counts[depth] != 0; +} + +static int gpu_affine_lock_host_next_depth(int current_depth, int max_depth) { + static const int curriculum_depths[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = curriculum_depths[i]; + if (depth > current_depth) { + return depth < max_depth ? depth : max_depth; + } + } + return max_depth; +} + +static void gpu_affine_lock_validate_curriculum( + const GpuAffineLockConfig* config) { + if (config->start_depth <= 0 || + config->max_depth < config->start_depth || + config->max_depth > MAX_SOLUTION_DEPTH) { + std::fprintf(stderr, + "Affine Lock CUDA: invalid curriculum range start=%d max=%d\n", + config->start_depth, config->max_depth); + std::exit(1); + } + int depth = config->start_depth; + for (int i = 0; i <= CURRICULUM_DEPTH_COUNT; i++) { + if (!gpu_affine_lock_host_has_depth(config, depth)) { + std::fprintf(stderr, + "Affine Lock CUDA: target table has no depth %d section\n", depth); + std::exit(1); + } + if (depth + config->step_grace <= 0) { + std::fprintf(stderr, + "Affine Lock CUDA: depth %d with step_grace=%d has no valid steps\n", + depth, config->step_grace); + std::exit(1); + } + if (depth == config->max_depth) { + return; + } + int next = gpu_affine_lock_host_next_depth(depth, config->max_depth); + if (next == depth) { + break; + } + depth = next; + } + std::fprintf(stderr, "Affine Lock CUDA: curriculum does not reach max depth %d\n", + config->max_depth); + std::exit(1); +} + +Env* puf_vec_create(int n, Dict* env_kwargs, + obs_t* observations, float* actions, + float* rewards, float* terminals) { + if (n <= 0) { + std::fprintf(stderr, "Affine Lock CUDA: vector size must be positive\n"); + std::exit(1); + } + if (g_gpu.envs != nullptr) { + std::fprintf(stderr, "Affine Lock CUDA: vector already exists\n"); + std::exit(1); + } + + VisibleTargetTable table = {}; + if (visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table) != 0) { + std::fprintf(stderr, + "Affine Lock CUDA: failed to load visible target table %s\n", + VISIBLE_TARGET_TABLE_PATH); + std::exit(1); + } + if (table.num_actions != NUM_ACTIONS) { + std::fprintf(stderr, + "Affine Lock CUDA: target table has %u actions, expected %d\n", + table.num_actions, NUM_ACTIONS); + visible_targets_free(&table); + std::exit(1); + } + + GpuAffineLockConfig config = {}; + config.start_depth = (int)dict_get(env_kwargs, "start_depth"); + config.max_depth = (int)dict_get(env_kwargs, "max_depth"); + config.step_grace = (int)dict_get(env_kwargs, "step_grace"); + config.perf_weighting = (int)dict_get(env_kwargs, "perf_weighting"); + for (uint32_t i = 0; i < table.depth_count; i++) { + const VisibleTargetDepth* depth = &table.depths[i]; + if (depth->depth >= AFFINE_LOCK_GPU_DEPTH_LUT_SIZE || + depth->stored_count == 0 || + config.depth_counts[depth->depth] != 0) { + std::fprintf(stderr, + "Affine Lock CUDA: invalid target-table depth section %u\n", + depth->depth); + visible_targets_free(&table); + std::exit(1); + } + config.depth_first[depth->depth] = depth->first_record; + config.depth_counts[depth->depth] = depth->stored_count; + for (uint32_t record_offset = 0; + record_offset < depth->stored_count; record_offset++) { + const VisibleTargetRecord* record = + &table.records[depth->first_record + record_offset]; + if (record->depth != depth->depth) { + std::fprintf(stderr, + "Affine Lock CUDA: record depth %u does not match section %u\n", + (unsigned int)record->depth, depth->depth); + visible_targets_free(&table); + std::exit(1); + } + } + } + gpu_affine_lock_validate_curriculum(&config); + + uint32_t* host_pairs = (uint32_t*)std::malloc( + (size_t)table.record_count * sizeof(uint32_t)); + if (host_pairs == nullptr) { + std::perror("malloc"); + visible_targets_free(&table); + std::exit(1); + } + for (uint32_t i = 0; i < table.record_count; i++) { + host_pairs[i] = (uint32_t)table.records[i].start + | ((uint32_t)table.records[i].target << 16); + } + + GpuAffineLockState* host_states = (GpuAffineLockState*)std::calloc( + (size_t)n, sizeof(GpuAffineLockState)); + if (host_states == nullptr) { + std::perror("calloc"); + std::free(host_pairs); + visible_targets_free(&table); + std::exit(1); + } + unsigned int running_seed = (unsigned int)dict_get(env_kwargs, "seed"); + for (int i = 0; i < n; i++) { + host_states[i].rng = (uint32_t)rand_r(&running_seed); + host_states[i].curriculum_depth = config.start_depth; + } + + Env* device_envs = nullptr; + GpuAffineLockState* device_states = nullptr; + uint32_t* device_pairs = nullptr; + gpu_affine_lock_check(cudaMalloc((void**)&device_envs, + (size_t)n * sizeof(Env)), "cudaMalloc envs"); + gpu_affine_lock_check(cudaMalloc((void**)&device_states, + (size_t)n * sizeof(GpuAffineLockState)), "cudaMalloc states"); + gpu_affine_lock_check(cudaMalloc((void**)&device_pairs, + (size_t)table.record_count * sizeof(uint32_t)), "cudaMalloc target pairs"); + gpu_affine_lock_check(cudaMemset(device_envs, 0, + (size_t)n * sizeof(Env)), "clear env logs"); + gpu_affine_lock_check(cudaMemcpy(device_states, host_states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), "copy states"); + gpu_affine_lock_check(cudaMemcpy(device_pairs, host_pairs, + (size_t)table.record_count * sizeof(uint32_t), cudaMemcpyHostToDevice), + "copy target pairs"); + gpu_affine_lock_check(cudaMemcpyToSymbol(d_affine_lock_config, + &config, sizeof(config)), "copy config"); + + std::free(host_pairs); + std::free(host_states); + visible_targets_free(&table); + + g_gpu.envs = device_envs; + g_gpu.states = device_states; + g_gpu.target_pairs = device_pairs; + g_gpu.n = n; + g_gpu.observations = observations; + g_gpu.actions = actions; + g_gpu.rewards = rewards; + g_gpu.terminals = terminals; + g_gpu.stream = nullptr; + g_gpu.config = config; + return device_envs; +} + +void puf_bind_stream(cudaStream_t stream) { + g_gpu.stream = stream; +} + +// GPU creation is vector-only; puf_init exists to satisfy the common API. +void puf_init(Env* env, Dict* kwargs) { + (void)env; + (void)kwargs; +} + +void puf_reset(Env* env) { + (void)env; +#if AFFINE_LOCK_GPU_SHARED_OBS + int blocks = (g_gpu.n + AFFINE_LOCK_GPU_SHARED_BLOCK - 1) + / AFFINE_LOCK_GPU_SHARED_BLOCK; + gpu_affine_lock_shared_reset_kernel<<< + blocks, AFFINE_LOCK_GPU_SHARED_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.observations, g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#else + int threads = g_gpu.n * AFFINE_LOCK_GPU_LANES; + gpu_affine_lock_reset_kernel<<< + gpu_affine_lock_grid(threads), AFFINE_LOCK_GPU_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.observations, g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#endif + gpu_affine_lock_check(cudaPeekAtLastError(), "launch reset kernel"); +} + +void puf_step(Env* env) { + (void)env; +#if AFFINE_LOCK_GPU_SHARED_OBS + int blocks = (g_gpu.n + AFFINE_LOCK_GPU_SHARED_BLOCK - 1) + / AFFINE_LOCK_GPU_SHARED_BLOCK; + gpu_affine_lock_shared_step_kernel<<< + blocks, AFFINE_LOCK_GPU_SHARED_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.actions, g_gpu.observations, + g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#else + int threads = g_gpu.n * AFFINE_LOCK_GPU_LANES; + gpu_affine_lock_step_kernel<<< + gpu_affine_lock_grid(threads), AFFINE_LOCK_GPU_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.actions, g_gpu.observations, + g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#endif + gpu_affine_lock_check(cudaPeekAtLastError(), "launch step kernel"); +} + +void puf_close(Env* env) { + (void)env; + if (IsWindowReady()) { + CloseWindow(); + } + if (g_gpu.envs != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.envs), "cudaFree envs"); + } + if (g_gpu.states != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.states), "cudaFree states"); + } + if (g_gpu.target_pairs != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.target_pairs), "cudaFree target pairs"); + } + g_gpu = {}; +} + +void puf_render(Env* env) { + (void)env; + if (g_gpu.envs == nullptr || g_gpu.n < 1) { + return; + } + if (IsWindowReady() && (WindowShouldClose() || IsKeyPressed(KEY_ESCAPE))) { + puf_close(g_gpu.envs); + std::exit(0); + } + if (!IsWindowReady()) { + InitWindow(780, 360, "PufferLib AffineLock CUDA"); + SetTargetFPS(30); + } + if (g_gpu.stream != nullptr) { + gpu_affine_lock_check(cudaStreamSynchronize(g_gpu.stream), + "synchronize render stream"); + } + GpuAffineLockState state; + float reward = 0.0f; + float terminal = 0.0f; + gpu_affine_lock_check(cudaMemcpy(&state, g_gpu.states, sizeof(state), + cudaMemcpyDeviceToHost), "copy render state"); + gpu_affine_lock_check(cudaMemcpy(&reward, g_gpu.rewards, sizeof(reward), + cudaMemcpyDeviceToHost), "copy render reward"); + gpu_affine_lock_check(cudaMemcpy(&terminal, g_gpu.terminals, sizeof(terminal), + cudaMemcpyDeviceToHost), "copy render terminal"); + + uint32_t mismatches = (state.state ^ state.target) & 0xffffu; + const char* status = terminal == 0.0f + ? "running" : (reward > 0.0f ? "solved" : "failed"); + Color status_color = terminal == 0.0f + ? (Color){190, 198, 206, 255} + : (reward > 0.0f + ? (Color){80, 210, 140, 255} + : (Color){238, 88, 88, 255}); + + BeginDrawing(); + ClearBackground((Color){6, 24, 24, 255}); + DrawText("Affine Lock CUDA", 30, 24, 28, RAYWHITE); + DrawText(TextFormat("depth %d/%d step %d/%d last reward %.2f", + state.scramble_depth, g_gpu.config.max_depth, + state.step_count, state.max_steps, reward), + 30, 62, 20, (Color){180, 190, 200, 255}); + DrawText(TextFormat("status %s mismatches 0x%04x", + status, mismatches), 30, 90, 20, status_color); + + const char* row_label[2] = {"current", "target"}; + uint32_t row_value[2] = {state.state, state.target}; + int row_y[2] = {138, 220}; + for (int row = 0; row < 2; row++) { + DrawText(row_label[row], 30, row_y[row] + 9, 20, RAYWHITE); + for (int bit = 0; bit < BITS; bit++) { + int x = 145 + bit * 34; + int on = (row_value[row] >> bit) & 1u; + int mismatch = ((state.state ^ state.target) >> bit) & 1u; + Color fill = on + ? (Color){80, 210, 140, 255} + : (Color){38, 48, 58, 255}; + Color border = mismatch + ? (Color){238, 88, 88, 255} + : (Color){182, 196, 205, 255}; + DrawRectangle(x, row_y[row], 24, 34, fill); + DrawRectangleLinesEx((Rectangle){(float)x, (float)row_y[row], 24, 34}, + mismatch ? 3 : 1, border); + DrawText(TextFormat("%d", bit), x + 5, row_y[row] + 40, 10, + (Color){128, 140, 150, 255}); + } + } + DrawText("GPU-resident environment", 30, 310, 16, + (Color){160, 170, 178, 255}); + EndDrawing(); + puf_web_vsync(); +} + +#endif diff --git a/ocean/affine_lock/tests/run_cuda.sh b/ocean/affine_lock/tests/run_cuda.sh new file mode 100755 index 0000000000..f7762be9a1 --- /dev/null +++ b/ocean/affine_lock/tests/run_cuda.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/../../.." && pwd)" +cd "$ROOT" + +OUT="${TMPDIR:-/tmp}/affine_lock_cuda_tests" +CUDA_ROOT="${CUDA_HOME:-${CUDA_PATH:-/usr/local/cuda}}" +NVCC_BIN="${NVCC:-$CUDA_ROOT/bin/nvcc}" +CUDA_ARCH="${NVCC_ARCH:-native}" + +RAYLIB_ROOT="$ROOT/raylib-5.5_linux_amd64" +if [ ! -d "$RAYLIB_ROOT/include" ]; then + echo "raylib-5.5_linux_amd64 not found" >&2 + exit 1 +fi + +"$NVCC_BIN" \ + -std=c++17 -O3 -lineinfo -arch="$CUDA_ARCH" \ + -Xcompiler=-Wall,-Wextra,-Werror,-Wno-unused-function,-Wno-unused-parameter,-Wno-missing-field-initializers \ + -Xcompiler=-ffunction-sections,-fdata-sections \ + -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" \ + -I"$ROOT/vendor" -I"$RAYLIB_ROOT/include" \ + "$ROOT/ocean/affine_lock/tests/test_affine_lock_cuda.cu" \ + "$RAYLIB_ROOT/lib/libraylib.a" \ + -Xlinker=--gc-sections \ + -lGL -lpthread -ldl -lrt -lm \ + -o "$OUT" + +"$OUT" diff --git a/ocean/affine_lock/tests/test_affine_lock_cuda.cu b/ocean/affine_lock/tests/test_affine_lock_cuda.cu new file mode 100644 index 0000000000..54cb06ad34 --- /dev/null +++ b/ocean/affine_lock/tests/test_affine_lock_cuda.cu @@ -0,0 +1,1109 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "../affine_lock.cu" + +#define EXPECT_TRUE(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: expected true: %s\n", \ + __FILE__, __LINE__, #condition); \ + std::exit(1); \ + } \ +} while (0) + +#define EXPECT_EQ(actual, expected) do { \ + auto actual_value = (actual); \ + auto expected_value = (expected); \ + if (actual_value != expected_value) { \ + std::fprintf(stderr, "%s:%d: expected %s == %s, got %lld != %lld\n", \ + __FILE__, __LINE__, #actual, #expected, \ + (long long)actual_value, (long long)expected_value); \ + std::exit(1); \ + } \ +} while (0) + +#define EXPECT_NEAR(actual, expected, tolerance) do { \ + float actual_value = (float)(actual); \ + float expected_value = (float)(expected); \ + if (!std::isfinite(actual_value) || !std::isfinite(expected_value) || \ + std::fabs(actual_value - expected_value) > (tolerance)) { \ + std::fprintf(stderr, "%s:%d: expected %s ~= %.9g, got %.9g\n", \ + __FILE__, __LINE__, #actual, expected_value, actual_value); \ + std::exit(1); \ + } \ +} while (0) + +static void check_cuda(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + std::fprintf(stderr, "%s failed: %s\n", operation, cudaGetErrorString(status)); + std::exit(1); + } +} + +typedef struct OracleState { + uint32_t rng; + uint16_t state; + uint16_t target; + int step_count; + int max_steps; + int scramble_depth; + int curriculum_depth; + int target_distance; + float episode_return; + Log log; +} OracleState; + +static uint32_t oracle_random_mixed_u32(OracleState* env) { + env->rng = env->rng * 1664525u + 1013904223u; + uint32_t x = env->rng; + x ^= x >> 16; + x *= 0x7feb352du; + x ^= x >> 15; + x *= 0x846ca68bu; + x ^= x >> 16; + return x; +} + +static int oracle_random_bounded(OracleState* env, int bound) { + uint32_t ubound = (uint32_t)bound; + uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; + uint32_t value = oracle_random_mixed_u32(env); + while (value >= limit) { + value = oracle_random_mixed_u32(env); + } + return (int)(value % ubound); +} + +static const VisibleTargetDepth* oracle_depth( + const VisibleTargetTable* table, int requested_depth) { + for (uint32_t i = 0; i < table->depth_count; i++) { + if ((int)table->depths[i].depth == requested_depth) { + return &table->depths[i]; + } + } + return nullptr; +} + +static void oracle_reset_state(OracleState* env, + const VisibleTargetTable* table, int step_grace) { + env->scramble_depth = env->curriculum_depth; + env->step_count = 0; + env->episode_return = 0.0f; + const VisibleTargetDepth* depth = oracle_depth(table, env->scramble_depth); + EXPECT_TRUE(depth != nullptr); + int choice = oracle_random_bounded(env, (int)depth->stored_count); + const VisibleTargetRecord* record = + &table->records[depth->first_record + (uint32_t)choice]; + env->state = record->start; + env->target = record->target; + env->target_distance = record->depth; + env->max_steps = env->target_distance + step_grace; +} + +static uint16_t oracle_apply_action(uint16_t state, int action) { + uint32_t next = state; + switch (action) { + case 0: next = (state >> 1) | ((state & 1u) << 15); break; + case 1: next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); break; + case 2: next = state ^ 0xfe00u; break; + case 3: next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); break; + case 4: next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); break; + case 5: next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); break; + case 6: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + break; + case 7: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); + break; + } + return (uint16_t)(next & 0xffffu); +} + +static int oracle_next_curriculum_depth(int current_depth, int max_depth) { + static const int curriculum_depths[] = {2, 4, 5, 6, 8, 16}; + for (int depth : curriculum_depths) { + if (depth > current_depth) { + return depth < max_depth ? depth : max_depth; + } + } + return max_depth; +} + +static void oracle_add_log(OracleState* env, int solved, + int max_depth, int perf_weighting) { + int log_depth = env->target_distance; + int at_max_depth = log_depth == max_depth; + float ratio = log_depth / (float)max_depth; + float solve_credit = 0.0f; + if (solved) { + solve_credit = perf_weighting == PERF_WEIGHTING_QUADRATIC + ? ratio * ratio : ratio; + } + env->log.perf += solve_credit; + env->log.score += solve_credit; + env->log.solve_rate += solved; + env->log.max_depth_solve += solved && at_max_depth; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->step_count; + env->log.solve_steps += solved ? env->step_count : 0; + env->log.timeout_rate += !solved; + env->log.solve_efficiency += solved + ? env->step_count / (float)log_depth : 0.0f; + env->log.target_distance += env->target_distance; + env->log.solved_target_distance += solved ? env->target_distance : 0; + env->log.d6_rate += log_depth == 6; + env->log.d6_solve_rate += solved && log_depth == 6; + env->log.d8_rate += log_depth == 8; + env->log.d8_solve_rate += solved && log_depth == 8; + env->log.d16_rate += log_depth == 16; + env->log.d16_solve_rate += solved && log_depth == 16; + env->log.n += 1; +} + +static void oracle_step(OracleState* env, float action, + const VisibleTargetTable* table, int start_depth, int max_depth, + int step_grace, int perf_weighting, float* reward, float* terminal) { + *reward = STEP_REWARD; + *terminal = 0.0f; + int solved = 0; + env->step_count += 1; + int invalid = !std::isfinite(action) || action < 0.0f || action > 7.0f; + if (invalid) { + *reward = -1.0f; + *terminal = 1.0f; + } else { + env->state = oracle_apply_action(env->state, (int)action); + if (env->state == env->target) { + *reward = 1.0f; + *terminal = 1.0f; + solved = 1; + } else if (env->step_count >= env->max_steps) { + *reward = -1.0f; + *terminal = 1.0f; + } + } + env->episode_return += *reward; + if (*terminal != 0.0f) { + oracle_add_log(env, solved, max_depth, perf_weighting); + env->curriculum_depth = solved + ? oracle_next_curriculum_depth(env->scramble_depth, max_depth) + : start_depth; + oracle_reset_state(env, table, step_grace); + } +} + +static uint16_t obs_bits(obs_t value) { + uint16_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +static uint32_t float_bits(float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +static void expect_log_equal(const Log& actual, const Log& expected) { + const float* a = (const float*)&actual; + const float* e = (const float*)&expected; + for (size_t i = 0; i < sizeof(Log) / sizeof(float); i++) { + EXPECT_EQ(float_bits(a[i]), float_bits(e[i])); + } +} + +static void expect_state_equal(const GpuAffineLockState& actual, + const OracleState& expected) { + EXPECT_EQ(actual.rng, expected.rng); + EXPECT_EQ(actual.state, expected.state); + EXPECT_EQ(actual.target, expected.target); + EXPECT_EQ(actual.step_count, expected.step_count); + EXPECT_EQ(actual.max_steps, expected.max_steps); + EXPECT_EQ(actual.scramble_depth, expected.scramble_depth); + EXPECT_EQ(actual.curriculum_depth, expected.curriculum_depth); + EXPECT_EQ(actual.target_distance, expected.target_distance); + EXPECT_EQ(float_bits(actual.episode_return), + float_bits(expected.episode_return)); +} + +static void expect_observation_equal(const obs_t* actual, + const OracleState& expected) { + uint32_t bits = (uint32_t)expected.state | ((uint32_t)expected.target << 16); + for (int bit = 0; bit < 32; bit++) { + float value = (bits & (1u << bit)) ? 1.0f : -1.0f; + EXPECT_EQ(obs_bits(actual[bit]), obs_bits(__float2bfloat16(value))); + } + float timer = expected.step_count / (float)expected.max_steps; + EXPECT_EQ(obs_bits(actual[TIMER_INDEX]), + obs_bits(__float2bfloat16(timer))); +} + +static void fill_kwargs(Dict* kwargs, int seed, int step_grace, + int perf_weighting) { + std::memset(kwargs, 0, sizeof(*kwargs)); + dict_set(kwargs, "seed", seed); + dict_set(kwargs, "start_depth", 2); + dict_set(kwargs, "max_depth", 16); + dict_set(kwargs, "step_grace", step_grace); + dict_set(kwargs, "perf_weighting", perf_weighting); +} + +static void test_deterministic_reset_and_step_parity() { + constexpr int n = 257; + constexpr int seed = 42; + constexpr int step_grace = 2; + constexpr int perf_weighting = PERF_WEIGHTING_QUADRATIC; + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), "cudaMalloc observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), "cudaMalloc actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), "cudaMalloc rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), "cudaMalloc terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, seed, step_grace, perf_weighting); + Env* envs = puf_vec_create(n, &kwargs, observations, actions, rewards, terminals); + EXPECT_TRUE(envs != nullptr); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "initial reset"); + + std::vector oracle(n); + unsigned int running_seed = seed; + for (int i = 0; i < n; i++) { + oracle[i].rng = (uint32_t)rand_r(&running_seed); + oracle[i].curriculum_depth = 2; + oracle_reset_state(&oracle[i], &table, step_grace); + } + + std::vector states(n); + std::vector host_envs(n); + std::vector host_obs((size_t)n * OBS_SIZE); + std::vector host_rewards(n), host_terminals(n), host_actions(n); + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), "copy reset states"); + check_cuda(cudaMemcpy(host_obs.data(), observations, + host_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), "copy reset observations"); + check_cuda(cudaMemcpy(host_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy reset rewards"); + check_cuda(cudaMemcpy(host_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy reset terminals"); + for (int i = 0; i < n; i++) { + expect_state_equal(states[i], oracle[i]); + expect_observation_equal(&host_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_NEAR(host_rewards[i], 0.0f, 0.0f); + EXPECT_NEAR(host_terminals[i], 0.0f, 0.0f); + } + + for (int step = 0; step < 96; step++) { + for (int i = 0; i < n; i++) { + int selector = (step * 17 + i * 13) % 41; + if (selector == 0) host_actions[i] = std::numeric_limits::quiet_NaN(); + else if (selector == 1) host_actions[i] = -0.25f; + else if (selector == 2) host_actions[i] = 8.0f; + else host_actions[i] = (float)((step + 3 * i) & 7) + (selector == 3 ? 0.75f : 0.0f); + } + check_cuda(cudaMemcpy(actions, host_actions.data(), + n * sizeof(float), cudaMemcpyHostToDevice), "copy actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "step"); + + for (int i = 0; i < n; i++) { + oracle_step(&oracle[i], host_actions[i], &table, + 2, 16, step_grace, perf_weighting, + &host_rewards[i], &host_terminals[i]); + } + + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), "copy states"); + check_cuda(cudaMemcpy(host_envs.data(), envs, + n * sizeof(Env), cudaMemcpyDeviceToHost), "copy logs"); + check_cuda(cudaMemcpy(host_obs.data(), observations, + host_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), "copy observations"); + std::vector actual_rewards(n), actual_terminals(n); + check_cuda(cudaMemcpy(actual_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy rewards"); + check_cuda(cudaMemcpy(actual_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy terminals"); + for (int i = 0; i < n; i++) { + expect_state_equal(states[i], oracle[i]); + expect_log_equal(host_envs[i].log, oracle[i].log); + expect_observation_equal(&host_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_NEAR(actual_rewards[i], host_rewards[i], 0.0f); + EXPECT_NEAR(actual_terminals[i], host_terminals[i], 0.0f); + } + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree observations"); + check_cuda(cudaFree(actions), "cudaFree actions"); + check_cuda(cudaFree(rewards), "cudaFree rewards"); + check_cuda(cudaFree(terminals), "cudaFree terminals"); +} + +static void test_reset_rejection_sampling() { + constexpr uint32_t rejection_seed = 24481u; + constexpr uint32_t expected_final_rng = 3424986747u; + static const int depths[] = {2, 16}; + static const int expected_choices[] = {44338, 15778}; + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc rejection observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), + "cudaMalloc rejection action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), + "cudaMalloc rejection reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), + "cudaMalloc rejection terminal"); + + Dict kwargs; + fill_kwargs(&kwargs, 1, 0, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + + for (int case_index = 0; case_index < 2; case_index++) { + int depth = depths[case_index]; + GpuAffineLockState injected = {}; + injected.rng = rejection_seed; + injected.curriculum_depth = depth; + check_cuda(cudaMemcpy(g_gpu.states, &injected, sizeof(injected), + cudaMemcpyHostToDevice), "inject rejection state"); + + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "rejection reset"); + + OracleState expected = {}; + expected.rng = rejection_seed; + expected.curriculum_depth = depth; + oracle_reset_state(&expected, &table, 0); + EXPECT_EQ(expected.rng, expected_final_rng); + const VisibleTargetDepth* table_depth = oracle_depth(&table, depth); + EXPECT_TRUE(table_depth != nullptr); + const VisibleTargetRecord* selected = &table.records[ + table_depth->first_record + (uint32_t)expected_choices[case_index]]; + EXPECT_EQ(expected.state, selected->start); + EXPECT_EQ(expected.target, selected->target); + + GpuAffineLockState actual = {}; + obs_t actual_obs[OBS_SIZE]; + float actual_reward = 123.0f; + float actual_terminal = 123.0f; + check_cuda(cudaMemcpy(&actual, g_gpu.states, sizeof(actual), + cudaMemcpyDeviceToHost), "copy rejection state"); + check_cuda(cudaMemcpy(actual_obs, observations, sizeof(actual_obs), + cudaMemcpyDeviceToHost), "copy rejection observations"); + check_cuda(cudaMemcpy(&actual_reward, rewards, sizeof(actual_reward), + cudaMemcpyDeviceToHost), "copy rejection reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, sizeof(actual_terminal), + cudaMemcpyDeviceToHost), "copy rejection terminal"); + expect_state_equal(actual, expected); + expect_observation_equal(actual_obs, expected); + EXPECT_EQ(float_bits(actual_reward), float_bits(0.0f)); + EXPECT_EQ(float_bits(actual_terminal), float_bits(0.0f)); + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree rejection observations"); + check_cuda(cudaFree(actions), "cudaFree rejection action"); + check_cuda(cudaFree(rewards), "cudaFree rejection reward"); + check_cuda(cudaFree(terminals), "cudaFree rejection terminal"); +} + +static void test_puf_log_exports_cpu_contract() { + Log log = {}; + log.perf = 1.25f; + log.score = 2.5f; + log.solve_rate = 2.0f; + log.max_depth_solve = 1.0f; + log.episode_return = 3.5f; + log.episode_length = 8.0f; + log.solve_steps = 5.0f; + log.timeout_rate = 1.0f; + log.solve_efficiency = 1.75f; + log.target_distance = 20.0f; + log.solved_target_distance = 12.0f; + log.d6_rate = 2.0f; + log.d6_solve_rate = 1.0f; + log.d8_rate = 4.0f; + log.d8_solve_rate = 3.0f; + log.d16_rate = 1.0f; + log.d16_solve_rate = 1.0f; + log.n = 3.0f; + Dict out = {}; + puf_log(&log, &out); + EXPECT_EQ(out.size, 15); + EXPECT_NEAR(dict_get(&out, "perf"), 1.25f, 0.0f); + EXPECT_NEAR(dict_get(&out, "score"), 2.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "solve_rate"), 2.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "max_depth_solve"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "episode_return"), 3.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "episode_length"), 8.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "timeout_rate"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "min_win_moves"), 20.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "solved_min_win_moves"), 6.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "conditional_solve_steps"), 2.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "conditional_solve_efficiency"), 0.875f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d6_solve_rate"), 0.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d8_solve_rate"), 0.75f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d16_solve_rate"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "n"), 3.0f, 0.0f); + dict_clear(&out); + + Log zero_denominators = {}; + zero_denominators.solved_target_distance = 12.0f; + zero_denominators.solve_steps = 5.0f; + zero_denominators.solve_efficiency = 1.75f; + zero_denominators.d6_solve_rate = 1.0f; + zero_denominators.d8_solve_rate = 1.0f; + zero_denominators.d16_solve_rate = 1.0f; + Dict zero_out = {}; + puf_log(&zero_denominators, &zero_out); + EXPECT_EQ(zero_out.size, 15); + EXPECT_NEAR(dict_get(&zero_out, "solved_min_win_moves"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "conditional_solve_steps"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "conditional_solve_efficiency"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d6_solve_rate"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d8_solve_rate"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d16_solve_rate"), 0.0f, 0.0f); + dict_clear(&zero_out); +} + +static void test_exhaustive_action_transforms() { + constexpr int n = 1 << BITS; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc exhaustive observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, 7, 100, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + std::vector states(n); + std::vector host_actions(n), host_rewards(n), host_terminals(n); + + for (int action = 0; action < NUM_ACTIONS; action++) { + for (int value = 0; value < n; value++) { + uint16_t expected = oracle_apply_action((uint16_t)value, action); + states[value] = {}; + states[value].rng = (uint32_t)(value + 1); + states[value].state = (uint16_t)value; + states[value].target = expected ^ 1u; + states[value].max_steps = 100; + states[value].scramble_depth = 16; + states[value].curriculum_depth = 16; + states[value].target_distance = 16; + host_actions[value] = (float)action; + } + check_cuda(cudaMemcpy(g_gpu.states, states.data(), + n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), + "copy exhaustive states"); + check_cuda(cudaMemcpy(actions, host_actions.data(), + n * sizeof(float), cudaMemcpyHostToDevice), + "copy exhaustive actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "exhaustive action step"); + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), + "copy exhaustive results"); + check_cuda(cudaMemcpy(host_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), + "copy exhaustive rewards"); + check_cuda(cudaMemcpy(host_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), + "copy exhaustive terminals"); + for (int value = 0; value < n; value++) { + EXPECT_EQ(states[value].state, + oracle_apply_action((uint16_t)value, action)); + EXPECT_EQ(states[value].step_count, 1); + EXPECT_NEAR(states[value].episode_return, STEP_REWARD, 0.0f); + EXPECT_NEAR(host_rewards[value], STEP_REWARD, 0.0f); + EXPECT_NEAR(host_terminals[value], 0.0f, 0.0f); + } + } + + puf_close(envs); + dict_clear(&kwargs); + check_cuda(cudaFree(observations), "cudaFree exhaustive observations"); + check_cuda(cudaFree(actions), "cudaFree exhaustive actions"); + check_cuda(cudaFree(rewards), "cudaFree exhaustive rewards"); + check_cuda(cudaFree(terminals), "cudaFree exhaustive terminals"); +} + +static void test_action_boundaries() { + const float infinity = std::numeric_limits::infinity(); + const float actions_under_test[] = { + -infinity, + -0.25f, + -std::numeric_limits::denorm_min(), + -0.0f, + 0.0f, + 0.5f, + 0.999f, + 1.0f, + 1.5f, + 6.999f, + 7.0f, + std::nextafter(7.0f, infinity), + 8.0f, + infinity, + std::numeric_limits::quiet_NaN(), + }; + constexpr int n = sizeof(actions_under_test) / sizeof(actions_under_test[0]); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc boundary observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc boundary actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc boundary rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc boundary terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, 11, 100, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + + std::vector injected(n); + std::vector expected(n); + std::vector expected_rewards(n), expected_terminals(n); + for (int i = 0; i < n; i++) { + injected[i] = {}; + injected[i].rng = (uint32_t)(1000 + i); + injected[i].state = 0x1234u; + injected[i].target = 0xbeefu; + injected[i].max_steps = 100; + injected[i].scramble_depth = 2; + injected[i].curriculum_depth = 2; + injected[i].target_distance = 2; + + expected[i].rng = injected[i].rng; + expected[i].state = injected[i].state; + expected[i].target = injected[i].target; + expected[i].max_steps = injected[i].max_steps; + expected[i].scramble_depth = injected[i].scramble_depth; + expected[i].curriculum_depth = injected[i].curriculum_depth; + expected[i].target_distance = injected[i].target_distance; + oracle_step(&expected[i], actions_under_test[i], &table, + 2, 16, 100, PERF_WEIGHTING_LINEAR, + &expected_rewards[i], &expected_terminals[i]); + } + check_cuda(cudaMemcpy(g_gpu.states, injected.data(), + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), + "copy boundary states"); + check_cuda(cudaMemcpy(actions, actions_under_test, + sizeof(actions_under_test), cudaMemcpyHostToDevice), + "copy boundary actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "boundary step"); + + std::vector actual_states(n); + std::vector actual_envs(n); + std::vector actual_obs((size_t)n * OBS_SIZE); + std::vector actual_rewards(n), actual_terminals(n); + check_cuda(cudaMemcpy(actual_states.data(), g_gpu.states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), + "copy boundary results"); + check_cuda(cudaMemcpy(actual_envs.data(), envs, + (size_t)n * sizeof(Env), cudaMemcpyDeviceToHost), + "copy boundary logs"); + check_cuda(cudaMemcpy(actual_obs.data(), observations, + actual_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), + "copy boundary observations"); + check_cuda(cudaMemcpy(actual_rewards.data(), rewards, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost), + "copy boundary rewards"); + check_cuda(cudaMemcpy(actual_terminals.data(), terminals, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost), + "copy boundary terminals"); + + for (int i = 0; i < n; i++) { + bool invalid = !std::isfinite(actions_under_test[i]) + || actions_under_test[i] < 0.0f || actions_under_test[i] > 7.0f; + expect_state_equal(actual_states[i], expected[i]); + expect_log_equal(actual_envs[i].log, expected[i].log); + expect_observation_equal( + &actual_obs[(size_t)i * OBS_SIZE], expected[i]); + EXPECT_EQ(float_bits(actual_rewards[i]), + float_bits(expected_rewards[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(expected_terminals[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(invalid ? 1.0f : 0.0f)); + EXPECT_EQ(float_bits(actual_envs[i].log.n), + float_bits(invalid ? 1.0f : 0.0f)); + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree boundary observations"); + check_cuda(cudaFree(actions), "cudaFree boundary actions"); + check_cuda(cudaFree(rewards), "cudaFree boundary rewards"); + check_cuda(cudaFree(terminals), "cudaFree boundary terminals"); +} + +static const VisibleTargetRecord* find_solution_record( + const VisibleTargetTable* table, const OracleState& state) { + const VisibleTargetDepth* depth = oracle_depth(table, state.scramble_depth); + EXPECT_TRUE(depth != nullptr); + for (uint32_t i = 0; i < depth->stored_count; i++) { + const VisibleTargetRecord* record = + &table->records[depth->first_record + i]; + if (record->start == state.state && record->target == state.target) { + return record; + } + } + return nullptr; +} + +static void test_solution_curriculum_and_logs() { + constexpr int seed = 69; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc solution observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), "cudaMalloc solution action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), "cudaMalloc solution reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), "cudaMalloc solution terminal"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + Dict kwargs; + fill_kwargs(&kwargs, seed, 0, PERF_WEIGHTING_QUADRATIC); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "solution reset"); + + OracleState oracle = {}; + unsigned int running_seed = seed; + oracle.rng = (uint32_t)rand_r(&running_seed); + oracle.curriculum_depth = 2; + oracle_reset_state(&oracle, &table, 0); + static const int expected_depths[] = {2, 4, 5, 6, 8, 16}; + for (int expected_depth : expected_depths) { + EXPECT_EQ(oracle.scramble_depth, expected_depth); + const VisibleTargetRecord* record = find_solution_record(&table, oracle); + EXPECT_TRUE(record != nullptr); + for (int move = 0; move < record->solution_length; move++) { + float action = (float)((record->packed_actions >> (3 * move)) & 7u); + check_cuda(cudaMemcpy(actions, &action, sizeof(float), + cudaMemcpyHostToDevice), "copy solution action"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "solution step"); + + float expected_reward = 0.0f; + float expected_terminal = 0.0f; + oracle_step(&oracle, action, &table, 2, 16, 0, + PERF_WEIGHTING_QUADRATIC, + &expected_reward, &expected_terminal); + GpuAffineLockState actual_state; + Env actual_env; + obs_t actual_obs[OBS_SIZE]; + float actual_reward = 0.0f; + float actual_terminal = 0.0f; + check_cuda(cudaMemcpy(&actual_state, g_gpu.states, + sizeof(actual_state), cudaMemcpyDeviceToHost), + "copy solution state"); + check_cuda(cudaMemcpy(&actual_env, envs, + sizeof(actual_env), cudaMemcpyDeviceToHost), + "copy solution log"); + check_cuda(cudaMemcpy(actual_obs, observations, + sizeof(actual_obs), cudaMemcpyDeviceToHost), + "copy solution observations"); + check_cuda(cudaMemcpy(&actual_reward, rewards, + sizeof(float), cudaMemcpyDeviceToHost), + "copy solution reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, + sizeof(float), cudaMemcpyDeviceToHost), + "copy solution terminal"); + expect_state_equal(actual_state, oracle); + expect_log_equal(actual_env.log, oracle.log); + expect_observation_equal(actual_obs, oracle); + EXPECT_NEAR(actual_reward, expected_reward, 0.0f); + EXPECT_NEAR(actual_terminal, expected_terminal, 0.0f); + EXPECT_NEAR(actual_terminal, + move + 1 == record->solution_length ? 1.0f : 0.0f, 0.0f); + } + } + EXPECT_NEAR(oracle.log.solve_rate, 6.0f, 0.0f); + EXPECT_NEAR(oracle.log.d6_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.d8_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.d16_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.max_depth_solve, 1.0f, 0.0f); + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree solution observations"); + check_cuda(cudaFree(actions), "cudaFree solution action"); + check_cuda(cudaFree(rewards), "cudaFree solution reward"); + check_cuda(cudaFree(terminals), "cudaFree solution terminal"); +} + +static void test_linear_solve_scoring() { + constexpr int seed = 17; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc linear observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), + "cudaMalloc linear action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), + "cudaMalloc linear reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), + "cudaMalloc linear terminal"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + Dict kwargs; + fill_kwargs(&kwargs, seed, 0, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "linear reset"); + + OracleState oracle = {}; + unsigned int running_seed = seed; + oracle.rng = (uint32_t)rand_r(&running_seed); + oracle.curriculum_depth = 2; + oracle_reset_state(&oracle, &table, 0); + const VisibleTargetRecord* record = find_solution_record(&table, oracle); + EXPECT_TRUE(record != nullptr); + for (int move = 0; move < record->solution_length; move++) { + float action = (float)((record->packed_actions >> (3 * move)) & 7u); + check_cuda(cudaMemcpy(actions, &action, sizeof(action), + cudaMemcpyHostToDevice), "copy linear solution action"); + puf_step(envs); + } + check_cuda(cudaDeviceSynchronize(), "linear solution"); + + Env actual_env = {}; + float actual_reward = 0.0f; + float actual_terminal = 0.0f; + check_cuda(cudaMemcpy(&actual_env, envs, sizeof(actual_env), + cudaMemcpyDeviceToHost), "copy linear log"); + check_cuda(cudaMemcpy(&actual_reward, rewards, sizeof(actual_reward), + cudaMemcpyDeviceToHost), "copy linear reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, sizeof(actual_terminal), + cudaMemcpyDeviceToHost), "copy linear terminal"); + EXPECT_NEAR(actual_env.log.perf, 2.0f / 16.0f, 0.0f); + EXPECT_NEAR(actual_env.log.score, 2.0f / 16.0f, 0.0f); + EXPECT_NEAR(actual_env.log.solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(actual_env.log.n, 1.0f, 0.0f); + EXPECT_NEAR(actual_reward, 1.0f, 0.0f); + EXPECT_NEAR(actual_terminal, 1.0f, 0.0f); + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree linear observations"); + check_cuda(cudaFree(actions), "cudaFree linear action"); + check_cuda(cudaFree(rewards), "cudaFree linear reward"); + check_cuda(cudaFree(terminals), "cudaFree linear terminal"); +} + +static void expect_device_canaries(const unsigned char* device_storage, + size_t prefix_bytes, size_t payload_bytes, size_t suffix_bytes, + unsigned char canary, const char* label) { + size_t total_bytes = prefix_bytes + payload_bytes + suffix_bytes; + std::vector host_storage(total_bytes); + check_cuda(cudaMemcpy(host_storage.data(), device_storage, total_bytes, + cudaMemcpyDeviceToHost), "copy canary storage"); + for (size_t i = 0; i < prefix_bytes; i++) { + if (host_storage[i] != canary) { + std::fprintf(stderr, + "%s prefix canary overwritten at byte %zu: 0x%02x != 0x%02x\n", + label, i, host_storage[i], canary); + std::exit(1); + } + } + size_t suffix_start = prefix_bytes + payload_bytes; + for (size_t i = suffix_start; i < total_bytes; i++) { + if (host_storage[i] != canary) { + std::fprintf(stderr, + "%s suffix canary overwritten at byte %zu: 0x%02x != 0x%02x\n", + label, i - suffix_start, host_storage[i], canary); + std::exit(1); + } + } +} + +static void run_io_canary_case(int n) { + constexpr size_t guard_bytes = 64; + constexpr size_t observation_prefix = guard_bytes + sizeof(obs_t); + constexpr unsigned char canary = 0xa5u; + size_t observation_bytes = (size_t)n * OBS_SIZE * sizeof(obs_t); + size_t scalar_bytes = (size_t)n * sizeof(float); + size_t observation_storage_bytes = observation_prefix + + observation_bytes + guard_bytes; + size_t scalar_storage_bytes = guard_bytes + scalar_bytes + guard_bytes; + + unsigned char* observation_storage = nullptr; + unsigned char* reward_storage = nullptr; + unsigned char* terminal_storage = nullptr; + float* actions = nullptr; + check_cuda(cudaMalloc(&observation_storage, observation_storage_bytes), + "cudaMalloc guarded observations"); + check_cuda(cudaMalloc(&reward_storage, scalar_storage_bytes), + "cudaMalloc guarded rewards"); + check_cuda(cudaMalloc(&terminal_storage, scalar_storage_bytes), + "cudaMalloc guarded terminals"); + check_cuda(cudaMalloc(&actions, scalar_bytes), + "cudaMalloc guarded actions"); + + obs_t* observations = reinterpret_cast( + observation_storage + observation_prefix); + float* rewards = reinterpret_cast(reward_storage + guard_bytes); + float* terminals = reinterpret_cast(terminal_storage + guard_bytes); + EXPECT_EQ((uintptr_t)observations & 1u, 0u); + EXPECT_EQ((uintptr_t)observations & 3u, 2u); + EXPECT_EQ((uintptr_t)rewards & 3u, 0u); + EXPECT_EQ((uintptr_t)terminals & 3u, 0u); + + check_cuda(cudaMemset(observation_storage, canary, + observation_storage_bytes), "initialize observation canaries"); + check_cuda(cudaMemset(reward_storage, canary, scalar_storage_bytes), + "initialize reward canaries"); + check_cuda(cudaMemset(terminal_storage, canary, scalar_storage_bytes), + "initialize terminal canaries"); + check_cuda(cudaMemset(actions, 0, scalar_bytes), + "initialize guarded actions"); + + Dict kwargs; + fill_kwargs(&kwargs, 31 + n, 1, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "guarded reset"); + expect_device_canaries(observation_storage, observation_prefix, + observation_bytes, guard_bytes, canary, "reset observations"); + expect_device_canaries(reward_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "reset rewards"); + expect_device_canaries(terminal_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "reset terminals"); + + check_cuda(cudaMemset(observation_storage, canary, + observation_storage_bytes), "reinitialize observation canaries"); + check_cuda(cudaMemset(reward_storage, canary, scalar_storage_bytes), + "reinitialize reward canaries"); + check_cuda(cudaMemset(terminal_storage, canary, scalar_storage_bytes), + "reinitialize terminal canaries"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "guarded step"); + expect_device_canaries(observation_storage, observation_prefix, + observation_bytes, guard_bytes, canary, "step observations"); + expect_device_canaries(reward_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "step rewards"); + expect_device_canaries(terminal_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "step terminals"); + + puf_close(envs); + dict_clear(&kwargs); + check_cuda(cudaFree(observation_storage), + "cudaFree guarded observations"); + check_cuda(cudaFree(reward_storage), "cudaFree guarded rewards"); + check_cuda(cudaFree(terminal_storage), "cudaFree guarded terminals"); + check_cuda(cudaFree(actions), "cudaFree guarded actions"); +} + +static void test_io_canaries_and_observation_alignment() { +#if AFFINE_LOCK_GPU_SHARED_OBS + constexpr int environments_per_block = AFFINE_LOCK_GPU_SHARED_BLOCK; +#else + constexpr int environments_per_block = + AFFINE_LOCK_GPU_BLOCK / AFFINE_LOCK_GPU_LANES; +#endif + run_io_canary_case(environments_per_block); + run_io_canary_case(environments_per_block + 1); +} + +static void test_nondefault_stream_and_cuda_graph() { + constexpr int n = 4099; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc graph observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc graph actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc graph rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc graph terminals"); + check_cuda(cudaMemset(actions, 0, (size_t)n * sizeof(float)), + "clear graph actions"); + + Dict kwargs; + fill_kwargs(&kwargs, 123, 3, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + cudaStream_t stream; + check_cuda(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), + "create graph stream"); + puf_bind_stream(stream); + puf_reset(envs); + check_cuda(cudaStreamSynchronize(stream), "graph reset"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + std::vector oracle(n); + unsigned int running_seed = 123; + for (int i = 0; i < n; i++) { + oracle[i].rng = (uint32_t)rand_r(&running_seed); + oracle[i].curriculum_depth = 2; + oracle_reset_state(&oracle[i], &table, 3); + } + + GpuAffineLockState* actual_states = nullptr; + Env* actual_envs = nullptr; + obs_t* actual_obs = nullptr; + float* actual_rewards = nullptr; + float* actual_terminals = nullptr; + check_cuda(cudaMallocHost((void**)&actual_states, + (size_t)n * sizeof(GpuAffineLockState)), + "cudaMallocHost graph states"); + check_cuda(cudaMallocHost((void**)&actual_envs, + (size_t)n * sizeof(Env)), "cudaMallocHost graph envs"); + check_cuda(cudaMallocHost((void**)&actual_obs, + (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMallocHost graph observations"); + check_cuda(cudaMallocHost((void**)&actual_rewards, + (size_t)n * sizeof(float)), "cudaMallocHost graph rewards"); + check_cuda(cudaMallocHost((void**)&actual_terminals, + (size_t)n * sizeof(float)), "cudaMallocHost graph terminals"); + + cudaGraph_t graph; + cudaGraphExec_t graph_exec; + check_cuda(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal), + "begin graph capture"); + for (int i = 0; i < 3; i++) { + puf_step(envs); + } + check_cuda(cudaStreamEndCapture(stream, &graph), "end graph capture"); + check_cuda(cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0), + "instantiate graph"); + check_cuda(cudaGraphLaunch(graph_exec, stream), "launch graph first"); + check_cuda(cudaGraphLaunch(graph_exec, stream), "launch graph second"); + + check_cuda(cudaMemcpyAsync(actual_states, g_gpu.states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost, stream), + "queue graph states D2H"); + check_cuda(cudaMemcpyAsync(actual_envs, envs, + (size_t)n * sizeof(Env), cudaMemcpyDeviceToHost, stream), + "queue graph logs D2H"); + check_cuda(cudaMemcpyAsync(actual_obs, observations, + (size_t)n * OBS_SIZE * sizeof(obs_t), cudaMemcpyDeviceToHost, stream), + "queue graph observations D2H"); + check_cuda(cudaMemcpyAsync(actual_rewards, rewards, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost, stream), + "queue graph rewards D2H"); + check_cuda(cudaMemcpyAsync(actual_terminals, terminals, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost, stream), + "queue graph terminals D2H"); + check_cuda(cudaStreamSynchronize(stream), "synchronize graph and D2H"); + + std::vector expected_rewards(n), expected_terminals(n); + for (int step = 0; step < 6; step++) { + for (int i = 0; i < n; i++) { + oracle_step(&oracle[i], 0.0f, &table, + 2, 16, 3, PERF_WEIGHTING_LINEAR, + &expected_rewards[i], &expected_terminals[i]); + } + } + for (int i = 0; i < n; i++) { + expect_state_equal(actual_states[i], oracle[i]); + expect_log_equal(actual_envs[i].log, oracle[i].log); + expect_observation_equal( + &actual_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_EQ(float_bits(actual_rewards[i]), + float_bits(expected_rewards[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(expected_terminals[i])); + } + + check_cuda(cudaGraphExecDestroy(graph_exec), "destroy graph exec"); + check_cuda(cudaGraphDestroy(graph), "destroy graph"); + check_cuda(cudaFreeHost(actual_states), "cudaFreeHost graph states"); + check_cuda(cudaFreeHost(actual_envs), "cudaFreeHost graph envs"); + check_cuda(cudaFreeHost(actual_obs), "cudaFreeHost graph observations"); + check_cuda(cudaFreeHost(actual_rewards), "cudaFreeHost graph rewards"); + check_cuda(cudaFreeHost(actual_terminals), "cudaFreeHost graph terminals"); + check_cuda(cudaStreamDestroy(stream), "destroy graph stream"); + puf_bind_stream(nullptr); + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree graph observations"); + check_cuda(cudaFree(actions), "cudaFree graph actions"); + check_cuda(cudaFree(rewards), "cudaFree graph rewards"); + check_cuda(cudaFree(terminals), "cudaFree graph terminals"); +} + +int main() { + test_deterministic_reset_and_step_parity(); + test_reset_rejection_sampling(); + test_exhaustive_action_transforms(); + test_action_boundaries(); + test_solution_curriculum_and_logs(); + test_linear_solve_scoring(); + test_io_canaries_and_observation_alignment(); + test_nondefault_stream_and_cuda_graph(); + test_puf_log_exports_cpu_contract(); + std::puts("affine_lock CUDA tests passed"); + return 0; +}