77 lines
2.2 KiB
Rust
77 lines
2.2 KiB
Rust
use candle_metal_kernels::{call_affine, Kernels};
|
|
use metal::objc::rc::autoreleasepool;
|
|
use metal::{Device, MTLResourceOptions};
|
|
use rand;
|
|
use std::any::type_name;
|
|
use std::time::Instant;
|
|
|
|
fn main() {
|
|
let device = Device::system_default().unwrap();
|
|
let kernels = Kernels::new();
|
|
|
|
let f32_1k = (0..1000).map(|_| rand::random::<f32>()).collect::<Vec<_>>();
|
|
let f32_10k = (0..10000)
|
|
.map(|_| rand::random::<f32>())
|
|
.collect::<Vec<_>>();
|
|
let f32_100k = (0..100000)
|
|
.map(|_| rand::random::<f32>())
|
|
.collect::<Vec<_>>();
|
|
|
|
println!(
|
|
"{0: <5} | {1: <19} | {2: <6} | {3: <5} | {4: <11} | {5: <11}",
|
|
"dtype", "kernel", "size", "runs", "total time", "avg time"
|
|
);
|
|
|
|
// f32
|
|
run_affine_bench(&device, &kernels, &f32_1k);
|
|
run_affine_bench(&device, &kernels, &f32_10k);
|
|
run_affine_bench(&device, &kernels, &f32_100k);
|
|
}
|
|
|
|
fn run_affine_bench<T: Clone>(device: &Device, kernels: &Kernels, v: &[T]) {
|
|
let command_queue = device.new_command_queue();
|
|
let options = MTLResourceOptions::StorageModeManaged;
|
|
|
|
let iterations = 10000;
|
|
let input = device.new_buffer_with_data(
|
|
v.as_ptr() as *const core::ffi::c_void,
|
|
core::mem::size_of_val(v) as u64,
|
|
options,
|
|
);
|
|
let mut output = device.new_buffer(core::mem::size_of_val(v) as u64, options);
|
|
|
|
let mul: f32 = 1.2345;
|
|
let add: f32 = 2.3456;
|
|
let total_time = autoreleasepool(|| {
|
|
let command_buffer = command_queue.new_command_buffer();
|
|
let start = Instant::now();
|
|
for _ in 0..iterations {
|
|
call_affine(
|
|
&device,
|
|
command_buffer,
|
|
&kernels,
|
|
"affine_float",
|
|
v.len(),
|
|
&input,
|
|
&mut output,
|
|
mul,
|
|
add,
|
|
)
|
|
.unwrap();
|
|
}
|
|
command_buffer.commit();
|
|
command_buffer.wait_until_completed();
|
|
|
|
start.elapsed()
|
|
});
|
|
println!(
|
|
"{0: <5} | {1: <19} | {2: <6} | {3: <5} | {4: <11?} | {5: <11?}",
|
|
type_name::<T>().split("::").last().unwrap(),
|
|
"affine",
|
|
v.len(),
|
|
iterations,
|
|
total_time,
|
|
total_time / iterations
|
|
);
|
|
}
|