From 6f7cd65be506149ff9edda2cc46d4d977bd0a42b Mon Sep 17 00:00:00 2001 From: pax-synetic Date: Sun, 5 Jul 2026 18:28:09 -0400 Subject: [PATCH] Add sparse COO weight upload for GPU CTRNN/IZNN eval NEAT nets are sparse, so the dense W [N,M,M] tensor uploaded per generation is mostly zeros. Add an opt-in sparse path (sparse_upload=False by default) that transfers a COO edge list (genome_idx, dst, src, weight) instead and scatters it into a device-side dense W via a RawKernel, isolating the transfer win from the existing matmul integration loop. - _padding.py: packers emit edge_index [E,3]/edge_weight [E] with W=None when sparse_upload=True; add num_genomes so backends no longer read W.shape. - _cupy_backend.py: scatter_edges kernel + _upload_W() (dense copy or zero-alloc + GPU scatter, 64-bit W offset for large N*M*M). - evaluator.py: both evaluators accept and forward sparse_upload. - tests: sparse packing scatters to the exact dense W; GPU sparse trajectories and fitness are bit-identical to dense (CTRNN, IZNN, evaluator). - benchmark: report W H2D bytes + dense/sparse timing; add 32-hidden CTRNN run. Transfer cut tracks density: ~1x at M=6, 3.2x at M=35. Co-Authored-By: Claude Fable 5 --- benchmarks/gpu_benchmark.py | 113 ++++++++++++++++------- neat/gpu/_cupy_backend.py | 80 ++++++++++++++-- neat/gpu/_padding.py | 144 ++++++++++++++++++++--------- neat/gpu/evaluator.py | 20 +++- tests/test_gpu.py | 177 ++++++++++++++++++++++++++++++++++++ 5 files changed, 450 insertions(+), 84 deletions(-) diff --git a/benchmarks/gpu_benchmark.py b/benchmarks/gpu_benchmark.py index 0d73567f..3690de5e 100644 --- a/benchmarks/gpu_benchmark.py +++ b/benchmarks/gpu_benchmark.py @@ -158,8 +158,37 @@ def make_iznn_genome(config, genome_id, num_hidden=0): # Benchmarks # --------------------------------------------------------------------------- +def fmt_bytes(n): + """Human-readable byte count.""" + if n >= 1 << 20: + return f"{n / (1 << 20):.2f} MB" + if n >= 1 << 10: + return f"{n / (1 << 10):.1f} KB" + return f"{n} B" + + +def w_h2d_bytes(packed): + """Bytes crossing the PCIe bus for the weight tensor of a packed pop.""" + if packed['W'] is not None: + return packed['W'].nbytes + return packed['edge_index'].nbytes + packed['edge_weight'].nbytes + + +_HEADER = (f"{'Pop':>8} {'MaxN':>6} {'CPU (s)':>9} {'GPU-d (s)':>10} " + f"{'GPU-s (s)':>10} {'W H2D dense':>12} {'W H2D sparse':>13} " + f"{'xfer cut':>9}") + + +def _print_row(pop_size, max_nodes, cpu_time, gpu_dense, gpu_sparse, + dense_bytes, sparse_bytes): + reduction = dense_bytes / sparse_bytes if sparse_bytes else float('inf') + print(f"{pop_size:>8d} {max_nodes:>6d} {cpu_time:>9.3f} {gpu_dense:>10.3f} " + f"{gpu_sparse:>10.3f} {fmt_bytes(dense_bytes):>12} " + f"{fmt_bytes(sparse_bytes):>13} {reduction:>8.1f}x") + + def benchmark_ctrnn(pop_sizes, num_hidden=3): - """Benchmark CTRNN CPU vs GPU at various population sizes.""" + """Benchmark CTRNN CPU vs GPU (dense and sparse W upload).""" from neat.gpu._padding import pack_ctrnn_population from neat.gpu._cupy_backend import evaluate_ctrnn_batch @@ -170,12 +199,12 @@ def benchmark_ctrnn(pop_sizes, num_hidden=3): input_vals = [0.5, -0.3] inputs_np = np.tile(np.array(input_vals, dtype=np.float32), (num_steps, 1)) - print(f"\n{'='*70}") + print(f"\n{'='*84}") print(f"CTRNN Benchmark: dt={dt}, t_max={t_max}, num_steps={num_steps}, " f"hidden_nodes={num_hidden}") - print(f"{'='*70}") - print(f"{'Pop Size':>10} {'Max Nodes':>10} {'CPU (s)':>10} {'GPU (s)':>10} {'Speedup':>10}") - print(f"{'-'*10:>10} {'-'*10:>10} {'-'*10:>10} {'-'*10:>10} {'-'*10:>10}") + print(f"{'='*84}") + print(_HEADER) + print('-' * 84) for pop_size in pop_sizes: np.random.seed(42) @@ -190,27 +219,36 @@ def benchmark_ctrnn(pop_sizes, num_hidden=3): net.advance(input_vals, dt, dt) cpu_time = time.perf_counter() - t0 - # GPU timing (include packing + transfer + compute). - # Warmup. - packed = pack_ctrnn_population(genomes, config) - _ = evaluate_ctrnn_batch(packed, inputs_np, dt) + # Warmup both paths (compiles activation + scatter kernels). + _ = evaluate_ctrnn_batch(pack_ctrnn_population(genomes, config), + inputs_np, dt) + _ = evaluate_ctrnn_batch( + pack_ctrnn_population(genomes, config, sparse_upload=True), + inputs_np, dt) cp.cuda.Stream.null.synchronize() + # GPU timing, dense upload (packing + transfer + compute). t0 = time.perf_counter() - packed = pack_ctrnn_population(genomes, config) - trajectory = evaluate_ctrnn_batch(packed, inputs_np, dt) + packed_dense = pack_ctrnn_population(genomes, config) + evaluate_ctrnn_batch(packed_dense, inputs_np, dt) cp.cuda.Stream.null.synchronize() - gpu_time = time.perf_counter() - t0 + gpu_dense_time = time.perf_counter() - t0 - max_nodes = packed['max_nodes'] - speedup = cpu_time / gpu_time if gpu_time > 0 else float('inf') + # GPU timing, sparse upload. + t0 = time.perf_counter() + packed_sparse = pack_ctrnn_population(genomes, config, + sparse_upload=True) + evaluate_ctrnn_batch(packed_sparse, inputs_np, dt) + cp.cuda.Stream.null.synchronize() + gpu_sparse_time = time.perf_counter() - t0 - print(f"{pop_size:>10d} {max_nodes:>10d} {cpu_time:>10.3f} {gpu_time:>10.3f} " - f"{speedup:>9.1f}x") + _print_row(pop_size, packed_dense['max_nodes'], cpu_time, + gpu_dense_time, gpu_sparse_time, + w_h2d_bytes(packed_dense), w_h2d_bytes(packed_sparse)) def benchmark_iznn(pop_sizes, num_hidden=3): - """Benchmark Izhikevich CPU vs GPU at various population sizes.""" + """Benchmark Izhikevich CPU vs GPU (dense and sparse W upload).""" from neat.gpu._padding import pack_iznn_population from neat.gpu._cupy_backend import evaluate_iznn_batch @@ -221,12 +259,12 @@ def benchmark_iznn(pop_sizes, num_hidden=3): input_vals = [1.0, 0.5] inputs_np = np.tile(np.array(input_vals, dtype=np.float32), (num_steps, 1)) - print(f"\n{'='*70}") + print(f"\n{'='*84}") print(f"Izhikevich Benchmark: dt={dt} ms, t_max={t_max} ms, " f"num_steps={num_steps}, hidden_nodes={num_hidden}") - print(f"{'='*70}") - print(f"{'Pop Size':>10} {'Max Nodes':>10} {'CPU (s)':>10} {'GPU (s)':>10} {'Speedup':>10}") - print(f"{'-'*10:>10} {'-'*10:>10} {'-'*10:>10} {'-'*10:>10} {'-'*10:>10}") + print(f"{'='*84}") + print(_HEADER) + print('-' * 84) for pop_size in pop_sizes: np.random.seed(42) @@ -242,25 +280,38 @@ def benchmark_iznn(pop_sizes, num_hidden=3): net.advance(dt) cpu_time = time.perf_counter() - t0 - # GPU timing. - packed = pack_iznn_population(genomes, config) - _ = evaluate_iznn_batch(packed, inputs_np, dt, num_steps) + # Warmup both paths. + _ = evaluate_iznn_batch(pack_iznn_population(genomes, config), + inputs_np, dt, num_steps) + _ = evaluate_iznn_batch( + pack_iznn_population(genomes, config, sparse_upload=True), + inputs_np, dt, num_steps) cp.cuda.Stream.null.synchronize() + # GPU timing, dense upload. t0 = time.perf_counter() - packed = pack_iznn_population(genomes, config) - trajectory = evaluate_iznn_batch(packed, inputs_np, dt, num_steps) + packed_dense = pack_iznn_population(genomes, config) + evaluate_iznn_batch(packed_dense, inputs_np, dt, num_steps) cp.cuda.Stream.null.synchronize() - gpu_time = time.perf_counter() - t0 + gpu_dense_time = time.perf_counter() - t0 - max_nodes = packed['max_nodes'] - speedup = cpu_time / gpu_time if gpu_time > 0 else float('inf') + # GPU timing, sparse upload. + t0 = time.perf_counter() + packed_sparse = pack_iznn_population(genomes, config, + sparse_upload=True) + evaluate_iznn_batch(packed_sparse, inputs_np, dt, num_steps) + cp.cuda.Stream.null.synchronize() + gpu_sparse_time = time.perf_counter() - t0 - print(f"{pop_size:>10d} {max_nodes:>10d} {cpu_time:>10.3f} {gpu_time:>10.3f} " - f"{speedup:>9.1f}x") + _print_row(pop_size, packed_dense['max_nodes'], cpu_time, + gpu_dense_time, gpu_sparse_time, + w_h2d_bytes(packed_dense), w_h2d_bytes(packed_sparse)) if __name__ == '__main__': pop_sizes = [100, 500, 1000] benchmark_ctrnn(pop_sizes) + # Larger networks: the dense W tensor scales with max_nodes^2, the sparse + # edge list only with connection count — this run shows the transfer win. + benchmark_ctrnn(pop_sizes, num_hidden=32) benchmark_iznn(pop_sizes) diff --git a/neat/gpu/_cupy_backend.py b/neat/gpu/_cupy_backend.py index 7a65f332..aadd11da 100644 --- a/neat/gpu/_cupy_backend.py +++ b/neat/gpu/_cupy_backend.py @@ -92,6 +92,68 @@ def _get_activation_kernel(): return cp.RawKernel(_ACTIVATION_KERNEL_CODE, 'apply_activation') +# --------------------------------------------------------------------------- +# COO-edge scatter kernel (sparse W upload) +# --------------------------------------------------------------------------- +# Scatters a COO edge list (genome_idx, dst, src, weight) into a +# zero-initialized dense W [N, M, M] on the device. Used when the population +# was packed with sparse_upload=True: only ~16 bytes per enabled connection +# cross the PCIe bus instead of the full N*M*M*4-byte dense tensor. + +_SCATTER_KERNEL_CODE = ''' +extern "C" __global__ +void scatter_edges( + const int* __restrict__ edge_index, // [E, 3] rows: (genome, dst, src) + const float* __restrict__ edge_weight, // [E] + float* __restrict__ W, // [N*M*M], zero-initialized + int num_edges, + int M +) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i >= num_edges) return; + + // 64-bit offset: N*M*M can exceed 2^31 for large populations. + long long g = edge_index[3 * i]; + long long dst = edge_index[3 * i + 1]; + long long src = edge_index[3 * i + 2]; + W[(g * M + dst) * M + src] = edge_weight[i]; +} +''' + + +def _get_scatter_kernel(): + """Compile and cache the edge scatter kernel.""" + cp = _import_cupy() + return cp.RawKernel(_SCATTER_KERNEL_CODE, 'scatter_edges') + + +def _upload_W(cp, packed): + """ + Materialize the dense weight tensor W [N, M, M] on the device. + + Dense-packed populations (packed['W'] is an ndarray) transfer it as-is. + Sparse-packed populations (packed['W'] is None) transfer only the COO + edge arrays and scatter them into a device-side zero tensor. + """ + if packed['W'] is not None: + return cp.asarray(packed['W']) + + N = packed['num_genomes'] + M = packed['max_nodes'] + W = cp.zeros((N, M, M), dtype=cp.float32) + + num_edges = packed['edge_index'].shape[0] + if num_edges: + edge_index = cp.asarray(packed['edge_index']) + edge_weight = cp.asarray(packed['edge_weight']) + kernel = _get_scatter_kernel() + block_size = 256 + grid_size = (num_edges + block_size - 1) // block_size + kernel((grid_size,), (block_size,), + (edge_index, edge_weight, W, num_edges, M)) + return W + + def evaluate_ctrnn_batch(packed, inputs_cpu, dt): """ Run batched CTRNN simulation on GPU using exponential Euler integration. @@ -99,7 +161,8 @@ def evaluate_ctrnn_batch(packed, inputs_cpu, dt): Parameters ---------- packed : dict - Output of pack_ctrnn_population(). NumPy arrays on CPU. + Output of pack_ctrnn_population(), dense or sparse_upload=True. + NumPy arrays on CPU. inputs_cpu : ndarray [num_steps, num_inputs] or [num_steps, N, num_inputs] Precomputed input trajectory. If 2-D, broadcast across population. dt : float @@ -113,14 +176,14 @@ def evaluate_ctrnn_batch(packed, inputs_cpu, dt): cp = _import_cupy() np = _import_numpy() - N = packed['W'].shape[0] + N = packed['num_genomes'] M = packed['max_nodes'] num_inputs = packed['num_inputs'] num_outputs = packed['num_outputs'] num_steps = inputs_cpu.shape[0] - # Transfer parameters to GPU. - W = cp.asarray(packed['W']) # [N, M, M] + # Transfer parameters to GPU (W: dense copy or sparse scatter). + W = _upload_W(cp, packed) # [N, M, M] bias = cp.asarray(packed['bias']) # [N, M] response = cp.asarray(packed['response']) # [N, M] tau = cp.asarray(packed['tau']) # [N, M] @@ -196,7 +259,8 @@ def evaluate_iznn_batch(packed, inputs_cpu, dt, num_steps): Parameters ---------- packed : dict - Output of pack_iznn_population(). NumPy arrays on CPU. + Output of pack_iznn_population(), dense or sparse_upload=True. + NumPy arrays on CPU. inputs_cpu : ndarray [num_steps, num_inputs] or [num_steps, N, num_inputs] Precomputed input trajectory. dt : float @@ -212,13 +276,13 @@ def evaluate_iznn_batch(packed, inputs_cpu, dt, num_steps): cp = _import_cupy() np = _import_numpy() - N = packed['W'].shape[0] + N = packed['num_genomes'] M = packed['max_nodes'] num_inputs = packed['num_inputs'] num_outputs = packed['num_outputs'] - # Transfer to GPU. - W = cp.asarray(packed['W']) # [N, M, M] + # Transfer to GPU (W: dense copy or sparse scatter). + W = _upload_W(cp, packed) # [N, M, M] bias = cp.asarray(packed['bias']) # [N, M] a = cp.asarray(packed['a']) # [N, M] b = cp.asarray(packed['b']) # [N, M] diff --git a/neat/gpu/_padding.py b/neat/gpu/_padding.py index 599cb0af..68b3da3a 100644 --- a/neat/gpu/_padding.py +++ b/neat/gpu/_padding.py @@ -35,6 +35,69 @@ ]) +def _collect_genome_edges(g_idx, genome, key_map, required, edges): + """ + Append (genome_idx, dst_idx, src_idx, weight) tuples for every enabled + connection of ``genome`` whose endpoints are packable, to ``edges``. + + The filtering rules match the original dense fill exactly: the connection + must be enabled, both endpoints must be in the key map, and the + destination must be a required (non-input) node. + """ + for cg in genome.connections.values(): + if not cg.enabled: + continue + + src_key, dst_key = cg.key + # Only include connections where both endpoints are in the key map. + if src_key not in key_map or dst_key not in key_map: + continue + # dst must be a required node (non-input). + if dst_key not in required: + continue + + edges.append((g_idx, key_map[dst_key], key_map[src_key], cg.weight)) + + +def _edges_to_arrays(np, edges): + """ + Convert an edge tuple list into COO arrays. + + Returns (edge_index [E, 3] int32, edge_weight [E] float32) where each + edge_index row is (genome_idx, dst_idx, src_idx). + """ + if edges: + edge_index = np.array([e[:3] for e in edges], dtype=np.int32) + edge_weight = np.array([e[3] for e in edges], dtype=np.float32) + else: + edge_index = np.zeros((0, 3), dtype=np.int32) + edge_weight = np.zeros((0,), dtype=np.float32) + return edge_index, edge_weight + + +def _finalize_weights(np, edges, N, M, sparse_upload): + """ + Turn the accumulated edge list into the weight entries of a packed dict: + either COO arrays for GPU-side scatter (sparse_upload=True) or a dense + CPU-built W [N, M, M] (sparse_upload=False). + """ + edge_index, edge_weight = _edges_to_arrays(np, edges) + + if sparse_upload: + return { + 'W': None, + 'edge_index': edge_index, + 'edge_weight': edge_weight, + } + + W = np.zeros((N, M, M), dtype=np.float32) + if len(edge_weight): + # Connection keys are unique per (src, dst), so no duplicate + # (genome, dst, src) triples exist and scatter order is irrelevant. + W[edge_index[:, 0], edge_index[:, 1], edge_index[:, 2]] = edge_weight + return {'W': W} + + def _build_node_key_map(genome, config, required_nodes): """ Build a mapping from neat-python node keys to dense indices. @@ -70,7 +133,7 @@ def _build_node_key_map(genome, config, required_nodes): return key_map, num_nodes -def pack_ctrnn_population(genomes, config): +def pack_ctrnn_population(genomes, config, sparse_upload=False): """ Convert a list of (genome_id, genome) pairs into padded NumPy arrays for GPU CTRNN evaluation. @@ -81,16 +144,28 @@ def pack_ctrnn_population(genomes, config): The population to pack. config : neat.Config The NEAT configuration object. + sparse_upload : bool + If True, connection weights are returned as a COO edge list + (``edge_index``/``edge_weight``) instead of a dense ``W``; the GPU + backend scatters them into a device-side dense W. This cuts the + host-to-device transfer from N*M*M*4 bytes to ~16 bytes per enabled + connection. If False (default), a dense ``W`` is built on the CPU. Returns ------- dict with keys: - W : ndarray [N, M, M] float32 — weight matrices + W : ndarray [N, M, M] float32 — weight matrices, or None when + sparse_upload=True + edge_index : ndarray [E, 3] int32 — (genome_idx, dst, src) rows; + only present when sparse_upload=True + edge_weight : ndarray [E] float32 — only present when + sparse_upload=True bias : ndarray [N, M] float32 response : ndarray [N, M] float32 tau : ndarray [N, M] float32 activation_id : ndarray [N, M] int32 node_mask : ndarray [N, M] bool + num_genomes : int num_inputs : int num_outputs : int max_nodes : int @@ -118,7 +193,6 @@ def pack_ctrnn_population(genomes, config): M = max_nodes # Allocate arrays. - W = np.zeros((N, M, M), dtype=np.float32) bias = np.zeros((N, M), dtype=np.float32) response = np.ones((N, M), dtype=np.float32) # default 1.0 tau = np.ones((N, M), dtype=np.float32) # default 1.0 (won't matter for masked-out nodes) @@ -130,6 +204,7 @@ def pack_ctrnn_population(genomes, config): node_mask[:, num_inputs:num_inputs + num_outputs] = True node_key_maps = [] + edges = [] # Second pass: fill arrays. for g_idx, (genome_id, genome, required, key_map, num_nodes) in enumerate(per_genome_info): @@ -167,38 +242,26 @@ def pack_ctrnn_population(genomes, config): f"'{agg_name}' is not supported on GPU. Only 'sum' aggregation " f"is supported (required for batched matrix-vector multiply).") - # Fill weight matrix from enabled connections. - for cg in genome.connections.values(): - if not cg.enabled: - continue - - src_key, dst_key = cg.key - # Only include connections where both endpoints are in the key map. - if src_key not in key_map or dst_key not in key_map: - continue - # dst must be a required node (non-input). - if dst_key not in required: - continue - - src_idx = key_map[src_key] - dst_idx = key_map[dst_key] - W[g_idx, dst_idx, src_idx] = cg.weight - - return { - 'W': W, + # Collect enabled connections as COO edges. + _collect_genome_edges(g_idx, genome, key_map, required, edges) + + packed = { 'bias': bias, 'response': response, 'tau': tau, 'activation_id': activation_id, 'node_mask': node_mask, + 'num_genomes': N, 'num_inputs': num_inputs, 'num_outputs': num_outputs, 'max_nodes': M, 'node_key_maps': node_key_maps, } + packed.update(_finalize_weights(np, edges, N, M, sparse_upload)) + return packed -def pack_iznn_population(genomes, config): +def pack_iznn_population(genomes, config, sparse_upload=False): """ Convert a list of (genome_id, genome) pairs into padded NumPy arrays for GPU Izhikevich spiking network evaluation. @@ -209,17 +272,26 @@ def pack_iznn_population(genomes, config): The population to pack. config : neat.Config The NEAT configuration object. + sparse_upload : bool + If True, return connection weights as a COO edge list instead of a + dense ``W`` (see ``pack_ctrnn_population``). Returns ------- dict with keys: - W : ndarray [N, M, M] float32 — weight matrices + W : ndarray [N, M, M] float32 — weight matrices, or None when + sparse_upload=True + edge_index : ndarray [E, 3] int32 — (genome_idx, dst, src) rows; + only present when sparse_upload=True + edge_weight : ndarray [E] float32 — only present when + sparse_upload=True bias : ndarray [N, M] float32 a : ndarray [N, M] float32 b : ndarray [N, M] float32 c : ndarray [N, M] float32 d : ndarray [N, M] float32 node_mask : ndarray [N, M] bool + num_genomes : int num_inputs : int num_outputs : int max_nodes : int @@ -247,7 +319,6 @@ def pack_iznn_population(genomes, config): M = max_nodes # Allocate arrays. - W = np.zeros((N, M, M), dtype=np.float32) bias_arr = np.zeros((N, M), dtype=np.float32) a_arr = np.zeros((N, M), dtype=np.float32) b_arr = np.zeros((N, M), dtype=np.float32) @@ -259,6 +330,7 @@ def pack_iznn_population(genomes, config): node_mask[:, num_inputs:num_inputs + num_outputs] = True node_key_maps = [] + edges = [] # Second pass: fill arrays. for g_idx, (genome_id, genome, required, key_map, num_nodes) in enumerate(per_genome_info): @@ -280,31 +352,21 @@ def pack_iznn_population(genomes, config): c_arr[g_idx, dense_idx] = node.c d_arr[g_idx, dense_idx] = node.d - # Fill weight matrix. - for cg in genome.connections.values(): - if not cg.enabled: - continue - - src_key, dst_key = cg.key - if src_key not in key_map or dst_key not in key_map: - continue - if dst_key not in required: - continue - - src_idx = key_map[src_key] - dst_idx = key_map[dst_key] - W[g_idx, dst_idx, src_idx] = cg.weight + # Collect enabled connections as COO edges. + _collect_genome_edges(g_idx, genome, key_map, required, edges) - return { - 'W': W, + packed = { 'bias': bias_arr, 'a': a_arr, 'b': b_arr, 'c': c_arr, 'd': d_arr, 'node_mask': node_mask, + 'num_genomes': N, 'num_inputs': num_inputs, 'num_outputs': num_outputs, 'max_nodes': M, 'node_key_maps': node_key_maps, } + packed.update(_finalize_weights(np, edges, N, M, sparse_upload)) + return packed diff --git a/neat/gpu/evaluator.py b/neat/gpu/evaluator.py index e88121bc..39bf5825 100644 --- a/neat/gpu/evaluator.py +++ b/neat/gpu/evaluator.py @@ -42,13 +42,19 @@ class GPUCTRNNEvaluator: ``fitness_fn(output_trajectory) -> float`` where output_trajectory is an ndarray of shape ``[num_steps, num_outputs]``. Called once per genome on CPU after GPU simulation. + sparse_upload : bool + If True, upload connection weights as a COO edge list and scatter + into the dense weight tensor on the GPU, instead of transferring the + full (mostly zero) dense tensor. Reduces host-to-device transfer + volume; results are identical. """ - def __init__(self, dt, t_max, input_fn, fitness_fn): + def __init__(self, dt, t_max, input_fn, fitness_fn, sparse_upload=False): self.dt = dt self.t_max = t_max self.input_fn = input_fn self.fitness_fn = fitness_fn + self.sparse_upload = sparse_upload def evaluate(self, genomes, config): """ @@ -76,7 +82,8 @@ def evaluate(self, genomes, config): self.input_fn(step * self.dt, self.dt), dtype=np.float32) # Pack genomes into padded arrays. - packed = pack_ctrnn_population(genomes, config) + packed = pack_ctrnn_population(genomes, config, + sparse_upload=self.sparse_upload) # Run GPU simulation. trajectory = evaluate_ctrnn_batch(packed, inputs, self.dt) @@ -104,13 +111,17 @@ class GPUIZNNEvaluator: ``fitness_fn(output_trajectory) -> float`` where output_trajectory is an ndarray of shape ``[num_steps, num_outputs]`` containing spike indicators (0.0 or 1.0). + sparse_upload : bool + If True, upload connection weights as a COO edge list and scatter + into the dense weight tensor on the GPU (see GPUCTRNNEvaluator). """ - def __init__(self, dt, t_max, input_fn, fitness_fn): + def __init__(self, dt, t_max, input_fn, fitness_fn, sparse_upload=False): self.dt = dt self.t_max = t_max self.input_fn = input_fn self.fitness_fn = fitness_fn + self.sparse_upload = sparse_upload def evaluate(self, genomes, config): """Evaluate all genomes on GPU. Same interface as NEAT fitness function.""" @@ -129,7 +140,8 @@ def evaluate(self, genomes, config): inputs[step] = np.asarray( self.input_fn(step * self.dt, self.dt), dtype=np.float32) - packed = pack_iznn_population(genomes, config) + packed = pack_iznn_population(genomes, config, + sparse_upload=self.sparse_upload) trajectory = evaluate_iznn_batch(packed, inputs, self.dt, num_steps) for i, (genome_id, genome) in enumerate(genomes): diff --git a/tests/test_gpu.py b/tests/test_gpu.py index 3d84dd07..5f43eb00 100644 --- a/tests/test_gpu.py +++ b/tests/test_gpu.py @@ -367,6 +367,92 @@ def test_izhikevich_parameters(self): print(f" Izhikevich params packed correctly") +class TestSparsePacking: + """Test COO edge-list packing (sparse_upload=True).""" + + @staticmethod + def _scatter_cpu(packed): + """Densify a sparse-packed population on CPU (reference scatter).""" + N = packed['num_genomes'] + M = packed['max_nodes'] + W = np.zeros((N, M, M), dtype=np.float32) + idx = packed['edge_index'] + if idx.shape[0]: + W[idx[:, 0], idx[:, 1], idx[:, 2]] = packed['edge_weight'] + return W + + def test_ctrnn_sparse_matches_dense(self): + """Sparse edges scattered on CPU must reproduce the dense W exactly.""" + from neat.gpu._padding import pack_ctrnn_population + + config = _make_ctrnn_config() + g1 = _make_simple_ctrnn_genome(config, genome_id=1, w_in1=3.0, w_in2=-1.5) + g2 = _make_simple_ctrnn_genome(config, genome_id=2, add_hidden=True) + g3 = _make_simple_ctrnn_genome(config, genome_id=3, w_in1=5.0) + g3.connections[(-2, 0)].enabled = False # disabled edge must be excluded + genomes = [(1, g1), (2, g2), (3, g3)] + + dense = pack_ctrnn_population(genomes, config) + sparse = pack_ctrnn_population(genomes, config, sparse_upload=True) + + assert sparse['W'] is None + assert sparse['edge_index'].dtype == np.int32 + assert sparse['edge_index'].shape[1] == 3 + assert sparse['edge_weight'].dtype == np.float32 + assert sparse['edge_index'].shape[0] == sparse['edge_weight'].shape[0] + # 2 (g1) + 3 (g2) + 1 (g3, one disabled) enabled connections. + assert sparse['edge_index'].shape[0] == 6 + + assert np.array_equal(self._scatter_cpu(sparse), dense['W']) + + # Non-weight arrays must be unaffected by the flag. + for key in ('bias', 'response', 'tau', 'activation_id', 'node_mask'): + assert np.array_equal(sparse[key], dense[key]) + assert sparse['num_genomes'] == dense['num_genomes'] == 3 + print(f" {sparse['edge_index'].shape[0]} edges reproduce dense W " + f"({dense['W'].nbytes} B dense vs " + f"{sparse['edge_index'].nbytes + sparse['edge_weight'].nbytes} B sparse)") + + def test_iznn_sparse_matches_dense(self): + from neat.gpu._padding import pack_iznn_population + + config = _make_iznn_config() + genomes = [ + (1, _make_simple_iznn_genome(config, genome_id=1, w_in1=15.0)), + (2, _make_simple_iznn_genome(config, genome_id=2, w_in1=5.0, w_in2=-2.0)), + ] + + dense = pack_iznn_population(genomes, config) + sparse = pack_iznn_population(genomes, config, sparse_upload=True) + + assert sparse['W'] is None + assert np.array_equal(self._scatter_cpu(sparse), dense['W']) + for key in ('bias', 'a', 'b', 'c', 'd', 'node_mask'): + assert np.array_equal(sparse[key], dense[key]) + print(f" {sparse['edge_index'].shape[0]} IZNN edges reproduce dense W") + + def test_sparse_no_connections(self): + """A connectionless genome packs to an empty edge list.""" + from neat.gpu._padding import pack_ctrnn_population + + config = _make_ctrnn_config() + genome = neat.DefaultGenome(1) + node0 = DefaultNodeGene(0) + node0.bias = 0.0 + node0.response = 1.0 + node0.activation = 'tanh' + node0.aggregation = 'sum' + node0.time_constant = 1.0 + genome.nodes[0] = node0 + genomes = [(1, genome)] + + packed = pack_ctrnn_population(genomes, config, sparse_upload=True) + assert packed['edge_index'].shape == (0, 3) + assert packed['edge_weight'].shape == (0,) + assert np.all(self._scatter_cpu(packed) == 0) + print(" No connections: empty edge arrays, zero W after scatter") + + # --------------------------------------------------------------------------- # GPU Numerical Equivalence Tests (require CuPy) # --------------------------------------------------------------------------- @@ -467,6 +553,38 @@ def test_multiple_genomes_same_as_individual(self): assert max_diff < 1e-6, ( f"Genome {gid}: batched result differs from individual by {max_diff}") + def test_sparse_upload_matches_dense(self): + """ + Sparse (COO + GPU scatter) and dense uploads must produce bit-identical + trajectories — the simulation kernels see the same device-side W. + """ + from neat.gpu._padding import pack_ctrnn_population + from neat.gpu._cupy_backend import evaluate_ctrnn_batch + + config = _make_ctrnn_config() + genomes = [ + (1, _make_simple_ctrnn_genome(config, genome_id=1, bias=0.3, w_in1=1.0)), + (2, _make_simple_ctrnn_genome(config, genome_id=2, bias=-0.5, w_in1=2.0)), + (3, _make_simple_ctrnn_genome(config, genome_id=3, bias=0.0, w_in1=-1.0, + add_hidden=True)), + ] + + dt = 0.005 + num_steps = 100 + inputs_np = np.tile(np.array([0.5, -0.3], dtype=np.float32), + (num_steps, 1)) + + traj_dense = evaluate_ctrnn_batch( + pack_ctrnn_population(genomes, config), inputs_np, dt) + traj_sparse = evaluate_ctrnn_batch( + pack_ctrnn_population(genomes, config, sparse_upload=True), + inputs_np, dt) + + assert np.array_equal(traj_dense, traj_sparse), ( + f"Sparse upload diverged from dense: max diff " + f"{np.max(np.abs(traj_dense - traj_sparse)):.2e}") + print(" Sparse and dense uploads produced identical trajectories.") + def test_response_parameter_effect(self): """ Two genomes with different response values but same weights should @@ -607,6 +725,31 @@ def test_multiple_genomes_independent(self): print(f" Genome {gid}: max batch vs individual diff = {max_diff:.2e}") assert max_diff < 1e-6 + def test_sparse_upload_matches_dense(self): + """Sparse and dense uploads must produce identical spike trains.""" + from neat.gpu._padding import pack_iznn_population + from neat.gpu._cupy_backend import evaluate_iznn_batch + + config = _make_iznn_config() + genomes = [ + (1, _make_simple_iznn_genome(config, genome_id=1, w_in1=15.0)), + (2, _make_simple_iznn_genome(config, genome_id=2, w_in1=5.0, bias=2.0)), + ] + + dt = 0.05 + num_steps = 200 + inputs_np = np.tile(np.array([1.0, 0.5], dtype=np.float32), + (num_steps, 1)) + + traj_dense = evaluate_iznn_batch( + pack_iznn_population(genomes, config), inputs_np, dt, num_steps) + traj_sparse = evaluate_iznn_batch( + pack_iznn_population(genomes, config, sparse_upload=True), + inputs_np, dt, num_steps) + + assert np.array_equal(traj_dense, traj_sparse) + print(" Sparse and dense IZNN uploads produced identical spike trains.") + def test_no_input_no_spikes(self): """With zero external input and zero bias, neurons should not spike.""" from neat.gpu._padding import pack_iznn_population @@ -663,6 +806,40 @@ def fitness_fn(trajectory): assert genome.fitness >= 0.0 print(f" Genome {gid}: fitness = {genome.fitness:.6f}") + def test_ctrnn_evaluator_sparse_matches_dense_fitness(self): + """sparse_upload=True must yield the same fitness values as dense.""" + from neat.gpu.evaluator import GPUCTRNNEvaluator + + config = _make_ctrnn_config() + + def make_genomes(): + return [ + (1, _make_simple_ctrnn_genome(config, genome_id=1, bias=0.3)), + (2, _make_simple_ctrnn_genome(config, genome_id=2, bias=-0.5, + add_hidden=True)), + ] + + def input_fn(t, dt): + return [math.sin(2 * math.pi * t), math.cos(2 * math.pi * t)] + + def fitness_fn(trajectory): + return float(np.mean(np.abs(trajectory))) + + genomes_dense = make_genomes() + GPUCTRNNEvaluator(dt=0.01, t_max=0.5, input_fn=input_fn, + fitness_fn=fitness_fn).evaluate(genomes_dense, config) + + genomes_sparse = make_genomes() + GPUCTRNNEvaluator(dt=0.01, t_max=0.5, input_fn=input_fn, + fitness_fn=fitness_fn, + sparse_upload=True).evaluate(genomes_sparse, config) + + for (gid, gd), (_, gs) in zip(genomes_dense, genomes_sparse): + assert gd.fitness == gs.fitness, ( + f"Genome {gid}: dense fitness {gd.fitness} != " + f"sparse fitness {gs.fitness}") + print(f" Genome {gid}: fitness {gd.fitness:.6f} (dense == sparse)") + def test_iznn_evaluator_assigns_fitness(self): """GPUIZNNEvaluator.evaluate should set genome.fitness for all genomes.""" from neat.gpu.evaluator import GPUIZNNEvaluator