hrx::gpu

Struct Stream

pub struct Stream { /* private fields */ }

Expand description

An ordered GPU stream with explicit prepared commands and bounded pools.

Implementations§

Source §

impl Stream

Source

pub fn access_graph(&self) -> Result<AccessGraph<'_>>

Begin an access-aware graph recording on this stream.

Source §

Allocate an owned copy of a nonempty byte slice and publish it to the GPU. Fresh backing is initialized directly, without a preliminary zero fill. This does not wait for unrelated stream work; no existing buffer changes.

Source

pub fn allocate_zeroed(&self, bytes: usize) -> Result<Buffer>

Allocate initialized storage containing zero bytes.

Source

pub fn allocate_shared(&self, bytes: usize) -> Result<Buffer>

Allocate GPU-coherent host storage for direct, synchronized host access.

Source

pub fn synchronize(&mut self) -> Result<()>

Wait for all preceding native work and reclaim transfer staging.

Examples found in repository ?

examples/queue_pool_bench.rs (line 22)

16    fn new(device: &Device, bytes: usize) -> Result<Self> {
17        let mut stream = device.stream()?;
18        let source = stream.allocate(bytes)?;
19        let destination = stream.allocate(bytes)?;
20        stream.fill(source.binding(), 0x3c)?;
21        stream.copy(destination.binding(), source.binding())?;
22        stream.synchronize()?;
23        Ok(Self {
24            stream,
25            source,
26            destination,
27        })
28    }
29    fn enqueue(&self) -> Result<()> {
30        self.stream
31            .copy(self.destination.binding(), self.source.binding())
32    }
33    fn check(&mut self) -> Result<()> {
34        let mut bytes = [0; 64];
35        self.stream
36            .read_blocking(self.destination.slice(0, 64), &mut bytes)?;
37        assert_eq!(bytes, [0x3c; 64]);
38        Ok(())
39    }
40}
41
42fn milliseconds(start: Instant) -> f64 {
43    start.elapsed().as_secs_f64() * 1_000.0
44}
45
46fn throughput(lanes: &mut [Lane], samples: usize, batch: usize) -> Result<Distribution> {
47    let mut timings = Vec::new();
48    for round in 0..samples + WARMUP {
49        let start = Instant::now();
50        for _ in 0..batch {
51            for lane in &*lanes {
52                lane.enqueue()?;
53            }
54        }
55        for lane in &mut *lanes {
56            lane.stream.synchronize()?;
57        }
58        if round >= WARMUP {
59            timings.push(milliseconds(start));
60        }
61    }
62    Distribution::from_samples(timings)
63}
64
65fn concurrent(lanes: &mut [Lane], samples: usize, batch: usize) -> Result<Distribution> {
66    let mut timings = Vec::new();
67    let barrier = Barrier::new(lanes.len());
68    for round in 0..samples + WARMUP {
69        // Creation and joining are excluded from each worker's elapsed time.
70        // The slowest worker determines the batch's throughput.
71        let elapsed = std::thread::scope(|scope| {
72            let workers: Vec<_> = lanes
73                .iter_mut()
74                .map(|lane| {
75                    let barrier = &barrier;
76                    scope.spawn(move || -> Result<f64> {
77                        barrier.wait();
78                        let start = Instant::now();
79                        for _ in 0..batch {
80                            lane.enqueue()?;
81                        }
82                        lane.stream.synchronize()?;
83                        Ok(milliseconds(start))
84                    })
85                })
86                .collect();
87            workers
88                .into_iter()
89                .map(|worker| worker.join().expect("benchmark worker panicked"))
90                .collect::<Result<Vec<_>>>()
91        })?;
92        if round >= WARMUP {
93            timings.push(elapsed.into_iter().fold(0.0, f64::max));
94        }
95    }
96    Distribution::from_samples(timings)
97}
98
99fn main() -> Result<()> {
100    let args: Vec<_> = std::env::args().skip(1).collect();
101    let parse = |value: &str| value.parse::<usize>().ok().filter(|&n| n > 0);
102    let queues = args.first().and_then(|arg| parse(arg));
103    let samples = args.get(1).map_or(Some(15), |arg| parse(arg));
104    let (queues, samples) = match (queues, samples, args.len()) {
105        (Some(queues), Some(samples), 1..=2) => (queues, samples),
106        _ => {
107            return Err(hrx::Error::Message(
108                "usage: queue_pool_bench QUEUES [SAMPLES] (positive integers)".into(),
109            ));
110        }
111    };
112    let device = Device::open(0)?;
113    device.set_stream_queue_count(queues)?;
114    let start = Instant::now();
115    let streams = (0..STREAMS)
116        .map(|_| device.stream())
117        .collect::<Result<Vec<_>>>()?;
118    let creation_ms = milliseconds(start);
119    // Keep these streams alive so every case uses the same already-created queues.
120    let mut small = (0..STREAMS)
121        .map(|_| Lane::new(&device, 4096))
122        .collect::<Result<Vec<_>>>()?;
123    let small_batch = 1024;
124    let small_times = throughput(&mut small, samples, small_batch)?;
125    let threaded_times = concurrent(&mut small, samples, small_batch)?;
126    let mut idle_latency = Vec::new();
127    for round in 0..samples + WARMUP {
128        for lane in &mut small {
129            let start = Instant::now();
130            lane.enqueue()?;
131            lane.stream.synchronize()?;
132            if round >= WARMUP {
133                idle_latency.push(milliseconds(start));
134            }
135        }
136    }
137    // A 1 GiB aggregate working set exceeds the on-device caches on gfx1151.
138    let bulk_bytes = 32 * 1024 * 1024;
139    let bulk_batch = 16;
140    let mut bulk = (0..STREAMS)
141        .map(|_| Lane::new(&device, bulk_bytes))
142        .collect::<Result<Vec<_>>>()?;
143    let bulk_times = throughput(&mut bulk, samples, bulk_batch)?;
144    let mut loaded_latency = Vec::new();
145    for round in 0..samples + WARMUP {
146        // Probe every queue position. Enqueue a bulk backlog on all streams,
147        // then measure one small copy's completion before draining the backlog.
148        for probe in &mut small {
149            for _ in 0..4 {
150                for lane in &bulk {
151                    lane.enqueue()?;
152                }
153            }
154            let start = Instant::now();
155            probe.enqueue()?;
156            probe.stream.synchronize()?;
157            if round >= WARMUP {
158                loaded_latency.push(milliseconds(start));
159            }
160            for lane in &mut bulk {
161                lane.stream.synchronize()?;
162            }
163        }
164    }
165    for lane in small.iter_mut().chain(bulk.iter_mut()) {
166        lane.check()?;
167    }
168    let idle = Distribution::from_samples(idle_latency)?;
169    let loaded = Distribution::from_samples(loaded_latency)?;
170    println!(
171        "{}",
172        json!({
173            "target": device.target().as_str(), "queues": queues, "streams": STREAMS,
174            "samples": samples, "warmup": WARMUP, "stream_creation_ms": creation_ms,
175            "small_commands_per_second": STREAMS as f64 * small_batch as f64 * 1_000.0 / small_times.median_ms,
176            "threaded_commands_per_second": STREAMS as f64 * small_batch as f64 * 1_000.0 / threaded_times.median_ms,
177            "bulk_copy_gb_s": 2.0 * bulk_bytes as f64 * STREAMS as f64 * bulk_batch as f64 / (bulk_times.median_ms * 1e6),
178            "small_batch_median_ms": small_times.median_ms, "small_batch_p95_ms": small_times.p95_ms,
179            "threaded_batch_median_ms": threaded_times.median_ms, "threaded_batch_p95_ms": threaded_times.p95_ms,
180            "bulk_batch_median_ms": bulk_times.median_ms, "bulk_batch_p95_ms": bulk_times.p95_ms,
181            "idle_latency_median_us": idle.median_ms * 1_000.0, "idle_latency_p95_us": idle.p95_ms * 1_000.0,
182            "loaded_latency_median_us": loaded.median_ms * 1_000.0, "loaded_latency_p95_us": loaded.p95_ms * 1_000.0,
183        })
184    );
185    drop(streams);
186    Ok(())
187}

