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>
62 lines
1.8 KiB
Rust
62 lines
1.8 KiB
Rust
use cudarc::driver::{CudaContext, DriverError, LaunchConfig, PushKernelArg};
|
|
use cudarc::nvrtc::compile_ptx;
|
|
|
|
const PTX_SRC: &str = "
|
|
extern \"C\" __global__ void matmul(float* A, float* B, float* C, int N) {
|
|
int ROW = blockIdx.y*blockDim.y+threadIdx.y;
|
|
int COL = blockIdx.x*blockDim.x+threadIdx.x;
|
|
|
|
float tmpSum = 0;
|
|
|
|
if (ROW < N && COL < N) {
|
|
// each thread computes one element of the block sub-matrix
|
|
for (int i = 0; i < N; i++) {
|
|
tmpSum += A[ROW * N + i] * B[i * N + COL];
|
|
}
|
|
}
|
|
// printf(\"pos, (%d, %d) - N %d - value %d\\n\", ROW, COL, N, tmpSum);
|
|
C[ROW * N + COL] = tmpSum;
|
|
}
|
|
";
|
|
|
|
fn main() -> Result<(), DriverError> {
|
|
let start = std::time::Instant::now();
|
|
|
|
let ptx = compile_ptx(PTX_SRC).unwrap();
|
|
println!("Compilation succeeded in {:?}", start.elapsed());
|
|
|
|
let ctx = CudaContext::new(0)?;
|
|
let stream = ctx.default_stream();
|
|
println!("Built in {:?}", start.elapsed());
|
|
|
|
let module = ctx.load_module(ptx)?;
|
|
let f = module.load_function("matmul")?;
|
|
println!("Loaded in {:?}", start.elapsed());
|
|
|
|
let a_host = [1.0f32, 2.0, 3.0, 4.0];
|
|
let b_host = [1.0f32, 2.0, 3.0, 4.0];
|
|
let mut c_host = [0.0f32; 4];
|
|
|
|
let a_dev = stream.clone_htod(&a_host)?;
|
|
let b_dev = stream.clone_htod(&b_host)?;
|
|
let mut c_dev = stream.clone_htod(&c_host)?;
|
|
|
|
println!("Copied in {:?}", start.elapsed());
|
|
|
|
let mut builder = stream.launch_builder(&f);
|
|
builder.arg(&a_dev);
|
|
builder.arg(&b_dev);
|
|
builder.arg(&mut c_dev);
|
|
builder.arg(&2i32);
|
|
let cfg = LaunchConfig {
|
|
block_dim: (2, 2, 1),
|
|
grid_dim: (1, 1, 1),
|
|
shared_mem_bytes: 0,
|
|
};
|
|
unsafe { builder.launch(cfg) }?;
|
|
|
|
stream.memcpy_dtoh(&c_dev, &mut c_host)?;
|
|
println!("Found {:?} in {:?}", c_host, start.elapsed());
|
|
Ok(())
|
|
}
|