Codegen Architecture#

Every GRiD algorithm is emitted in three layers (_host / _kernel / _device). Knowing the layering helps when you want to compose generated functions, call kernels from your own CUDA host code, or read the emitter source in grid_codegen/algorithms/.

Note

New here (human or agent)? Read Algorithm Design Principles & Best Practices first — it is the shared mental model (smart inners / thin wrappers, the inner owns its memory placement, the spill ladder, validation discipline, and the anti-patterns to avoid). This page covers the mechanics of the three layers; that page covers the ethos behind them.

The three emission layers#

For each algorithm X (e.g. inverse_dynamics, forward_dynamics, crba, fdsva_so), the codegen emits:

X_device — a __device__ entry point. Scratch arguments are algorithm-specific: simple wrappers such as RNEA declare their helper storage internally, while composing orchestrators accept caller-supplied buffers and workspace. Placement belongs to the owning device routine; do not infer one algorithm’s signature from another’s. See Algorithm Design Principles & Best Practices and the generated header for the selected model.
X_kernel — a __global__ entry point that handles batch scheduling (blockIdx.x loops over timesteps) and global ↔ shared memory transfer. It allocates __shared__ smem for inputs/outputs and the s_temp pool, loads inputs, calls X_device, and writes outputs back. Per-tier dispatch ( RESOURCE_TIER template ) picks the spill flags; the kernel never decides placement itself.
X (no suffix) — a host function that wraps X_kernel and handles H↔D copies for inputs and outputs.

Inner helpers (X_inner, sub-step helpers like fdsva_so_contract) still exist where useful, but they are internal to ``X_device`` — not part of the external surface. Sub-algorithm composition routes through other algorithms’ _inner helpers when they are placement-free building blocks (e.g. fdsva_so_device calls minv_inner and forward_dynamics_inner to reuse one XImats load across all of them).

Scratch ownership#

Composing algorithms reuse placement-free _inner helpers. Public device entry points perform the required setup; their buffer and resource-tier contracts vary. The shared emitter helper gen_device_wrapper implements the simple device-wrapper pattern. Inspect the generated signature rather than relying on historical _full_inner names.

Nested composition#

The split exists for nested composition. Higher-level algorithms call other algorithms’ _inner (the placement-free math) directly to reuse one expensive XImats load.

For example, second-order forward dynamics (_fdsva_so.py) needs both forward dynamics and direct inverse-mass-matrix outputs internally. fdsva_so_device loads XImats once at the top, then calls the placement-free _inner variants:

The final -Minv contraction of fdsva_so_contract (iL,Ljk->ijk over the four n³ tensors) is emitted register-tiled along i: each thread produces R outputs that share every strided arena load, where R is the largest divisor of n that is at most 8 and leaves at least two tiles (35 → 7, 18 → 6, 7 → 1). The per-cell dot is accumulated in exactly the untiled order, so the outputs are bit-identical to the one-output-per-thread loop; measured on g1-floating the kernel is 7–11% faster (2026-09-24). GRID_FDSVA_SO_MINV_TILE=1 at generation time forces the untiled loop (an A/B knob; it is part of every header cache key).

The world-frame second-order inner (idsva_so_world_frame_inner, the floating-base and mimic path) runs its triple ancestor walk with ONE pair of block barriers per body column: every ancestor-or-self velocity column (j, t) of body i is built into a per-pair scratch slab (72 × WF_MAX_PAIRS floats, a codegen constant) in one parallel loop, and one flattened (pair, k, r) loop then performs every contraction. The write set was enumerated per topology to prove the cells written by different pairs are disjoint, except the symmetric dM_dq pair on a multi-column body (the floating root), whose sequential last writer (vel_j > vel_k) is made the sole writer — so the outputs are bit-identical to the former per-pair loop while the kernel’s barrier count falls by an order of magnitude (go2 264 → 36; g1 712 → 70 per launch). Measured 2026-09-24: idsva_so −20% and fdsva_so −18% of compute on g1-floating, −13…−52% on go2.

fdsva_so_device
  ├── [s_temp repoint based on SCRATCH_IN_SMEM]
  ├── load_update_XImats()       # once
  ├── minv_inner()               # reuses s_XImats
  ├── forward_dynamics_inner()   # reuses s_XImats
  ├── fd_gradient_inline()       # may surgically spill to d_fd_grad_spill
  ├── idsva_so_{world,body}_inner()
  └── fdsva_so_inner()           # the rank-3 contraction step

If the _inner building blocks were collapsed into their owning _device, the compositional algorithms would pay the XImats load three times instead of once. The separation is a performance contract, not a stylistic preference.

Schematic layers (RNEA / inverse_dynamics)#