More examples

Hide additional examples

examples/stream_bench.rs (line 14)

5fn measure(
6    stream: &mut Stream,
7    samples: usize,
8    operations: usize,
9    mut run: impl FnMut(&mut Stream) -> Result<()>,
10) -> Result<(f64, f64)> {
11    let mut enqueue = Vec::new();
12    let mut complete = Vec::new();
13    for sample in 0..samples + 3 {
14        stream.synchronize()?;
15        let start = Instant::now();
16        run(stream)?;
17        let recorded = start.elapsed();
18        stream.synchronize()?;
19        let finished = start.elapsed();
20        if sample >= 3 {
21            enqueue.push(recorded.as_secs_f64() * 1e9 / operations as f64);
22            complete.push(finished.as_secs_f64() * 1e9 / operations as f64);
23        }
24    }
25    enqueue.sort_by(f64::total_cmp);
26    complete.sort_by(f64::total_cmp);
27    Ok((enqueue[samples / 2], complete[samples / 2]))
28}
29
30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

examples/gpu_load.rs (line 103)

88fn main() -> Result<()> {
89    let options = match parse() {
90        Ok(options) => options,
91        Err(message) => {
92            eprintln!("gpu_load: {message}");
93            std::process::exit(2);
94        }
95    };
96    let size = options.mib * 1024 * 1024;
97    let mut stream = Stream::open()?;
98    let source = stream.allocate(size)?;
99    let destination = stream.allocate(size)?;
100    // Touch both allocations so the steady-state loop never pays first-use faults.
101    stream.fill(source.binding(), 0x3c)?;
102    stream.fill(destination.binding(), 0)?;
103    stream.synchronize()?;
104
105    let run_batch = |stream: &mut Stream| -> Result<()> {
106        for _ in 0..options.batch {
107            match options.mode {
108                Mode::Copy => stream.copy(destination.binding(), source.binding())?,
109                Mode::Fill => stream.fill(destination.binding(), 0x5a)?,
110            }
111        }
112        stream.synchronize()
113    };
114
115    // Warm up outside the measured window: the first batch pays queue and paging costs.
116    run_batch(&mut stream)?;
117
118    let window = Duration::from_secs_f64(options.seconds);
119    let start = Instant::now();
120    let mut iterations = 0usize;
121    // Per-batch rates, so the harness can see whether contention was steady or bursty.
122    let mut batch_rates = Vec::new();
123    while start.elapsed() < window {
124        let batch_start = Instant::now();
125        run_batch(&mut stream)?;
126        let seconds = batch_start.elapsed().as_secs_f64();
127        iterations += options.batch;
128        let bytes = options.mode.traffic(size) * options.batch;
129        batch_rates.push(bytes as f64 / seconds / 1e9);
130    }
131    let elapsed = start.elapsed().as_secs_f64();
132    let total = options.mode.traffic(size) as f64 * iterations as f64;
133    let mean = total / elapsed / 1e9;
134    batch_rates.sort_by(f64::total_cmp);
135    let median = batch_rates[batch_rates.len() / 2];
136    let low = batch_rates[batch_rates.len() / 20];
137
138    println!(
139        concat!(
140            r#"{{"device":"gpu","mode":"{}","buffer_mib":{},"batch":{},"#,
141            r#""seconds":{:.4},"iterations":{},"gb_s_mean":{:.3},"#,
142            r#""gb_s_median":{:.3},"gb_s_p5":{:.3}}}"#
143        ),
144        options.mode.name(),
145        options.mib,
146        options.batch,
147        elapsed,
148        iterations,
149        mean,
150        median,
151        low
152    );
153    eprintln!(
154        "gpu_load: {} {:.1} GB/s mean over {:.1}s ({} iterations)",
155        options.mode.name(),
156        mean,
157        elapsed,
158        iterations
159    );
160    Ok(())
161}

examples/patch_bench.rs (line 403)

