Method load

cake-core/src/backends/rocm/ffi.rs:65–109

()

    /// Searches LD_LIBRARY_PATH, then default system paths.
    #[allow(clippy::missing_transmute_annotations)]
    pub fn load() -> std::result::Result<Self, String> {
        unsafe {
            let hip_lib = libloading::Library::new("libamdhip64.so")
                .or_else(|_| libloading::Library::new("libamdhip64.so.5"))
                .map_err(|e| format!("cannot load libamdhip64: {e}"))?;

            let blas_lib = libloading::Library::new("librocblas.so")
                .or_else(|_| libloading::Library::new("librocblas.so.0"))
                .map_err(|e| format!("cannot load librocblas: {e}"))?;

            macro_rules! sym {
                ($lib:expr, $name:expr) => {{
                    let raw: libloading::Symbol<*const std::ffi::c_void> = $lib.get($name)
                        .map_err(|e| format!("symbol {}: {e}", String::from_utf8_lossy($name)))?;
                    std::mem::transmute(*raw)
                }};
            }

            Ok(Self {
                hip_init: sym!(hip_lib, b"hipInit\0"),
                hip_get_device_count: sym!(hip_lib, b"hipGetDeviceCount\0"),
                hip_set_device: sym!(hip_lib, b"hipSetDevice\0"),
                hip_device_get_name: sym!(hip_lib, b"hipDeviceGetName\0"),
                hip_malloc: sym!(hip_lib, b"hipMalloc\0"),
                hip_free: sym!(hip_lib, b"hipFree\0"),
                hip_host_malloc: sym!(hip_lib, b"hipHostMalloc\0"),
                hip_host_free: sym!(hip_lib, b"hipHostFree\0"),
                hip_memcpy: sym!(hip_lib, b"hipMemcpy\0"),
                hip_mem_get_info: sym!(hip_lib, b"hipMemGetInfo\0"),
                hip_device_synchronize: sym!(hip_lib, b"hipDeviceSynchronize\0"),
                hip_stream_create: sym!(hip_lib, b"hipStreamCreate\0"),
                hip_stream_synchronize: sym!(hip_lib, b"hipStreamSynchronize\0"),
                hip_memcpy_async: sym!(hip_lib, b"hipMemcpyAsync\0"),

                rocblas_create_handle: sym!(blas_lib, b"rocblas_create_handle\0"),
                rocblas_destroy_handle: sym!(blas_lib, b"rocblas_destroy_handle\0"),
                rocblas_set_stream: sym!(blas_lib, b"rocblas_set_stream\0"),
                rocblas_sgemm: sym!(blas_lib, b"rocblas_sgemm\0"),
                rocblas_sgemm_strided_batched: sym!(blas_lib, b"rocblas_sgemm_strided_batched\0"),

                _hip_lib: hip_lib,
                _blas_lib: blas_lib,
            })
        }
    }
}