added dynamic-load branch

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:

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<PathBuf> {
     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<PathBuf> {
         .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<PathBuf> {
     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");
 }

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,38 +42,122 @@ 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<T: Copy>(&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<f32> {
     assert!(n >= 3 && (0.0..=0.5).contains(&alpha));
     let mut host = initial_grid(n);
     let bytes = n * size_of::<f32>();
     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<f32> {
     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;
-    }
-}
-