348fn main() -> Result<()> {
349    let args: Vec<_> = std::env::args().collect();
350    if args.len() < 3 || args.len() > 5 {
351        return Err(hrx::Error::Message(
352            "usage: patch_bench BASELINE_LIB CANDIDATE_LIB [pairs=30] [default|cu|wgp]".into(),
353        ));
354    }
355    let pairs = args
356        .get(3)
357        .map(|s| s.parse::<usize>())
358        .transpose()
359        .map_err(|e| hrx::Error::Message(e.to_string()))?
360        .unwrap_or(30);
361    if pairs < 30 {
362        return Err(hrx::Error::Message("at least 30 pairs are required".into()));
363    }
364    let mode = match args.get(4).map(String::as_str).unwrap_or("default") {
365        "default" => ProcessorMode::Default,
366        "cu" => ProcessorMode::ComputeUnit,
367        "wgp" => ProcessorMode::WorkgroupProcessor,
368        _ => {
369            return Err(hrx::Error::Message(
370                "mode must be default, cu or wgp".into(),
371            ));
372        }
373    };
374    let mut stream = Stream::open()?;
375    let mut options = CompilerOptions::for_target(stream.target());
376    options.processor_mode = mode;
377    let compilers = [
378        Compiler::with_options(Some(Path::new(&args[1])), options.clone())?,
379        Compiler::with_options(Some(Path::new(&args[2])), options)?,
380    ];
381    for case in [
382        attention(128, false),
383        attention(512, false),
384        gemm(false, false),
385        gemm(true, false),
386        gemm(false, true),
387        gemm(true, true),
388        conv(56, 64, 64),
389        conv(14, 256, 256),
390        conv(7, 512, 512),
391    ] {
392        let baseline = prepare(&mut stream, &compilers[0], &case, None)?;
393        let candidate = prepare(
394            &mut stream,
395            &compilers[1],
396            &case,
397            Some(baseline.buffers.clone()),
398        )?;
399        let mut arms = [baseline, candidate];
400        for arm in &mut arms {
401            for _ in 0..5 {
402                stream.launch(&mut arm.graph)?;
403                stream.synchronize()?;
404            }
405            check(&mut stream, arm, &case)?;
406        }
407        let errors = [
408            check(&mut stream, &arms[0], &case)?,
409            check(&mut stream, &arms[1], &case)?,
410        ];
411        println!(
412            "{}",
413            serde_json::json!({"case":case.name,"mode":format!("{mode:?}"),"baseline":arms[0].identity,"candidate":arms[1].identity,"max_absolute_error":errors,"tolerance":case.tolerance,"replays":32})
414        );
415        for pair in 0..pairs {
416            let mut ns = [0.0; 2];
417            for index in if pair % 2 == 0 { [0, 1] } else { [1, 0] } {
418                let start = Instant::now();
419                stream.launch(&mut arms[index].graph)?;
420                stream.synchronize()?;
421                ns[index] = start.elapsed().as_secs_f64() * 1e9 / 32.0;
422            }
423            println!(
424                "{}",
425                serde_json::json!({"case":case.name,"pair":pair,"baseline_ns":ns[0],"candidate_ns":ns[1]})
426            );
427        }
428        for arm in &mut arms {
429            stream.launch(&mut arm.graph)?;
430            stream.synchronize()?;
431            check(&mut stream, arm, &case)?;
432        }
433    }
434    Ok(())
435}

Source

pub fn record_event(&mut self) -> Result<Event>

Snapshot preceding published work without a host wait. A submission still waiting for shared queue capacity is not included.

Source

pub fn wait_event(&mut self, event: &Event) -> Result<()>

Order subsequent stream work after an immutable event. Events refer only to already-published native commands; they cannot name future work from a scheduler callback or an unlaunched graph.

Source

pub fn submit(&mut self) -> Result<Submission<'_>>

Observe the current native submission prefix.

Source

