diff --git a/examples/heat-equation-ffi/README.md b/examples/heat-equation-ffi/README.md index 06d2395..4886d03 100644 --- a/examples/heat-equation-ffi/README.md +++ b/examples/heat-equation-ffi/README.md @@ -28,9 +28,17 @@ LIBRARY_PATH=/usr/local/cuda/lib64:/opt/rocm/lib cargo build --release The platform description must identify the backend available on the machine, for example `acc_backend = ["CUDA"]` or `acc_backend = ["ROCM"]` in the -accelerator section of `Platform.toml`. The same executable contains both -implementations and selects the matching version when it starts; no Cargo -feature or recompilation is involved in that selection. +accelerator section of `Platform.toml`. The build creates the executable and +two backend plugins beside it: `libheat_cuda.so` and `libheat_rocm.so`. The +executable selects the matching version and dynamically opens only that +version's plugin, so it has no startup dependency on either vendor runtime. +No Cargo feature or recompilation is involved in that selection. + +Only the selected plugin and its vendor runtime need to be installed. For +example, an AMD installation can contain the executable and +`libheat_rocm.so`, without `libheat_cuda.so` or the CUDA runtime. By default, +plugins are found beside the executable. `HEAT_CUDA_LIBRARY` and +`HEAT_ROCM_LIBRARY` can override their respective full paths. Run that same executable on either platform after placing the appropriate `Platform.toml` in the working directory: diff --git a/examples/heat-equation-ffi/build.rs b/examples/heat-equation-ffi/build.rs index 4f1b079..1e0b0bf 100644 --- a/examples/heat-equation-ffi/build.rs +++ b/examples/heat-equation-ffi/build.rs @@ -1,4 +1,8 @@ -use std::{env, path::PathBuf, process::Command}; +use std::{ + env, + path::{Path, PathBuf}, + process::Command, +}; fn run(mut command: Command, tool: &str) { let status = command.status().unwrap_or_else(|e| panic!("cannot execute {tool}: {e}")); @@ -9,13 +13,11 @@ fn hip_gcc_install_dir() -> Option { if let Some(path) = env::var_os("HIP_GCC_INSTALL_DIR") { return Some(path.into()); } - let cxx = env::var_os("CXX").unwrap_or_else(|| "c++".into()); let output = Command::new(cxx).arg("-print-libgcc-file-name").output().ok()?; if !output.status.success() { return None; } - let libgcc = PathBuf::from(String::from_utf8(output.stdout).ok()?.trim()); let directory = libgcc.parent()?; directory @@ -24,72 +26,72 @@ fn hip_gcc_install_dir() -> Option { .then(|| directory.into()) } -fn add_sdk_library_path(environment: &str, default_root: &str, directories: &[&str]) { +fn sdk_library_dir(environment: &str, default_root: &str, directories: &[&str]) -> Option { let root = env::var_os(environment) .map(PathBuf::from) .unwrap_or_else(|| PathBuf::from(default_root)); - - if let Some(directory) = directories + directories .iter() .map(|directory| root.join(directory)) .find(|directory| directory.is_dir()) - { - println!("cargo:rustc-link-search=native={}", directory.display()); - } +} + +fn profile_dir(out: &Path) -> &Path { + // target/{profile}/build/{package-hash}/out + out.ancestors().nth(3).expect("unexpected Cargo OUT_DIR") } fn main() { let out = PathBuf::from(env::var_os("OUT_DIR").unwrap()); + let destination = profile_dir(&out); println!("cargo:rerun-if-env-changed=CXX"); println!("cargo:rerun-if-env-changed=HIP_GCC_INSTALL_DIR"); println!("cargo:rerun-if-env-changed=CUDA_PATH"); println!("cargo:rerun-if-env-changed=ROCM_PATH"); - add_sdk_library_path( + let cuda_object = out.join("heat_cuda.o"); + let mut nvcc = Command::new("nvcc"); + nvcc.args(["-c", "kernels/heat_cuda.cu", "-o"]).arg(&cuda_object).args([ + "-O2", + "--compiler-options", + "-fPIC", + ]); + run(nvcc, "nvcc"); + let mut cuda_linker = Command::new(env::var_os("CXX").unwrap_or_else(|| "c++".into())); + cuda_linker + .args(["-shared", "-o"]) + .arg(destination.join("libheat_cuda.so")) + .arg(&cuda_object) + .arg("-lcudart"); + if let Some(directory) = sdk_library_dir( "CUDA_PATH", "/usr/local/cuda", &["lib64", "targets/x86_64-linux/lib"], - ); - add_sdk_library_path("ROCM_PATH", "/opt/rocm", &["lib", "lib64"]); - - { - let object = out.join("heat_cuda.o"); - let library = out.join("libheat_cuda.a"); - let mut nvcc = Command::new("nvcc"); - nvcc.args(["-c", "kernels/heat_cuda.cu", "-o"]).arg(&object).args([ - "-O2", - "--compiler-options", - "-fPIC", - ]); - run(nvcc, "nvcc"); - let mut ar = Command::new("ar"); - ar.arg("crus").arg(&library).arg(&object); - run(ar, "ar"); - println!("cargo:rustc-link-search=native={}", out.display()); - println!("cargo:rustc-link-lib=static=heat_cuda"); - println!("cargo:rustc-link-lib=cudart"); - println!("cargo:rerun-if-changed=kernels/heat_cuda.cu"); + ) { + cuda_linker.arg(format!("-L{}", directory.display())); } + run(cuda_linker, "C++ linker for CUDA plugin"); + println!("cargo:rerun-if-changed=kernels/heat_cuda.cu"); - { - let object = out.join("heat_rocm.o"); - let library = out.join("libheat_rocm.a"); - let mut hipcc = Command::new("hipcc"); - if let Some(directory) = hip_gcc_install_dir() { - hipcc.arg(format!("--gcc-install-dir={}", directory.display())); - } - hipcc - .args(["-c", "kernels/heat_rocm.hip", "-o"]) - .arg(&object) - .args(["-O2", "-fPIC"]); - run(hipcc, "hipcc"); - let mut ar = Command::new("ar"); - ar.arg("crus").arg(&library).arg(&object); - run(ar, "ar"); - println!("cargo:rustc-link-search=native={}", out.display()); - println!("cargo:rustc-link-lib=static=heat_rocm"); - println!("cargo:rustc-link-lib=amdhip64"); - println!("cargo:rerun-if-changed=kernels/heat_rocm.hip"); + let mut hipcc = Command::new("hipcc"); + if let Some(directory) = hip_gcc_install_dir() { + hipcc.arg(format!("--gcc-install-dir={}", directory.display())); + } + let rocm_object = out.join("heat_rocm.o"); + hipcc + .args(["-c", "kernels/heat_rocm.hip", "-o"]) + .arg(&rocm_object) + .args(["-O2", "-fPIC"]); + run(hipcc, "hipcc"); + let mut rocm_linker = Command::new(env::var_os("CXX").unwrap_or_else(|| "c++".into())); + rocm_linker + .args(["-shared", "-o"]) + .arg(destination.join("libheat_rocm.so")) + .arg(&rocm_object) + .args(["-lamdhip64", "-lstdc++"]); + if let Some(directory) = sdk_library_dir("ROCM_PATH", "/opt/rocm", &["lib", "lib64"]) { + rocm_linker.arg(format!("-L{}", directory.display())); } - println!("cargo:rustc-link-lib=stdc++"); + run(rocm_linker, "C++ linker for ROCm plugin"); + println!("cargo:rerun-if-changed=kernels/heat_rocm.hip"); } diff --git a/examples/heat-equation-ffi/src/main.rs b/examples/heat-equation-ffi/src/main.rs index e132552..fc247cc 100644 --- a/examples/heat-equation-ffi/src/main.rs +++ b/examples/heat-equation-ffi/src/main.rs @@ -21,8 +21,13 @@ fn main() { println!("centre after 100 steps: {}", answer[512]); } - -use std::{mem::size_of, ptr}; +use std::{ + env, + ffi::{c_char, c_int, c_void, CStr, CString}, + mem::size_of, + path::{Path, PathBuf}, + ptr, +}; const HOST_TO_DEVICE: i32 = 1; const DEVICE_TO_HOST: i32 = 2; @@ -37,6 +42,83 @@ fn check(code: i32, operation: &str) { assert_eq!(code, 0, "{operation} failed with GPU error {code}"); } +const RTLD_NOW: c_int = 2; +const RTLD_LOCAL: c_int = 0; + +#[link(name = "dl")] +extern "C" { + fn dlopen(filename: *const c_char, flags: c_int) -> *mut c_void; + fn dlsym(handle: *mut c_void, symbol: *const c_char) -> *mut c_void; + fn dlclose(handle: *mut c_void) -> c_int; + fn dlerror() -> *const c_char; +} + +struct DynamicLibrary(*mut c_void); + +impl DynamicLibrary { + fn open(environment: &str, filename: &str) -> Self { + let path = env::var_os(environment) + .map(PathBuf::from) + .unwrap_or_else(|| sibling_library(filename)); + let path = CString::new(path.as_os_str().as_encoded_bytes()) + .unwrap_or_else(|_| panic!("invalid plugin path for {filename}")); + let handle = unsafe { dlopen(path.as_ptr(), RTLD_NOW | RTLD_LOCAL) }; + if handle.is_null() { + panic!("cannot load {filename}: {}", dynamic_loader_error()); + } + Self(handle) + } + + unsafe fn symbol(&self, name: &str) -> T { + let name = CString::new(name).expect("symbol contains a NUL byte"); + dlerror(); + let address = dlsym(self.0, name.as_ptr()); + let error = dlerror(); + if !error.is_null() { + panic!( + "cannot load symbol {}: {}", + name.to_string_lossy(), + CStr::from_ptr(error).to_string_lossy() + ); + } + std::mem::transmute_copy(&address) + } +} + +impl Drop for DynamicLibrary { + fn drop(&mut self) { + unsafe { + dlclose(self.0); + } + } +} + +fn sibling_library(filename: &str) -> PathBuf { + env::current_exe() + .ok() + .and_then(|executable| executable.parent().map(Path::to_path_buf)) + .unwrap_or_default() + .join(filename) +} + +fn dynamic_loader_error() -> String { + unsafe { + let error = dlerror(); + if error.is_null() { + "unknown dynamic loader error".into() + } else { + CStr::from_ptr(error).to_string_lossy().into_owned() + } + } +} + +type Malloc = unsafe extern "C" fn(*mut *mut c_void, usize) -> i32; +type Free = unsafe extern "C" fn(*mut c_void) -> i32; +type Memcpy = unsafe extern "C" fn(*mut c_void, *const c_void, usize, i32) -> i32; +type Memset = unsafe extern "C" fn(*mut c_void, i32, usize) -> i32; +type Synchronize = unsafe extern "C" fn() -> i32; +type Launch = unsafe extern "C" fn(*const f32, *mut f32, i32, f32, u32, u32) -> i32; + fn solve_heat_cuda_impl(n: usize, steps: usize, alpha: f32) -> Vec { assert!(n >= 3 && (0.0..=0.5).contains(&alpha)); let mut host = initial_grid(n); @@ -44,31 +126,38 @@ fn solve_heat_cuda_impl(n: usize, steps: usize, alpha: f32) -> Vec { let (mut u, mut next) = (ptr::null_mut(), ptr::null_mut()); unsafe { - check(cuda::cudaMalloc(&mut u, bytes), "cudaMalloc(u)"); - check(cuda::cudaMalloc(&mut next, bytes), "cudaMalloc(next)"); + let library = DynamicLibrary::open("HEAT_CUDA_LIBRARY", "libheat_cuda.so"); + let malloc: Malloc = library.symbol("cudaMalloc"); + let free: Free = library.symbol("cudaFree"); + let memcpy: Memcpy = library.symbol("cudaMemcpy"); + let memset: Memset = library.symbol("cudaMemset"); + let synchronize: Synchronize = library.symbol("cudaDeviceSynchronize"); + let launch: Launch = library.symbol("heat_cuda_launch"); + check(malloc(&mut u, bytes), "cudaMalloc(u)"); + check(malloc(&mut next, bytes), "cudaMalloc(next)"); check( - cuda::cudaMemcpy(u, host.as_ptr().cast(), bytes, HOST_TO_DEVICE), + memcpy(u, host.as_ptr().cast(), bytes, HOST_TO_DEVICE), "cudaMemcpy(u)", ); - check(cuda::cudaMemset(next, 0, bytes), "cudaMemset(next)"); + check(memset(next, 0, bytes), "cudaMemset(next)"); let threads = 256_u32; let blocks = ((n as u32 - 2) + threads - 1) / threads; for _ in 0..steps { // Dirichlet boundary values u[0] = u[n-1] = 0 remain untouched. check( - cuda::heat_cuda_launch(u.cast(), next.cast(), n as i32, alpha, blocks, threads), + launch(u.cast(), next.cast(), n as i32, alpha, blocks, threads), "heat_cuda_launch", ); std::mem::swap(&mut u, &mut next); } - check(cuda::cudaDeviceSynchronize(), "cudaDeviceSynchronize"); + check(synchronize(), "cudaDeviceSynchronize"); check( - cuda::cudaMemcpy(host.as_mut_ptr().cast(), u, bytes, DEVICE_TO_HOST), + memcpy(host.as_mut_ptr().cast(), u, bytes, DEVICE_TO_HOST), "cudaMemcpy(result)", ); - check(cuda::cudaFree(u), "cudaFree(u)"); - check(cuda::cudaFree(next), "cudaFree(next)"); + check(free(u), "cudaFree(u)"); + check(free(next), "cudaFree(next)"); } host } @@ -80,69 +169,37 @@ fn solve_heat_rocm_impl(n: usize, steps: usize, alpha: f32) -> Vec { let (mut u, mut next) = (ptr::null_mut(), ptr::null_mut()); unsafe { - check(rocm::hipMalloc(&mut u, bytes), "hipMalloc(u)"); - check(rocm::hipMalloc(&mut next, bytes), "hipMalloc(next)"); + let library = DynamicLibrary::open("HEAT_ROCM_LIBRARY", "libheat_rocm.so"); + let malloc: Malloc = library.symbol("hipMalloc"); + let free: Free = library.symbol("hipFree"); + let memcpy: Memcpy = library.symbol("hipMemcpy"); + let memset: Memset = library.symbol("hipMemset"); + let synchronize: Synchronize = library.symbol("hipDeviceSynchronize"); + let launch: Launch = library.symbol("heat_rocm_launch"); + check(malloc(&mut u, bytes), "hipMalloc(u)"); + check(malloc(&mut next, bytes), "hipMalloc(next)"); check( - rocm::hipMemcpy(u, host.as_ptr().cast(), bytes, HOST_TO_DEVICE), + memcpy(u, host.as_ptr().cast(), bytes, HOST_TO_DEVICE), "hipMemcpy(u)", ); - check(rocm::hipMemset(next, 0, bytes), "hipMemset(next)"); + check(memset(next, 0, bytes), "hipMemset(next)"); let threads = 256_u32; let blocks = ((n as u32 - 2) + threads - 1) / threads; for _ in 0..steps { check( - rocm::heat_rocm_launch(u.cast(), next.cast(), n as i32, alpha, blocks, threads), + launch(u.cast(), next.cast(), n as i32, alpha, blocks, threads), "heat_rocm_launch", ); std::mem::swap(&mut u, &mut next); } - check(rocm::hipDeviceSynchronize(), "hipDeviceSynchronize"); + check(synchronize(), "hipDeviceSynchronize"); check( - rocm::hipMemcpy(host.as_mut_ptr().cast(), u, bytes, DEVICE_TO_HOST), + memcpy(host.as_mut_ptr().cast(), u, bytes, DEVICE_TO_HOST), "hipMemcpy(result)", ); - check(rocm::hipFree(u), "hipFree(u)"); - check(rocm::hipFree(next), "hipFree(next)"); + check(free(u), "hipFree(u)"); + check(free(next), "hipFree(next)"); } host } - -mod cuda { - use std::ffi::c_void; - extern "C" { - pub fn heat_cuda_launch( - u: *const f32, - next: *mut f32, - n: i32, - alpha: f32, - blocks: u32, - threads: u32, - ) -> i32; - pub fn cudaMalloc(p: *mut *mut c_void, bytes: usize) -> i32; - pub fn cudaFree(p: *mut c_void) -> i32; - pub fn cudaMemcpy(dst: *mut c_void, src: *const c_void, bytes: usize, kind: i32) -> i32; - pub fn cudaMemset(dst: *mut c_void, value: i32, bytes: usize) -> i32; - pub fn cudaDeviceSynchronize() -> i32; - } -} - -mod rocm { - use std::ffi::c_void; - extern "C" { - pub fn heat_rocm_launch( - u: *const f32, - next: *mut f32, - n: i32, - alpha: f32, - blocks: u32, - threads: u32, - ) -> i32; - pub fn hipMalloc(p: *mut *mut c_void, bytes: usize) -> i32; - pub fn hipFree(p: *mut c_void) -> i32; - pub fn hipMemcpy(dst: *mut c_void, src: *const c_void, bytes: usize, kind: i32) -> i32; - pub fn hipMemset(dst: *mut c_void, value: i32, bytes: usize) -> i32; - pub fn hipDeviceSynchronize() -> i32; - } -} -