The following sketches explain the layers; they are not callable API declarations. Exact signatures depend on generated options, including external-force and acceleration inputs. Use the generated grid.cuh or the runnable walkthrough in examples/cuda/ for integration.

// pure math; assumes s_XImats already loaded; placement-free
template <typename T>
__device__ void inverse_dynamics_inner(
    T *s_vaf, const T *s_q, const T *s_qd,
    /* topology + scratch */ T *s_temp,
    const T gravity);

// simple device wrapper: declares helper scratch internally
template <typename T>
__device__ void inverse_dynamics_device(
    T *s_c, const T *s_q, const T *s_qd,
    const robotModel<T> *d_robotModel, T *d_f_ext, const T gravity);

// global entry, batched over timesteps; per-tier RESOURCE_TIER dispatch
template <typename T, int RESOURCE_TIER = GRID_DEFAULT_RESOURCE_TIER>
__global__ void inverse_dynamics_kernel(
    T *d_c, const T *d_q_qd, const int stride_q_qd,
    const robotModel<T> *d_robotModel, const T gravity,
    const int NUM_TIMESTEPS);

// CPU launcher with H↔D copies
template <typename T, bool USE_QDD_FLAG=false, bool USE_COMPRESSED_MEM=false>
__host__ void inverse_dynamics(
    gridData<T> *hd_data, const robotModel<T> *d_robotModel,
    const T gravity, const int num_timesteps,
    dim3 block_dimms, dim3 thread_dimms, cudaStream_t *streams);

Orchestrator signature (fdsva_so / inverse_dynamics_gradient / forward_dynamics_gradient / integrator_gradient)#

The orchestrators take the placement flags + a d_workspace pointer in addition to the standard inputs:

template <typename T,
          bool SCRATCH_IN_SMEM = true,        // s_temp pool location
          bool FD_GRAD_USE_SPILL = false,     // selective spill bits
          bool CONTRACT_IN_SMEM = true>       // (fdsva_so only)
__device__ __forceinline__
void fdsva_so_device(
    T *s_df2, T *s_idsva_so, T *s_Minv, T *s_df_du, T *s_qdd,
    const T *s_q, const T *s_qd, const T *s_u,
    /* XImats helpers */
    T *s_temp,              // smem pool; ignored when !SCRATCH_IN_SMEM
    T *d_workspace,         // global pool; ignored when SCRATCH_IN_SMEM
    T *d_fd_grad_spill,     // band-spill region; nullptr unless FD_GRAD_USE_SPILL
    T *s_fdsva_temp,        // contraction scratch; per CONTRACT_IN_SMEM
    const robotModel<T> *d_robotModel, const T gravity);

The kernel emitter passes per-tier true/false literals for the template flags and threads the right d_workspace offsets in. Inline-CUDA users size their smem from *_DEVICE_INLINE_SMEM_BYTES<T, TIER>() and d_workspace from *_DEVICE_INLINE_WORKSPACE_BYTES<T, TIER>().

Thread-count assumptions#

GRiD emits a MAX_PERF_LEVEL_THREADS constant per generated header, computed from the robot’s DVA parallelism (rounded up to a warp, capped at 512) — for iiwa14 it is 352, for go2_fixed it is 288, and so on. The host wrappers default to launching with dim3(MAX_PERF_LEVEL_THREADS, 1, 1).

After the v2.0 cuBLASDx removal, MAX_PERF_LEVEL_THREADS is a hint, not an enforced floor. Every emitted X_inner does block-cooperative compute on one timestep — threads within a block split work via block-stride loops (the gen_add_parallel_loop helper emits for (int i = threadIdx.x + threadIdx.y*blockDim.x; i < max_val; i += blockDim.x*blockDim.y)). Any block size that fits per-block covers the work correctly. Batching across timesteps is handled by the outer X_kernel via a grid-stride loop over blockIdx — each block processes one or more timesteps.

CUDA-inline users can launch GRiD kernels with any block size that suits their outer kernel. See cuBLASDx Removal & Any-Thread-Count Library Functions for the rationale.

Design sweet spot: tens-to-hundreds of parallel computations#

The one-timestep-per-block layout — block-cooperative compute inside, grid-stride over batch outside — is tuned for batch sizes in the tens to hundreds. Many batch robotics workloads (MPC shooting nodes, trajectory optimization horizons, behavioral cloning rollouts, real-time control with N robots) sit squarely in that range, which is the design target.

For very small batches (N=1-8) the per-block fixed overhead dominates, and a design that packed multiple timesteps per block could be faster; for very large batches (N=10k+) a design that splits one timestep across multiple blocks could expose more parallelism. Neither is what GRiD optimizes for. If your workload sits at one of those extremes, a codegen layered on a different parallelism map (or an entirely different library) will likely beat GRiD; for the tens-to-hundreds range, GRiD’s layout is the right tool.