pub fn upload_blocking(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()>

Write a checked range after draining this stream.

Examples found in repository ?

examples/patch_bench.rs (line 296)

279fn prepare(
280    stream: &mut Stream,
281    compiler: &Compiler,
282    case: &Case,
283    shared: Option<std::sync::Arc<Vec<Buffer>>>,
284) -> Result<Arm> {
285    let artifact = compiler.module(case.source).compile(&case.spec)?;
286    let kernel = unsafe { stream.load_artifact(&artifact) }?;
287    let constants = Constants::indices(&kernel, &case.indices)?;
288    let buffers = if let Some(buffers) = shared {
289        buffers
290    } else {
291        std::sync::Arc::new(
292            case.data
293                .iter()
294                .map(|data| {
295                    let buffer = stream.allocate(data.len())?;
296                    stream.upload_blocking(buffer.binding(), data)?;
297                    Ok(buffer)
298                })
299                .collect::<Result<Vec<_>>>()?,
300        )
301    };
302    let bindings: Vec<_> = buffers.iter().map(Buffer::binding).collect();
303    let mut graph = stream.graph()?;
304    let mut prior = None;
305    for _ in 0..32 {
306        prior = Some(unsafe {
307            graph.dispatch(
308                prior.as_slice(),
309                &kernel,
310                case.grid,
311                case.block,
312                &constants,
313                &bindings,
314            )
315        }?);
316    }
317    let identity = serde_json::json!({"compiler":artifact.compiler_identity(), "artifact":hrx::bundle::digest(artifact.bytes()),
318        "report":artifact.report().map(|report| report.json())});
319    Ok(Arm {
320        graph: graph.finish()?,
321        buffers,
322        identity,
323    })
324}

More examples

Hide additional examples

examples/prepare_weights.rs (line 40)

6fn main() -> Result<(), Box<dyn std::error::Error>> {
7    let mut args: Vec<_> = std::env::args().skip(1).collect();
8    let upload = args.last().is_some_and(|arg| arg == "--upload");
9    if upload {
10        args.pop();
11    }
12    if args.len() != 3 || !matches!(args[2].as_str(), "none" | "advice" | "populate") {
13        return Err("usage: prepare_weights FILE TENSOR none|advice|populate [--upload]".into());
14    }
15    // The command requires an immutable checkpoint for the mapping's lifetime.
16    let file = unsafe { FileView::map(&args[0]) }?;
17    let tensor = file.get(&args[1])?;
18    let mut stream = if upload {
19        Some(hrx::Stream::open()?)
20    } else {
21        None
22    };
23    let buffer = stream
24        .as_ref()
25        .map(|stream| stream.allocate(tensor.bytes.len()))
26        .transpose()?;
27    let start = Instant::now();
28    let prepared = match args[2].as_str() {
29        "populate" => format!("{:?}", file.prepare_bytes(tensor.bytes)?),
30        "advice" => {
31            file.will_need(tensor);
32            "Advised".into()
33        }
34        _ => "None".into(),
35    };
36    let preparation_ms = start.elapsed().as_secs_f64() * 1000.;
37    let checksum: u64;
38    let total_ms;
39    if let (Some(stream), Some(buffer)) = (&mut stream, &buffer) {
40        stream.upload_blocking(buffer.binding(), tensor.bytes)?;
41        total_ms = start.elapsed().as_secs_f64() * 1000.;
42        // Validation is outside timing, after the upload fence has retired.
43        let mut actual = vec![0; tensor.bytes.len()];
44        stream.read_blocking(buffer.binding(), &mut actual)?;
45        if actual != tensor.bytes {
46            return Err("upload bytes differ from checkpoint".into());
47        }
48        checksum = actual.iter().map(|&v| u64::from(v)).sum();
49    } else {
50        checksum = tensor.bytes.iter().map(|&v| u64::from(v)).sum();
51        total_ms = start.elapsed().as_secs_f64() * 1000.;
52    }
53    println!(
54        "{}",
55        serde_json::json!({"mode":args[2], "preparation":prepared,
56        "bytes":tensor.bytes.len(), "preparation_ms":preparation_ms,
57        "total_ms":total_ms, "checksum":checksum, "upload":upload,
58        "allocation_policy":"preallocated destination; fresh process",
59        "timed_scope":if upload { "page preparation plus completed upload" } else { "page preparation plus complete host read" }, "cache_state":"externally controlled"})
60    );
61    Ok(())
62}

examples/stream_bench.rs (line 87)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn upload_blocking_at( &mut self, dst: &Buffer, offset: usize, bytes: &[u8], ) -> Result<()>

Write at an allocation-relative byte offset.

Source

pub fn read_blocking(&mut self, src: View<'_>, bytes: &mut [u8]) -> Result<()>

Read a checked range after draining this stream.

Examples found in repository ?

examples/queue_pool_bench.rs (line 36)

33    fn check(&mut self) -> Result<()> {
34        let mut bytes = [0; 64];
35        self.stream
36            .read_blocking(self.destination.slice(0, 64), &mut bytes)?;
37        assert_eq!(bytes, [0x3c; 64]);
38        Ok(())
39    }

More examples

Hide additional examples

examples/patch_bench.rs (line 327)

325fn check(stream: &mut Stream, arm: &Arm, case: &Case) -> Result<f64> {
326    let mut bytes = vec![0; case.expected.len() * 2];
327    stream.read_blocking(arm.buffers[case.output].binding(), &mut bytes)?;
328    let mut maximum = 0f64;
329    for (index, (value, expected)) in bytes
330        .as_chunks::<2>()
331        .0
332        .iter()
333        .zip(&case.expected)
334        .enumerate()
335    {
336        let value = float(u16::from_le_bytes(*value));
337        let error = (value - expected).abs();
338        if !value.is_finite() || error > case.tolerance {
339            return Err(hrx::Error::Message(format!(
340                "{} output {index}: {value} vs {expected}, error {error}, limit {}",
341                case.name, case.tolerance
342            )));
343        }
344        maximum = maximum.max(error);
345    }
346    Ok(maximum)
347}

examples/prepare_weights.rs (line 44)

6fn main() -> Result<(), Box<dyn std::error::Error>> {
7    let mut args: Vec<_> = std::env::args().skip(1).collect();
8    let upload = args.last().is_some_and(|arg| arg == "--upload");
9    if upload {
10        args.pop();
11    }
12    if args.len() != 3 || !matches!(args[2].as_str(), "none" | "advice" | "populate") {
13        return Err("usage: prepare_weights FILE TENSOR none|advice|populate [--upload]".into());
14    }
15    // The command requires an immutable checkpoint for the mapping's lifetime.
16    let file = unsafe { FileView::map(&args[0]) }?;
17    let tensor = file.get(&args[1])?;
18    let mut stream = if upload {
19        Some(hrx::Stream::open()?)
20    } else {
21        None
22    };
23    let buffer = stream
24        .as_ref()
25        .map(|stream| stream.allocate(tensor.bytes.len()))
26        .transpose()?;
27    let start = Instant::now();
28    let prepared = match args[2].as_str() {
29        "populate" => format!("{:?}", file.prepare_bytes(tensor.bytes)?),
30        "advice" => {
31            file.will_need(tensor);
32            "Advised".into()
33        }
34        _ => "None".into(),
35    };
36    let preparation_ms = start.elapsed().as_secs_f64() * 1000.;
37    let checksum: u64;
38    let total_ms;
39    if let (Some(stream), Some(buffer)) = (&mut stream, &buffer) {
40        stream.upload_blocking(buffer.binding(), tensor.bytes)?;
41        total_ms = start.elapsed().as_secs_f64() * 1000.;
42        // Validation is outside timing, after the upload fence has retired.
43        let mut actual = vec![0; tensor.bytes.len()];
44        stream.read_blocking(buffer.binding(), &mut actual)?;
45        if actual != tensor.bytes {
46            return Err("upload bytes differ from checkpoint".into());
47        }
48        checksum = actual.iter().map(|&v| u64::from(v)).sum();
49    } else {
50        checksum = tensor.bytes.iter().map(|&v| u64::from(v)).sum();
51        total_ms = start.elapsed().as_secs_f64() * 1000.;
52    }
53    println!(
54        "{}",
55        serde_json::json!({"mode":args[2], "preparation":prepared,
56        "bytes":tensor.bytes.len(), "preparation_ms":preparation_ms,
57        "total_ms":total_ms, "checksum":checksum, "upload":upload,
58        "allocation_policy":"preallocated destination; fresh process",
59        "timed_scope":if upload { "page preparation plus completed upload" } else { "page preparation plus complete host read" }, "cache_state":"externally controlled"})
60    );
61    Ok(())
62}

examples/stream_bench.rs (lines 69-72)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn read_blocking_at( &mut self, src: &Buffer, offset: usize, bytes: &mut [u8], ) -> Result<()>

Read at an allocation-relative byte offset.

Source

pub fn fill(&self, dst: View<'_>, value: u8) -> Result<()>

Enqueue a native byte-pattern fill.

Examples found in repository ?

examples/queue_pool_bench.rs (line 20)

16    fn new(device: &Device, bytes: usize) -> Result<Self> {
17        let mut stream = device.stream()?;
18        let source = stream.allocate(bytes)?;
19        let destination = stream.allocate(bytes)?;
20        stream.fill(source.binding(), 0x3c)?;
21        stream.copy(destination.binding(), source.binding())?;
22        stream.synchronize()?;
23        Ok(Self {
24            stream,
25            source,
26            destination,
27        })
28    }

More examples

Hide additional examples

examples/gpu_load.rs (line 101)

88fn main() -> Result<()> {
89    let options = match parse() {
90        Ok(options) => options,
91        Err(message) => {
92            eprintln!("gpu_load: {message}");
93            std::process::exit(2);
94        }
95    };
96    let size = options.mib * 1024 * 1024;
97    let mut stream = Stream::open()?;
98    let source = stream.allocate(size)?;
99    let destination = stream.allocate(size)?;
100    // Touch both allocations so the steady-state loop never pays first-use faults.
101    stream.fill(source.binding(), 0x3c)?;
102    stream.fill(destination.binding(), 0)?;
103    stream.synchronize()?;
104
105    let run_batch = |stream: &mut Stream| -> Result<()> {
106        for _ in 0..options.batch {
107            match options.mode {
108                Mode::Copy => stream.copy(destination.binding(), source.binding())?,
109                Mode::Fill => stream.fill(destination.binding(), 0x5a)?,
110            }
111        }
112        stream.synchronize()
113    };
114
115    // Warm up outside the measured window: the first batch pays queue and paging costs.
116    run_batch(&mut stream)?;
117
118    let window = Duration::from_secs_f64(options.seconds);
119    let start = Instant::now();
120    let mut iterations = 0usize;
121    // Per-batch rates, so the harness can see whether contention was steady or bursty.
122    let mut batch_rates = Vec::new();
123    while start.elapsed() < window {
124        let batch_start = Instant::now();
125        run_batch(&mut stream)?;
126        let seconds = batch_start.elapsed().as_secs_f64();
127        iterations += options.batch;
128        let bytes = options.mode.traffic(size) * options.batch;
129        batch_rates.push(bytes as f64 / seconds / 1e9);
130    }
131    let elapsed = start.elapsed().as_secs_f64();
132    let total = options.mode.traffic(size) as f64 * iterations as f64;
133    let mean = total / elapsed / 1e9;
134    batch_rates.sort_by(f64::total_cmp);
135    let median = batch_rates[batch_rates.len() / 2];
136    let low = batch_rates[batch_rates.len() / 20];
137
138    println!(
139        concat!(
140            r#"{{"device":"gpu","mode":"{}","buffer_mib":{},"batch":{},"#,
141            r#""seconds":{:.4},"iterations":{},"gb_s_mean":{:.3},"#,
142            r#""gb_s_median":{:.3},"gb_s_p5":{:.3}}}"#
143        ),
144        options.mode.name(),
145        options.mib,
146        options.batch,
147        elapsed,
148        iterations,
149        mean,
150        median,
151        low
152    );
153    eprintln!(
154        "gpu_load: {} {:.1} GB/s mean over {:.1}s ({} iterations)",
155        options.mode.name(),
156        mean,
157        elapsed,
158        iterations
159    );
160    Ok(())
161}

Source

pub fn copy(&self, dst: View<'_>, src: View<'_>) -> Result<()>

Enqueue a native copy between equal, non-overlapping ranges.

Examples found in repository ?

examples/queue_pool_bench.rs (line 21)

16    fn new(device: &Device, bytes: usize) -> Result<Self> {
17        let mut stream = device.stream()?;
18        let source = stream.allocate(bytes)?;
19        let destination = stream.allocate(bytes)?;
20        stream.fill(source.binding(), 0x3c)?;
21        stream.copy(destination.binding(), source.binding())?;
22        stream.synchronize()?;
23        Ok(Self {
24            stream,
25            source,
26            destination,
27        })
28    }
29    fn enqueue(&self) -> Result<()> {
30        self.stream
31            .copy(self.destination.binding(), self.source.binding())
32    }

More examples

Hide additional examples

examples/gpu_load.rs (line 108)

88fn main() -> Result<()> {
89    let options = match parse() {
90        Ok(options) => options,
91        Err(message) => {
92            eprintln!("gpu_load: {message}");
93            std::process::exit(2);
94        }
95    };
96    let size = options.mib * 1024 * 1024;
97    let mut stream = Stream::open()?;
98    let source = stream.allocate(size)?;
99    let destination = stream.allocate(size)?;
100    // Touch both allocations so the steady-state loop never pays first-use faults.
101    stream.fill(source.binding(), 0x3c)?;
102    stream.fill(destination.binding(), 0)?;
103    stream.synchronize()?;
104
105    let run_batch = |stream: &mut Stream| -> Result<()> {
106        for _ in 0..options.batch {
107            match options.mode {
108                Mode::Copy => stream.copy(destination.binding(), source.binding())?,
109                Mode::Fill => stream.fill(destination.binding(), 0x5a)?,
110            }
111        }
112        stream.synchronize()
113    };
114
115    // Warm up outside the measured window: the first batch pays queue and paging costs.
116    run_batch(&mut stream)?;
117
118    let window = Duration::from_secs_f64(options.seconds);
119    let start = Instant::now();
120    let mut iterations = 0usize;
121    // Per-batch rates, so the harness can see whether contention was steady or bursty.
122    let mut batch_rates = Vec::new();
123    while start.elapsed() < window {
124        let batch_start = Instant::now();
125        run_batch(&mut stream)?;
126        let seconds = batch_start.elapsed().as_secs_f64();
127        iterations += options.batch;
128        let bytes = options.mode.traffic(size) * options.batch;
129        batch_rates.push(bytes as f64 / seconds / 1e9);
130    }
131    let elapsed = start.elapsed().as_secs_f64();
132    let total = options.mode.traffic(size) as f64 * iterations as f64;
133    let mean = total / elapsed / 1e9;
134    batch_rates.sort_by(f64::total_cmp);
135    let median = batch_rates[batch_rates.len() / 2];
136    let low = batch_rates[batch_rates.len() / 20];
137
138    println!(
139        concat!(
140            r#"{{"device":"gpu","mode":"{}","buffer_mib":{},"batch":{},"#,
141            r#""seconds":{:.4},"iterations":{},"gb_s_mean":{:.3},"#,
142            r#""gb_s_median":{:.3},"gb_s_p5":{:.3}}}"#
143        ),
144        options.mode.name(),
145        options.mib,
146        options.batch,
147        elapsed,
148        iterations,
149        mean,
150        median,
151        low
152    );
153    eprintln!(
154        "gpu_load: {} {:.1} GB/s mean over {:.1}s ({} iterations)",
155        options.mode.name(),
156        mean,
157        elapsed,
158        iterations
159    );
160    Ok(())
161}

Source

pub fn upload(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()>

Enqueue a host upload using bounded, reusable staging chunks. The full destination is checked before any chunk is submitted.

Examples found in repository ?

examples/stream_bench.rs (line 60)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn upload_at( &mut self, dst: &Buffer, offset: usize, bytes: &[u8], ) -> Result<()>

Enqueue a checked upload at an allocation-relative offset.

Source

pub unsafe fn load(&self, path: &Path, symbol: &str) -> Result<Kernel>

Load a trusted native code object from disk.

§Safety

Native code must obey its declared memory and argument contract.

Source

pub unsafe fn load_artifact(&self, artifact: &Artifact) -> Result<Kernel>

Load a compiler artifact for this exact device target.

§Safety

Native code must obey its declared memory and argument contract.

Examples found in repository ?

examples/gemm_pipeline.rs (line 39)

23fn gpu_kernel(
24    runtime: &Runtime,
25    source: &str,
26    symbol: &str,
27    config: &[(&str, String)],
28    elements: usize,
29    bindings: Vec<BindingContract>,
30) -> Result<GpuKernel> {
31    let compiler = hrx::loom::Compiler::for_target(None, runtime.gpu()?.target())?;
32    let module = compiler.module(source);
33    let mut spec = hrx::loom::Specialization::new(symbol);
34    for (key, value) in config {
35        spec.set_config(*key, value.clone());
36    }
37    let artifact = module.compile(&spec)?;
38    let stream = hrx::gpu::Stream::open()?;
39    let raw = unsafe { stream.load_artifact(&artifact) }?;
40    let mut constants = hrx::gpu::Constants::new();
41    match raw.info().constant_byte_length {
42        8 => constants.push(elements as u32)?,
43        12 => constants.push(elements as u64)?,
44        _ => return Err(hrx::Error::Message("unexpected scalar ABI".into())),
45    }
46    constants.push(1f32)?;
47    unsafe {
48        runtime.load_gpu_kernel(
49            artifact.path(),
50            symbol,
51            [elements.div_ceil(256) as u32, 1, 1],
52            [256, 1, 1],
53            KernelContract {
54                bindings,
55                constants: constants.as_bytes().to_vec(),
56            },
57        )
58    }
59}

More examples

Hide additional examples

examples/patch_bench.rs (line 286)

279fn prepare(
280    stream: &mut Stream,
281    compiler: &Compiler,
282    case: &Case,
283    shared: Option<std::sync::Arc<Vec<Buffer>>>,
284) -> Result<Arm> {
285    let artifact = compiler.module(case.source).compile(&case.spec)?;
286    let kernel = unsafe { stream.load_artifact(&artifact) }?;
287    let constants = Constants::indices(&kernel, &case.indices)?;
288    let buffers = if let Some(buffers) = shared {
289        buffers
290    } else {
291        std::sync::Arc::new(
292            case.data
293                .iter()
294                .map(|data| {
295                    let buffer = stream.allocate(data.len())?;
296                    stream.upload_blocking(buffer.binding(), data)?;
297                    Ok(buffer)
298                })
299                .collect::<Result<Vec<_>>>()?,
300        )
301    };
302    let bindings: Vec<_> = buffers.iter().map(Buffer::binding).collect();
303    let mut graph = stream.graph()?;
304    let mut prior = None;
305    for _ in 0..32 {
306        prior = Some(unsafe {
307            graph.dispatch(
308                prior.as_slice(),
309                &kernel,
310                case.grid,
311                case.block,
312                &constants,
313                &bindings,
314            )
315        }?);
316    }
317    let identity = serde_json::json!({"compiler":artifact.compiler_identity(), "artifact":hrx::bundle::digest(artifact.bytes()),
318        "report":artifact.report().map(|report| report.json())});
319    Ok(Arm {
320        graph: graph.finish()?,
321        buffers,
322        identity,
323    })
324}

examples/stream_bench.rs (line 83)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub unsafe fn load_artifact_for_workload( &self, artifact: &Artifact, workload_argument_bits: &[u64], ) -> Result<(Kernel, LaunchConfig)>

Load an artifact and evaluate its compiler-authored launch configuration. All evaluation and allocation happens before dispatch preparation.

§Safety

Native code must obey the caller’s buffer-access contract for this workload.

Source

pub unsafe fn dispatch( &self, kernel: &Kernel, grid: [u32; 3], block: [u32; 3], constants: &Constants, bindings: &[View<'_>], ) -> Result<()>

Enqueue a kernel invocation with declaration-order constants and buffers.

§Safety

Dimensions, constants, and actual kernel accesses must obey each binding’s range.

Examples found in repository ?

examples/stream_bench.rs (line 101)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn scratch(&mut self, bytes: usize) -> Result<Buffer>

Reuse scratch storage under this queue’s execution order.

Examples found in repository ?

examples/stream_bench.rs (line 40)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn recycle(&mut self, buffer: Buffer) -> Result<()>

Return this stream’s allocation after recording its last use.

Examples found in repository ?

examples/stream_bench.rs (line 42)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn set_scratch_limit(&mut self, bytes: usize)

Bound retained scratch storage, evicting large entries first.

Source

pub fn read(&mut self, source: View<'_>) -> Result<Readback>

Queue an owned readback.

Source §

impl Stream

Source

pub fn graph(&self) -> Result<Graph<'_>>

Begin a dependency-checked recording borrowing this stream.

Examples found in repository ?

examples/graph_prepare.rs (line 18)

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/patch_bench.rs (line 303)

279fn prepare(
280    stream: &mut Stream,
281    compiler: &Compiler,
282    case: &Case,
283    shared: Option<std::sync::Arc<Vec<Buffer>>>,
284) -> Result<Arm> {
285    let artifact = compiler.module(case.source).compile(&case.spec)?;
286    let kernel = unsafe { stream.load_artifact(&artifact) }?;
287    let constants = Constants::indices(&kernel, &case.indices)?;
288    let buffers = if let Some(buffers) = shared {
289        buffers
290    } else {
291        std::sync::Arc::new(
292            case.data
293                .iter()
294                .map(|data| {
295                    let buffer = stream.allocate(data.len())?;
296                    stream.upload_blocking(buffer.binding(), data)?;
297                    Ok(buffer)
298                })
299                .collect::<Result<Vec<_>>>()?,
300        )
301    };
302    let bindings: Vec<_> = buffers.iter().map(Buffer::binding).collect();
303    let mut graph = stream.graph()?;
304    let mut prior = None;
305    for _ in 0..32 {
306        prior = Some(unsafe {
307            graph.dispatch(
308                prior.as_slice(),
309                &kernel,
310                case.grid,
311                case.block,
312                &constants,
313                &bindings,
314            )
315        }?);
316    }
317    let identity = serde_json::json!({"compiler":artifact.compiler_identity(), "artifact":hrx::bundle::digest(artifact.bytes()),
318        "report":artifact.report().map(|report| report.json())});
319    Ok(Arm {
320        graph: graph.finish()?,
321        buffers,
322        identity,
323    })
324}

examples/stream_bench.rs (line 120)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Source

pub fn owned_graph(&self) -> Result<Graph<'static>>

Begin a recording that owns its execution domain and each prepared command. Buffers and kernels need only live through their recording call. Native allocations and residency charges remain held by the graph. Pooled leases still require caller-controlled reuse and ordering through final replay.

Source

pub fn launch_profiled( &mut self, graph: &mut GraphExec, ) -> Result<DeviceProfile>

Replay an instrumented graph, wait for retirement, and read GPU times. Ordinary graphs return an error; construct with Graph::finish_profiled.

Source

pub fn launch(&mut self, graph: &mut GraphExec) -> Result<()>

Replay fixed native commands on the originating stream.

Examples found in repository ?

examples/patch_bench.rs (line 402)

348fn main() -> Result<()> {
349    let args: Vec<_> = std::env::args().collect();
350    if args.len() < 3 || args.len() > 5 {
351        return Err(hrx::Error::Message(
352            "usage: patch_bench BASELINE_LIB CANDIDATE_LIB [pairs=30] [default|cu|wgp]".into(),
353        ));
354    }
355    let pairs = args
356        .get(3)
357        .map(|s| s.parse::<usize>())
358        .transpose()
359        .map_err(|e| hrx::Error::Message(e.to_string()))?
360        .unwrap_or(30);
361    if pairs < 30 {
362        return Err(hrx::Error::Message("at least 30 pairs are required".into()));
363    }
364    let mode = match args.get(4).map(String::as_str).unwrap_or("default") {
365        "default" => ProcessorMode::Default,
366        "cu" => ProcessorMode::ComputeUnit,
367        "wgp" => ProcessorMode::WorkgroupProcessor,
368        _ => {
369            return Err(hrx::Error::Message(
370                "mode must be default, cu or wgp".into(),
371            ));
372        }
373    };
374    let mut stream = Stream::open()?;
375    let mut options = CompilerOptions::for_target(stream.target());
376    options.processor_mode = mode;
377    let compilers = [
378        Compiler::with_options(Some(Path::new(&args[1])), options.clone())?,
379        Compiler::with_options(Some(Path::new(&args[2])), options)?,
380    ];
381    for case in [
382        attention(128, false),
383        attention(512, false),
384        gemm(false, false),
385        gemm(true, false),
386        gemm(false, true),
387        gemm(true, true),
388        conv(56, 64, 64),
389        conv(14, 256, 256),
390        conv(7, 512, 512),
391    ] {
392        let baseline = prepare(&mut stream, &compilers[0], &case, None)?;
393        let candidate = prepare(
394            &mut stream,
395            &compilers[1],
396            &case,
397            Some(baseline.buffers.clone()),
398        )?;
399        let mut arms = [baseline, candidate];
400        for arm in &mut arms {
401            for _ in 0..5 {
402                stream.launch(&mut arm.graph)?;
403                stream.synchronize()?;
404            }
405            check(&mut stream, arm, &case)?;
406        }
407        let errors = [
408            check(&mut stream, &arms[0], &case)?,
409            check(&mut stream, &arms[1], &case)?,
410        ];
411        println!(
412            "{}",
413            serde_json::json!({"case":case.name,"mode":format!("{mode:?}"),"baseline":arms[0].identity,"candidate":arms[1].identity,"max_absolute_error":errors,"tolerance":case.tolerance,"replays":32})
414        );
415        for pair in 0..pairs {
416            let mut ns = [0.0; 2];
417            for index in if pair % 2 == 0 { [0, 1] } else { [1, 0] } {
418                let start = Instant::now();
419                stream.launch(&mut arms[index].graph)?;
420                stream.synchronize()?;
421                ns[index] = start.elapsed().as_secs_f64() * 1e9 / 32.0;
422            }
423            println!(
424                "{}",
425                serde_json::json!({"case":case.name,"pair":pair,"baseline_ns":ns[0],"candidate_ns":ns[1]})
426            );
427        }
428        for arm in &mut arms {
429            stream.launch(&mut arm.graph)?;
430            stream.synchronize()?;
431            check(&mut stream, arm, &case)?;
432        }
433    }
434    Ok(())
435}

More examples

Hide additional examples

examples/stream_bench.rs (line 160)

30fn main() -> Result<()> {
31    let mut stream = Stream::open()?;
32    let mut metrics = BTreeMap::new();
33    let samples = std::env::var("HRX_BENCH_SAMPLES")
34        .ok()
35        .and_then(|value| value.parse::<usize>().ok())
36        .filter(|&value| value > 0)
37        .unwrap_or(9);
38    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
39        for _ in 0..50_000 {
40            let buffer = stream.scratch(1024 * 1024)?;
41            black_box(&buffer);
42            stream.recycle(buffer)?;
43        }
44        Ok(())
45    })?;
46    metrics.insert("scratch_reuse_ns", ns);
47    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
48        for _ in 0..2_000 {
49            black_box(stream.allocate(1024 * 1024)?);
50        }
51        Ok(())
52    })?;
53    metrics.insert("allocate_drop_ns", ns);
54
55    let bytes = 256 * 1024 * 1024;
56    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
57    let weights = stream.allocate(bytes)?;
58    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
59        for offset in (0..bytes).step_by(chunk.len()) {
60            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
61        }
62        Ok(())
63    })?;
64    metrics.insert(
65        "queued_upload_gib_s",
66        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
67    );
68    let mut check = [0u8; 16];
69    stream.read_blocking(
70        weights.try_slice(bytes - check.len(), check.len())?,
71        &mut check,
72    )?;
73    assert_eq!(check, [0x3c; 16]);
74    drop(weights);
75
76    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
77    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
78    let mut spec = hrx::loom::Specialization::new("krea2_euler");
79    spec.set_config("krea2.euler.grid_x", "1");
80    spec.set_config("krea2.euler.grid_y", "1");
81    let artifact = module.compile(&spec)?;
82    // This repository owns the source; use its exact dimensions and scalar layout.
83    let kernel = unsafe { stream.load_artifact(&artifact)? };
84    let sample = stream.allocate(512)?;
85    let velocity = stream.allocate(512)?;
86    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
87    stream.upload_blocking(sample.binding(), &ones)?;
88    stream.upload_blocking(velocity.binding(), &ones)?;
89    let mut constants = Constants::new();
90    match kernel.info().constant_byte_length {
91        8 => constants.push(256u32)?,
92        12 => constants.push(256u64)?,
93        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
94    }
95    constants.push(0f32)?;
96    let bindings = [sample.binding(), velocity.binding()];
97    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
98        for _ in 0..2048 {
99            // The source accesses exactly the 256 bf16 elements in each buffer.
100            unsafe {
101                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
102            }
103        }
104        Ok(())
105    })?;
106    metrics.insert("dispatch_enqueue_ns", enqueue);
107    metrics.insert("dispatch_complete_ns", complete);
108
109    // Use disjoint allocations in both graphs. Even an idempotent store is a
110    // write: concurrent Euler dispatches cannot share their in/out buffer.
111    let mut graph_buffers = Vec::new();
112    for _ in 0..32 {
113        let sample = stream.allocate(512)?;
114        let velocity = stream.allocate(512)?;
115        stream.upload(sample.binding(), &ones)?;
116        stream.upload(velocity.binding(), &ones)?;
117        graph_buffers.push((sample, velocity));
118    }
119    stream.synchronize()?;
120    let mut builder = stream.graph()?;
121    let mut previous = None;
122    for (sample, velocity) in &graph_buffers {
123        let after = previous.as_slice();
124        previous = Some(unsafe {
125            builder.dispatch(
126                after,
127                &kernel,
128                [1; 3],
129                [256, 1, 1],
130                &constants,
131                &[sample.binding(), velocity.binding()],
132            )?
133        });
134    }
135    let chained = builder.finish()?;
136
137    // The same disjoint workloads, this time allowing concurrent scheduling.
138    let mut builder = stream.graph()?;
139    for (sample, velocity) in &graph_buffers {
140        unsafe {
141            builder.dispatch(
142                &[],
143                &kernel,
144                [1; 3],
145                [256, 1, 1],
146                &constants,
147                &[sample.binding(), velocity.binding()],
148            )?;
149        }
150    }
151    let independent = builder.finish()?;
152    let mut graphs = [chained, independent];
153    let mut enqueue = [Vec::new(), Vec::new()];
154    let mut complete = [Vec::new(), Vec::new()];
155    for round in 0..samples + 3 {
156        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
157            stream.synchronize()?;
158            let start = Instant::now();
159            for _ in 0..256 {
160                stream.launch(&mut graphs[arm])?;
161            }
162            let submitted = start.elapsed();
163            stream.synchronize()?;
164            let elapsed = start.elapsed();
165            if round >= 3 {
166                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
167                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
168            }
169        }
170    }
171    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
172        values.sort_by(f64::total_cmp);
173    }
174    let ns = complete[0][samples / 2];
175    let ns_free = complete[1][samples / 2];
176    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
177    metrics.insert("graph_complete_ns_per_kernel", ns);
178    metrics.insert("graph_independent_ns_per_kernel", ns_free);
179    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
180    for (sample, _) in &graph_buffers {
181        let mut output = vec![0; 512];
182        stream.read_blocking(sample.binding(), &mut output)?;
183        assert_eq!(output, ones);
184    }
185    let mut output = vec![0; 512];
186    stream.read_blocking(sample.binding(), &mut output)?;
187    assert_eq!(output, ones);
188    println!("{}", serde_json::to_string(&metrics)?);
189    Ok(())
190}

