Skip to main content

Buffer

hrx::execution

Struct Buffer 

Source 
pub struct Buffer { /* private fields */ }
Expand description

A cloneable, tracked allocation. Every alias shares lifetime and access state.

Implementations§

Source§

impl Buffer

Source

pub fn len(&self) -> usize

Allocation size in bytes.

Source

pub fn is_empty(&self) -> bool

Whether the allocation is empty. Allocations reject zero lengths.

Source

pub fn view(&self) -> BufferView

Borrow all bytes as an owned device binding.

Examples found in repository?
examples/graph_prepare.rs (line 13)
4fn main() -> hrx::Result<()> {
5    let runtime = Runtime::new()?;
6    let buffer = runtime.allocate(4096, MemoryPlacement::GpuLocal)?;
7    let stream = Stream::open()?;
8    let raw = stream.allocate(4096)?;
9    for count in [128, 512, 1024, 2048] {
10        let begin = Instant::now();
11        let mut graph = runtime.graph();
12        for _ in 0..count {
13            graph.fill(buffer.view(), 0)?;
14        }
15        let graph = graph.prepare()?;
16        let inferred = begin.elapsed();
17        let begin = Instant::now();
18        let mut native = stream.graph()?;
19        let mut last = None;
20        for _ in 0..count {
21            last = Some(native.fill(last.as_slice(), raw.binding(), 0)?);
22        }
23        let _native = native.finish()?;
24        println!(
25            "nodes={count} inferred_us={} explicit_us={}",
26            inferred.as_micros(),
27            begin.elapsed().as_micros()
28        );
29        graph.submit()?.wait()?;
30    }
31    Ok(())
32}
More examples
Hide additional examples
examples/shared_roundtrip.rs (line 49)
8fn main() -> Result<()> {
9    let args: Vec<_> = std::env::args().collect();
10    if args.len() > 2 {
11        return Err(hrx::Error::Message(
12            "usage: shared_roundtrip [bytes]".into(),
13        ));
14    }
15    let bytes = args
16        .get(1)
17        .map(String::as_str)
18        .unwrap_or("1048576")
19        .parse::<usize>()
20        .map_err(|e| hrx::Error::Message(e.to_string()))?;
21    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
22        return Err(hrx::Error::Message(
23            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
24        ));
25    }
26    let runtime = Runtime::new()?;
27    let device = runtime.npu(0)?;
28    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
29        .module(include_str!("../tests/kernels/copy.xdna.loom"))
30        .compile(
31            &hrx::loom::Specialization::new("copy")
32                .with_config("copy.packets", (bytes / 4096).to_string()),
33        )?;
34    let binding = |bytes, access| BindingContract {
35        bytes,
36        access,
37        alignment: 4,
38        layout: "u32 passthrough".into(),
39    };
40    let contract = KernelContract {
41        bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
42        constants: vec![],
43    };
44    let kernel = unsafe { device.load_artifact(&artifact, 1, contract) }?;
45    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
46    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
47    let result = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
48    let mut graph = runtime.graph();
49    graph.fill(a.view(), 0x5a)?;
50    graph.npu(&kernel, &[a.view(), c.view()])?;
51    graph.copy(result.view(), c.view())?;
52    let graph = graph.prepare()?;
53    for _ in 0..10 {
54        graph.submit()?.wait()?;
55        let mapped = result.map_read()?;
56        if mapped.iter().any(|&b| b != 0x5a) {
57            return Err(hrx::Error::Message("GPU/NPU roundtrip mismatch".into()));
58        }
59    }
60    println!(
61        "GPU -> NPU -> GPU: all {bytes} bytes correct, 10 executions, {} prepared regions",
62        graph.regions()
63    );
64    Ok(())
65}
examples/shared_bench.rs (line 52)
9fn main() -> Result<()> {
10    let args: Vec<_> = std::env::args().collect();
11    if args.len() > 2 {
12        return Err(hrx::Error::Message("usage: shared_bench [bytes]".into()));
13    }
14    let bytes = args
15        .get(1)
16        .map(String::as_str)
17        .unwrap_or("1048576")
18        .parse::<usize>()
19        .map_err(|e| hrx::Error::Message(e.to_string()))?;
20    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
21        return Err(hrx::Error::Message(
22            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
23        ));
24    }
25    let runtime = Runtime::new()?;
26    let device = runtime.npu(0)?;
27    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
28        .module(include_str!("../tests/kernels/copy.xdna.loom"))
29        .compile(
30            &hrx::loom::Specialization::new("copy")
31                .with_config("copy.packets", (bytes / 4096).to_string()),
32        )?;
33    let binding = |bytes, access| BindingContract {
34        bytes,
35        access,
36        alignment: 4,
37        layout: "passthrough".into(),
38    };
39    let kernel = unsafe {
40        device.load_artifact(
41            &artifact,
42            1,
43            KernelContract {
44                bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
45                constants: vec![],
46            },
47        )
48    }?;
49    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
50    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
51    let mut graph = runtime.graph();
52    graph.fill(a.view(), 0x35)?;
53    graph.npu(&kernel, &[a.view(), c.view()])?;
54    let graph = graph.prepare()?;
55    for _ in 0..5 {
56        graph.submit()?.wait()?;
57    }
58    let mut samples = Vec::new();
59    for _ in 0..100 {
60        let start = Instant::now();
61        graph.submit()?.wait()?;
62        samples.push(start.elapsed().as_secs_f64());
63    }
64    if c.map_read()?.iter().any(|&b| b != 0x35) {
65        return Err(hrx::Error::Message("incorrect benchmark output".into()));
66    }
67    samples.sort_by(f64::total_cmp);
68    println!(
69        "bytes={bytes} p50_us={:.2} p95_us={:.2}",
70        percentile(&samples, 50) * 1e6,
71        percentile(&samples, 95) * 1e6
72    );
73    Ok(())
74}
examples/gemm_pipeline.rs (line 159)
60fn main() -> Result<()> {
61    let args: Vec<_> = std::env::args().collect();
62    let args = if args.len() == 1 {
63        vec![
64            args[0].clone(),
65            String::new(),
66            "gemm".into(),
67            "8".into(),
68            "8".into(),
69            "8".into(),
70        ]
71    } else {
72        args
73    };
74    if args.len() != 6 {
75        return Err(hrx::Error::Message(
76            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
77        ));
78    }
79    let parse = |i: usize| {
80        args[i]
81            .parse::<usize>()
82            .map_err(|e| hrx::Error::Message(e.to_string()))
83    };
84    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
85    let extent = |a: usize, b: usize, width: usize| {
86        a.checked_mul(b)
87            .filter(|&count| count > 0 && count <= (1 << 30))
88            .and_then(|count| count.checked_mul(width))
89            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
90    };
91    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
92    let runtime = Runtime::new()?;
93    let device = runtime.npu(0)?;
94    let (source, spec) = if args[1].is_empty() {
95        (
96            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
97            hrx::loom::Specialization::new("gemm"),
98        )
99    } else {
100        (
101            std::fs::read_to_string(&args[1])?,
102            hrx::loom::Specialization::new(&args[2])
103                .with_config("gemm.m", m.to_string())
104                .with_config("gemm.k", k.to_string())
105                .with_config("gemm.n", n.to_string()),
106        )
107    };
108    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
109        .module(&source)
110        .compile(&spec)?;
111    let npu = unsafe {
112        device.load_artifact(
113            &artifact,
114            1,
115            KernelContract {
116                bindings: vec![
117                    binding(a_bytes, Access::Read),
118                    binding(b_bytes, Access::Read),
119                    binding(c_bytes, Access::Write),
120                ],
121                constants: vec![],
122            },
123        )
124    }?;
125    let prep = gpu_kernel(
126        &runtime,
127        include_str!("../tests/kernels/euler.loom"),
128        "krea2_euler",
129        &[
130            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
131            ("krea2.euler.grid_y", "1".into()),
132        ],
133        m * k,
134        vec![
135            binding(a_bytes, Access::ReadWrite),
136            binding(a_bytes, Access::Read),
137        ],
138    )?;
139    let epilogue = gpu_kernel(
140        &runtime,
141        include_str!("../tests/kernels/add_f32.loom"),
142        "add_f32",
143        &[("add.grid", (m * n).div_ceil(256).to_string())],
144        m * n,
145        vec![binding(c_bytes, Access::ReadWrite)],
146    )?;
147    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
148    let weights = allocate(b_bytes)?;
149    let ones = allocate(a_bytes)?;
150    for buffer in [&weights, &ones] {
151        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
152            word.copy_from_slice(&0x3f80u16.to_le_bytes());
153        }
154    }
155    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
156        let a = allocate(a_bytes)?;
157        let c = allocate(c_bytes)?;
158        let mut graph = runtime.graph();
159        graph.fill(a.view(), 0)?;
160        graph.gpu(&prep, &[a.view(), ones.view()])?;
161        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
162        graph.gpu(&epilogue, &[c.view()])?;
163        Ok((graph.prepare()?, c))
164    };
165    let (first, c0) = prepare()?;
166    let (second, c1) = prepare()?;
167    for _ in 0..5 {
168        first.submit()?.wait()?;
169        second.submit()?.wait()?;
170    }
171    let before = runtime.statistics();
172    let mut latency = Vec::with_capacity(50);
173    for _ in 0..50 {
174        let start = Instant::now();
175        first.submit()?.wait()?;
176        latency.push(start.elapsed().as_secs_f64());
177    }
178    let start = Instant::now();
179    for _ in 0..50 {
180        let a = first.submit()?;
181        let b = second.submit()?;
182        a.wait()?;
183        b.wait()?;
184    }
185    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
186    for c in [&c0, &c1] {
187        if c.map_read()?
188            .as_chunks::<4>()
189            .0
190            .iter()
191            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
192        {
193            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
194        }
195    }
196    let after = runtime.statistics();
197    assert_eq!(before.allocations, after.allocations);
198    assert_eq!(before.imports, after.imports);
199    assert_eq!(after.copied_bytes, 0);
200    latency.sort_by(f64::total_cmp);
201    println!(
202        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
203        percentile(&latency, 50) * 1e3,
204        percentile(&latency, 95) * 1e3
205    );
206    println!("{}", serde_json::to_string(&after)?);
207    Ok(())
208}
examples/gpu_npu_parallel.rs (line 190)
104fn main() -> Result<()> {
105    let args: Vec<_> = std::env::args().skip(1).collect();
106    if args.len() > 4 || args.iter().any(|arg| arg == "--help") {
107        println!(
108            "usage: gpu_npu_parallel [gpu-elements=16777216] [npu-matrices=1024] [samples=21] [trace.json]"
109        );
110        return if args.len() <= 4 {
111            Ok(())
112        } else {
113            Err(Error::Message("too many arguments".into()))
114        };
115    }
116    let number = |index: usize, default: usize| -> Result<usize> {
117        args.get(index).map_or(Ok(default), |value| {
118            value
119                .parse()
120                .map_err(|_| Error::Message(format!("invalid integer: {value}")))
121        })
122    };
123    let elements = number(0, 1 << 24)?;
124    let matrices = number(1, 1024)?;
125    let samples = number(2, 21)?;
126    if !(256..=1 << 26).contains(&elements)
127        || !elements.is_multiple_of(256)
128        || !(1..=16384).contains(&matrices)
129        || !(3..=1000).contains(&samples)
130    {
131        return Err(Error::Message("gpu-elements must be a multiple of 256 in 256..=67108864; npu-matrices in 1..=16384; samples in 3..=1000".into()));
132    }
133    let trace_path = args.get(3).map(PathBuf::from);
134    let runtime = Runtime::new()?;
135    let gpu = runtime.gpu()?;
136    let npu = runtime.npu(0)?;
137    println!(
138        "GPU: {}; NPU: {}",
139        gpu.target().as_str(),
140        npu.target().as_str()
141    );
142    let gpu_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn slice(&self, range: Range<usize>) -> Result<BufferView>

Create a checked subregion.

Source

pub fn map_read(&self) -> Result<ReadGuard<'_>>

Wait for device writers and map initialized bytes for host reading. Live conflicting host guards return Busy instead of waiting.

Examples found in repository?
examples/gpu_npu_parallel.rs (line 63)
62fn validate(gpu: &Buffer, matrices: &Buffer, pass: usize) -> Result<()> {
63    for (index, word) in gpu.map_read()?.as_chunks::<4>().0.iter().enumerate() {
64        let actual = i32::from_le_bytes(*word);
65        if actual != gpu_value(index, pass) * 3 + 7 {
66            return Err(Error::Message(format!("GPU mismatch at element {index}")));
67        }
68    }
69    for (index, word) in matrices.map_read()?.as_chunks::<4>().0.iter().enumerate() {
70        let matrix = index / 64;
71        let row = index % 64 / 8;
72        let column = index % 8;
73        let expected: f32 = (0..8)
74            .map(|k| {
75                matrix_value(matrix, row * 8 + k, 0, pass)
76                    * matrix_value(matrix, k * 8 + column, 1, pass)
77            })
78            .sum();
79        if f32::from_le_bytes(*word) != expected {
80            return Err(Error::Message(format!(
81                "NPU mismatch at matrix {matrix}, ({row}, {column})"
82            )));
83        }
84    }
85    Ok(())
86}
More examples
Hide additional examples
examples/shared_roundtrip.rs (line 55)
8fn main() -> Result<()> {
9    let args: Vec<_> = std::env::args().collect();
10    if args.len() > 2 {
11        return Err(hrx::Error::Message(
12            "usage: shared_roundtrip [bytes]".into(),
13        ));
14    }
15    let bytes = args
16        .get(1)
17        .map(String::as_str)
18        .unwrap_or("1048576")
19        .parse::<usize>()
20        .map_err(|e| hrx::Error::Message(e.to_string()))?;
21    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
22        return Err(hrx::Error::Message(
23            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
24        ));
25    }
26    let runtime = Runtime::new()?;
27    let device = runtime.npu(0)?;
28    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
29        .module(include_str!("../tests/kernels/copy.xdna.loom"))
30        .compile(
31            &hrx::loom::Specialization::new("copy")
32                .with_config("copy.packets", (bytes / 4096).to_string()),
33        )?;
34    let binding = |bytes, access| BindingContract {
35        bytes,
36        access,
37        alignment: 4,
38        layout: "u32 passthrough".into(),
39    };
40    let contract = KernelContract {
41        bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
42        constants: vec![],
43    };
44    let kernel = unsafe { device.load_artifact(&artifact, 1, contract) }?;
45    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
46    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
47    let result = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
48    let mut graph = runtime.graph();
49    graph.fill(a.view(), 0x5a)?;
50    graph.npu(&kernel, &[a.view(), c.view()])?;
51    graph.copy(result.view(), c.view())?;
52    let graph = graph.prepare()?;
53    for _ in 0..10 {
54        graph.submit()?.wait()?;
55        let mapped = result.map_read()?;
56        if mapped.iter().any(|&b| b != 0x5a) {
57            return Err(hrx::Error::Message("GPU/NPU roundtrip mismatch".into()));
58        }
59    }
60    println!(
61        "GPU -> NPU -> GPU: all {bytes} bytes correct, 10 executions, {} prepared regions",
62        graph.regions()
63    );
64    Ok(())
65}
examples/shared_bench.rs (line 64)
9fn main() -> Result<()> {
10    let args: Vec<_> = std::env::args().collect();
11    if args.len() > 2 {
12        return Err(hrx::Error::Message("usage: shared_bench [bytes]".into()));
13    }
14    let bytes = args
15        .get(1)
16        .map(String::as_str)
17        .unwrap_or("1048576")
18        .parse::<usize>()
19        .map_err(|e| hrx::Error::Message(e.to_string()))?;
20    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
21        return Err(hrx::Error::Message(
22            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
23        ));
24    }
25    let runtime = Runtime::new()?;
26    let device = runtime.npu(0)?;
27    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
28        .module(include_str!("../tests/kernels/copy.xdna.loom"))
29        .compile(
30            &hrx::loom::Specialization::new("copy")
31                .with_config("copy.packets", (bytes / 4096).to_string()),
32        )?;
33    let binding = |bytes, access| BindingContract {
34        bytes,
35        access,
36        alignment: 4,
37        layout: "passthrough".into(),
38    };
39    let kernel = unsafe {
40        device.load_artifact(
41            &artifact,
42            1,
43            KernelContract {
44                bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
45                constants: vec![],
46            },
47        )
48    }?;
49    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
50    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
51    let mut graph = runtime.graph();
52    graph.fill(a.view(), 0x35)?;
53    graph.npu(&kernel, &[a.view(), c.view()])?;
54    let graph = graph.prepare()?;
55    for _ in 0..5 {
56        graph.submit()?.wait()?;
57    }
58    let mut samples = Vec::new();
59    for _ in 0..100 {
60        let start = Instant::now();
61        graph.submit()?.wait()?;
62        samples.push(start.elapsed().as_secs_f64());
63    }
64    if c.map_read()?.iter().any(|&b| b != 0x35) {
65        return Err(hrx::Error::Message("incorrect benchmark output".into()));
66    }
67    samples.sort_by(f64::total_cmp);
68    println!(
69        "bytes={bytes} p50_us={:.2} p95_us={:.2}",
70        percentile(&samples, 50) * 1e6,
71        percentile(&samples, 95) * 1e6
72    );
73    Ok(())
74}
examples/gemm_pipeline.rs (line 187)
60fn main() -> Result<()> {
61    let args: Vec<_> = std::env::args().collect();
62    let args = if args.len() == 1 {
63        vec![
64            args[0].clone(),
65            String::new(),
66            "gemm".into(),
67            "8".into(),
68            "8".into(),
69            "8".into(),
70        ]
71    } else {
72        args
73    };
74    if args.len() != 6 {
75        return Err(hrx::Error::Message(
76            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
77        ));
78    }
79    let parse = |i: usize| {
80        args[i]
81            .parse::<usize>()
82            .map_err(|e| hrx::Error::Message(e.to_string()))
83    };
84    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
85    let extent = |a: usize, b: usize, width: usize| {
86        a.checked_mul(b)
87            .filter(|&count| count > 0 && count <= (1 << 30))
88            .and_then(|count| count.checked_mul(width))
89            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
90    };
91    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
92    let runtime = Runtime::new()?;
93    let device = runtime.npu(0)?;
94    let (source, spec) = if args[1].is_empty() {
95        (
96            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
97            hrx::loom::Specialization::new("gemm"),
98        )
99    } else {
100        (
101            std::fs::read_to_string(&args[1])?,
102            hrx::loom::Specialization::new(&args[2])
103                .with_config("gemm.m", m.to_string())
104                .with_config("gemm.k", k.to_string())
105                .with_config("gemm.n", n.to_string()),
106        )
107    };
108    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
109        .module(&source)
110        .compile(&spec)?;
111    let npu = unsafe {
112        device.load_artifact(
113            &artifact,
114            1,
115            KernelContract {
116                bindings: vec![
117                    binding(a_bytes, Access::Read),
118                    binding(b_bytes, Access::Read),
119                    binding(c_bytes, Access::Write),
120                ],
121                constants: vec![],
122            },
123        )
124    }?;
125    let prep = gpu_kernel(
126        &runtime,
127        include_str!("../tests/kernels/euler.loom"),
128        "krea2_euler",
129        &[
130            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
131            ("krea2.euler.grid_y", "1".into()),
132        ],
133        m * k,
134        vec![
135            binding(a_bytes, Access::ReadWrite),
136            binding(a_bytes, Access::Read),
137        ],
138    )?;
139    let epilogue = gpu_kernel(
140        &runtime,
141        include_str!("../tests/kernels/add_f32.loom"),
142        "add_f32",
143        &[("add.grid", (m * n).div_ceil(256).to_string())],
144        m * n,
145        vec![binding(c_bytes, Access::ReadWrite)],
146    )?;
147    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
148    let weights = allocate(b_bytes)?;
149    let ones = allocate(a_bytes)?;
150    for buffer in [&weights, &ones] {
151        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
152            word.copy_from_slice(&0x3f80u16.to_le_bytes());
153        }
154    }
155    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
156        let a = allocate(a_bytes)?;
157        let c = allocate(c_bytes)?;
158        let mut graph = runtime.graph();
159        graph.fill(a.view(), 0)?;
160        graph.gpu(&prep, &[a.view(), ones.view()])?;
161        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
162        graph.gpu(&epilogue, &[c.view()])?;
163        Ok((graph.prepare()?, c))
164    };
165    let (first, c0) = prepare()?;
166    let (second, c1) = prepare()?;
167    for _ in 0..5 {
168        first.submit()?.wait()?;
169        second.submit()?.wait()?;
170    }
171    let before = runtime.statistics();
172    let mut latency = Vec::with_capacity(50);
173    for _ in 0..50 {
174        let start = Instant::now();
175        first.submit()?.wait()?;
176        latency.push(start.elapsed().as_secs_f64());
177    }
178    let start = Instant::now();
179    for _ in 0..50 {
180        let a = first.submit()?;
181        let b = second.submit()?;
182        a.wait()?;
183        b.wait()?;
184    }
185    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
186    for c in [&c0, &c1] {
187        if c.map_read()?
188            .as_chunks::<4>()
189            .0
190            .iter()
191            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
192        {
193            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
194        }
195    }
196    let after = runtime.statistics();
197    assert_eq!(before.allocations, after.allocations);
198    assert_eq!(before.imports, after.imports);
199    assert_eq!(after.copied_bytes, 0);
200    latency.sort_by(f64::total_cmp);
201    println!(
202        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
203        percentile(&latency, 50) * 1e3,
204        percentile(&latency, 95) * 1e3
205    );
206    println!("{}", serde_json::to_string(&after)?);
207    Ok(())
208}
Source

