Skip to content

shaders: panic-free Mat3/Mat2 column access (ColumnIndex::col_at) #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sebcrozet merged 2 commits into from
Jul 16, 2026
Merged

shaders: panic-free Mat3/Mat2 column access (ColumnIndex::col_at) #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a24bb9c
shaders: panic-free Mat3 column access for the cuda-oxide device build
haixuanTao Jul 8, 2026
b02e035
review: ColBranchless::col_b -> ColumnIndex::col_at; direct axis acce…
haixuanTao Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 24 additions & 24 deletions src_rbd_shaders/dynamics/joint_constraint_builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use crate::Rotation;
#[cfg(feature = "dim2")]
use crate::rotation_angle;
use crate::utils::{Slice, SliceMut};
use crate::{AngVector, MAX_FLT, Pose, Vector, gdot, rotation_to_matrix};
use crate::{AngVector, ColumnIndex, MAX_FLT, Pose, Vector, gdot, rotation_to_matrix};
use khal_std::index::MaybeIndexUnchecked;

#[cfg(feature = "dim2")]
Expand Down Expand Up @@ -122,7 +122,7 @@ pub fn new_helper(
// Then snap the locked ones.
for i in 0..DIM {
if (locked_lin_axes & (1u32 << i)) != 0 {
let axis = basis.col(i);
let axis = basis.col_at(i);
new_center1 -= axis * lin_err.dot(axis);
}
}
Expand All @@ -141,8 +141,8 @@ pub fn new_helper(
basis,
// In 2D, cmat is a Vec2 representing [-r.y, r.x], and we need the dot product
// with each column of basis to get the angular Jacobian components.
cmat1_basis: [cmat1.dot(basis.col(0)), cmat1.dot(basis.col(1))],
cmat2_basis: [cmat2.dot(basis.col(0)), cmat2.dot(basis.col(1))],
cmat1_basis: [cmat1.dot(basis.x_axis), cmat1.dot(basis.y_axis)],
cmat2_basis: [cmat2.dot(basis.x_axis), cmat2.dot(basis.y_axis)],
lin_err,
ang_err,
}
Expand Down Expand Up @@ -170,7 +170,7 @@ pub fn new_helper(
// Then snap the locked ones.
for i in 0..DIM {
if (locked_lin_axes & (1u32 << i)) != 0 {
let axis = basis.col(i);
let axis = basis.col_at(i);
new_center1 -= axis * lin_err.dot(axis);
}
}
Expand Down Expand Up @@ -241,19 +241,19 @@ impl JointConstraintHelper {
params: &RbdSimParams,
) -> JointConstraintElement {
#[cfg(feature = "dim2")]
let lin_jac = self.basis.col(locked_axis);
let lin_jac = self.basis.col_at(locked_axis);
#[cfg(feature = "dim3")]
let lin_jac = self.basis.col(locked_axis);
let lin_jac = self.basis.col_at(locked_axis);

#[cfg(feature = "dim2")]
let ang_jac1 = self.cmat1_basis.read(locked_axis);
#[cfg(feature = "dim2")]
let ang_jac2 = self.cmat2_basis.read(locked_axis);

#[cfg(feature = "dim3")]
let ang_jac1 = self.cmat1_basis.col(locked_axis);
let ang_jac1 = self.cmat1_basis.col_at(locked_axis);
#[cfg(feature = "dim3")]
let ang_jac2 = self.cmat2_basis.col(locked_axis);
let ang_jac2 = self.cmat2_basis.col_at(locked_axis);

let rhs_wo_bias = 0.0;
let erp_inv_dt = params.joint_erp_inv_dt();
Expand Down Expand Up @@ -294,7 +294,7 @@ impl JointConstraintHelper {
#[cfg(feature = "dim2")]
let ang_jac = 1.0;
#[cfg(feature = "dim3")]
let ang_jac = self.ang_basis.col(_locked_axis);
let ang_jac = self.ang_basis.col_at(_locked_axis);

let rhs_wo_bias = 0.0;
let erp_inv_dt = params.joint_erp_inv_dt();
Expand Down Expand Up @@ -407,8 +407,8 @@ impl JointConstraintHelper {
#[cfg(feature = "dim2")]
for i in 0..DIM {
if (coupled_axes & (1u32 << i)) != 0 {
let coeff = self.basis.col(i).dot(self.lin_err);
lin_jac += self.basis.col(i) * coeff;
let coeff = self.basis.col_at(i).dot(self.lin_err);
lin_jac += self.basis.col_at(i) * coeff;
ang_jac1 += self.cmat1_basis.read(i) * coeff;
ang_jac2 += self.cmat2_basis.read(i) * coeff;
}
Expand All @@ -417,10 +417,10 @@ impl JointConstraintHelper {
#[cfg(feature = "dim3")]
for i in 0..DIM {
if (coupled_axes & (1u32 << i)) != 0 {
let coeff = self.basis.col(i).dot(self.lin_err);
lin_jac += self.basis.col(i) * coeff;
ang_jac1 += self.cmat1_basis.col(i) * coeff;
ang_jac2 += self.cmat2_basis.col(i) * coeff;
let coeff = self.basis.col_at(i).dot(self.lin_err);
lin_jac += self.basis.col_at(i) * coeff;
ang_jac1 += self.cmat1_basis.col_at(i) * coeff;
ang_jac2 += self.cmat2_basis.col_at(i) * coeff;
}
}

Expand Down Expand Up @@ -517,8 +517,8 @@ impl JointConstraintHelper {
#[cfg(feature = "dim2")]
for i in 0..DIM {
if (coupled_axes & (1u32 << i)) != 0 {
let coeff = self.basis.col(i).dot(self.lin_err);
lin_jac += self.basis.col(i) * coeff;
let coeff = self.basis.col_at(i).dot(self.lin_err);
lin_jac += self.basis.col_at(i) * coeff;
ang_jac1 += self.cmat1_basis.read(i) * coeff;
ang_jac2 += self.cmat2_basis.read(i) * coeff;
}
Expand All @@ -527,10 +527,10 @@ impl JointConstraintHelper {
#[cfg(feature = "dim3")]
for i in 0..DIM {
if (coupled_axes & (1u32 << i)) != 0 {
let coeff = self.basis.col(i).dot(self.lin_err);
lin_jac += self.basis.col(i) * coeff;
ang_jac1 += self.cmat1_basis.col(i) * coeff;
ang_jac2 += self.cmat2_basis.col(i) * coeff;
let coeff = self.basis.col_at(i).dot(self.lin_err);
lin_jac += self.basis.col_at(i) * coeff;
ang_jac1 += self.cmat1_basis.col_at(i) * coeff;
ang_jac2 += self.cmat2_basis.col_at(i) * coeff;
}
}

Expand Down Expand Up @@ -602,7 +602,7 @@ impl JointConstraintHelper {
#[cfg(feature = "dim2")]
let ang_jac = 1.0;
#[cfg(feature = "dim3")]
let ang_jac = self.ang_basis.col(_limited_axis);
let ang_jac = self.ang_basis.col_at(_limited_axis);

let rhs_wo_bias = 0.0;
let erp_inv_dt = params.joint_erp_inv_dt();
Expand Down Expand Up @@ -643,7 +643,7 @@ impl JointConstraintHelper {
#[cfg(feature = "dim2")]
let ang_jac = 1.0;
#[cfg(feature = "dim3")]
let ang_jac = self.basis.col(_motor_axis);
let ang_jac = self.basis.col_at(_motor_axis);

let mut rhs_wo_bias = 0.0;
if motor_params.erp_inv_dt != 0.0 {
Expand Down
42 changes: 42 additions & 0 deletions src_rbd_shaders/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,48 @@ pub fn abs(x: f32) -> f32 {
Float::abs(x)
}

/// Panic-free matrix column access for GPU shader code.
///
/// glam's `col` panics with a message on an out-of-range index, which GPU
/// backends that cannot lower panic-message formatting (e.g. cuda-oxide /
/// PTX) reject at device-codegen time — even when the panic branch is
/// unreachable at runtime.
///
/// `index` **must** be in range (`0..2` for `Mat2`, `0..3` for `Mat3`); every
/// caller in this crate derives it from an axis mask, so this always holds.
/// An out-of-range index is a bug: debug (host) builds catch it with a
/// `debug_assert!`, release/device builds fold it to the last column rather
/// than carrying panic machinery into the kernel.
pub trait ColumnIndex {
type Column;
fn col_at(&self, index: usize) -> Self::Column;
}

impl ColumnIndex for glamx::Mat3 {
type Column = glamx::Vec3;
#[inline(always)]
fn col_at(&self, index: usize) -> glamx::Vec3 {
debug_assert!(index < 3, "Mat3 column index out of range");
match index {
0 => self.x_axis,
1 => self.y_axis,
_ => self.z_axis,
}
}
}

impl ColumnIndex for glamx::Mat2 {
type Column = glamx::Vec2;
#[inline(always)]
fn col_at(&self, index: usize) -> glamx::Vec2 {
debug_assert!(index < 2, "Mat2 column index out of range");
match index {
0 => self.x_axis,
_ => self.y_axis,
}
}
}

//
// Modules
//
Expand Down
Loading