T1: scaffold repo + Rust/CUDA build chain (vecadd smoke test)

Stand up the xtrain project skeleton: a Cargo workspace mirroring xserv's
csrc/ + crates/ layout, with a single xtrain-cuda crate that wraps the CUDA
Runtime over hand-written extern "C" FFI. build.rs compiles csrc/test/vecadd.cu
via the cc crate targeting sm_120 (RTX 5090) and links cudart.

A gated integration test runs the vector-add kernel on the GPU and asserts the
result. When nvcc is absent (local GPU-less machine), build.rs skips CUDA
compilation and sets a `no_cuda` cfg so host-side cargo check still works.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 14:42:43 +08:00
commit 92acf9f413
13 changed files with 376 additions and 0 deletions

View File

@@ -0,0 +1,16 @@
use crate::error::{self, Result};
use crate::ffi;
pub fn device_count() -> Result<i32> {
let mut count = 0;
error::check(unsafe { ffi::cudaGetDeviceCount(&mut count) })?;
Ok(count)
}
pub fn set_device(device: u32) -> Result<()> {
error::check(unsafe { ffi::cudaSetDevice(device as i32) })
}
pub fn synchronize() -> Result<()> {
error::check(unsafe { ffi::cudaDeviceSynchronize() })
}

View File

@@ -0,0 +1,40 @@
use crate::ffi;
use std::ffi::CStr;
use std::fmt;
#[derive(Debug)]
pub enum CudaError {
OutOfMemory,
Raw { code: i32, message: String },
}
impl fmt::Display for CudaError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
CudaError::OutOfMemory => write!(f, "CUDA out of memory"),
CudaError::Raw { code, message } => write!(f, "CUDA error {code}: {message}"),
}
}
}
impl std::error::Error for CudaError {}
pub type Result<T> = std::result::Result<T, CudaError>;
pub fn check(code: i32) -> Result<()> {
if code == ffi::CUDA_SUCCESS {
return Ok(());
}
let message = unsafe {
let ptr = ffi::cudaGetErrorString(code);
if ptr.is_null() {
"unknown error".to_string()
} else {
CStr::from_ptr(ptr).to_string_lossy().into_owned()
}
};
Err(match code {
ffi::CUDA_ERROR_OUT_OF_MEMORY => CudaError::OutOfMemory,
_ => CudaError::Raw { code, message },
})
}

View File

@@ -0,0 +1,32 @@
use std::ffi::c_void;
use std::os::raw::c_char;
pub type CudaStream = *mut c_void;
pub const CUDA_MEMCPY_H2D: i32 = 1;
pub const CUDA_MEMCPY_D2H: i32 = 2;
pub const CUDA_SUCCESS: i32 = 0;
pub const CUDA_ERROR_OUT_OF_MEMORY: i32 = 2;
unsafe extern "C" {
// --- Device ---
pub fn cudaGetDeviceCount(count: *mut i32) -> i32;
pub fn cudaSetDevice(device: i32) -> i32;
pub fn cudaDeviceSynchronize() -> i32;
// --- Memory ---
pub fn cudaMalloc(devptr: *mut *mut u8, size: usize) -> i32;
pub fn cudaFree(devptr: *mut u8) -> i32;
pub fn cudaMemcpy(dst: *mut u8, src: *const u8, count: usize, kind: i32) -> i32;
// --- Error ---
pub fn cudaGetErrorString(error: i32) -> *const c_char;
}
// The vector-add smoke-test kernel, compiled from csrc/test/vecadd.cu by build.rs.
// Only linked when CUDA is actually compiled (i.e. nvcc was present).
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream);
}

View File

@@ -0,0 +1,7 @@
pub mod device;
pub mod error;
pub mod ffi;
pub mod memory;
pub use error::{CudaError, Result};
pub use memory::GpuBuffer;

View File

@@ -0,0 +1,59 @@
use crate::error::{self, Result};
use crate::ffi;
/// RAII wrapper around a GPU memory allocation. Dropping frees the memory.
pub struct GpuBuffer {
ptr: *mut u8,
len: usize,
}
impl GpuBuffer {
pub fn alloc(len: usize) -> Result<Self> {
assert!(len > 0, "cannot allocate 0 bytes on GPU");
let mut ptr = std::ptr::null_mut();
error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?;
Ok(Self { ptr, len })
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn as_ptr(&self) -> *const u8 {
self.ptr
}
pub fn as_mut_ptr(&mut self) -> *mut u8 {
self.ptr
}
/// Copy data from a host (CPU) slice to this GPU buffer.
pub fn copy_from_host(&mut self, src: &[u8]) -> Result<()> {
assert!(src.len() <= self.len, "source larger than buffer");
error::check(unsafe {
ffi::cudaMemcpy(self.ptr, src.as_ptr(), src.len(), ffi::CUDA_MEMCPY_H2D)
})
}
/// Copy data from this GPU buffer to a host (CPU) slice.
pub fn copy_to_host(&self, dst: &mut [u8]) -> Result<()> {
assert!(dst.len() <= self.len, "destination larger than buffer");
error::check(unsafe {
ffi::cudaMemcpy(dst.as_mut_ptr(), self.ptr, dst.len(), ffi::CUDA_MEMCPY_D2H)
})
}
}
impl Drop for GpuBuffer {
fn drop(&mut self) {
if !self.ptr.is_null() {
unsafe { ffi::cudaFree(self.ptr) };
}
}
}
unsafe impl Send for GpuBuffer {}