use std::collections::BTreeMap;
use std::time::{Duration, Instant};
use lampshade::{
ArgminByKey, Compactor, Context, GpuProfile, Histogram, KeyValue, KeyValueCompactor,
KeyValueSorter, MaskGenerator, Reducer, Scanner, Sorter, U32Predicate, U32Reduction,
};
use wgpu::util::DeviceExt;
const DEFAULT_INPUT_SIZES: [usize; 3] = [1_000_000, 10_000_000, 100_000_000];
const DEFAULT_SAMPLES: usize = 5;
const DEFAULT_WARMUP_MS: u64 = 1_000;
const DEFAULT_CASES: [ProfileCase; 10] = [
ProfileCase::Predicate(50),
ProfileCase::Histogram256,
ProfileCase::ReductionSum,
ProfileCase::ArgminByKey,
ProfileCase::Scan,
ProfileCase::Compact(50),
ProfileCase::KeyValueCompact(50),
ProfileCase::KeySort,
ProfileCase::KeyValueBounded16,
ProfileCase::KeyValueFullWidth,
];
#[derive(Clone, Copy)]
enum ProfileCase {
Predicate(u32),
Histogram256,
ReductionSum,
ArgminByKey,
Scan,
Compact(u32),
KeyValueCompact(u32),
KeySort,
KeyValueBounded16,
KeyValueFullWidth,
}
struct ProfileConfig {
samples: usize,
warmup: Duration,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let context = Context::init().await?;
if !context
.device
.features()
.contains(wgpu::Features::TIMESTAMP_QUERY)
{
return Err("selected adapter does not support timestamp queries".into());
}
let sizes = input_sizes()?;
let cases = profile_cases()?;
let config = profile_config()?;
println!(
"adapter={:?} vendor={} device={} device_type={:?} backend={:?} driver={:?} driver_info={:?} subgroup_min={} subgroup_max={} samples={} warmup_ms={}",
context.adapter_info.name,
context.adapter_info.vendor,
context.adapter_info.device,
context.adapter_info.device_type,
context.adapter_info.backend,
context.adapter_info.driver,
context.adapter_info.driver_info,
context.adapter_info.subgroup_min_size,
context.adapter_info.subgroup_max_size,
config.samples,
config.warmup.as_millis(),
);
println!(
"primitive,items,resident_wall_median_ms,gpu_elapsed_median_ms,dispatch_median_ms,inter_pass_gap_ms,wall_minus_gpu_ms"
);
for item_count in sizes {
for case in &cases {
match case {
ProfileCase::Predicate(selectivity) => {
profile_predicate(&context, item_count, *selectivity, &config).await?
}
ProfileCase::Histogram256 => {
profile_histogram(&context, item_count, &config).await?
}
ProfileCase::ReductionSum => {
profile_reduction_sum(&context, item_count, &config).await?
}
ProfileCase::ArgminByKey => {
profile_argmin_by_key(&context, item_count, &config).await?
}
ProfileCase::Scan => profile_scan(&context, item_count, &config).await?,
ProfileCase::Compact(selectivity) => {
profile_compaction(&context, item_count, *selectivity, &config).await?
}
ProfileCase::KeyValueCompact(selectivity) => {
profile_key_value_compaction(&context, item_count, *selectivity, &config)
.await?
}
ProfileCase::KeySort => profile_key_sort(&context, item_count, &config).await?,
ProfileCase::KeyValueBounded16 => {
profile_key_value_sort(&context, item_count, &config, false).await?
}
ProfileCase::KeyValueFullWidth => {
profile_key_value_sort(&context, item_count, &config, true).await?
}
}
}
}
Ok(())
}
async fn profile_histogram(
context: &Context,
item_count: usize,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
const BIN_COUNT: u32 = 256;
let input: Vec<_> = deterministic_keys(item_count)
.into_iter()
.map(|value| value & 0xff)
.collect();
let gpu_input = create_input(context, "Profile Histogram Input", &input);
let gpu_output = create_output(
context,
"Profile Histogram Output",
Histogram::output_buffer_size(BIN_COUNT)?,
);
let histogram = Histogram::from_context(context);
warm_up(
config.warmup,
|| histogram.histogram_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32, BIN_COUNT),
context,
)?;
let wall = measure_wall(
config.samples,
|| histogram.histogram_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32, BIN_COUNT),
context,
)?;
let _ = histogram
.profile_histogram_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32, BIN_COUNT)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
histogram
.profile_histogram_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32, BIN_COUNT)
.await?,
);
}
if profile_validation_enabled() {
let actual = histogram.histogram(&input, BIN_COUNT).await?;
let mut expected = vec![0_u32; BIN_COUNT as usize];
for value in &input {
expected[*value as usize] += 1;
}
if actual != expected {
return Err("histogram_256 validation failed".into());
}
println!("validation,histogram_256,{item_count},passed,{BIN_COUNT}");
}
report("histogram_256", item_count, wall, &profiles);
Ok(())
}
async fn profile_reduction_sum(
context: &Context,
item_count: usize,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input = deterministic_keys(item_count);
let gpu_input = create_input(context, "Profile Reduction Input", &input);
let gpu_output = create_output(
context,
"Profile Reduction Output",
Reducer::output_buffer_size(),
);
let mut reducer = Reducer::from_context(context);
warm_up(
config.warmup,
|| {
reducer.reduce_gpu_to_gpu(
&gpu_input,
&gpu_output,
item_count as u32,
U32Reduction::Sum,
)
},
context,
)?;
let wall = measure_wall(
config.samples,
|| {
reducer.reduce_gpu_to_gpu(
&gpu_input,
&gpu_output,
item_count as u32,
U32Reduction::Sum,
)
},
context,
)?;
let _ = reducer
.profile_reduce_gpu_to_gpu(
&gpu_input,
&gpu_output,
item_count as u32,
U32Reduction::Sum,
)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
reducer
.profile_reduce_gpu_to_gpu(
&gpu_input,
&gpu_output,
item_count as u32,
U32Reduction::Sum,
)
.await?,
);
}
report("reduction_sum", item_count, wall, &profiles);
if profile_validation_enabled() {
let actual = reducer.sum(&input).await?;
let expected = input
.iter()
.fold(0_u32, |sum, value| sum.wrapping_add(*value));
if actual != expected {
return Err(format!(
"reduction_sum validation failed: expected {expected}, got {actual}"
)
.into());
}
println!("validation,reduction_sum,{item_count},passed,1");
}
Ok(())
}
async fn profile_argmin_by_key(
context: &Context,
item_count: usize,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input: Vec<_> = deterministic_keys(item_count)
.into_iter()
.enumerate()
.map(|(index, key)| KeyValue::new(key, index as u32))
.collect();
let gpu_input = create_input(context, "Profile Argmin Input", &input);
let gpu_output = create_output(
context,
"Profile Argmin Output",
ArgminByKey::output_buffer_size(),
);
let mut selector = ArgminByKey::from_context(context);
warm_up(
config.warmup,
|| selector.argmin_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32),
context,
)?;
let wall = measure_wall(
config.samples,
|| selector.argmin_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32),
context,
)?;
let _ = selector
.profile_argmin_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
selector
.profile_argmin_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32)
.await?,
);
}
report("argmin_by_key", item_count, wall, &profiles);
if profile_validation_enabled() {
let actual = selector.argmin(&input).await?;
let expected = input
.iter()
.copied()
.min_by_key(|item| (item.key, item.value))
.ok_or("argmin validation input is empty")?;
if actual != expected {
return Err(format!(
"argmin_by_key validation failed: expected {expected:?}, got {actual:?}"
)
.into());
}
println!("validation,argmin_by_key,{item_count},passed,1");
}
Ok(())
}
async fn profile_predicate(
context: &Context,
item_count: usize,
selectivity: u32,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input = deterministic_keys(item_count);
let predicate = predicate_for_selectivity(selectivity);
let gpu_input = create_input(context, "Profile Predicate Input", &input);
let gpu_mask = create_output(
context,
"Profile Predicate Mask",
MaskGenerator::mask_buffer_size(item_count as u32)?,
);
let generator = MaskGenerator::from_context(context);
warm_up(
config.warmup,
|| generator.mask_gpu_to_gpu(&gpu_input, &gpu_mask, item_count as u32, predicate),
context,
)?;
let wall = measure_wall(
config.samples,
|| generator.mask_gpu_to_gpu(&gpu_input, &gpu_mask, item_count as u32, predicate),
context,
)?;
let _ = generator
.profile_mask_gpu_to_gpu(&gpu_input, &gpu_mask, item_count as u32, predicate)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
generator
.profile_mask_gpu_to_gpu(&gpu_input, &gpu_mask, item_count as u32, predicate)
.await?,
);
}
let primitive = format!("predicate_{selectivity}");
report(&primitive, item_count, wall, &profiles);
if profile_validation_enabled() {
let actual = generator.mask(&input, predicate).await?;
let expected: Vec<_> = input
.iter()
.map(|&value| u32::from(predicate_matches(value, predicate)))
.collect();
if actual != expected {
return Err(format!("{primitive} validation failed").into());
}
let selected = actual.iter().map(|&flag| flag as usize).sum::<usize>();
println!("validation,{primitive},{item_count},passed,{selected}");
}
Ok(())
}
async fn profile_compaction(
context: &Context,
item_count: usize,
selectivity: u32,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input = deterministic_keys(item_count);
let mask: Vec<u32> = (0..item_count)
.map(|index| u32::from(index % 100 < selectivity as usize))
.collect();
let gpu_input = create_input(context, "Profile Compaction Input", &input);
let gpu_mask = create_input(context, "Profile Compaction Mask", &mask);
let gpu_output = create_output(context, "Profile Compaction Output", gpu_input.size());
let gpu_count = create_output(
context,
"Profile Compaction Output Count",
size_of::<u32>() as u64,
);
let mut compactor = Compactor::from_context(context);
warm_up(
config.warmup,
|| {
compactor.compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
},
context,
)?;
let wall = measure_wall(
config.samples,
|| {
compactor.compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
},
context,
)?;
let _ = compactor
.profile_compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
compactor
.profile_compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
.await?,
);
}
let primitive = format!("compact_{selectivity}");
report(&primitive, item_count, wall, &profiles);
if profile_validation_enabled() {
let actual = compactor.compact(&input, &mask).await?;
let expected: Vec<_> = input
.iter()
.zip(&mask)
.filter_map(|(&value, &keep)| (keep == 1).then_some(value))
.collect();
if actual != expected {
return Err(format!(
"{primitive} validation failed: expected {} items, got {}",
expected.len(),
actual.len()
)
.into());
}
println!(
"validation,{primitive},{item_count},passed,{}",
actual.len()
);
}
Ok(())
}
async fn profile_key_value_compaction(
context: &Context,
item_count: usize,
selectivity: u32,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input: Vec<_> = deterministic_keys(item_count)
.into_iter()
.enumerate()
.map(|(index, key)| KeyValue::new(key, index as u32))
.collect();
let mask: Vec<u32> = (0..item_count)
.map(|index| u32::from(index % 100 < selectivity as usize))
.collect();
let gpu_input = create_input(context, "Profile Key-Value Compaction Input", &input);
let gpu_mask = create_input(context, "Profile Key-Value Compaction Mask", &mask);
let gpu_output = create_output(
context,
"Profile Key-Value Compaction Output",
gpu_input.size(),
);
let gpu_count = create_output(
context,
"Profile Key-Value Compaction Output Count",
size_of::<u32>() as u64,
);
let mut compactor = KeyValueCompactor::from_context(context);
warm_up(
config.warmup,
|| {
compactor.compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
},
context,
)?;
let wall = measure_wall(
config.samples,
|| {
compactor.compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
},
context,
)?;
let _ = compactor
.profile_compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
compactor
.profile_compact_gpu_to_gpu(
&gpu_input,
&gpu_mask,
&gpu_output,
&gpu_count,
item_count as u32,
)
.await?,
);
}
let primitive = format!("key_value_compact_{selectivity}");
report(&primitive, item_count, wall, &profiles);
if profile_validation_enabled() {
let actual = compactor.compact(&input, &mask).await?;
let expected: Vec<_> = input
.iter()
.zip(&mask)
.filter_map(|(&item, &keep)| (keep == 1).then_some(item))
.collect();
if actual != expected {
return Err(format!(
"{primitive} validation failed: expected {} items, got {}",
expected.len(),
actual.len()
)
.into());
}
println!(
"validation,{primitive},{item_count},passed,{}",
actual.len()
);
}
Ok(())
}
async fn profile_scan(
context: &Context,
item_count: usize,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input: Vec<u32> = (0..item_count as u32)
.map(|value| value ^ 0xA5A5_A5A5)
.collect();
let gpu_input = create_input(context, "Profile Scan Input", &input);
let gpu_output = create_output(context, "Profile Scan Output", gpu_input.size());
let mut scanner = Scanner::from_context(context);
warm_up(
config.warmup,
|| scanner.scan_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32),
context,
)?;
let wall = measure_wall(
config.samples,
|| scanner.scan_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32),
context,
)?;
let _ = scanner
.profile_scan_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
scanner
.profile_scan_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32)
.await?,
);
}
report("scan", item_count, wall, &profiles);
Ok(())
}
async fn profile_key_sort(
context: &Context,
item_count: usize,
config: &ProfileConfig,
) -> Result<(), Box<dyn std::error::Error>> {
let input = deterministic_keys(item_count);
let gpu_input = create_input(context, "Profile Sort Input", &input);
let gpu_output = create_output(context, "Profile Sort Output", gpu_input.size());
let mut sorter = Sorter::from_context(context);
warm_up(
config.warmup,
|| sorter.sort_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32),
context,
)?;
let wall = measure_wall(
config.samples,
|| sorter.sort_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32),
context,
)?;
let _ = sorter
.profile_sort_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
sorter
.profile_sort_gpu_to_gpu(&gpu_input, &gpu_output, item_count as u32)
.await?,
);
}
report("key_sort", item_count, wall, &profiles);
Ok(())
}
async fn profile_key_value_sort(
context: &Context,
item_count: usize,
config: &ProfileConfig,
full_width: bool,
) -> Result<(), Box<dyn std::error::Error>> {
let input: Vec<_> = deterministic_keys(item_count)
.into_iter()
.enumerate()
.map(|(index, key)| {
let key = if full_width { key } else { key & 0xffff };
KeyValue::new(key, index as u32)
})
.collect();
let gpu_input = create_input(context, "Profile Key-Value Input", &input);
let gpu_output = create_output(context, "Profile Key-Value Output", gpu_input.size());
let mut sorter = KeyValueSorter::from_context(context);
let key_bits = if full_width { u32::BITS } else { 16 };
warm_up(
config.warmup,
|| {
sorter.sort_gpu_to_gpu_with_key_bits(
&gpu_input,
&gpu_output,
item_count as u32,
key_bits,
)
},
context,
)?;
let wall = measure_wall(
config.samples,
|| {
sorter.sort_gpu_to_gpu_with_key_bits(
&gpu_input,
&gpu_output,
item_count as u32,
key_bits,
)
},
context,
)?;
let _ = sorter
.profile_sort_gpu_to_gpu_with_key_bits(&gpu_input, &gpu_output, item_count as u32, key_bits)
.await?;
let mut profiles = Vec::with_capacity(config.samples);
for _ in 0..config.samples {
profiles.push(
sorter
.profile_sort_gpu_to_gpu_with_key_bits(
&gpu_input,
&gpu_output,
item_count as u32,
key_bits,
)
.await?,
);
}
let primitive = if full_width {
"key_value_sort_full_width"
} else {
"key_value_sort_bounded16"
};
report(primitive, item_count, wall, &profiles);
Ok(())
}
fn measure_wall(
samples: usize,
mut submit: impl FnMut() -> Result<(), lampshade::Error>,
context: &Context,
) -> Result<Duration, lampshade::Error> {
let mut durations = Vec::with_capacity(samples);
for _ in 0..samples {
let started = Instant::now();
submit()?;
wait_for_gpu(&context.device)?;
durations.push(started.elapsed());
}
Ok(median(durations))
}
fn warm_up(
minimum_duration: Duration,
mut submit: impl FnMut() -> Result<(), lampshade::Error>,
context: &Context,
) -> Result<(), lampshade::Error> {
let started = Instant::now();
loop {
submit()?;
wait_for_gpu(&context.device)?;
if started.elapsed() >= minimum_duration {
break;
}
}
Ok(())
}
fn report(primitive: &str, item_count: usize, wall: Duration, profiles: &[GpuProfile]) {
let gpu_elapsed = median(profiles.iter().map(|profile| profile.gpu_elapsed).collect());
let dispatch_time = median(
profiles
.iter()
.map(|profile| profile.dispatch_time)
.collect(),
);
let inter_pass_gap = gpu_elapsed.saturating_sub(dispatch_time);
let wall_minus_gpu = wall.saturating_sub(gpu_elapsed);
println!(
"{primitive},{item_count},{:.3},{:.3},{:.3},{:.3},{:.3}",
milliseconds(wall),
milliseconds(gpu_elapsed),
milliseconds(dispatch_time),
milliseconds(inter_pass_gap),
milliseconds(wall_minus_gpu),
);
let mut stages: BTreeMap<&str, Vec<Duration>> = BTreeMap::new();
let mut spans: BTreeMap<String, Vec<Duration>> = BTreeMap::new();
for profile in profiles {
let mut sample_stages: BTreeMap<&str, Duration> = BTreeMap::new();
for span in &profile.spans {
*sample_stages.entry(stage(&span.label)).or_default() += span.duration;
spans
.entry(span.label.clone())
.or_default()
.push(span.duration);
}
for (stage, duration) in sample_stages {
stages.entry(stage).or_default().push(duration);
}
}
let stage_medians: Vec<_> = stages
.into_iter()
.map(|(stage, durations)| (stage, median(durations)))
.collect();
let stage_total = stage_medians
.iter()
.map(|(_, duration)| *duration)
.sum::<Duration>();
for (stage, median_duration) in stage_medians {
let percent = median_duration.as_secs_f64() / stage_total.as_secs_f64() * 100.0;
println!(
"stage,{primitive},{item_count},{stage},{:.3},{percent:.1}%",
milliseconds(median_duration)
);
}
for (label, durations) in spans {
println!(
"span,{primitive},{item_count},{label},{:.3}",
milliseconds(median(durations))
);
}
}
fn stage(label: &str) -> &'static str {
if label == "predicate.mask" {
"predicate"
} else if label == "histogram.count" || label.ends_with(".histogram") {
"histogram"
} else if label.ends_with(".prefix") {
"prefix"
} else if label.ends_with(".reduce") {
"reduce"
} else if label.ends_with(".scatter") {
"scatter"
} else if label.starts_with("compact.scan.") {
"scan"
} else if label.starts_with("radix.") && label.contains(".scan.") {
"histogram_scan"
} else if label.starts_with("reduction.") {
"reduce"
} else if label.contains(".level.") {
"scan"
} else {
"add"
}
}
fn median(mut durations: Vec<Duration>) -> Duration {
durations.sort_unstable();
let middle = durations.len() / 2;
if durations.len().is_multiple_of(2) {
(durations[middle - 1] + durations[middle]) / 2
} else {
durations[middle]
}
}
fn milliseconds(duration: Duration) -> f64 {
duration.as_secs_f64() * 1_000.0
}
fn deterministic_keys(item_count: usize) -> Vec<u32> {
let mut state = 0x9E37_79B9_u32 ^ item_count as u32;
(0..item_count)
.map(|_| {
state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
state
})
.collect()
}
fn predicate_for_selectivity(selectivity: u32) -> U32Predicate {
match selectivity {
0 => U32Predicate::LessThan(0),
100 => U32Predicate::LessThanOrEqual(u32::MAX),
_ => U32Predicate::LessThan(((1_u64 << 32) * u64::from(selectivity) / 100) as u32),
}
}
fn predicate_matches(value: u32, predicate: U32Predicate) -> bool {
match predicate {
U32Predicate::Equal(target) => value == target,
U32Predicate::NotEqual(target) => value != target,
U32Predicate::LessThan(target) => value < target,
U32Predicate::LessThanOrEqual(target) => value <= target,
U32Predicate::GreaterThan(target) => value > target,
U32Predicate::GreaterThanOrEqual(target) => value >= target,
U32Predicate::BetweenInclusive { min, max } => value >= min && value <= max,
}
}
fn create_input<T: bytemuck::Pod>(
context: &Context,
label: &'static str,
input: &[T],
) -> wgpu::Buffer {
context
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some(label),
contents: bytemuck::cast_slice(input),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
})
}
fn create_output(context: &Context, label: &'static str, size: u64) -> wgpu::Buffer {
context.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(label),
size,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_DST
| wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
})
}
fn wait_for_gpu(device: &wgpu::Device) -> Result<(), lampshade::Error> {
device.poll(wgpu::PollType::Wait {
submission_index: None,
timeout: None,
})?;
Ok(())
}
fn input_sizes() -> Result<Vec<usize>, Box<dyn std::error::Error>> {
let Some(raw) = std::env::var_os("WGPU_PRIMITIVES_PROFILE_ITEMS") else {
return Ok(DEFAULT_INPUT_SIZES.to_vec());
};
raw.to_string_lossy()
.split(',')
.map(|value| Ok(value.trim().replace('_', "").parse()?))
.collect()
}
fn profile_cases() -> Result<Vec<ProfileCase>, Box<dyn std::error::Error>> {
let Some(raw) = std::env::var_os("WGPU_PRIMITIVES_PROFILE_CASES") else {
return Ok(DEFAULT_CASES.to_vec());
};
let cases: Vec<_> = raw
.to_string_lossy()
.split(',')
.map(|value| {
let value = value.trim();
match value {
"histogram_256" => Ok(ProfileCase::Histogram256),
"reduction_sum" => Ok(ProfileCase::ReductionSum),
"argmin_by_key" => Ok(ProfileCase::ArgminByKey),
"scan" => Ok(ProfileCase::Scan),
"key_sort" => Ok(ProfileCase::KeySort),
"key_value_bounded16" => Ok(ProfileCase::KeyValueBounded16),
"key_value_full_width" => Ok(ProfileCase::KeyValueFullWidth),
value if value.starts_with("predicate_") => {
let selectivity: u32 = value["predicate_".len()..]
.parse()
.map_err(|error| format!("invalid predicate selectivity: {error}"))?;
if selectivity > 100 {
return Err(format!(
"predicate selectivity must be at most 100, got {selectivity}"
));
}
Ok(ProfileCase::Predicate(selectivity))
}
value if value.starts_with("key_value_compact_") => {
let selectivity: u32 = value["key_value_compact_".len()..]
.parse()
.map_err(|error| {
format!("invalid key-value compaction selectivity: {error}")
})?;
if selectivity > 100 {
return Err(format!(
"key-value compaction selectivity must be at most 100, got {selectivity}"
));
}
Ok(ProfileCase::KeyValueCompact(selectivity))
}
value if value.starts_with("compact_") => {
let selectivity: u32 = value["compact_".len()..]
.parse()
.map_err(|error| format!("invalid compaction selectivity: {error}"))?;
if selectivity > 100 {
return Err(format!(
"compaction selectivity must be at most 100, got {selectivity}"
));
}
Ok(ProfileCase::Compact(selectivity))
}
value => Err(format!(
"unknown WGPU_PRIMITIVES_PROFILE_CASES value {value:?}"
)),
}
})
.collect::<Result<_, _>>()?;
if cases.is_empty() {
return Err("WGPU_PRIMITIVES_PROFILE_CASES must not be empty".into());
}
Ok(cases)
}
fn profile_validation_enabled() -> bool {
std::env::var("WGPU_PRIMITIVES_PROFILE_VALIDATE")
.is_ok_and(|value| matches!(value.as_str(), "1" | "true" | "yes"))
}
fn profile_config() -> Result<ProfileConfig, Box<dyn std::error::Error>> {
let samples = std::env::var("WGPU_PRIMITIVES_PROFILE_SAMPLES")
.unwrap_or_else(|_| DEFAULT_SAMPLES.to_string())
.parse()?;
if samples == 0 {
return Err("WGPU_PRIMITIVES_PROFILE_SAMPLES must be greater than zero".into());
}
let warmup_ms = std::env::var("WGPU_PRIMITIVES_PROFILE_WARMUP_MS")
.unwrap_or_else(|_| DEFAULT_WARMUP_MS.to_string())
.parse()?;
Ok(ProfileConfig {
samples,
warmup: Duration::from_millis(warmup_ms),
})
}