07d0e60fe4
- Fork cudarc locally (vendor/cudarc): add CudaContext::load_cubin() that calls cuModuleLoadData directly — zero nvrtc dependency - Remove "nvrtc" feature from ml-core, ml-dqn, ml-ppo Cargo.toml - Replace all 89 Ptx::from_binary + load_module calls with load_cubin - ml-core cuda_autograd: wire 9 stub constructors to precompiled cubins (activation, elementwise, linear, loss, reduction, dropout, layer_norm, optimizer) - ml-core build.rs: compile 8 BF16-native CUDA kernels via nvcc - cubin_loader.rs: thin wrapper around CudaContext::load_cubin() - Fix size_of::<f32> in gpu_tensor.rs, stream_ops.rs, layer_norm.rs - Fix test data: Vec<f32> → Vec<half::bf16> for memcpy_htod - Stub ml-ppo/ml-dqn runtime compile_ptx calls (dead code) - backtest_metrics_kernel.cu: full native BF16 rewrite (no float) - backtest_env_kernel.cu: shared memory → __nv_bfloat16 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
55 lines
1.3 KiB
Rust
55 lines
1.3 KiB
Rust
use cudarc::{driver::*, nvrtc::compile_ptx};
|
|
|
|
/// Here's the struct in rust, note that we have #[repr(C)]
|
|
/// here which allows us to pass it to cuda.
|
|
#[repr(C)]
|
|
struct MyCoolRustStruct {
|
|
a: f32,
|
|
b: f64,
|
|
c: u32,
|
|
d: usize,
|
|
}
|
|
|
|
/// We have to implement this to send it to cuda!
|
|
unsafe impl DeviceRepr for MyCoolRustStruct {}
|
|
|
|
const PTX_SRC: &str = "
|
|
// here's the same struct in cuda
|
|
struct MyCoolStruct {
|
|
float a;
|
|
double b;
|
|
unsigned int c;
|
|
size_t d;
|
|
};
|
|
extern \"C\" __global__ void my_custom_kernel(MyCoolStruct thing) {
|
|
assert(thing.a == 1.0);
|
|
assert(thing.b == 2.34);
|
|
assert(thing.c == 57);
|
|
assert(thing.d == 420);
|
|
}
|
|
";
|
|
|
|
fn main() -> Result<(), DriverError> {
|
|
let ctx = CudaContext::new(0)?;
|
|
let stream = ctx.default_stream();
|
|
|
|
let ptx = compile_ptx(PTX_SRC).unwrap();
|
|
let module = ctx.load_module(ptx)?;
|
|
let f = module.load_function("my_custom_kernel")?;
|
|
|
|
// try changing some of these values to see a device assert
|
|
let thing = MyCoolRustStruct {
|
|
a: 1.0,
|
|
b: 2.34,
|
|
c: 57,
|
|
d: 420,
|
|
};
|
|
|
|
let mut builder = stream.launch_builder(&f);
|
|
// since MyCoolRustStruct implements DeviceRepr, we can pass it to launch.
|
|
builder.arg(&thing);
|
|
unsafe { builder.launch(LaunchConfig::for_num_elems(1)) }?;
|
|
|
|
Ok(())
|
|
}
|