Merging VAMP bindings + robogpu collision checker into main + cleanup

.gitignore

@@ -18,5 +18,8 @@ build
 HJCD-IK
 resources/learned_ik/*
 pRRTC
-src/pyroffi/cuda_kernels/*.so
-src/pyroffi/cuda_kernels/*.o
+src/pyroffi/cuda_kernels/**/*.so
+src/pyroffi/cuda_kernels/**/*.o
+src/pyroffi/cuda_kernels/**/*.ptx
+NVIDIA*.sh
+NVIDIA-OptiX-SDK-9.1.0-linux64-x86_64

src/pyroffi/cuda_kernels/build_all.sh → build_kernels/build_all.sh

@@ -2,11 +2,11 @@
 # Build all CUDA kernels.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_all.sh
-#   bash src/pyroffi/cuda_kernels/build_all.sh --debug
+#   bash build_kernels/build_all.sh
+#   bash build_kernels/build_all.sh --debug
 #
 # Override GPU arch for all kernels:
-#   GPU_ARCH=-arch=sm_80 bash src/pyroffi/cuda_kernels/build_all.sh
+#   GPU_ARCH=-arch=sm_80 bash build_kernels/build_all.sh
 
 set -euo pipefail
 
@@ -65,5 +65,7 @@ bash "${SCRIPT_DIR}/build_sco_trajopt_cuda.sh" "${BUILD_ARGS[@]}"
 bash "${SCRIPT_DIR}/build_ls_trajopt_cuda.sh" "${BUILD_ARGS[@]}"
 bash "${SCRIPT_DIR}/build_chomp_trajopt_cuda.sh" "${BUILD_ARGS[@]}"
 bash "${SCRIPT_DIR}/build_stomp_trajopt_cuda.sh" "${BUILD_ARGS[@]}"
+bash "${SCRIPT_DIR}/build_robogpu_collision.sh" "${BUILD_ARGS[@]}"
+bash "${SCRIPT_DIR}/build_cricket_jit.sh" 
 
 echo "All CUDA kernels built successfully."

src/pyroffi/cuda_kernels/build_brownian_motion_ik_cuda.sh → build_kernels/build_brownian_motion_ik_cuda.sh

@@ -2,8 +2,8 @@
 # Build _brownian_motion_ik_cuda_lib.so from _brownian_motion_ik_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_brownian_motion_ik_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_brownian_motion_ik_cuda.sh --debug
+#   bash build_kernels/build_brownian_motion_ik_cuda.sh
+#   bash build_kernels/build_brownian_motion_ik_cuda.sh --debug
 
 set -euo pipefail
 
@@ -45,8 +45,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_brownian_motion_ik_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_brownian_motion_ik_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/region_ik/_brownian_motion_ik_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/region_ik/_brownian_motion_ik_cuda_lib.so"
 
 JAXLIB_INC="$(python -c "import os, jaxlib; print(os.path.join(os.path.dirname(jaxlib.__file__), 'include'))")"
 
@@ -71,8 +72,9 @@ nvcc \
   ${GPU_ARCH} \
   --shared \
   --compiler-options "-fPIC" \
-  -I"${SCRIPT_DIR}" \
+  -I"${KERNELS_DIR}" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_chomp_trajopt_cuda.sh → build_kernels/build_chomp_trajopt_cuda.sh

@@ -2,8 +2,8 @@
 # Build _chomp_trajopt_cuda_lib.so from _chomp_trajopt_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_chomp_trajopt_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_chomp_trajopt_cuda.sh --debug
+#   bash build_kernels/build_chomp_trajopt_cuda.sh
+#   bash build_kernels/build_chomp_trajopt_cuda.sh --debug
 
 set -euo pipefail
 
@@ -45,8 +45,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_chomp_trajopt_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_chomp_trajopt_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/trajopt/_chomp_trajopt_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/trajopt/_chomp_trajopt_cuda_lib.so"
 
 JAXLIB_INC="$(python -c \
   "import os, jaxlib; print(os.path.join(os.path.dirname(jaxlib.__file__), 'include'))")"
@@ -73,6 +74,7 @@ nvcc \
   --shared \
   --compiler-options "-fPIC" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_collision_binary_cuda.sh → build_kernels/build_collision_binary_cuda.sh

@@ -2,8 +2,8 @@
 # Build _collision_binary_cuda_lib.so from _collision_binary_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_collision_binary_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_collision_binary_cuda.sh --debug
+#   bash build_kernels/build_collision_binary_cuda.sh
+#   bash build_kernels/build_collision_binary_cuda.sh --debug
 #
 # Requirements:
 #   - nvcc (CUDA toolkit)
