dmabuf_probe/
dmabuf_probe.rs

1//! Is the GPU->NPU dma-buf bridge blocked by the driver, or only by libhrx's API surface?
2//!
3//! `hsa_amd_portable_export_dmabuf` refuses locked host pages
4//! (`HSA_STATUS_ERROR_INVALID_ALLOCATION`); it exports allocations from an HSA memory pool.
5//! libhrx will not hand out a device pointer for a device-local buffer, so that path cannot
6//! be driven through the GPU runtime today.
7//!
8//! This goes around libhrx and allocates from an HSA pool directly, purely to answer the
9//! roadmap question: if the export and the XRT import both succeed here, zero-copy needs a
10//! small libhrx addition. If the driver refuses, it needs a driver change.
11//!
12//! Usage: dmabuf_probe <xclbin> [--mib N]
13
14use hrx::npu::raw as dvxrt;
15
16use std::ffi::c_void;
17
18const ARG_A: i32 = 3;
19
20type Status = u32;
21const HSA_STATUS_SUCCESS: Status = 0;
22/// hsa_device_type_t
23const HSA_DEVICE_TYPE_GPU: u32 = 1;
24/// hsa_agent_info_t::HSA_AGENT_INFO_DEVICE
25const HSA_AGENT_INFO_DEVICE: u32 = 17;
26/// hsa_amd_memory_pool_info_t
27const POOL_INFO_SEGMENT: u32 = 0;
28const POOL_INFO_SIZE: u32 = 2;
29const POOL_INFO_ALLOC_ALLOWED: u32 = 5;
30/// hsa_amd_segment_t::HSA_AMD_SEGMENT_GLOBAL
31const SEGMENT_GLOBAL: u32 = 0;
32
33struct Hsa {
34    library: libloading::Library,
35}
36
37impl Hsa {
38    fn open() -> Result<Self, String> {
39        let cache = std::env::var_os("XDG_CACHE_HOME")
40            .map(std::path::PathBuf::from)
41            .filter(|path| path.is_absolute())
42            .or_else(|| {
43                std::env::var_os("HOME").map(|home| std::path::PathBuf::from(home).join(".cache"))
44            })
45            .ok_or("no cache directory")?;
46        let root = cache.join("hrx/runtime");
47        let entry = std::fs::read_dir(&root)
48            .map_err(|e| format!("{}: {e}", root.display()))?
49            .flatten()
50            .map(|entry| entry.path().join("libhsa-runtime64.so.1"))
51            .find(|path| path.is_file())
52            .ok_or_else(|| format!("no libhsa-runtime64.so.1 under {}", root.display()))?;
53        let library =
54            unsafe { libloading::Library::new(&entry) }.map_err(|e| format!("{entry:?}: {e}"))?;
55        Ok(Hsa { library })
56    }
57
58    unsafe fn symbol<T>(&self, name: &[u8]) -> Result<libloading::Symbol<'_, T>, String> {
59        unsafe { self.library.get(name) }
60            .map_err(|e| format!("{}: {e}", String::from_utf8_lossy(name)))
61    }
62}
63
64/// Collected by the agent/pool iteration callbacks below.
65#[derive(Default)]
66struct Found {
67    agent: u64,
68    pool: u64,
69    pool_bytes: usize,
70}
71
72fn main() -> Result<(), String> {
73    let args: Vec<String> = std::env::args().skip(1).collect();
74    let (xclbin, mib) = match args.as_slice() {
75        [xclbin] => (xclbin.clone(), 4usize),
76        [xclbin, flag, value] if flag == "--mib" => (
77            xclbin.clone(),
78            value.parse().map_err(|_| "--mib expects an integer")?,
79        ),
80        _ => return Err("usage: dmabuf_probe <xclbin> [--mib N]".into()),
81    };
82    let bytes = mib * 1024 * 1024;
83    let hsa = Hsa::open()?;
84
85    unsafe {
86        let init: libloading::Symbol<unsafe extern "C" fn() -> Status> =
87            hsa.symbol(b"hsa_init\0")?;
88        let status = init();
89        if status != HSA_STATUS_SUCCESS {
90            return Err(format!("hsa_init returned {status}"));
91        }
92
93        // --- find a GPU agent -------------------------------------------------
94        let iterate_agents: libloading::Symbol<
95            unsafe extern "C" fn(
96                unsafe extern "C" fn(u64, *mut c_void) -> Status,
97                *mut c_void,
98            ) -> Status,
99        > = hsa.symbol(b"hsa_iterate_agents\0")?;
100        let agent_info: libloading::Symbol<unsafe extern "C" fn(u64, u32, *mut c_void) -> Status> =
101            hsa.symbol(b"hsa_agent_get_info\0")?;
102
103        // The callback needs the symbol; pass it through the user-data pointer.
104        struct AgentScan<'a> {
105            info: &'a dyn Fn(u64, u32, *mut c_void) -> Status,
106            found: Found,
107        }
108        unsafe extern "C" fn on_agent(agent: u64, data: *mut c_void) -> Status {
109            let scan = unsafe { &mut *data.cast::<AgentScan>() };
110            let mut kind: u32 = 0;
111            if (scan.info)(agent, HSA_AGENT_INFO_DEVICE, (&raw mut kind).cast())
112                == HSA_STATUS_SUCCESS
113                && kind == HSA_DEVICE_TYPE_GPU
114                && scan.found.agent == 0
115            {
116                scan.found.agent = agent;
117            }
118            HSA_STATUS_SUCCESS
119        }
120        let info_fn = |agent, key, out| agent_info(agent, key, out);
121        let mut scan = AgentScan {
122            info: &info_fn,
123            found: Found::default(),
124        };
125        iterate_agents(on_agent, (&raw mut scan).cast());
126        if scan.found.agent == 0 {
127            return Err("no GPU agent found".into());
128        }
129        println!("gpu agent: 0x{:x}", scan.found.agent);
130
131        // --- find an allocatable global pool ----------------------------------
132        let iterate_pools: libloading::Symbol<
133            unsafe extern "C" fn(
134                u64,
135                unsafe extern "C" fn(u64, *mut c_void) -> Status,
136                *mut c_void,
137            ) -> Status,
138        > = hsa.symbol(b"hsa_amd_agent_iterate_memory_pools\0")?;
139        let pool_info: libloading::Symbol<unsafe extern "C" fn(u64, u32, *mut c_void) -> Status> =
140            hsa.symbol(b"hsa_amd_memory_pool_get_info\0")?;
141
142        struct PoolScan<'a> {
143            info: &'a dyn Fn(u64, u32, *mut c_void) -> Status,
144            found: Found,
145        }
146        unsafe extern "C" fn on_pool(pool: u64, data: *mut c_void) -> Status {
147            let scan = unsafe { &mut *data.cast::<PoolScan>() };
148            if scan.found.pool != 0 {
149                return HSA_STATUS_SUCCESS;
150            }
151            let mut segment: u32 = u32::MAX;
152            let mut allowed: bool = false;
153            let mut size: usize = 0;
154            (scan.info)(pool, POOL_INFO_SEGMENT, (&raw mut segment).cast());
155            (scan.info)(pool, POOL_INFO_ALLOC_ALLOWED, (&raw mut allowed).cast());
156            (scan.info)(pool, POOL_INFO_SIZE, (&raw mut size).cast());
157            if segment == SEGMENT_GLOBAL && allowed && size > 0 {
158                scan.found.pool = pool;
159                scan.found.pool_bytes = size;
160            }
161            HSA_STATUS_SUCCESS
162        }
163        let pool_fn = |pool, key, out| pool_info(pool, key, out);
164        let mut pools = PoolScan {
165            info: &pool_fn,
166            found: Found::default(),
167        };
168        iterate_pools(scan.found.agent, on_pool, (&raw mut pools).cast());
169        if pools.found.pool == 0 {
170            return Err("no allocatable global memory pool on the GPU agent".into());
171        }
172        println!(
173            "gpu memory pool: 0x{:x} ({} MiB)",
174            pools.found.pool,
175            pools.found.pool_bytes / (1024 * 1024)
176        );
177
178        // --- allocate from it and export a dma-buf ----------------------------
179        let allocate: libloading::Symbol<
180            unsafe extern "C" fn(u64, usize, u32, *mut *mut c_void) -> Status,
181        > = hsa.symbol(b"hsa_amd_memory_pool_allocate\0")?;
182        let mut pointer: *mut c_void = std::ptr::null_mut();
183        let status = allocate(pools.found.pool, bytes, 0, &raw mut pointer);
184        if status != HSA_STATUS_SUCCESS {
185            return Err(format!("hsa_amd_memory_pool_allocate returned {status}"));
186        }
187        println!("pool allocation: {mib} MiB at {pointer:p}");
188
189        // The route that would make a shared buffer a first-class hrx buffer: allocate it
190        // through the GPU runtime, take its device pointer, and export that.
191        match hrx::Stream::open() {
192            Ok(mut stream) => match stream.allocate_shared(bytes) {
193                Ok(buffer) => {
194                    stream.fill(buffer.binding(), 0x11).ok();
195                    stream.synchronize().ok();
196                    match buffer.device_ptr() {
197                        Ok(hrx_pointer) => {
198                            println!("hrx shared allocation: ok, device ptr {hrx_pointer:p}");
199                            let mut fd: i32 = -1;
200                            let mut offset: u64 = 0;
201                            let export: libloading::Symbol<
202                                unsafe extern "C" fn(
203                                    *const c_void,
204                                    usize,
205                                    *mut i32,
206                                    *mut u64,
207                                ) -> Status,
208                            > = hsa.symbol(b"hsa_amd_portable_export_dmabuf\0")?;
209                            let status = export(hrx_pointer, bytes, &raw mut fd, &raw mut offset);
210                            if status != HSA_STATUS_SUCCESS {
211                                println!("  dma-buf export of an hrx buffer: NO (status {status})");
212                            } else {
213                                println!(
214                                    "  dma-buf export of an hrx buffer: YES (fd {fd}, offset {offset})"
215                                );
216                                match dvxrt::Context::new(0, &xclbin) {
217                                    Ok(context) => match context.import_dmabuf(fd, bytes) {
218                                        Ok(bo) => {
219                                            let mut seen = vec![0u8; 4096];
220                                            bo.read(&mut seen).ok();
221                                            let wrong = seen.iter().filter(|&&b| b != 0x11).count();
222                                            println!(
223                                                "  npu reads the hrx buffer: {}",
224                                                if wrong == 0 {
225                                                    "YES -- one buffer, both engines".into()
226                                                } else {
227                                                    format!("NO ({wrong}/4096 differ)")
228                                                }
229                                            );
230                                        }
231                                        Err(error) => println!("  npu import: NO ({error})"),
232                                    },
233                                    Err(error) => println!("  npu context: {error}"),
234                                }
235                            }
236                        }
237                        Err(error) => println!("hrx shared allocation device ptr: NO ({error})"),
238                    }
239                }
240                Err(error) => println!("hrx shared allocation: NO ({error})"),
241            },
242            Err(error) => println!("hrx stream unavailable: {error}"),
243        }
244
245        let export: libloading::Symbol<
246            unsafe extern "C" fn(*const c_void, usize, *mut i32, *mut u64) -> Status,
247        > = hsa.symbol(b"hsa_amd_portable_export_dmabuf\0")?;
248        let mut fd: i32 = -1;
249        let mut offset: u64 = 0;
250        let status = export(pointer, bytes, &raw mut fd, &raw mut offset);
251        if status != HSA_STATUS_SUCCESS {
252            println!("hsa dma-buf export of pool memory: NO (status {status})");
253            return Ok(());
254        }
255        println!("hsa dma-buf export of pool memory: YES (fd {fd}, offset {offset})");
256
257        // --- and hand that descriptor to the NPU ------------------------------
258        match dvxrt::Context::new(0, &xclbin) {
259            Ok(context) => {
260                let _group = context.group_id(ARG_A).map_err(|e| e.to_string())?;
261                match context.import_dmabuf(fd, bytes) {
262                    Ok(imported) => {
263                        println!("npu import of the GPU dma-buf: YES");
264                        // Import alone only proves the descriptor was accepted. Write with
265                        // the GPU and read through the NPU's BO to prove it is one buffer.
266                        let fill: libloading::Symbol<
267                            unsafe extern "C" fn(*mut c_void, u32, usize) -> Status,
268                        > = hsa.symbol(b"hsa_amd_memory_fill\0")?;
269                        let pattern = 0xa5a5_a5a5u32;
270                        let status = fill(pointer, pattern, bytes / 4);
271                        if status != HSA_STATUS_SUCCESS {
272                            println!("  gpu fill of pool memory failed (status {status})");
273                            return Ok(());
274                        }
275                        let mut seen = vec![0u8; 4096];
276                        match imported.read(&mut seen) {
277                            Ok(()) => {
278                                let wrong = seen.iter().filter(|&&b| b != 0xa5).count();
279                                if wrong == 0 {
280                                    println!(
281                                        "  gpu write -> npu read over dma-buf: YES -- zero-copy confirmed"
282                                    );
283                                } else {
284                                    println!(
285                                        "  gpu write -> npu read over dma-buf: NO ({wrong}/4096 bytes differ)"
286                                    );
287                                }
288                            }
289                            Err(error) => {
290                                println!("  read through the imported BO failed: {error}")
291                            }
292                        }
293                    }
294                    Err(error) => println!("npu import of the GPU dma-buf: NO ({error})"),
295                }
296            }
297            Err(error) => println!("npu context: UNAVAILABLE ({error})"),
298        }
299    }
300    Ok(())
301}
dmabuf_probe/dmabuf_probe.rs

dmabuf_probe/
dmabuf_probe.rs
