diff --git a/crates/xtrain-cuda/build.rs b/crates/xtrain-cuda/build.rs index 1fb0fb0..1d27f02 100644 --- a/crates/xtrain-cuda/build.rs +++ b/crates/xtrain-cuda/build.rs @@ -28,6 +28,7 @@ fn main() { .cudart("shared") .flag("-gencode=arch=compute_120,code=sm_120") .file("../../csrc/test/vecadd.cu") + .file("../../csrc/ops/elementwise.cu") .compile("xtrain_cuda_kernels"); } diff --git a/crates/xtrain-cuda/src/ffi.rs b/crates/xtrain-cuda/src/ffi.rs index 8bb0b0f..1f6bfc7 100644 --- a/crates/xtrain-cuda/src/ffi.rs +++ b/crates/xtrain-cuda/src/ffi.rs @@ -24,9 +24,19 @@ unsafe extern "C" { 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). +// GPU kernels compiled from csrc/ by build.rs. Only linked when CUDA is +// actually compiled (i.e. nvcc was present). #[cfg(not(no_cuda))] unsafe extern "C" { + // Vector-add smoke test (csrc/test/vecadd.cu). pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream); + + // Elementwise scale: out[i] = in[i] * alpha (csrc/ops/elementwise.cu). + pub fn launch_scale_f32( + input: *const f32, + out: *mut f32, + alpha: f32, + n: i32, + stream: CudaStream, + ); } diff --git a/csrc/ops/elementwise.cu b/csrc/ops/elementwise.cu new file mode 100644 index 0000000..eec217f --- /dev/null +++ b/csrc/ops/elementwise.cu @@ -0,0 +1,17 @@ +extern "C" { + +// out[i] = in[i] * alpha (in-place safe: out may alias in) +__global__ void scale_f32(const float* in, float* out, float alpha, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + out[idx] = in[idx] * alpha; + } +} + +void launch_scale_f32(const float* in, float* out, float alpha, int n, void* stream) { + int block = 256; + int grid = (n + block - 1) / block; + scale_f32<<>>(in, out, alpha, n); +} + +}