#[repr(C)]
struct Runner {
uid_counter: u64,
n_agents: u64,
agents_up_to_date: bool,
agents: *mut AgentContainer,
positions: *mut f32,
velocities: *mut f32,
forces: *mut f32,
}
impl Drop for Runner {
fn drop(&mut self) {
unsafe {
drop_runner(self);
}
}
}
type UID = u64;
#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct Position {
p: [f32; 3],
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct Velocity {
v: [f32; 3],
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct Force {
f: [f32; 3],
}
#[repr(C)]
struct Agent {
id: UID,
position: Position,
velocity: Velocity,
}
#[repr(C)]
struct AgentContainer {
agent: Agent,
}
unsafe extern "C" {
fn new_runner(agents: *const Agent, n_agents: usize) -> *mut Runner;
fn drop_runner(runner: *mut Runner);
fn do_compute();
fn update_positions(runner: &mut Runner);
fn print_positions(runner: &Runner);
}
pub fn run() {
let n_agents = 10;
let agents: Vec<_> = (0..n_agents)
.map(|n| Agent {
id: n,
position: Position { p: [n as f32; 3] },
velocity: Velocity { v: [0.; 3] },
})
.collect();
let runner = unsafe { new_runner(agents.as_ptr(), agents.len()) };
unsafe {
update_positions(&mut *runner);
do_compute();
}
}
#[test]
fn test_main() {
run();
}
#[test]
fn test_cudarc() -> Result<(), Box<dyn std::error::Error>> {
use cudarc::driver::*;
const SIN_KERNEL: &str = include_str!("kernel.cu");
let ctx = cudarc::driver::CudaContext::new(0)?;
let stream = ctx.default_stream();
let inp = stream.memcpy_stod(&[1.0f32; 100])?;
let mut out = stream.alloc_zeros::<f32>(100)?;
let ptx = cudarc::nvrtc::compile_ptx(SIN_KERNEL)?;
let module = ctx.load_module(ptx)?;
let sin_kernel = module.load_function("sin_kernel")?;
let mut builder = stream.launch_builder(&sin_kernel);
builder.arg(&mut out);
builder.arg(&inp);
builder.arg(&100usize);
unsafe { builder.launch(cudarc::driver::LaunchConfig::for_num_elems(100)) }?;
let out_host: Vec<f32> = stream.memcpy_dtov(&out)?;
for q in out_host.into_iter() {
assert!((q - f32::sin(1.).abs() < 0.001));
}
Ok(())
}