Trait Implementations§

Source §

impl Debug for Stream

Source §

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more

Source §

impl Drop for Stream

Source §

fn drop(&mut self)

Executes the destructor for this type. Read more

Source §

fn pin_drop(self: Pin<&mut Self>)

🔬This is a nightly-only experimental API. (pin_ergonomics)

Execute the destructor for this type, but different to Drop::drop, it requires self to be pinned. Read more

Auto Trait Implementations§

§

impl UnwindSafe for Stream

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> 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, 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 Stream Copy item path

Implementations§

impl Stream

pub fn access_graph(&self) -> Result<AccessGraph<'_>>

impl Stream

pub fn open() -> Result<Self>

pub fn device_id(&self) -> usize

pub fn id(&self) -> usize

pub fn target(&self) -> &Target

pub fn with_memory_budget(self, budget: MemoryBudget) -> Self

pub fn memory_budget(&self) -> Option<&MemoryBudget>

pub fn allocate(&self, bytes: usize) -> Result<Buffer>

pub fn allocate_from(&self, data: &[u8]) -> Result<Buffer>

pub fn allocate_zeroed(&self, bytes: usize) -> Result<Buffer>

pub fn allocate_shared(&self, bytes: usize) -> Result<Buffer>

pub fn synchronize(&mut self) -> Result<()>

