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 189)

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_compiler = Compiler::for_target(None, gpu.target())?;
143    let gpu_artifact = gpu_compiler
144        .import_cxx(CxxSource::new(
145            "affine.cpp",
146            format!(
147                r#"
148#include <hip/hip_runtime.h>
149__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
150void affine(const int* input, int* output) {{
151  unsigned i = blockIdx.x * 256u + threadIdx.x;
152  output[i] = input[i] * 3 + 7;
153}}
154"#,
155                elements / 256
156            ),
157        ))?
158        .compile(&Specialization::new("affine"))?;
159    // SAFETY: the fixed launch covers exactly elements i32 values in each
160    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
161    let gpu_kernel = unsafe {
162        runtime.load_gpu_artifact(
163            &gpu_artifact,
164            &[],
165            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
166        )
167    }?;
168    let npu_artifact = Compiler::for_target(None, npu.target())?
169        .module(include_str!("kernels/batched_gemm.xdna.loom"))
170        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
171    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
172    // produces the same number of 8x8 FP32 outputs on one NPU column.
173    let npu_kernel = unsafe {
174        npu.load_artifact(
175            &npu_artifact,
176            1,
177            contract(&[
178                (matrices * 128, Access::Read),
179                (matrices * 128, Access::Read),
180                (matrices * 256, Access::Write),
181            ]),
182        )
183    }?;
184    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
185    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
187    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let mut upload = runtime.graph();
189    upload.copy(input.view(), host_input.view())?;
190    let upload = upload.prepare()?;
191    let mut download = runtime.graph();
192    download.copy(host_output.view(), output.view())?;
193    let download = download.prepare()?;
194    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
195    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
197
198    let mut gpu_only = runtime.graph();
199    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
200    let gpu_only = gpu_only.prepare()?;
201    let mut npu_only = runtime.graph();
202    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
203    let npu_only = npu_only.prepare()?;
204
205    // Disjoint buffers imply no dependency edge. One submit admits both
206    // branches to separate GPU/NPU workers; one wait joins their completions.
207    let mut parallel = runtime.graph();
208    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
209    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
210    let parallel = parallel.prepare()?;
211    for pass in 0..3 {
212        initialize(&host_input, &a, &b, pass * 2)?;
213        upload.submit()?.wait()?;
214        gpu_only.submit()?.wait()?;
215        npu_only.submit()?.wait()?;
216        download.submit()?.wait()?;
217        validate(&host_output, &c, pass * 2)?;
218        // Change both inputs so validation cannot accept the sequential run's
219        // old output if either parallel branch fails to write its result.
220        initialize(&host_input, &a, &b, pass * 2 + 1)?;
221        upload.submit()?.wait()?;
222        parallel.submit()?.wait()?;
223        download.submit()?.wait()?;
224        validate(&host_output, &c, pass * 2 + 1)?;
225    }
226
227    let before = runtime.statistics();
228    let mut sequential_ms = Vec::with_capacity(samples);
229    let mut parallel_ms = Vec::with_capacity(samples);
230    for sample in 0..samples {
231        // Alternate order to reduce systematic clock/thermal ordering bias.
232        for concurrent in if sample % 2 == 0 {
233            [false, true]
234        } else {
235            [true, false]
236        } {
237            if concurrent {
238                parallel_ms.push(completed_ms(&parallel)?);
239            } else {
240                let start = Instant::now();
241                gpu_only.submit()?.wait()?;
242                npu_only.submit()?.wait()?;
243                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
244            }
245        }
246    }
247    let after = runtime.statistics();
248    if before.allocations != after.allocations || before.imports != after.imports {
249        return Err(Error::Message(
250            "replay unexpectedly allocated or imported storage".into(),
251        ));
252    }
253    download.submit()?.wait()?;
254    validate(&host_output, &c, 5)?;
255    let sequential = median(&mut sequential_ms);
256    let concurrent = median(&mut parallel_ms);
257    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
258    println!(
259        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
260        sequential / concurrent
261    );
262    println!(
263        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
264    );
265
266    // Capture separately so tracing does not affect the benchmark samples.
267    runtime.start_trace(32)?;
268    for _ in 0..4 {
269        parallel.submit()?.wait()?;
270    }
271    let trace = runtime
272        .finish_trace()
273        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
274    let overlap: f64 = trace
275        .events
276        .iter()
277        .filter(|e| e.lane == 1)
278        .flat_map(|gpu| {
279            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
280                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
281                    - gpu.start_us.max(npu.start_us))
282                .max(0.0)
283            })
284        })
285        .sum();
286    println!(
287        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
288    );
289    if let Some(path) = trace_path {
290        std::fs::write(&path, trace.to_json()?)?;
291        println!("Trace: {}", path.display());
292    }
293    Ok(())
294}

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§

§

impl UnwindSafe for Buffer

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 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 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 for T
where U: TryFrom<T>,

Source §

type Error = >::Error

The type returned in the event of a conversion error.

Source §

fn try_into(self) -> Result<U, >::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

Struct Buffer Copy item path

Implementations§

impl Buffer

pub fn len(&self) -> usize

pub fn is_empty(&self) -> bool

pub fn view(&self) -> BufferView

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

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

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

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

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

Trait Implementations§

impl Clone for Buffer

fn clone(&self) -> Self

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

Auto Trait Implementations§

impl Freeze for Buffer

impl RefUnwindSafe for Buffer

impl Send for Buffer

impl Sync for Buffer

impl Unpin for Buffer

impl UnsafeUnpin for Buffer

impl UnwindSafe for Buffer

Blanket Implementations§

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

fn type_id(&self) -> TypeId

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

fn borrow(&self) -> &T

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

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

impl<T> CloneToUninit for Twhere T: Clone,

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

impl<T> DropFlavorWrapper<T> for T

type Flavor = MayDrop

impl<T> From<T> for T

fn from(t: T) -> T

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

const WITNESS: W = W::MAKE

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

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

type Type = T

impl<T> Instrument for T

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

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

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

fn into(self) -> U

impl<T> IntoEither for T

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

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

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

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

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

impl<T> Same for T

type Output = T

impl<T> ToOwned for Twhere T: Clone,

type Owned = T

fn to_owned(&self) -> T

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

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

type Error = !

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

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

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

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

impl<T> WithSubscriber for T

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

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

Struct Buffer

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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