Skip to main content

max / everycycle

15.5 KB · 430 lines History Blame Raw
1 //! Memory-bandwidth probe — BabelStream-style triad.
2 //!
3 //! Allocates three device-local buffers (A, B, C), dispatches a
4 //! compute shader that computes `c[i] = a[i] + s * b[i]` for `i` in
5 //! `[0, N)`, times the dispatches, and reports observed read and
6 //! write bandwidth from the triad:
7 //!
8 //! ```text
9 //! read_bps = (2 * N * sizeof(f32)) * iterations / elapsed_seconds
10 //! write_bps = (1 * N * sizeof(f32)) * iterations / elapsed_seconds
11 //! ```
12 //!
13 //! Two reads per element (A and B), one write (C). The shader is a
14 //! pure streaming kernel so the achieved figure tracks the device's
15 //! DRAM throughput, not its cache or compute path.
16 //!
17 //! # Timing
18 //!
19 //! The dispatch is timed **on the device**, with a timestamp query
20 //! written either side of it inside the command buffer, and only the
21 //! deltas are accumulated. Queue submission and the fence wait cost
22 //! tens of microseconds against a dispatch of a few hundred, so timing
23 //! the host-side loop instead charges that overhead to the memory
24 //! system and understates bandwidth by low double-digit percent —
25 //! measured at ~13% on an RTX 5070 Ti.
26 //!
27 //! A queue family that reports no valid timestamp bits falls back to
28 //! host wall-clock timing. The result records which clock produced it,
29 //! because the two are not comparable and a report must not imply they
30 //! are.
31
32 use std::time::Instant;
33
34 use ash::vk;
35 use naga::back::spv;
36 use naga::front::wgsl;
37 use naga::valid;
38
39 use super::vulkan_ctx::VulkanContext;
40
41 // Naga's WGSL frontend does not accept `var<push_constant>`, so the
42 // triad params (which are compile-time constants anyway) are inlined
43 // as module-scope `const`s in the shader. The N literal below must
44 // equal ELEMENTS on the Rust side; a const_assert below guards it.
45 const TRIAD_WGSL: &str = r"
46 const SCALAR: f32 = 2.0;
47 const N: u32 = 16777216u;
48
49 @group(0) @binding(0) var<storage, read> a: array<f32>;
50 @group(0) @binding(1) var<storage, read> b: array<f32>;
51 @group(0) @binding(2) var<storage, read_write> c: array<f32>;
52
53 @compute @workgroup_size(256)
54 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
55 let i = gid.x;
56 if (i < N) {
57 c[i] = a[i] + SCALAR * b[i];
58 }
59 }
60 ";
61
62 /// Element count per buffer. 16 Mi f32 → 64 MiB per buffer, 192 MiB
63 /// total. Comfortably larger than any L2/LLC so the triad is DRAM-bound.
64 const ELEMENTS: u32 = 16 * 1024 * 1024;
65 const _: () = assert!(
66 ELEMENTS == 16_777_216,
67 "ELEMENTS must equal the N literal in TRIAD_WGSL"
68 );
69 const ELEMENT_BYTES: u32 = 4;
70 const BUFFER_BYTES: vk::DeviceSize =
71 (ELEMENTS as vk::DeviceSize) * (ELEMENT_BYTES as vk::DeviceSize);
72 const WORKGROUP_SIZE: u32 = 256;
73 const WARMUP_ITERATIONS: u32 = 4;
74 const MEASURE_ITERATIONS: u32 = 32;
75
76 /// Which clock produced a measurement.
77 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
78 pub enum TimingSource {
79 /// Timestamp queries around the dispatch. Excludes submit and
80 /// fence-wait overhead.
81 DeviceTimestamps,
82 /// Host wall clock around submit-and-wait. Includes per-dispatch
83 /// submission overhead, so the figure is a lower bound.
84 HostWallClock,
85 }
86
87 #[derive(Clone, Copy, Debug)]
88 pub struct BandwidthResult {
89 pub read_bps: f64,
90 pub write_bps: f64,
91 pub timing: TimingSource,
92 }
93
94 #[derive(Debug)]
95 pub enum BandwidthError {
96 ShaderParse(String),
97 ShaderValidate(String),
98 ShaderEmit(String),
99 Vulkan(vk::Result),
100 NoDeviceLocalMemory,
101 }
102
103 impl core::fmt::Display for BandwidthError {
104 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
105 match self {
106 Self::ShaderParse(s) => write!(f, "WGSL parse failed: {s}"),
107 Self::ShaderValidate(s) => write!(f, "WGSL validation failed: {s}"),
108 Self::ShaderEmit(s) => write!(f, "SPIR-V emission failed: {s}"),
109 Self::Vulkan(r) => write!(f, "Vulkan call failed: {r:?}"),
110 Self::NoDeviceLocalMemory => write!(f, "no DEVICE_LOCAL memory type available"),
111 }
112 }
113 }
114
115 impl std::error::Error for BandwidthError {}
116
117 /// Compile the triad shader and run the measurement against `ctx`.
118 pub fn measure_bandwidth(ctx: &VulkanContext) -> Result<BandwidthResult, BandwidthError> {
119 let spirv = compile_triad_shader()?;
120 let device = &ctx.device;
121
122 let memory_type = ctx
123 .find_memory_type(u32::MAX, vk::MemoryPropertyFlags::DEVICE_LOCAL)
124 .ok_or(BandwidthError::NoDeviceLocalMemory)?;
125
126 // Three device-local buffers, used as storage buffers.
127 let buffers = create_buffers(ctx, memory_type)?;
128
129 // SAFETY: Vulkan handle lifetimes are managed manually below. The
130 // teardown loop at the end of `measure_bandwidth` releases every
131 // handle in reverse creation order.
132 let result = unsafe { dispatch_and_time(ctx, &spirv, &buffers) };
133
134 unsafe { destroy_buffers(device, &buffers) };
135
136 let (elapsed, timing) = result?;
137 Ok(compute_bandwidth(elapsed, MEASURE_ITERATIONS, timing))
138 }
139
140 fn compile_triad_shader() -> Result<Vec<u32>, BandwidthError> {
141 let module =
142 wgsl::parse_str(TRIAD_WGSL).map_err(|e| BandwidthError::ShaderParse(format!("{e:?}")))?;
143 let info = valid::Validator::new(valid::ValidationFlags::all(), valid::Capabilities::all())
144 .validate(&module)
145 .map_err(|e| BandwidthError::ShaderValidate(format!("{e:?}")))?;
146 let options = spv::Options::default();
147 spv::write_vec(&module, &info, &options, None)
148 .map_err(|e| BandwidthError::ShaderEmit(format!("{e:?}")))
149 }
150
151 struct TriadBuffers {
152 buffers: [vk::Buffer; 3],
153 memories: [vk::DeviceMemory; 3],
154 }
155
156 fn create_buffers(ctx: &VulkanContext, memory_type: u32) -> Result<TriadBuffers, BandwidthError> {
157 let device = &ctx.device;
158 let mut buffers = [vk::Buffer::null(); 3];
159 let mut memories = [vk::DeviceMemory::null(); 3];
160
161 for slot in 0..3 {
162 let info = vk::BufferCreateInfo::default()
163 .size(BUFFER_BYTES)
164 .usage(vk::BufferUsageFlags::STORAGE_BUFFER)
165 .sharing_mode(vk::SharingMode::EXCLUSIVE);
166 let buffer =
167 unsafe { device.create_buffer(&info, None) }.map_err(BandwidthError::Vulkan)?;
168 let reqs = unsafe { device.get_buffer_memory_requirements(buffer) };
169 let alloc = vk::MemoryAllocateInfo::default()
170 .allocation_size(reqs.size)
171 .memory_type_index(memory_type);
172 let memory = match unsafe { device.allocate_memory(&alloc, None) } {
173 Ok(m) => m,
174 Err(e) => {
175 unsafe { device.destroy_buffer(buffer, None) };
176 return Err(BandwidthError::Vulkan(e));
177 }
178 };
179 if let Err(e) = unsafe { device.bind_buffer_memory(buffer, memory, 0) } {
180 unsafe {
181 device.free_memory(memory, None);
182 device.destroy_buffer(buffer, None);
183 }
184 return Err(BandwidthError::Vulkan(e));
185 }
186 buffers[slot] = buffer;
187 memories[slot] = memory;
188 }
189
190 Ok(TriadBuffers { buffers, memories })
191 }
192
193 unsafe fn destroy_buffers(device: &ash::Device, b: &TriadBuffers) {
194 for i in 0..3 {
195 unsafe {
196 device.destroy_buffer(b.buffers[i], None);
197 device.free_memory(b.memories[i], None);
198 }
199 }
200 }
201
202 #[allow(clippy::too_many_lines)]
203 unsafe fn dispatch_and_time(
204 ctx: &VulkanContext,
205 spirv: &[u32],
206 buffers: &TriadBuffers,
207 ) -> Result<(f64, TimingSource), BandwidthError> {
208 let device = &ctx.device;
209
210 // Descriptor set layout: three storage-buffer bindings.
211 let bindings: [vk::DescriptorSetLayoutBinding; 3] = std::array::from_fn(|i| {
212 #[allow(clippy::cast_possible_truncation)]
213 let binding_idx = i as u32;
214 vk::DescriptorSetLayoutBinding::default()
215 .binding(binding_idx)
216 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
217 .descriptor_count(1)
218 .stage_flags(vk::ShaderStageFlags::COMPUTE)
219 });
220 let dsl_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings);
221 let dsl = unsafe { device.create_descriptor_set_layout(&dsl_info, None) }
222 .map_err(BandwidthError::Vulkan)?;
223
224 // Pipeline layout: no push constants — the triad params are baked
225 // into the shader as module-scope consts.
226 let dsls = [dsl];
227 let pl_info = vk::PipelineLayoutCreateInfo::default().set_layouts(&dsls);
228 let pipeline_layout =
229 unsafe { device.create_pipeline_layout(&pl_info, None) }.map_err(BandwidthError::Vulkan)?;
230
231 // Shader module + compute pipeline.
232 let shader_info = vk::ShaderModuleCreateInfo::default().code(spirv);
233 let shader = unsafe { device.create_shader_module(&shader_info, None) }
234 .map_err(BandwidthError::Vulkan)?;
235 let stage = vk::PipelineShaderStageCreateInfo::default()
236 .stage(vk::ShaderStageFlags::COMPUTE)
237 .module(shader)
238 .name(c"main");
239 let pipeline_info = [vk::ComputePipelineCreateInfo::default()
240 .stage(stage)
241 .layout(pipeline_layout)];
242 let pipelines =
243 unsafe { device.create_compute_pipelines(vk::PipelineCache::null(), &pipeline_info, None) }
244 .map_err(|(_, r)| BandwidthError::Vulkan(r))?;
245 let pipeline = pipelines[0];
246
247 // Descriptor pool + descriptor set, wired to our three buffers.
248 let pool_sizes = [vk::DescriptorPoolSize::default()
249 .ty(vk::DescriptorType::STORAGE_BUFFER)
250 .descriptor_count(3)];
251 let pool_info = vk::DescriptorPoolCreateInfo::default()
252 .max_sets(1)
253 .pool_sizes(&pool_sizes);
254 let descriptor_pool = unsafe { device.create_descriptor_pool(&pool_info, None) }
255 .map_err(BandwidthError::Vulkan)?;
256 let alloc_info = vk::DescriptorSetAllocateInfo::default()
257 .descriptor_pool(descriptor_pool)
258 .set_layouts(&dsls);
259 let descriptor_set =
260 unsafe { device.allocate_descriptor_sets(&alloc_info) }.map_err(BandwidthError::Vulkan)?[0];
261
262 let buf_infos: [vk::DescriptorBufferInfo; 3] = std::array::from_fn(|i| {
263 vk::DescriptorBufferInfo::default()
264 .buffer(buffers.buffers[i])
265 .offset(0)
266 .range(BUFFER_BYTES)
267 });
268 let writes: [vk::WriteDescriptorSet; 3] = std::array::from_fn(|i| {
269 #[allow(clippy::cast_possible_truncation)]
270 let binding_idx = i as u32;
271 vk::WriteDescriptorSet::default()
272 .dst_set(descriptor_set)
273 .dst_binding(binding_idx)
274 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
275 .buffer_info(std::slice::from_ref(&buf_infos[i]))
276 });
277 unsafe { device.update_descriptor_sets(&writes, &[]) };
278
279 // Command buffer that runs one triad dispatch.
280 let cmd_info = vk::CommandBufferAllocateInfo::default()
281 .command_pool(ctx.command_pool)
282 .level(vk::CommandBufferLevel::PRIMARY)
283 .command_buffer_count(1);
284 let cmd =
285 unsafe { device.allocate_command_buffers(&cmd_info) }.map_err(BandwidthError::Vulkan)?[0];
286
287 // A queue that cannot write timestamps leaves only the host clock.
288 let timing = if ctx.timestamp_valid_bits > 0 {
289 TimingSource::DeviceTimestamps
290 } else {
291 TimingSource::HostWallClock
292 };
293
294 let query_pool = if timing == TimingSource::DeviceTimestamps {
295 let info = vk::QueryPoolCreateInfo::default()
296 .query_type(vk::QueryType::TIMESTAMP)
297 .query_count(2);
298 Some(unsafe { device.create_query_pool(&info, None) }.map_err(BandwidthError::Vulkan)?)
299 } else {
300 None
301 };
302
303 let begin = vk::CommandBufferBeginInfo::default();
304 let groups = ELEMENTS.div_ceil(WORKGROUP_SIZE);
305
306 // Recorded once and resubmitted: the buffer carries no
307 // ONE_TIME_SUBMIT flag, so re-recording per iteration bought
308 // nothing and put another host-side cost inside the timed region.
309 // The query-pool reset lives in the command buffer so each
310 // submission overwrites the previous pair.
311 unsafe {
312 device
313 .begin_command_buffer(cmd, &begin)
314 .map_err(BandwidthError::Vulkan)?;
315 if let Some(pool) = query_pool {
316 device.cmd_reset_query_pool(cmd, pool, 0, 2);
317 device.cmd_write_timestamp(cmd, vk::PipelineStageFlags::TOP_OF_PIPE, pool, 0);
318 }
319 device.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::COMPUTE, pipeline);
320 device.cmd_bind_descriptor_sets(
321 cmd,
322 vk::PipelineBindPoint::COMPUTE,
323 pipeline_layout,
324 0,
325 &[descriptor_set],
326 &[],
327 );
328 device.cmd_dispatch(cmd, groups, 1, 1);
329 if let Some(pool) = query_pool {
330 device.cmd_write_timestamp(cmd, vk::PipelineStageFlags::BOTTOM_OF_PIPE, pool, 1);
331 }
332 device
333 .end_command_buffer(cmd)
334 .map_err(BandwidthError::Vulkan)?;
335 }
336
337 let submit = || -> Result<(), BandwidthError> {
338 let info = vk::SubmitInfo::default().command_buffers(std::slice::from_ref(&cmd));
339 unsafe {
340 device
341 .queue_submit(ctx.queue, &[info], vk::Fence::null())
342 .map_err(BandwidthError::Vulkan)?;
343 device
344 .queue_wait_idle(ctx.queue)
345 .map_err(BandwidthError::Vulkan)?;
346 }
347 Ok(())
348 };
349
350 // Ticks are only meaningful in the low `timestamp_valid_bits`, so
351 // mask before subtracting or a wrapped counter reads as a huge
352 // interval.
353 let tick_mask: u64 = if ctx.timestamp_valid_bits >= 64 {
354 u64::MAX
355 } else {
356 (1_u64 << ctx.timestamp_valid_bits) - 1
357 };
358
359 let read_dispatch_ns = |pool: vk::QueryPool| -> Result<f64, BandwidthError> {
360 let mut ticks = [0_u64; 2];
361 unsafe {
362 device
363 .get_query_pool_results(
364 pool,
365 0,
366 &mut ticks,
367 vk::QueryResultFlags::TYPE_64 | vk::QueryResultFlags::WAIT,
368 )
369 .map_err(BandwidthError::Vulkan)?;
370 }
371 let delta = (ticks[1] & tick_mask).wrapping_sub(ticks[0] & tick_mask) & tick_mask;
372 #[allow(clippy::cast_precision_loss)]
373 Ok(delta as f64 * f64::from(ctx.timestamp_period_ns))
374 };
375
376 // Warm-up: drive the device into steady state so the measured pass
377 // reflects thermal-stable behavior, not first-launch overhead.
378 for _ in 0..WARMUP_ITERATIONS {
379 submit()?;
380 }
381
382 let result = match query_pool {
383 Some(pool) => {
384 let mut total_ns = 0.0_f64;
385 for _ in 0..MEASURE_ITERATIONS {
386 submit()?;
387 total_ns += read_dispatch_ns(pool)?;
388 }
389 total_ns / 1e9
390 }
391 None => {
392 let start = Instant::now();
393 for _ in 0..MEASURE_ITERATIONS {
394 submit()?;
395 }
396 start.elapsed().as_secs_f64()
397 }
398 };
399
400 unsafe {
401 if let Some(pool) = query_pool {
402 device.destroy_query_pool(pool, None);
403 }
404 device.free_command_buffers(ctx.command_pool, &[cmd]);
405 device.destroy_descriptor_pool(descriptor_pool, None);
406 device.destroy_pipeline(pipeline, None);
407 device.destroy_shader_module(shader, None);
408 device.destroy_pipeline_layout(pipeline_layout, None);
409 device.destroy_descriptor_set_layout(dsl, None);
410 }
411
412 Ok((result, timing))
413 }
414
415 #[allow(clippy::cast_precision_loss)]
416 fn compute_bandwidth(
417 elapsed_seconds: f64,
418 iterations: u32,
419 timing: TimingSource,
420 ) -> BandwidthResult {
421 let bytes_per_iter_read = 2.0 * (BUFFER_BYTES as f64);
422 let bytes_per_iter_write = BUFFER_BYTES as f64;
423 let iters = f64::from(iterations);
424 BandwidthResult {
425 read_bps: bytes_per_iter_read * iters / elapsed_seconds,
426 write_bps: bytes_per_iter_write * iters / elapsed_seconds,
427 timing,
428 }
429 }
430