pub fn record_event(&mut self) -> Result<Event>

pub fn wait_event(&mut self, event: &Event) -> Result<()>

pub fn submit(&mut self) -> Result<Submission<'_>>

pub fn upload_blocking(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()>

pub fn upload_blocking_at( &mut self, dst: &Buffer, offset: usize, bytes: &[u8], ) -> Result<()>

pub fn read_blocking(&mut self, src: View<'_>, bytes: &mut [u8]) -> Result<()>

pub fn read_blocking_at( &mut self, src: &Buffer, offset: usize, bytes: &mut [u8], ) -> Result<()>

pub fn fill(&self, dst: View<'_>, value: u8) -> Result<()>

pub fn copy(&self, dst: View<'_>, src: View<'_>) -> Result<()>

pub fn upload(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()>

pub fn upload_at( &mut self, dst: &Buffer, offset: usize, bytes: &[u8], ) -> Result<()>

pub unsafe fn load(&self, path: &Path, symbol: &str) -> Result<Kernel>

§Safety

pub unsafe fn load_artifact(&self, artifact: &Artifact) -> Result<Kernel>

§Safety

pub unsafe fn load_artifact_for_workload( &self, artifact: &Artifact, workload_argument_bits: &[u64], ) -> Result<(Kernel, LaunchConfig)>

§Safety

pub unsafe fn dispatch( &self, kernel: &Kernel, grid: [u32; 3], block: [u32; 3], constants: &Constants, bindings: &[View<'_>], ) -> Result<()>

§Safety

pub fn scratch(&mut self, bytes: usize) -> Result<Buffer>

pub fn recycle(&mut self, buffer: Buffer) -> Result<()>

pub fn set_scratch_limit(&mut self, bytes: usize)

pub fn read(&mut self, source: View<'_>) -> Result<Readback>

impl Stream

pub fn graph(&self) -> Result<Graph<'_>>

pub fn owned_graph(&self) -> Result<Graph<'static>>

pub fn launch_profiled( &mut self, graph: &mut GraphExec, ) -> Result<DeviceProfile>

pub fn launch(&mut self, graph: &mut GraphExec) -> Result<()>

Trait Implementations§

impl Debug for Stream

fn fmt(&self, f: &mut Formatter<'_>) -> Result

impl Drop for Stream

fn drop(&mut self)

fn pin_drop(self: Pin<&mut Self>)

Auto Trait Implementations§

impl !Freeze for Stream

impl !RefUnwindSafe for Stream

impl !Sync for Stream

impl Send for Stream

impl Unpin for Stream

impl UnsafeUnpin for Stream

impl UnwindSafe for Stream

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

Struct Stream

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, 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, 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>,