@@ -51,8 +51,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_collision_binary_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_collision_binary_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/collision/_collision_binary_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/collision/_collision_binary_cuda_lib.so"
 
 # Locate the jaxlib include directory that ships xla/ffi/api/ffi.h.
 JAXLIB_INC="$(python -c \
@@ -81,6 +82,7 @@ nvcc \
   --shared \
   --compiler-options "-fPIC" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_collision_cuda.sh → build_kernels/build_collision_cuda.sh

@@ -2,8 +2,8 @@
 # Build _collision_cuda_lib.so from _collision_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_collision_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_collision_cuda.sh --debug
+#   bash build_kernels/build_collision_cuda.sh
+#   bash build_kernels/build_collision_cuda.sh --debug
 #
 # Requirements:
 #   - nvcc (CUDA toolkit)
@@ -51,8 +51,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_collision_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_collision_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/collision/_collision_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/collision/_collision_cuda_lib.so"
 
 # Locate the jaxlib include directory that ships xla/ffi/api/ffi.h.
 JAXLIB_INC="$(python -c \
@@ -81,6 +82,7 @@ nvcc \
   --shared \
   --compiler-options "-fPIC" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

build_kernels/build_cricket_jit.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Build + install cricket with the Python extension AND the runtime JIT enabled,
+# so pyroffi's VAMPCPUCollisionChecker can JIT-compile robot-specialised VAMP
+# collision checkers at runtime.
+#
+# This mirrors the exact, verified setup (conda-forge deps + scikit-build-core).
+# Run it inside the target conda env, e.g.:
+#
+#   conda activate pyroffi
+#   bash build_kernels/build_cricket_jit.sh
+#
+# Dependencies are taken from cricket/environment.yaml.  The JIT additionally
+# needs a `clang` binary on PATH at *runtime* (the JIT driver shells out to it to
+# discover system headers) — clangdev provides it inside the env.
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+CRICKET_DIR="${REPO_ROOT}/external/cricket"
+
+if [[ -z "${CONDA_PREFIX:-}" ]]; then
+  echo "ERROR: activate the target conda env first (conda activate pyroffi)." >&2
+  exit 1
+fi
+
+# 1. Install build/runtime dependencies (no-op if already present).
+conda install -c conda-forge --solver=libmamba -y \
+  pinocchio cppad eigen cgal-cpp nlohmann_json fmt \
+  llvmdev clangdev lld cxx-compiler ninja pkg-config patch \
+  nanobind scikit-build-core
+
+if ! command -v clang >/dev/null 2>&1; then
+  echo "ERROR: clang still not on PATH after install." >&2
+  exit 1
+fi
+
+# 2. Build + install the cricket Python extension (JIT + Python both ON).
+export CMAKE_PREFIX_PATH="${CONDA_PREFIX}:${CMAKE_PREFIX_PATH:-}"
+export CMAKE_ARGS="-DCMAKE_PREFIX_PATH=${CONDA_PREFIX} -DCRICKET_BUILD_JIT=ON -DCRICKET_BUILD_PYTHON=ON"
+pip install -e "${CRICKET_DIR}" --no-build-isolation
+
+echo
+echo "Done. Verify with:"
+echo "  python -c 'from cricket import _core_ext as e; print(e.jit.JitSession)'"
+echo "  python -m pytest tests/test_vamp_cpu_collision.py -s"

src/pyroffi/cuda_kernels/build_fk_cuda.sh → build_kernels/build_fk_cuda.sh

@@ -2,8 +2,8 @@
 # Build _fk_cuda.so from _fk_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_fk_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_fk_cuda.sh --debug
+#   bash build_kernels/build_fk_cuda.sh
+#   bash build_kernels/build_fk_cuda.sh --debug
 #
 # Requirements:
 #   - nvcc (CUDA toolkit)
@@ -50,8 +50,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_fk_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_fk_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/fk/_fk_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/fk/_fk_cuda_lib.so"
 
 # Locate the jaxlib include directory that ships xla/ffi/api/ffi.h.
 JAXLIB_INC="$(python -c \
@@ -82,6 +83,7 @@ nvcc \
   --shared \
   --compiler-options "-fPIC" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_hit_and_run_ik_cuda.sh → build_kernels/build_hit_and_run_ik_cuda.sh

@@ -2,8 +2,8 @@
 # Build _hit_and_run_ik_cuda_lib.so from _hit_and_run_ik_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_hit_and_run_ik_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_hit_and_run_ik_cuda.sh --debug
+#   bash build_kernels/build_hit_and_run_ik_cuda.sh
+#   bash build_kernels/build_hit_and_run_ik_cuda.sh --debug
 
 set -euo pipefail
 
