Skip to content

Commit 45e6d7e

Browse files
KeavonKeavon
committed
Isolate wgpu's loop shape and time variants interleaved after a sustained warmup
1 parent 601371f commit 45e6d7e

1 file changed

Lines changed: 273 additions & 54 deletions

File tree

‎.github/mac-gpu-probe/src/metal_direct.rs‎

Lines changed: 273 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
// Runs the same workloads as hand-written Metal, timed by each command buffer's own GPU start and end times,
1+
// Runs the probe's workloads through Metal directly, timed by each command buffer's own GPU start and end times,
22
// which Metal reports even where counter sampling (and therefore wgpu's timestamp queries) is unavailable
33

44
use std::ptr::NonNull;
5-
use std::time::Instant;
5+
use std::time::{Duration, Instant};
66

77
use objc2::rc::Retained;
88
use objc2::runtime::ProtocolObject;
@@ -42,6 +42,199 @@ kernel void bandwidth(device const float4 *source [[buffer(0)]], device float4 *
4242
}
4343
"#;
4444

45+
// What wgpu generated for the probe's WGSL compute shader (copied from its debug log), with its default checks
46+
const NAGA_CHECKED: &str = r#"
47+
#include <metal_stdlib>
48+
#include <simd/simd.h>
49+
50+
using metal::uint;
51+
52+
struct _mslBufferSizes {
53+
uint size0;
54+
};
55+
56+
typedef metal::float4 type_1[1];
57+
58+
[[max_total_threads_per_threadgroup(256)]] kernel void main_(
59+
metal::uint3 id [[thread_position_in_grid]]
60+
, device type_1& output [[buffer(0)]]
61+
, constant _mslBufferSizes& _buffer_sizes [[buffer(1)]]
62+
) {
63+
metal::float4 a0_ = {};
64+
metal::float4 a1_ = {};
65+
metal::float4 a2_ = {};
66+
metal::float4 a3_ = {};
67+
uint i = 0u;
68+
a0_ = metal::float4(static_cast<float>(id.x) * 0.0000001) + metal::float4(0.0, 1.0, 2.0, 3.0);
69+
metal::float4 _e13 = a0_;
70+
a1_ = _e13 + metal::float4(0.1);
71+
metal::float4 _e18 = a0_;
72+
a2_ = _e18 + metal::float4(0.2);
73+
metal::float4 _e23 = a0_;
74+
a3_ = _e23 + metal::float4(0.3);
75+
metal::float4 b0_ = metal::float4(0.9999, 0.9998, 0.9997, 0.9996);
76+
metal::float4 b1_ = metal::float4(0.9995, 0.9994, 0.9993, 0.9992);
77+
metal::float4 b2_ = metal::float4(0.9991, 0.999, 0.9989, 0.9988);
78+
metal::float4 b3_ = metal::float4(0.9987, 0.9986, 0.9985, 0.9984);
79+
metal::float4 c = metal::float4(0.0001, 0.0002, 0.0003, 0.0004);
80+
uint2 loop_bound = uint2(4294967295u);
81+
bool loop_init = true;
82+
while(true) {
83+
if (metal::all(loop_bound == uint2(0u))) { break; }
84+
loop_bound -= uint2(loop_bound.y == 0u, 1u);
85+
if (!loop_init) {
86+
uint _e67 = i;
87+
i = _e67 + 1u;
88+
}
89+
loop_init = false;
90+
uint _e55 = i;
91+
if (_e55 < 256u) {
92+
} else {
93+
break;
94+
}
95+
{
96+
metal::float4 _e58 = a0_;
97+
a0_ = metal::fma(_e58, b0_, c);
98+
metal::float4 _e60 = a1_;
99+
a1_ = metal::fma(_e60, b1_, c);
100+
metal::float4 _e62 = a2_;
101+
a2_ = metal::fma(_e62, b2_, c);
102+
metal::float4 _e64 = a3_;
103+
a3_ = metal::fma(_e64, b3_, c);
104+
}
105+
}
106+
metal::float4 _e72 = a0_;
107+
metal::float4 _e73 = a1_;
108+
metal::float4 _e75 = a2_;
109+
metal::float4 _e77 = a3_;
110+
output[metal::min(unsigned(id.x), (_buffer_sizes.size0 - 0 - 16) / 16)] = ((_e72 + _e73) + _e75) + _e77;
111+
return;
112+
}
113+
"#;
114+
115+
// What wgpu generated with its checks off, which still emits the loop as `while (true)` with a `loop_init` flag
116+
const NAGA_UNCHECKED: &str = r#"
117+
#include <metal_stdlib>
118+
#include <simd/simd.h>
119+
120+
using metal::uint;
121+
122+
struct _mslBufferSizes {
123+
uint size0;
124+
};
125+
126+
typedef metal::float4 type_1[1];
127+
128+
[[max_total_threads_per_threadgroup(256)]] kernel void main_(
129+
metal::uint3 id [[thread_position_in_grid]]
130+
, device type_1& output [[buffer(0)]]
131+
, constant _mslBufferSizes& _buffer_sizes [[buffer(1)]]
132+
) {
133+
metal::float4 a0_ = {};
134+
metal::float4 a1_ = {};
135+
metal::float4 a2_ = {};
136+
metal::float4 a3_ = {};
137+
uint i = 0u;
138+
a0_ = metal::float4(static_cast<float>(id.x) * 0.0000001) + metal::float4(0.0, 1.0, 2.0, 3.0);
139+
metal::float4 _e13 = a0_;
140+
a1_ = _e13 + metal::float4(0.1);
141+
metal::float4 _e18 = a0_;
142+
a2_ = _e18 + metal::float4(0.2);
143+
metal::float4 _e23 = a0_;
144+
a3_ = _e23 + metal::float4(0.3);
145+
metal::float4 b0_ = metal::float4(0.9999, 0.9998, 0.9997, 0.9996);
146+
metal::float4 b1_ = metal::float4(0.9995, 0.9994, 0.9993, 0.9992);
147+
metal::float4 b2_ = metal::float4(0.9991, 0.999, 0.9989, 0.9988);
148+
metal::float4 b3_ = metal::float4(0.9987, 0.9986, 0.9985, 0.9984);
149+
metal::float4 c = metal::float4(0.0001, 0.0002, 0.0003, 0.0004);
150+
bool loop_init = true;
151+
while(true) {
152+
if (!loop_init) {
153+
uint _e67 = i;
154+
i = _e67 + 1u;
155+
}
156+
loop_init = false;
157+
uint _e55 = i;
158+
if (_e55 < 256u) {
159+
} else {
160+
break;
161+
}
162+
{
163+
metal::float4 _e58 = a0_;
164+
a0_ = metal::fma(_e58, b0_, c);
165+
metal::float4 _e60 = a1_;
166+
a1_ = metal::fma(_e60, b1_, c);
167+
metal::float4 _e62 = a2_;
168+
a2_ = metal::fma(_e62, b2_, c);
169+
metal::float4 _e64 = a3_;
170+
a3_ = metal::fma(_e64, b3_, c);
171+
}
172+
}
173+
metal::float4 _e72 = a0_;
174+
metal::float4 _e73 = a1_;
175+
metal::float4 _e75 = a2_;
176+
metal::float4 _e77 = a3_;
177+
output[id.x] = ((_e72 + _e73) + _e75) + _e77;
178+
return;
179+
}
180+
"#;
181+
182+
// The unchecked version with only its loop rewritten as a counted `for` loop
183+
const NAGA_UNCHECKED_FOR_LOOP: &str = r#"
184+
#include <metal_stdlib>
185+
#include <simd/simd.h>
186+
187+
using metal::uint;
188+
189+
struct _mslBufferSizes {
190+
uint size0;
191+
};
192+
193+
typedef metal::float4 type_1[1];
194+
195+
[[max_total_threads_per_threadgroup(256)]] kernel void main_(
196+
metal::uint3 id [[thread_position_in_grid]]
197+
, device type_1& output [[buffer(0)]]
198+
, constant _mslBufferSizes& _buffer_sizes [[buffer(1)]]
199+
) {
200+
metal::float4 a0_ = {};
201+
metal::float4 a1_ = {};
202+
metal::float4 a2_ = {};
203+
metal::float4 a3_ = {};
204+
a0_ = metal::float4(static_cast<float>(id.x) * 0.0000001) + metal::float4(0.0, 1.0, 2.0, 3.0);
205+
metal::float4 _e13 = a0_;
206+
a1_ = _e13 + metal::float4(0.1);
207+
metal::float4 _e18 = a0_;
208+
a2_ = _e18 + metal::float4(0.2);
209+
metal::float4 _e23 = a0_;
210+
a3_ = _e23 + metal::float4(0.3);
211+
metal::float4 b0_ = metal::float4(0.9999, 0.9998, 0.9997, 0.9996);
212+
metal::float4 b1_ = metal::float4(0.9995, 0.9994, 0.9993, 0.9992);
213+
metal::float4 b2_ = metal::float4(0.9991, 0.999, 0.9989, 0.9988);
214+
metal::float4 b3_ = metal::float4(0.9987, 0.9986, 0.9985, 0.9984);
215+
metal::float4 c = metal::float4(0.0001, 0.0002, 0.0003, 0.0004);
216+
for (uint i = 0u; i < 256u; i++) {
217+
metal::float4 _e58 = a0_;
218+
a0_ = metal::fma(_e58, b0_, c);
219+
metal::float4 _e60 = a1_;
220+
a1_ = metal::fma(_e60, b1_, c);
221+
metal::float4 _e62 = a2_;
222+
a2_ = metal::fma(_e62, b2_, c);
223+
metal::float4 _e64 = a3_;
224+
a3_ = metal::fma(_e64, b3_, c);
225+
}
226+
metal::float4 _e72 = a0_;
227+
metal::float4 _e73 = a1_;
228+
metal::float4 _e75 = a2_;
229+
metal::float4 _e77 = a3_;
230+
output[id.x] = ((_e72 + _e73) + _e75) + _e77;
231+
return;
232+
}
233+
"#;
234+
235+
const SUSTAINED_WARMUP: Duration = Duration::from_secs(1);
236+
const INTERLEAVED_ROUNDS: usize = 5;
237+
45238
type Device = ProtocolObject<dyn MTLDevice>;
46239
type Queue = ProtocolObject<dyn MTLCommandQueue>;
47240
type Pipeline = Retained<ProtocolObject<dyn MTLComputePipelineState>>;
@@ -59,17 +252,12 @@ pub fn run() {
59252
println!("RESULT: no Metal command queue");
60253
return;
61254
};
62-
// The same source compiled with Metal's defaults, with the invariance setting wgpu uses, and without fast math
63-
let default_options = MTLCompileOptions::new();
64-
let invariance_options = MTLCompileOptions::new();
65-
invariance_options.setPreserveInvariance(true);
66-
let precise_options = MTLCompileOptions::new();
67-
#[allow(deprecated)]
68-
precise_options.setFastMathEnabled(false);
69-
let compile_variants = [("Metal's default options", &default_options), ("preserveInvariance, as wgpu sets", &invariance_options), ("fast math disabled", &precise_options)];
70-
71-
let pipeline = |name: &str, options: &MTLCompileOptions| {
72-
let library = match device.newLibraryWithSource_options_error(&NSString::from_str(SHADERS), Some(options)) {
255+
256+
// Compiled with wgpu's own Metal compile options, so only the source differs between variants
257+
let options = MTLCompileOptions::new();
258+
options.setPreserveInvariance(true);
259+
let pipeline = |source: &str, name: &str| {
260+
let library = match device.newLibraryWithSource_options_error(&NSString::from_str(source), Some(&options)) {
73261
Ok(library) => library,
74262
Err(error) => {
75263
println!("RESULT: Metal shader compilation failed: {error}");
@@ -80,34 +268,77 @@ pub fn run() {
80268
device.newComputePipelineStateWithFunction_error(&function).ok()
81269
};
82270

83-
// FP32 compute
271+
let variants = [
272+
("hand-written", pipeline(SHADERS, "alu")),
273+
("wgpu's output, default checks", pipeline(NAGA_CHECKED, "main_")),
274+
("wgpu's output, no checks", pipeline(NAGA_UNCHECKED, "main_")),
275+
("wgpu's output, no checks, `for` loop", pipeline(NAGA_UNCHECKED_FOR_LOOP, "main_")),
276+
];
277+
let variants: Vec<_> = variants.into_iter().filter_map(|(name, pipeline)| Some((name, pipeline?))).collect();
278+
84279
let invocations = ALU_WORKGROUPS as u64 * 256;
85-
let Some(output) = device.newBufferWithLength_options(invocations as usize * 16, MTLResourceOptions::StorageModePrivate) else { return };
280+
let output_bytes = invocations as usize * 16;
281+
let Some(output) = device.newBufferWithLength_options(output_bytes, MTLResourceOptions::StorageModePrivate) else { return };
282+
let buffer_sizes = (output_bytes as u32).to_le_bytes();
283+
let Some(pointer) = NonNull::new(buffer_sizes.as_ptr() as *mut std::ffi::c_void) else { return };
284+
let Some(sizes) = (unsafe { device.newBufferWithBytes_length_options(pointer, buffer_sizes.len(), MTLResourceOptions::StorageModeShared) }) else { return };
285+
let buffers = [&output, &sizes];
286+
287+
// Keeps the GPU busy long enough for its clocks to ramp up before anything is timed
288+
let warmup_start = Instant::now();
289+
let mut warmup_dispatches = 0;
290+
while warmup_start.elapsed() < SUSTAINED_WARMUP {
291+
let _ = time_once(&queue, &variants[0].1, &buffers, ALU_WORKGROUPS as usize);
292+
warmup_dispatches += 1;
293+
}
294+
295+
// Each round times every variant in turn, so drift over time (clocks, other VMs on the host) affects them all alike
296+
let mut gpu_milliseconds = vec![Vec::new(); variants.len()];
297+
for _ in 0..INTERLEAVED_ROUNDS {
298+
for (index, (_, pipeline)) in variants.iter().enumerate() {
299+
for _ in 0..(TIMED_RUNS / INTERLEAVED_ROUNDS) {
300+
if let Some((gpu, _)) = time_once(&queue, pipeline, &buffers, ALU_WORKGROUPS as usize) {
301+
gpu_milliseconds[index].push(gpu);
302+
}
303+
}
304+
}
305+
}
306+
86307
let gflops = |milliseconds: f64| invocations as f64 * ALU_FLOPS_PER_INVOCATION / (milliseconds * 1e6);
87-
println!("\n== FP32 compute, direct Metal");
88-
for (variant, options) in compile_variants {
89-
let Some(alu_pipeline) = pipeline("alu", options) else { continue };
90-
let (gpu_milliseconds, wall_milliseconds) = time_dispatch(&queue, &alu_pipeline, &[&output], ALU_WORKGROUPS as usize);
91-
println!("-- Compiled with {variant}");
92-
crate::report(&gpu_milliseconds, &wall_milliseconds, "GFLOPS", gflops);
308+
println!("\n== FP32 compute, direct Metal, after {warmup_dispatches} warmup dispatches over {:?}, interleaved in {INTERLEAVED_ROUNDS} rounds", SUSTAINED_WARMUP);
309+
for ((name, _), times) in variants.iter().zip(&gpu_milliseconds) {
310+
if times.is_empty() {
311+
println!("-- {name}: GPU time unavailable");
312+
} else {
313+
println!("-- {name}: {:.1} GFLOPS from median GPU time; GPU time (ms): {}", gflops(median(times)), summarize(times));
314+
}
93315
}
94-
let Some(bandwidth_pipeline) = pipeline("bandwidth", &default_options) else { return };
95316

96317
// Memory bandwidth, with nonzero contents in case the hardware compresses zeros
318+
let Some(bandwidth_pipeline) = pipeline(SHADERS, "bandwidth") else { return };
97319
let pattern: Vec<u8> = (0..BANDWIDTH_BUFFER_BYTES).map(|i| (i * 2_654_435_761 >> 24) as u8).collect();
98320
let Some(pointer) = NonNull::new(pattern.as_ptr() as *mut std::ffi::c_void) else { return };
99321
let source = unsafe { device.newBufferWithBytes_length_options(pointer, pattern.len(), MTLResourceOptions::StorageModeShared) };
100322
let destination = device.newBufferWithLength_options(BANDWIDTH_BUFFER_BYTES as usize, MTLResourceOptions::StorageModePrivate);
101323
let (Some(source), Some(destination)) = (source, destination) else { return };
102324
let workgroups = (BANDWIDTH_BUFFER_BYTES / 16 / 256) as usize;
103-
let (gpu_milliseconds, wall_milliseconds) = time_dispatch(&queue, &bandwidth_pipeline, &[&source, &destination], workgroups);
325+
326+
let mut gpu_milliseconds = Vec::new();
327+
let mut wall_milliseconds = Vec::new();
328+
for i in 0..(WARMUP_RUNS + TIMED_RUNS) {
329+
let Some((gpu, wall)) = time_once(&queue, &bandwidth_pipeline, &[&source, &destination], workgroups) else { continue };
330+
if i >= WARMUP_RUNS {
331+
gpu_milliseconds.push(gpu);
332+
wall_milliseconds.push(wall);
333+
}
334+
}
104335
let gigabytes_per_second = |milliseconds: f64| 2. * BANDWIDTH_BUFFER_BYTES as f64 / (milliseconds * 1e6);
105336
println!("\n== Memory bandwidth, direct Metal");
106337
crate::report(&gpu_milliseconds, &wall_milliseconds, "GB/s", gigabytes_per_second);
107338

108-
if !gpu_milliseconds.is_empty() {
109-
let overhead: Vec<f64> = wall_milliseconds.iter().zip(&gpu_milliseconds).map(|(wall, gpu)| (wall - gpu) * 1e3).collect();
110-
println!("Wall minus GPU time per dispatch (µs): {}, median {:.1}", summarize(&overhead), median(&overhead));
339+
let overhead: Vec<f64> = wall_milliseconds.iter().zip(&gpu_milliseconds).map(|(wall, gpu)| (wall - gpu) * 1e3).collect();
340+
if !overhead.is_empty() {
341+
println!("Wall minus GPU time per dispatch (µs): {}", summarize(&overhead));
111342
}
112343
}
113344

@@ -128,34 +359,22 @@ fn describe(device: &Device) {
128359
println!("counter sets: {}", if sets.is_empty() { "none".to_string() } else { sets.join(", ") });
129360
}
130361

131-
// Returns per-run GPU times from the command buffer (empty if Metal reports none) and wall-clock times around commit-and-wait
132-
fn time_dispatch(queue: &Queue, pipeline: &Pipeline, buffers: &[&Buffer], workgroups: usize) -> (Vec<f64>, Vec<f64>) {
133-
let mut gpu_milliseconds = Vec::new();
134-
let mut wall_milliseconds = Vec::new();
135-
136-
for i in 0..(WARMUP_RUNS + TIMED_RUNS) {
137-
let Some(command_buffer) = queue.commandBuffer() else { break };
138-
let Some(encoder) = command_buffer.computeCommandEncoder() else { break };
139-
encoder.setComputePipelineState(pipeline);
140-
for (index, buffer) in buffers.iter().enumerate() {
141-
unsafe { encoder.setBuffer_offset_atIndex(Some(buffer), 0, index) };
142-
}
143-
encoder.dispatchThreadgroups_threadsPerThreadgroup(MTLSize { width: workgroups, height: 1, depth: 1 }, MTLSize { width: 256, height: 1, depth: 1 });
144-
encoder.endEncoding();
145-
146-
let start = Instant::now();
147-
command_buffer.commit();
148-
command_buffer.waitUntilCompleted();
149-
let wall = start.elapsed().as_secs_f64() * 1e3;
150-
151-
let gpu = (command_buffer.GPUEndTime() - command_buffer.GPUStartTime()) * 1e3;
152-
if i >= WARMUP_RUNS {
153-
wall_milliseconds.push(wall);
154-
if command_buffer.GPUStartTime() > 0. && gpu > 0. {
155-
gpu_milliseconds.push(gpu);
156-
}
157-
}
362+
// One dispatch, returning its GPU time from the command buffer (if Metal reports one) and the wall-clock time around commit-and-wait
363+
fn time_once(queue: &Queue, pipeline: &Pipeline, buffers: &[&Buffer], workgroups: usize) -> Option<(f64, f64)> {
364+
let command_buffer = queue.commandBuffer()?;
365+
let encoder = command_buffer.computeCommandEncoder()?;
366+
encoder.setComputePipelineState(pipeline);
367+
for (index, buffer) in buffers.iter().enumerate() {
368+
unsafe { encoder.setBuffer_offset_atIndex(Some(buffer), 0, index) };
158369
}
370+
encoder.dispatchThreadgroups_threadsPerThreadgroup(MTLSize { width: workgroups, height: 1, depth: 1 }, MTLSize { width: 256, height: 1, depth: 1 });
371+
encoder.endEncoding();
372+
373+
let start = Instant::now();
374+
command_buffer.commit();
375+
command_buffer.waitUntilCompleted();
376+
let wall = start.elapsed().as_secs_f64() * 1e3;
159377

160-
(gpu_milliseconds, wall_milliseconds)
378+
let gpu = (command_buffer.GPUEndTime() - command_buffer.GPUStartTime()) * 1e3;
379+
(command_buffer.GPUStartTime() > 0. && gpu > 0.).then_some((gpu, wall))
161380
}

0 commit comments

Comments
 (0)