pub fn try_map_read(&self) -> Result<ReadGuard<'_>>

Map for reading without waiting for device work.

Source

pub fn map_write(&self) -> Result<WriteGuard<'_>>

Wait for device access and map initialized bytes for host mutation.

Examples found in repository?
examples/gpu_npu_parallel.rs (line 38)
36fn initialize(gpu: &Buffer, a: &Buffer, b: &Buffer, pass: usize) -> Result<()> {
37    for (index, word) in gpu
38        .map_write()?
39        .as_chunks_mut::<4>()
40        .0
41        .iter_mut()
42        .enumerate()
43    {
44        word.copy_from_slice(&gpu_value(index, pass).to_le_bytes());
45    }
46    for (operand, buffer) in [a, b].into_iter().enumerate() {
47        for (index, word) in buffer
48            .map_write()?
49            .as_chunks_mut::<2>()
50            .0
51            .iter_mut()
52            .enumerate()
53        {
54            let value = matrix_value(index / 64, index % 64, operand, pass);
55            // Integers in -3..=3 are exactly representable in BF16 and BFP16.
56            word.copy_from_slice(&((value.to_bits() >> 16) as u16).to_le_bytes());
57        }
58    }
59    Ok(())
60}
More examples
Hide additional examples
examples/gemm_pipeline.rs (line 151)
60fn main() -> Result<()> {
61    let args: Vec<_> = std::env::args().collect();
62    let args = if args.len() == 1 {
63        vec![
64            args[0].clone(),
65            String::new(),
66            "gemm".into(),
67            "8".into(),
68            "8".into(),
69            "8".into(),
70        ]
71    } else {
72        args
73    };
74    if args.len() != 6 {
75        return Err(hrx::Error::Message(
76            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
77        ));
78    }
79    let parse = |i: usize| {
80        args[i]
81            .parse::<usize>()
82            .map_err(|e| hrx::Error::Message(e.to_string()))
83    };
84    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
85    let extent = |a: usize, b: usize, width: usize| {
86        a.checked_mul(b)
87            .filter(|&count| count > 0 && count <= (1 << 30))
88            .and_then(|count| count.checked_mul(width))
89            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
90    };
91    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
92    let runtime = Runtime::new()?;
93    let device = runtime.npu(0)?;
94    let (source, spec) = if args[1].is_empty() {
95        (
96            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
97            hrx::loom::Specialization::new("gemm"),
98        )
99    } else {
100        (
101            std::fs::read_to_string(&args[1])?,
102            hrx::loom::Specialization::new(&args[2])
103                .with_config("gemm.m", m.to_string())
104                .with_config("gemm.k", k.to_string())
105                .with_config("gemm.n", n.to_string()),
106        )
107    };
108    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
109        .module(&source)
110        .compile(&spec)?;
111    let npu = unsafe {
112        device.load_artifact(
113            &artifact,
114            1,
115            KernelContract {
116                bindings: vec![
117                    binding(a_bytes, Access::Read),
118                    binding(b_bytes, Access::Read),
119                    binding(c_bytes, Access::Write),
120                ],
121                constants: vec![],
122            },
123        )
124    }?;
125    let prep = gpu_kernel(
126        &runtime,
127        include_str!("../tests/kernels/euler.loom"),
128        "krea2_euler",
129        &[
130            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
131            ("krea2.euler.grid_y", "1".into()),
132        ],
133        m * k,
134        vec![
135            binding(a_bytes, Access::ReadWrite),
136            binding(a_bytes, Access::Read),
137        ],
138    )?;
139    let epilogue = gpu_kernel(
140        &runtime,
141        include_str!("../tests/kernels/add_f32.loom"),
142        "add_f32",
143        &[("add.grid", (m * n).div_ceil(256).to_string())],
144        m * n,
145        vec![binding(c_bytes, Access::ReadWrite)],
146    )?;
147    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
148    let weights = allocate(b_bytes)?;
149    let ones = allocate(a_bytes)?;
150    for buffer in [&weights, &ones] {
151        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
152            word.copy_from_slice(&0x3f80u16.to_le_bytes());
153        }
154    }
155    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
156        let a = allocate(a_bytes)?;
157        let c = allocate(c_bytes)?;
158        let mut graph = runtime.graph();
159        graph.fill(a.view(), 0)?;
160        graph.gpu(&prep, &[a.view(), ones.view()])?;
161        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
162        graph.gpu(&epilogue, &[c.view()])?;
163        Ok((graph.prepare()?, c))
164    };
165    let (first, c0) = prepare()?;
166    let (second, c1) = prepare()?;
167    for _ in 0..5 {
168        first.submit()?.wait()?;
169        second.submit()?.wait()?;
170    }
171    let before = runtime.statistics();
172    let mut latency = Vec::with_capacity(50);
173    for _ in 0..50 {
174        let start = Instant::now();
175        first.submit()?.wait()?;
176        latency.push(start.elapsed().as_secs_f64());
177    }
178    let start = Instant::now();
179    for _ in 0..50 {
180        let a = first.submit()?;
181        let b = second.submit()?;
182        a.wait()?;
183        b.wait()?;
184    }
185    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
186    for c in [&c0, &c1] {
187        if c.map_read()?
188            .as_chunks::<4>()
189            .0
190            .iter()
191            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
192        {
193            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
194        }
195    }
196    let after = runtime.statistics();
197    assert_eq!(before.allocations, after.allocations);
198    assert_eq!(before.imports, after.imports);
199    assert_eq!(after.copied_bytes, 0);
200    latency.sort_by(f64::total_cmp);
201    println!(
202        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
203        percentile(&latency, 50) * 1e3,
204        percentile(&latency, 95) * 1e3
205    );
206    println!("{}", serde_json::to_string(&after)?);
207    Ok(())
208}
Source