The MAX_PERF_LEVEL_THREADS constant is what the codegen picks as the best block-cooperative thread count for this robot (DOF-aware, warp-rounded). External callers are free to override (see grid_rbd.RobotHandle.set_threads_per_block() or grid_rbd_set_threads_per_block in the C ABI), but smaller block sizes will be slower at the same batch size (work-per-block stays constant; fewer threads cover it).

Each X_kernel is emitted with __launch_bounds__(tier_max_threads<RESOURCE_TIER>()) — tier-templated and load-bearing (it is what grid_rbd_kernel_max_threads introspects and what the baked-THREADS clamp validates against). Only the integrator family still uses MAX_PERF_LEVEL_THREADS. (An older revision here said the attribute was “about to drop in phase B1” — that plan was superseded; “B1/B2” now name the tier-matrix/kernel-attr phases in the live plan.)

Per-algo metadata: the descriptor table#

Each algorithm also carries a small amount of irregular per-algo metadata: which autotune launch-config key(s) it uses, the cudaFuncSetAttribute opt-in gate, the dynamic-shared-memory bytes-macro stem, whether it has an mjx (MUJOCO_OUTPUT) twin, and so on. This lives as one AlgoDescriptor row per algorithm in grid_codegen/algo_registry.py (the ALGO_DESCRIPTORS tuple) — the single source of truth from which the generator derives:

the GridAlgo enum and the launch-config symbol map (build_launch_config_algo_to_symbol),
the KERNEL_ATTR_MANIFEST and the mjx (floating-twin) manifest heads.

Previously these were several hand-maintained module-level dicts that had to be kept in lockstep by hand; the descriptor table removed that duplication. Adding an algorithm is now (metadata-wise) one row. test/test_algo_descriptor_parity.py locks the table to the generated output so a mismatch fails CPU-only in CI. The per-algo arena/spill t_count math is now folded into the table too (the ArenaRegion / SpillRung / ArenaCtx machinery in algo_registry.py, composed by compose_arena_full / compose_arena_rungs): every select_shared_tier_3way site is driven from the composer, so arena sizes are no longer hand-written in GRiDCodeGenerator.py. The arena’s correctness is guarded independently by test/test_shared_arena_covers_carve.py, which checks each kernel’s launch-sizing macro against the regions the kernel actually carves. To change an algorithm’s arena, edit its closure in algo_registry.py — do NOT hand-edit t_count expressions in the generator.

The binding surface: generated from `abi_specs.py`#

The descriptor table has a sibling: grid_codegen/abi_specs.py (the ABI_SPECS rows) transcribes every C-ABI body of the Python binding — signature, packing, qdd/f_ext routing, launch template shape, output buffer/size, mjx-twin variance — and grid_codegen/wrapper_body_gen.py EMITS six checked-in generated regions of bindings/grid_rbd/wrapper_template.cu from those rows:

the extern "C" algorithm bodies (30+ functions),
the grid_rbd_kernel_max_threads introspection branch table,
all 30 grid_rbd_<algo>_mujoco twin bodies,
the JAX FFI handlers,
the torch op bodies,
the torch op table.

grid_codegen/wrapper_plant_gen.py emits the remaining two generated regions — the JAX plant tail and the torch plant tail — for eight BEGIN/END GENERATED regions in all. The regions live between those markers in the checked-in file — never hand-edit inside them. To change a generated body, edit the spec row (or the emitter) and regenerate:

.venv/bin/python -m grid_codegen.wrapper_body_gen          # rewrite
.venv/bin/python -m grid_codegen.wrapper_body_gen --check  # CI drift gate

test/test_wrapper_generated_block.py runs --check in CI, and test/test_abi_spec_crosscheck.py validates every spec field against the template text — so a hand-edit inside a marker region, or a stale block after a table edit, fails a plain pytest -q. Only tool_fext, contact_fext (wrapper_template.cu ~:902-967, outside the generated C-ABI region), fk_batched, the four plant cost twins, and the plant/FFI/torch/pybind scaffolding outside the marker regions remain hand-written (bespoke by design).

The same table also drives the numpy Runner’s pybind surface: grid_codegen/core_body_gen.py emits the generated region of bindings/src/_core.cpp — the 61 spec-backed value/gradient/second-order methods plus their has_*_mujoco accessors — from each row’s inputs/py_out_dims/py_rc3_msg/py_twin_guard fields (-m grid_codegen.core_body_gen to regenerate, --check gated by test/test_core_generated_block.py). Adding an algorithm’s spec row therefore produces its wrapper C-ABI body, its mjx twin, AND its pybind method in one regeneration. The pybind .def list (with its user-facing docstrings), the plant family, and the metadata/overlay surface stay hand-written.