Skip to main content

ABA_SIMPLE_SHADER

Constant ABA_SIMPLE_SHADER 

Source
pub const ABA_SIMPLE_SHADER: &str = r#"
struct SimParams {
    nworld: u32,
    nv: u32,
    dt: f32,
    _padding: u32,
}

struct BodyParams {
    mass: f32,
    inertia: f32,
    com_y: f32,
    damping: f32,
    gravity_y: f32,
    _padding0: f32,
    _padding1: f32,
    _padding2: f32,
}

@group(0) @binding(0) var<uniform> params: SimParams;
@group(0) @binding(1) var<uniform> body: BodyParams;
@group(0) @binding(2) var<storage, read> q: array<f32>;
@group(0) @binding(3) var<storage, read> v: array<f32>;
@group(0) @binding(4) var<storage, read> ctrl: array<f32>;
@group(0) @binding(5) var<storage, read_write> qdd: array<f32>;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let world_idx = gid.x;

    if (world_idx >= params.nworld) {
        return;
    }

    // For single revolute joint: qdd = (tau - damping*v - m*g*L*sin(q)) / I
    let idx = world_idx;
    let q_val = q[idx];
    let v_val = v[idx];
    let tau = ctrl[idx];

    // Gravity torque: m * g * L * sin(q)
    // Note: gravity_y is magnitude (positive), com_y is typically negative
    // ABA uses base acceleration trick, so we need positive sign here
    let gravity_torque = body.mass * body.gravity_y * body.com_y * sin(q_val);

    // Total torque: applied torque + gravity torque - damping torque
    let total_torque = tau + gravity_torque - body.damping * v_val;

    // Total inertia for pendulum: I = m*L²/3 (parallel axis theorem)
    // For simplicity, we pass the computed inertia from CPU
    let total_inertia = body.inertia;

    // qdd = torque / inertia
    qdd[idx] = total_torque / total_inertia;
}
"#;
Expand description

WGSL shader for simplified ABA (single revolute joint systems).

This is a simplified version that handles pendulum-like systems with single revolute joints. For multi-body systems, we’d need a more complex shader with tree traversal.