pub fn try_map_write(&self) -> Result<WriteGuard<'_>>

Map for mutation without waiting for device work.

Trait Implementations§

Source§

impl Clone for Buffer

Source§

fn clone(&self) -> Self

Returns a duplicate of the value. Read more
1.0.0 (const: unstable) · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> CloneToUninit for T
where T: Clone,

Source§

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)
Performs copy-assignment from self to dest. Read more
Source§

impl<T> DropFlavorWrapper<T> for T

Source§

type Flavor = MayDrop

The DropFlavor that wraps T into Self
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, W> HasTypeWitness<W> for T
where W: MakeTypeWitness<Arg = T>, T: ?Sized,

Source§

const WITNESS: W = W::MAKE

A constant of the type witness
Source§

impl<T> Identity for T
where T: ?Sized,

Source§

const TYPE_EQ: TypeEq<T, <T as Identity>::Type> = TypeEq::NEW

Proof that Self is the same type as Self::Type, provides methods for casting between Self and Self::Type.
Source§

type Type = T

The same type as Self, used to emulate type equality bounds (T == U) with associated type equality constraints (T: Identity<Type = U>).
Source§

impl<T> Instrument for T

Source§

fn instrument(self, span: Span) -> Instrumented<Self> ⓘ

Instruments this type with the provided Span, returning an Instrumented wrapper. Read more
Source§

