1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.
//! End-to-end test for the **raw-GPU** CUDA custom-kernel seam
//! (`CudaGpuKernel`): a downstream `Op::Custom` that NVRTC-compiles a CUDA-C
//! kernel and launches it straight against the arena device buffer — no D2H/H2D
//! host roundtrip (contrast the host-delegate `Step::CustomHost` path). Needs a
//! real NVIDIA GPU (built with the `cuda` feature); runs on the msi rig.
#![cfg(feature = "cuda")]
use std::sync::Arc;
use rlx_cuda::cuda_gpu_kernels::{CudaGpuKernel, register_cuda_gpu_kernel};
use rlx_ir::{DType, Graph, OpExtension, Shape, register_op};
use rlx_runtime::{Device, Session, is_available};
/// IR-level shape inference for `test.times_three_cuda` (identity).
struct TimesThreeIr;
impl OpExtension for TimesThreeIr {
fn name(&self) -> &str {
"test.times_three_cuda"
}
fn num_inputs(&self) -> usize {
1
}
fn infer_shape(&self, inputs: &[&Shape], _: &[u8]) -> Shape {
inputs[0].clone()
}
}
/// CUDA-C with the fixed launch signature (`float* arena`, out/in element
/// offsets). `y = x * 3`.
const CU: &str = r#"
extern "C" __global__ void rlx_custom(
float* arena,
unsigned out_off, unsigned out_len, unsigned n_inputs,
unsigned in0_off, unsigned in0_len,
unsigned in1_off, unsigned in1_len,
unsigned in2_off, unsigned in2_len,
unsigned in3_off, unsigned in3_len)
{
unsigned i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < out_len) {
arena[out_off + i] = arena[in0_off + i] * 3.0f;
}
}
"#;
/// Raw-GPU CUDA kernel: `y = x * 3`, no host roundtrip.
#[derive(Debug)]
struct TimesThreeCuda;
impl CudaGpuKernel for TimesThreeCuda {
fn name(&self) -> &str {
"test.times_three_cuda"
}
fn cuda_c(&self) -> &str {
CU
}
}
#[test]
fn times_three_runs_on_cuda_via_raw_gpu_kernel() {
if !is_available(Device::Cuda) {
eprintln!("skip: no CUDA device");
return;
}
register_op(Arc::new(TimesThreeIr));
register_cuda_gpu_kernel(Arc::new(TimesThreeCuda));
let n = 8usize;
let mut g = Graph::new("cuda_gpu_custom");
let x = g.input("x", Shape::new(&[n], DType::F32));
let y = g.custom_op("test.times_three_cuda", vec![], vec![x]);
g.set_outputs(vec![y]);
// Device::Cuda → the lower path emits Step::CudaGpuKernel (an NVRTC kernel
// launch against the arena, no host staging).
let mut c = Session::new(Device::Cuda).compile(g);
let x0: Vec<f32> = (0..n).map(|i| i as f32 - 3.0).collect();
let outs = c.run(&[("x", &x0)]);
let out = &outs[0];
let expect: Vec<f32> = x0.iter().map(|v| v * 3.0).collect();
assert_eq!(out.len(), expect.len(), "out={out:?}");
for (a, b) in out.iter().zip(expect.iter()) {
assert!((a - b).abs() < 1e-4, "out={out:?} expect={expect:?}");
}
}