@@ -45,8 +45,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_hit_and_run_ik_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_hit_and_run_ik_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/region_ik/_hit_and_run_ik_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/region_ik/_hit_and_run_ik_cuda_lib.so"
 
 JAXLIB_INC="$(python -c "import os, jaxlib; print(os.path.join(os.path.dirname(jaxlib.__file__), 'include'))")"
 
@@ -71,8 +72,9 @@ nvcc \
   ${GPU_ARCH} \
   --shared \
   --compiler-options "-fPIC" \
-  -I"${SCRIPT_DIR}" \
+  -I"${KERNELS_DIR}" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_hjcd_ik_cuda.sh → build_kernels/build_hjcd_ik_cuda.sh

@@ -2,8 +2,8 @@
 # Build _hjcd_ik_cuda_lib.so from _hjcd_ik_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_hjcd_ik_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_hjcd_ik_cuda.sh --debug
+#   bash build_kernels/build_hjcd_ik_cuda.sh
+#   bash build_kernels/build_hjcd_ik_cuda.sh --debug
 #
 # Requirements:
 #   - nvcc (CUDA toolkit)
@@ -50,8 +50,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_hjcd_ik_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_hjcd_ik_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/ik/_hjcd_ik_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/ik/_hjcd_ik_cuda_lib.so"
 
 # Locate the jaxlib include directory that ships xla/ffi/api/ffi.h.
 JAXLIB_INC="$(python -c \
@@ -81,8 +82,9 @@ nvcc \
   ${GPU_ARCH} \
   --shared \
   --compiler-options "-fPIC" \
-  -I"${SCRIPT_DIR}" \
+  -I"${KERNELS_DIR}" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_ls_ik_cuda.sh → build_kernels/build_ls_ik_cuda.sh

@@ -2,8 +2,8 @@
 # Build _ls_ik_cuda_lib.so from _ls_ik_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_ls_ik_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_ls_ik_cuda.sh --debug
+#   bash build_kernels/build_ls_ik_cuda.sh
+#   bash build_kernels/build_ls_ik_cuda.sh --debug
 #
 # Requirements:
 #   - nvcc (CUDA toolkit)
@@ -50,8 +50,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_ls_ik_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_ls_ik_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/ik/_ls_ik_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/ik/_ls_ik_cuda_lib.so"
 
 JAXLIB_INC="$(python -c \
   "import os, jaxlib; print(os.path.join(os.path.dirname(jaxlib.__file__), 'include'))")"
@@ -77,8 +78,9 @@ nvcc \
   ${GPU_ARCH} \
   --shared \
   --compiler-options "-fPIC" \
-  -I"${SCRIPT_DIR}" \
+  -I"${KERNELS_DIR}" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"
 

src/pyroffi/cuda_kernels/build_ls_trajopt_cuda.sh → build_kernels/build_ls_trajopt_cuda.sh

@@ -2,8 +2,8 @@
 # Build _ls_trajopt_cuda_lib.so from _ls_trajopt_cuda_kernel.cu.
 #
 # Usage (from repo root):
-#   bash src/pyroffi/cuda_kernels/build_ls_trajopt_cuda.sh
-#   bash src/pyroffi/cuda_kernels/build_ls_trajopt_cuda.sh --debug
+#   bash build_kernels/build_ls_trajopt_cuda.sh
+#   bash build_kernels/build_ls_trajopt_cuda.sh --debug
 
 set -euo pipefail
 
@@ -45,8 +45,9 @@ if [[ -n "${MAX_JOINTS_OVERRIDE}" ]]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SRC="${SCRIPT_DIR}/_ls_trajopt_cuda_kernel.cu"
-OUT="${SCRIPT_DIR}/_ls_trajopt_cuda_lib.so"
+KERNELS_DIR="$(cd "${SCRIPT_DIR}/../src/pyroffi/cuda_kernels" && pwd)"
+SRC="${KERNELS_DIR}/trajopt/_ls_trajopt_cuda_kernel.cu"
+OUT="${KERNELS_DIR}/trajopt/_ls_trajopt_cuda_lib.so"
 
 JAXLIB_INC="$(python -c \
   "import os, jaxlib; print(os.path.join(os.path.dirname(jaxlib.__file__), 'include'))")"
@@ -72,8 +73,9 @@ nvcc \
   ${GPU_ARCH} \
   --shared \
   --compiler-options "-fPIC" \
-  -I"${SCRIPT_DIR}" \
+  -I"${KERNELS_DIR}" \
   -I"${JAXLIB_INC}" \
+  -I"${KERNELS_DIR}" \
   -o "${OUT}" \
   "${SRC}"