casadi-on-gpu

KO | EN

gitlite — search

#javascript #nodejs #swift #python #ios #chrome #macos #aws #vue #webpack #android #opengl

★ 61 Open GitHub ↗

turn CasADi function into GPU ready kernels so large batches of function evaluations can run in parallel at high speed [𝐏𝐲𝐭𝐨𝐫𝐜𝐡 & 𝐂𝐮𝐩𝐲 𝐜𝐨𝐦𝐩𝐚𝐭𝐢𝐛𝐥𝐞]

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# **casadi on gpu**

---

This library lets you run CasADi functions on NVIDIA GPUs from Python in large batches at high speed.

You provide CasADi `.casadi` function files, generate CUDA kernel files (`.cu/.cuh`), and then call those kernels directly from `PyTorch` or `CuPy` using GPU pointers. The package provides a single Python module (`casadi_on_gpu`) that exposes kernel discovery and launch APIs, so large batched evaluations can be executed on GPU with very little glue code.

The repository includes two example assets:

1) `fk_eval.casadi`: forward kinematics for a 4-DoF manipulator.
2) `dynamics_eval.casadi`: stochastic forward dynamics model with sample parameters in `examples/assets/posterior.bin`.

## Requirements

You need:

- An NVIDIA GPU and CUDA toolkit supported by your driver.
- Python 3.10 or newer.
- CasADi built from the CUDA codegen branch in [casadi/casadi#4291](https://github.com/casadi/casadi/pull/4291).
- A `.casadi` file containing one `ca.Function` that you want to evaluate on GPU.

Build the patched CasADi first. These commands are for the upstream `casadi` repository, not this repository:

```bash
git clone -b cuda_codegen https://github.com/edxmorgan/casadi.git casadi
cd casadi
cmake -S . -B build -DWITH_PYTHON=ON -DWITH_PYTHON3=ON
cmake --build build -j8
sudo cmake --install build
```

Clone this repository after CasADi imports correctly from Python:

```bash
git clone https://github.com/edxmorgan/casadi-on-gpu.git
cd casadi-on-gpu
python3 -c "import casadi as ca; print(ca.__version__)"
```

## Custom Functions

Create or load a CasADi function in Python, then save it as a `.casadi` file:

```python
import casadi as ca

x = ca.MX.sym("x", 4)
p = ca.MX.sym("p", 2)
y = ca.vertcat(x[0] + p[0], x[1] * p[1])

f = ca.Function("my_eval", [x, p], [y], ["x", "p"], ["y"])
f.save("my_eval.casadi")
```

The function name (`my_eval` above) is the name you pass to `casadi_on_gpu.launch(...)` after building.

## Generate kernels

Generate CUDA source, a manifest, and a Python-callable kernel registry from one or more `.casadi` files:

```bash
./tools/generate_manifest_and_registry.py \
  --casadi path/to/my_eval.casadi \
  --batch-inputs 0
```

This emits:

- `src/generated/*.cu` and `src/generated/*.cuh`
- `src/generated/kernels_manifest.json`
- `src/python/casadi_on_gpu_kernel_registry.cu`

`--batch-inputs` is a comma-separated list of CasADi input indices that vary per GPU thread. In the example above, input `0` (`x`) is batched and input `1` (`p`) is shared by every thread.

For a batch size `N`, each batched input pointer must contain `N * nnz(input)` `float32` values on the GPU. Each non-batched input pointer must contain only `nnz(input)` values. Every output is batched by the runtime and must contain `N * nnz(output)` values.

Generate several kernels in one module by repeating `--casadi` and assigning per-file batch inputs:

```bash
./tools/generate_manifest_and_registry.py \
  --casadi examples/assets/fk_eval.casadi \
  --casadi examples/assets/dynamics_eval.casadi \
  --batch-inputs-for examples/assets/fk_eval.casadi=0 \
  --batch-inputs-for examples/assets/dynamics_eval.casadi=2
```

Use `--entry` when you need explicit generated file or symbol names:

```bash
./tools/generate_manifest_and_registry.py \
  --entry path/to/my_eval.casadi:my_eval_unit:0:my_eval_kernel:device_my_eval
```

Inspect the generated manifest if you are unsure about input order, output order, shapes, or `nnz` sizes:

```bash
python3 -m json.tool src/generated/kernels_manifest.json
```

## Build casadi-on-gpu

Choose the CUDA toolkit and architecture for the GPU that will run the kernels. First confirm that `nvcc` is new enough for your target GPU:

```bash
which nvcc
nvcc --version
nvidia-smi --query-gpu=name,compute_cap --format=csv,noheader
```

For example, RTX 5090 uses compute capability `12.0`, so the CMake architecture value is `120`. CUDA 12.0 does not support `sm_120`; use a newer NVIDIA toolkit such as CUDA 13.x for that GPU.

Build with CMake:

```bash
cmake -S . -B build \
  -DBUILD_PYTHON=ON \
  -DCMAKE_INSTALL_PREFIX=$HOME/.local \
  -DCMAKE_CUDA_ARCHITECTURES=120
cmake --build build -j
cmake --install build
```

If CMake finds an older toolkit, pass the compiler explicitly:

```bash
cmake -S . -B build \
  -DBUILD_PYTHON=ON \
  -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
  -DCMAKE_CUDA_ARCHITECTURES=120
```

For older GPUs, replace `120` with the matching architecture, such as `86` for Ampere or `89` for Ada. If you omit `CMAKE_CUDA_ARCHITECTURES`, CMake lets the CUDA compiler choose its default target.

Build with pip:

```bash
CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=120" python3 -m pip install .
```

If pip or CMake finds an older toolkit, include the compiler in `CMAKE_ARGS`:

```bash
CMAKE_ARGS="-DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc -DCMAKE_CUDA_ARCHITECTURES=120" python3 -m pip install .
```

If CMake already cached an old CUDA compiler or architecture, delete the build directory before rebuilding.

## Launch from Python

The runtime accepts raw GPU pointers from PyTorch or CuPy. Pointer order must match the CasADi function input and output order.

### PyTorch

```python
import torch
import casadi_on_gpu as cog

print(cog.list_kernels())

N = 80000
x = torch.zeros((N, 4), device="cuda", dtype=torch.float32)  # batched input 0
p = torch.zeros((2,), device="cuda", dtype=torch.float32)     # shared input 1
y = torch.empty((N, 2), device="cuda", dtype=torch.float32)   # batched output 0

stream = torch.cuda.current_stream().cuda_stream
cog.launch(
    "my_eval",
    [x.data_ptr(), p.data_ptr()],
    [y.data_ptr()],
    N,
    stream_ptr=stream,
    sync=True,
)
```

### CuPy

```python
import cupy as cp
import casadi_on_gpu as cog

N = 80000
x = cp.zeros((N, 4), dtype=cp.float32)
p = cp.zeros((2,), dtype=cp.float32)
y = cp.empty((N, 2), dtype=cp.float32)

stream = cp.cuda.get_current_stream().ptr
cog.launch(
    "my_eval",
    [x.data.ptr, p.data.ptr],
    [y.data.ptr],
    N,
    stream_ptr=stream,
    sync=False,
)
```

Notes:

- Inputs and outputs must be CUDA device memory, not host NumPy arrays.
- The generated code defaults to `casadi_real=float`, so use `float32` tensors unless you generated with `--casadi-real double`.
- Use `sync=True` if you are not coordinating with your own CUDA streams.
- The first launch can be slower while CUDA loads and JITs the module; benchmark after warmup launches.
- Large CasADi work arrays use a cached device workspace per CUDA device and stream, so repeated launches do not allocate and free temporary GPU memory each time.

## Performance & Benchmarks

This project targets high-throughput batched evaluation of CasADi-generated CUDA kernels. Benchmarks below reflect realistic large-batch workloads and were measured using CUDA event timing with warmup and repeated launches.

---

## Benchmark Environment

All measurements were collected on:

**Hardware**

- NVIDIA GeForce RTX 5090 Laptop GPU  
- 24 GB VRAM

**Software**

- NVIDIA driver: 580.105.08  
- CUDA runtime: 13.0

**Methodology**

- Batched evaluation with N = 80,000 samples  
- Warmup runs performed before timing  
- CUDA events used for accurate GPU timing  
- Hundreds of repeated launches averaged  
- Asynchronous stream execution

Performance will vary with GPU architecture, clock behavior, thermals, and batch size.

---

## Forward kinematics kernel  
4-DOF manipulator model

Batch size: **N = 80,000**

| Interface | Time / batch | Throughput |
|----------|--------------|------------|
| C++      | 0.00820 ms   | 9.7568 × 10⁹ eval/s |
| PyTorch  | 0.00830 ms   | 9.6766 × 10⁹ eval/s |
| CuPy     | 0.00830 ms   | 9.6922 × 10⁹ eval/s |

This kernel is extremely lightweight, so runtime is dominated by kernel launch overhead. Throughput increases further with larger batches.

---

## Forward dynamics kernel  
12-state, 33-parameter stochastic model

Batch size: **N = 80,000**

| Interface | Time / batch | Throughput |
|----------|--------------|------------|
| C++      | 6.6391 ms    | 1.2050 × 10⁷ eval/s |
| PyTorch  | 6.6731 ms    | 1.1988 × 10⁷ eval/s |
| CuPy     | 6.6523 ms    | 1.2026 × 10⁷ eval/s |

This kernel is compute-heavy and represents realistic stochastic forward dynamics workloads.

---

## Observations

- GPU acceleration is most effective for **large batched evaluations**
- PyTorch and CuPy wrappers introduce **negligible overhead**
- Performance is close to raw CUDA launch speed
- Lightweight kernels approach launch-bound limits
- Heavy kernels scale with compute intensity
- Throughput improves with larger batch sizes

---

## Practical implications

These results demonstrate that CasADi-generated CUDA kernels can support:

- Massive batched simulation  
- Stochastic parameter sampling  
- Real-time trajectory rollout  
- Monte Carlo dynamics evaluation  
- Differentiable robotics pipelines

Performance scales with GPU architecture and problem size, making this approach well suited for large-scale robotics and control workloads.

casadi-on-gpu

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README