fn in_current_span(self) -> Instrumented<Self> ⓘ

Instruments this type with the current Span, returning an Instrumented wrapper. Read more
Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> IntoEither for T

Source§

fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ

Converts self into a Left variant of Either<Self, Self> if into_left is true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
where F: FnOnce(&Self) -> bool,

Converts self into a Left variant of Either<Self, Self> if into_left(&self) returns true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

impl<T> PolicyExt for T
where T: ?Sized,

Source§

fn and<P, B, E>(self, other: P) -> And<T, P>
where T: Sized + Policy<B, E>, P: Policy<B, E>,

Create a new Policy that returns Action::Follow only if self and other return Action::Follow. Read more
Source§

fn or<P, B, E>(self, other: P) -> Or<T, P>
where T: Sized + Policy<B, E>, P: Policy<B, E>,

Create a new Policy that returns Action::Follow if either self or other returns Action::Follow. Read more
Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T> ToOwned for T
where T: Clone,

Source§

type Owned = T

The resulting type after obtaining ownership.
Source§

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more
Source§

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<T> WithSubscriber for T

Source§

fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self> ⓘ
where S: Into<Dispatch>,

Attaches the provided Subscriber to this type, returning a WithDispatch wrapper. Read more
Source§

fn with_current_subscriber(self) -> WithDispatch<Self> ⓘ

Attaches the current default Subscriber to this type, returning a WithDispatch wrapper. Read more