import_probe/
import_probe.rs

1//! Does the GPU write straight into host pages we own? The zero-copy prerequisite.
2//!
3//! Allocates page-aligned host memory, imports it into the runtime with no copy, has the
4//! GPU fill it, and then reads the *original host pointer* back. If the pattern is there,
5//! the device wrote into our pages in place -- which is what makes the same allocation
6//! usable as an XRT BO for the NPU at the same time.
7//!
8//! Usage: import_probe [--mib N]
9
10use hrx::{Result, Stream};
11use std::alloc::{Layout, alloc_zeroed, dealloc};
12
13const PAGE: usize = 4096;
14const PATTERN: u8 = 0xa5;
15
16fn main() -> Result<()> {
17    let mib: usize = match std::env::args().collect::<Vec<_>>().as_slice() {
18        [_] => 16,
19        [_, flag, value] if flag == "--mib" => value.parse().expect("--mib expects an integer"),
20        _ => {
21            eprintln!("usage: import_probe [--mib N]");
22            std::process::exit(2);
23        }
24    };
25    let bytes = mib * 1024 * 1024;
26    let layout = Layout::from_size_align(bytes, PAGE).expect("page-aligned layout");
27
28    // Owned for the whole scope below; the imported buffer never frees it.
29    let host = unsafe { alloc_zeroed(layout) };
30    assert!(!host.is_null(), "host allocation failed");
31    let result = (|| -> Result<()> {
32        let mut stream = Stream::open()?;
33        let buffer = match unsafe { stream.import_host(host.cast(), bytes) } {
34            Ok(buffer) => buffer,
35            Err(error) => {
36                println!("import: UNSUPPORTED ({error})");
37                return Ok(());
38            }
39        };
40        println!("import: ok ({mib} MiB at {host:p})");
41
42        // Prove the host pages start clear, so a match afterwards can only be the GPU.
43        let before = unsafe { std::slice::from_raw_parts(host, bytes) };
44        assert!(before.iter().all(|&b| b == 0), "host memory was not zeroed");
45
46        stream.fill(buffer.binding(), PATTERN)?;
47        stream.synchronize()?;
48
49        // Read the ORIGINAL host pointer -- no download, no staging buffer.
50        let after = unsafe { std::slice::from_raw_parts(host, bytes) };
51        let wrong = after.iter().filter(|&&b| b != PATTERN).count();
52        if wrong == 0 {
53            println!("gpu fill visible on host pointer: YES (zero-copy confirmed)");
54        } else {
55            println!("gpu fill visible on host pointer: NO ({wrong}/{bytes} bytes unwritten)");
56        }
57
58        // A device-side read of the same buffer must also see host writes.
59        let mirror = stream.allocate(bytes)?;
60        unsafe { std::slice::from_raw_parts_mut(host, bytes) }.fill(0x3c);
61        stream.copy(mirror.binding(), buffer.binding())?;
62        let mut check = [0u8; 64];
63        stream.read_blocking(mirror.binding().slice(0, check.len())?, &mut check)?;
64        let coherent = check.iter().all(|&b| b == 0x3c);
65        println!(
66            "host write visible to device: {}",
67            if coherent { "YES" } else { "NO" }
68        );
69        Ok(())
70    })();
71    unsafe { dealloc(host, layout) };
72    result
73}
import_probe/import_probe.rs

import_probe/
import_probe.rs
