tensor: add scale elementwise CUDA kernel + FFI

New csrc/ops/elementwise.cu (out[i]=in[i]*alpha), compiled by
xtrain-cuda/build.rs and exposed via launch_scale_f32 FFI, gated behind
not(no_cuda) like the existing vecadd smoke test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 15:13:06 +08:00
parent 8557a289a2
commit 63dc05fd10
3 changed files with 30 additions and 2 deletions

View File

@@ -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,
);
}