JIT Compilation

Overview

One of the most powerful features of cuda-mppi is its support for Just-In-Time (JIT) Compilation of dynamics and cost functions. This allows users to define their system models in C++ strings (e.g., generated from symbolic libraries like SymPy) and compile them into optimized CUDA kernels at runtime.

This is powered by NVRTC (NVIDIA Runtime Compilation).

How it Works

The JITMPPIController takes two strings as input:

Dynamics Code: Source code defining a UserDynamics struct.
Cost Code: Source code defining a UserCost struct.

The controller wraps these definitions into a complete CUDA kernel source file, compiles it to PTX (Parallel Thread Execution) assembly using NVRTC, and loads it into the CUDA context.

Writing JIT-Compatible Code

Your code must define specific C++ structs with __device__ member functions.

1. Dynamics Interface

You must define a struct named UserDynamics.

struct UserDynamics {
    // Constant parameters can be hardcoded or passed via constructor/members 
    // (if supported by your generation pipeline).
    
    __device__ void step(const float* x, const float* u, float* x_next, float dt) const {
        // x: [x, y, theta, v]
        // u: [v_cmd, w_cmd]
        
        float v = x[3];
        float theta = x[2];
        
        x_next[0] = x[0] + v * cosf(theta) * dt;
        x_next[1] = x[1] + v * sinf(theta) * dt;
        x_next[2] = x[2] + u[1] * dt;
        x_next[3] = x[3] + (u[0] - v) * dt; // Simple P-control on velocity
    }
};

2. Cost Interface

You must define a struct named UserCost.

struct UserCost {
    __device__ float compute(const float* x, const float* u, int t) const {
        float cost = 0.0f;
        // Penalize control effort
        cost += 0.1f * (u[0]*u[0] + u[1]*u[1]);
        // Penalize deviation from origin
        cost += 1.0f * (x[0]*x[0] + x[1]*x[1]);
        return cost;
    }
    
    __device__ float terminal_cost(const float* x) const {
        return 10.0f * (x[0]*x[0] + x[1]*x[1]);
    }
};

Python Usage

The JIT controller is exposed to Python via cuda_mppi.JITMPPIController.

import cuda_mppi

dynamics_code = """
struct UserDynamics {
    __device__ void step(const float* x, const float* u, float* x_next, float dt) const {
        x_next[0] = x[0] + u[0] * dt;
    }
};
"""

cost_code = """
struct UserCost {
    __device__ float compute(const float* x, const float* u, int t) const { return x[0]*x[0]; }
    __device__ float terminal_cost(const float* x) const { return 0; }
};
"""

config = cuda_mppi.MPPIConfig(...)
include_paths = ["/path/to/cuda_mppi/include"]

controller = cuda_mppi.JITMPPIController(
    config, 
    dynamics_code, 
    cost_code, 
    include_paths
)