unified_probe/
unified_probe.rs

1//! Can one host allocation back both engines at once, with no copy between them?
2//!
3//! Strix Halo puts the gfx1151 iGPU and the XDNA2 NPU on the same physical memory, but
4//! each driver pins pages into its own IOMMU domain: the GPU through KFD, the NPU through
5//! amdxdna's DRM node. Physical sharing is not the question -- whether both drivers will
6//! map the *same* user pages is.
7//!
8//! This allocates page-aligned host memory once, imports it into the GPU runtime
9//! (`hrx_allocator_import_buffer`) and into XRT as a userptr BO, then writes with the GPU
10//! and reads through the NPU's BO. If the pattern survives, a GPU->NPU handoff costs a
11//! cache flush instead of a copy through system memory.
12//!
13//! Usage: unified_probe <xclbin> [--mib N]
14
15use hrx::npu::raw as dvxrt;
16
17use hrx::{Result, Stream};
18use std::alloc::{Layout, alloc_zeroed, dealloc};
19
20const PAGE: usize = 4096;
21/// config3's A argument, the binding a real GEMM would stream from.
22const ARG_A: i32 = 3;
23const GPU_PATTERN: u8 = 0xa5;
24const NPU_PATTERN: u8 = 0x5c;
25
26/// Allocate a GPU buffer, fill it, and export its allocation as a dma-buf descriptor.
27///
28/// Returns the descriptor, its length, and the buffer itself -- which must outlive the fd.
29fn export_dmabuf(
30    pointer: *mut std::ffi::c_void,
31    bytes: usize,
32) -> std::result::Result<(i32, usize), String> {
33    // The bundle's HSA runtime is already resident; this resolves the symbol in it rather
34    // than loading a second copy.
35    let library = hsa_library()?;
36    let export: libloading::Symbol<
37        unsafe extern "C" fn(*const std::ffi::c_void, usize, *mut i32, *mut u64) -> u32,
38    > = unsafe { library.get(b"hsa_amd_portable_export_dmabuf\0") }
39        .map_err(|e| format!("hsa_amd_portable_export_dmabuf: {e}"))?;
40
41    let mut fd = -1i32;
42    let mut offset = 0u64;
43    let status = unsafe { export(pointer.cast_const(), bytes, &mut fd, &mut offset) };
44    if status != 0 {
45        return Err(format!("hsa_amd_portable_export_dmabuf returned {status}"));
46    }
47    if offset != 0 {
48        return Err(format!(
49            "dma-buf carries a nonzero offset ({offset}); unsupported here"
50        ));
51    }
52    Ok((fd, bytes))
53}
54
55fn hsa_library() -> std::result::Result<libloading::Library, String> {
56    let cache = dirs_cache().ok_or("no cache directory")?;
57    let pattern = cache.join("hrx/runtime");
58    let entries = std::fs::read_dir(&pattern).map_err(|e| format!("{}: {e}", pattern.display()))?;
59    for entry in entries.flatten() {
60        let candidate = entry.path().join("libhsa-runtime64.so.1");
61        if candidate.is_file() {
62            return unsafe { libloading::Library::new(&candidate) }
63                .map_err(|e| format!("{}: {e}", candidate.display()));
64        }
65    }
66    Err(format!(
67        "no libhsa-runtime64.so.1 under {}",
68        pattern.display()
69    ))
70}
71
72fn dirs_cache() -> Option<std::path::PathBuf> {
73    std::env::var_os("XDG_CACHE_HOME")
74        .map(std::path::PathBuf::from)
75        .filter(|path| path.is_absolute())
76        .or_else(|| {
77            std::env::var_os("HOME").map(|home| std::path::PathBuf::from(home).join(".cache"))
78        })
79}
80
81fn main() -> Result<()> {
82    let args: Vec<String> = std::env::args().skip(1).collect();
83    let (xclbin, mib) = match args.as_slice() {
84        [xclbin] => (xclbin.clone(), 4usize),
85        [xclbin, flag, value] if flag == "--mib" => (
86            xclbin.clone(),
87            value.parse().expect("--mib expects an integer"),
88        ),
89        _ => {
90            eprintln!("usage: unified_probe <xclbin> [--mib N]");
91            std::process::exit(2);
92        }
93    };
94    let bytes = mib * 1024 * 1024;
95    let layout = Layout::from_size_align(bytes, PAGE).expect("page-aligned layout");
96    let host = unsafe { alloc_zeroed(layout) };
97    assert!(!host.is_null(), "host allocation failed");
98
99    let result = (|| -> Result<()> {
100        println!("host allocation: {mib} MiB at {host:p}");
101
102        let context = match unsafe { dvxrt::Context::new(0, &xclbin) } {
103            Ok(context) => context,
104            Err(error) => {
105                println!("npu context: UNAVAILABLE ({error})");
106                return Ok(());
107            }
108        };
109        let group = context.group_id(ARG_A).expect("group id for A");
110
111        // 1. Can amdxdna pin pages we allocated? Reported either way -- if it can, either
112        //    runtime may own the allocation; if it cannot, only the NPU's allocator may.
113        match unsafe { context.import_bo(host.cast(), bytes, dvxrt::BoKind::HostOnly, group) } {
114            Ok(_) => println!("npu userptr import (host-owned pages): YES"),
115            Err(error) => println!("npu userptr import (host-owned pages): NO ({error})"),
116        }
117
118        // 2. The direction that matters: let the NPU's allocator own the pages -- it is the
119        //    more constrained of the two -- and import its host mapping into the GPU runtime.
120        let bo = context
121            .alloc_bo(bytes, dvxrt::BoKind::HostOnly, group)
122            .expect("allocate an NPU BO");
123        let shared = bo.map().expect("map the NPU BO");
124        println!("npu allocation mapped at {shared:p}");
125        assert_eq!(shared as usize % PAGE, 0, "NPU mapping is not page-aligned");
126
127        let mut stream = Stream::open()?;
128        let imported = unsafe { stream.import_host(shared.cast(), bytes) };
129        match &imported {
130            Ok(_) => println!("gpu import of the NPU mapping: YES"),
131            Err(error) => println!("gpu import of the NPU mapping: NO ({error})"),
132        }
133
134        // 3. If the shared mapping took, check that a GPU write lands in the NPU's BO.
135        if let Ok(buffer) = &imported {
136            stream.fill(buffer.binding(), GPU_PATTERN)?;
137            stream.synchronize()?;
138            bo.sync(false, bytes).expect("invalidate for device reads");
139            let mut seen = vec![0u8; 4096];
140            bo.read(&mut seen).expect("read through the NPU BO");
141            let wrong = seen.iter().filter(|&&b| b != GPU_PATTERN).count();
142            println!(
143                "gpu write -> npu read: {}",
144                if wrong == 0 {
145                    "YES".into()
146                } else {
147                    format!("NO ({wrong} bytes differ)")
148                }
149            );
150        }
151
152        // 3b. The dma-buf route: the GPU allocates, HSA exports the allocation as a
153        //     dma-buf, and XRT imports that descriptor. Neither runtime copies, and unlike
154        //     the host-pointer routes above neither driver has to pin the other's pages.
155        // The GPU-side import above HSA-locked these pages, so they are exportable.
156        let locked = unsafe { stream.import_host(host.cast(), bytes) };
157        let export = match &locked {
158            Ok(buffer) => {
159                stream.fill(buffer.binding(), GPU_PATTERN)?;
160                stream.synchronize()?;
161                export_dmabuf(host.cast(), bytes)
162            }
163            Err(error) => Err(format!("host pages are not GPU-resident: {error}")),
164        };
165        match export {
166            Ok((fd, size)) => {
167                println!("hsa dma-buf export of a GPU buffer: YES (fd {fd})");
168                match unsafe { context.import_dmabuf(fd, size) } {
169                    Ok(imported) => {
170                        println!("npu import of the GPU dma-buf: YES");
171                        let mut seen = vec![0u8; 4096];
172                        match imported.read(&mut seen) {
173                            Ok(()) => {
174                                let wrong = seen.iter().filter(|&&b| b != GPU_PATTERN).count();
175                                println!(
176                                    "  gpu write -> npu read over dma-buf: {}",
177                                    if wrong == 0 {
178                                        "YES".into()
179                                    } else {
180                                        format!("NO ({wrong} bytes differ)")
181                                    }
182                                );
183                            }
184                            Err(error) => {
185                                println!("  read through the imported BO failed: {error}")
186                            }
187                        }
188                    }
189                    Err(error) => println!("npu import of the GPU dma-buf: NO ({error})"),
190                }
191            }
192            Err(error) => println!("hsa dma-buf export of a GPU buffer: NO ({error})"),
193        }
194
195        // 4. And the reverse: the NPU's BO writes, the GPU reads.
196        if let Ok(buffer) = &imported {
197            bo.write(&vec![NPU_PATTERN; bytes])
198                .expect("write through the NPU BO");
199            let mirror = stream.allocate(bytes)?;
200            stream.copy(mirror.binding(), buffer.binding())?;
201            let mut check = [0u8; 64];
202            stream.read_blocking(mirror.binding().slice(0, check.len())?, &mut check)?;
203            let back = check.iter().all(|&b| b == NPU_PATTERN);
204            println!("npu write -> gpu read: {}", if back { "YES" } else { "NO" });
205        }
206        Ok(())
207    })();
208    unsafe { dealloc(host, layout) };
209    result
210}
unified_probe/unified_probe.rs

unified_probe/
unified_probe.rs
