Enable parallel CUDA simulations for grids#1
Open
faezs wants to merge 6 commits into
Open
Conversation
- Add accelerate and CUDA dependencies to package.yaml and stack.yaml - Create CUDA.Types module with GPU-compatible data types - Create CUDA.Kernels module with parallel computation kernels - Create CUDA.Simulation module for high-level CUDA orchestration - Create Grid/CUDA.hs for grid-level CUDA integration - Update Main.hs with --cuda and --cuda-steps CLI options This enables parallel execution of household simulations (battery, PV, consumption) across all nodes in any grid generated by the library.
Architecture refactor to support batched simulation across many environments: - CUDA.Types: Add VecEnvConfig, VecEnvState, StepResult for vectorized envs - CUDA.Kernels: Add vec* kernels operating on [num_envs, num_nodes] arrays - vecGridStepKernel: Steps all envs × all nodes in single kernel launch - vecInitStatesKernel: Initialize states for all environments - vecResetKernel: Auto-reset done environments - vecComputeRewardsKernel: Per-env and per-node reward computation - vecExtractObsKernel: Extract observations from states - vecCheckDoneKernel: Check done/truncated conditions - CUDA.Simulation: Add VecEnv interface - makeVecEnv: Create N parallel environments from single grid - makeVecEnvFromGrids: Create from multiple different grids - vecStep: Step all envs with batched actions [num_envs, num_nodes] - vecReset: Reset all environments - Main.hs: Add --vecenv and --vecenv-bench commands for testing This enables high-throughput RL training by parallelizing across environments (not just nodes), similar to PufferLib/EnvPool/Isaac Gym. Example: 4096 envs × 100 nodes = 409,600 parallel ops per step
Change voltage update from incorrect absolute formula (vNom * SoC/100) to correct proportional formula matching CPU physics (v * SoC_next/SoC_current). Reference: Sandia module temperature model docs at https://pvpmc.sandia.gov/modeling-steps/2-dc-module-iv/module-temperature/sandia-module-temperature-model/
Implements scatter/gather based message passing using Accelerate's permute and gather primitives: - gatherNodeFeatures: read neighbor values - scatterAddMessages: aggregate messages at nodes - computePowerFlowMessages: compute edge power flows - messagePassingRound: single-env message passing - vecMessagePassingKernel: batched message passing [num_envs, num_nodes] - vecApplyPowerExchange: update HH states with grid power flow This enables proper simulation of power transmission across grid topology.
Previous implementation used the same flow value for both source and destination, which didn't model transmission losses correctly. Now properly implements: - Source loses: P_sent (full amount) - Destination gains: η * P_sent (reduced by efficiency) - Lost to transmission: (1-η) * P_sent This ensures total system energy decreases appropriately when power flows through lossy transmission lines.
Previous implementation was O(N × E) - each node iterated over all edges. Now O(E + N) using proper scatter/gather pattern: 1. GATHER: O(E) parallel reads per env 2. COMPUTE: O(E) parallel edge flow calculations 3. SCATTER: O(E) atomic adds via permute(+) Flattens [num_envs, num_edges] → [num_envs * num_edges] for scatter, uses (env * numNodes + node) indexing, then reshapes back. Total work: O(num_envs × num_edges) fully parallel GPU operations.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This enables parallel execution of household simulations (battery, PV, consumption) across all nodes in any grid generated by the library.