From 6eb44d1bcef00052e2a56424af82bd14d65f2df8 Mon Sep 17 00:00:00 2001 From: Ivar Flakstad <69173633+ivarflakstad@users.noreply.github.com> Date: Mon, 1 Jan 2024 20:22:44 +0100 Subject: [PATCH] Added fill bench --- candle-core/Cargo.toml | 4 ++ candle-core/benches/fill.rs | 57 ++++++++++++++++ candle-core/src/metal_backend.rs | 10 +-- candle-metal-kernels/Cargo.toml | 5 ++ candle-metal-kernels/src/lib.rs | 104 ++++++++++++------------------ candle-metal-kernels/src/tests.rs | 8 +-- 6 files changed, 117 insertions(+), 71 deletions(-) create mode 100644 candle-core/benches/fill.rs diff --git a/candle-core/Cargo.toml b/candle-core/Cargo.toml index 91655f57..6bd12589 100644 --- a/candle-core/Cargo.toml +++ b/candle-core/Cargo.toml @@ -49,3 +49,7 @@ metal = ["dep:metal", "dep:candle-metal-kernels"] name = "matmul" harness = false +[[bench]] +name = "fill" +harness = false + diff --git a/candle-core/benches/fill.rs b/candle-core/benches/fill.rs new file mode 100644 index 00000000..9bcb4775 --- /dev/null +++ b/candle-core/benches/fill.rs @@ -0,0 +1,57 @@ +use candle_core::{DType, Device, Tensor}; +use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; +use std::time::Instant; + +fn run(shape: (usize, usize, usize), dtype: DType, device: &Device) { + Tensor::ones(shape, dtype, device).unwrap(); +} + +fn criterion_benchmark(c: &mut Criterion) { + let b = 1; + let rows = 4096; + let columns = 4096; + + let flops = b * rows * columns; + + let device1 = Device::new_metal(0).unwrap(); + let device2 = device1.clone(); + + let mut group = c.benchmark_group("fill_metal_u8"); + group.throughput(Throughput::Bytes(flops as u64)); + group.bench_function("iter", move |bencher| { + bencher.iter_custom(|iters| { + let start = Instant::now(); + for _i in 0..iters { + run(black_box((b, rows, columns)), black_box(DType::U8), black_box(&device1)); + } + if let Device::Metal(device) = &device1 { + device.wait_until_completed().unwrap(); + } else { + panic!("Expected metal device"); + } + start.elapsed() + }) + }); + group.finish(); + + let mut group = c.benchmark_group("fill_metal_f32"); + group.throughput(Throughput::Bytes((flops * DType::F32.size_in_bytes()) as u64)); + group.bench_function("iter", move |bencher| { + bencher.iter_custom(|iters| { + let start = Instant::now(); + for _i in 0..iters { + run(black_box((b, rows, columns)), black_box(DType::F32), black_box(&device2)); + } + if let Device::Metal(device) = &device2 { + device.wait_until_completed().unwrap(); + } else { + panic!("Expected metal device"); + } + start.elapsed() + }) + }); + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/candle-core/src/metal_backend.rs b/candle-core/src/metal_backend.rs index f2b55d4e..21eb1336 100644 --- a/candle-core/src/metal_backend.rs +++ b/candle-core/src/metal_backend.rs @@ -3,7 +3,7 @@ use crate::conv::{ParamsConv1D, ParamsConv2D, ParamsConvTranspose1D, ParamsConvT use crate::op::{BinaryOpT, CmpOp, ReduceOp, UnaryOpT}; use crate::{CpuStorage, DType, Layout, Result, Shape}; use candle_metal_kernels; -use candle_metal_kernels::{CallFill, Fill, Kernels}; +use candle_metal_kernels::{FillOp, Unary, Kernels}; use half::{bf16, f16}; use metal; use metal::{Buffer, CommandBuffer, CommandQueue, MTLResourceOptions, NSUInteger}; @@ -1405,9 +1405,9 @@ impl BackendDevice for MetalDevice { let command_buffer = self.command_buffer()?; command_buffer.set_label("zeros"); - // This assumes the specific zero type DType is equal to 0x00u8 + // This assumes the zero value of this DType is equal to 0x00u8 // (which is true for all current types) - Fill::call_fill( + Unary::fill( &self.device, &command_buffer, &self.kernels, @@ -1421,13 +1421,13 @@ impl BackendDevice for MetalDevice { } fn ones_impl(&self, shape: &Shape, dtype: DType) -> Result { - let buffer = self.new_buffer(shape.elem_count(), dtype, "zeros")?; + let buffer = self.new_buffer(shape.elem_count(), dtype, "ones")?; let command_buffer = self.command_buffer()?; command_buffer.set_label("ones"); macro_rules! fill { ($value:expr) => { - Fill::call_fill( + Unary::fill( &self.device, &command_buffer, &self.kernels, diff --git a/candle-metal-kernels/Cargo.toml b/candle-metal-kernels/Cargo.toml index 6c64a8e5..25446d29 100644 --- a/candle-metal-kernels/Cargo.toml +++ b/candle-metal-kernels/Cargo.toml @@ -19,3 +19,8 @@ num-traits = "0.2.17" [dev-dependencies] rand = "0.8.5" +criterion = "0.5.1" + +[[bench]] +name = "fill" +harness = false diff --git a/candle-metal-kernels/src/lib.rs b/candle-metal-kernels/src/lib.rs index e0985d94..f5b0653b 100644 --- a/candle-metal-kernels/src/lib.rs +++ b/candle-metal-kernels/src/lib.rs @@ -1573,19 +1573,19 @@ pub fn call_upsample_nearest_2d( Ok(()) } -#[inline] +#[inline(always)] fn divide(m: usize, b: usize) -> NSUInteger { ((m + b - 1) / b) as NSUInteger } -pub struct Fill { +pub struct Unary { _marker: PhantomData, } -pub trait CallFill { - const KERNEL_NAME: &'static str; +pub trait FillOp { + const FILL_KERNEL: &'static str; - fn call_fill( + fn fill( device: &Device, command_buffer: &CommandBufferRef, kernels: &Kernels, @@ -1598,11 +1598,26 @@ pub trait CallFill { macro_rules ! impl_call_fill { ($($t:ty),*) => { $( - impl CallFill<$t> for Fill<$t> { - const KERNEL_NAME: &'static str = concat!("fill_", stringify!($t)); + impl FillOp<$t> for Unary<$t> { + const FILL_KERNEL: &'static str = concat!("fill_", stringify!($t)); - fn call_fill(device: &Device, command_buffer: &CommandBufferRef, kernels: &Kernels, elem_count: usize, buffer: &Buffer, value: $t) -> Result<(), MetalKernelError> { - _call_fill(device, command_buffer, kernels, Self::KERNEL_NAME, elem_count, buffer, value) + #[inline(always)] + fn fill(device: &Device, command_buffer: &CommandBufferRef, kernels: &Kernels, elem_count: usize, buffer: &Buffer, value: $t) -> Result<(), MetalKernelError> { + let pipeline = kernels.load_pipeline(device, Source::Fill, Self::FILL_KERNEL)?; + let encoder = command_buffer.new_compute_command_encoder(); + encoder.wait_for_fence(&kernels.fence); + encoder.set_compute_pipeline_state(&pipeline); + encoder.set_threadgroup_memory_length(0, elem_count as NSUInteger); + + set_params!(encoder, (buffer, value, elem_count)); + + let (thread_group_count, thread_group_size) = linear_split(&pipeline, elem_count); + encoder.dispatch_thread_groups(thread_group_count, thread_group_size); + encoder.use_resource(buffer, metal::MTLResourceUsage::Write); + encoder.update_fence(&kernels.fence); + encoder.end_encoding(); + + Ok(()) } } )* @@ -1610,10 +1625,11 @@ macro_rules ! impl_call_fill { } impl_call_fill!(u32, i64, f16, bf16, f32); -impl CallFill for Fill { - const KERNEL_NAME: &'static str = ""; +impl FillOp for Unary { + const FILL_KERNEL: &'static str = ""; - fn call_fill( + #[inline(always)] + fn fill( _: &Device, command_buffer: &CommandBufferRef, kernels: &Kernels, @@ -1621,58 +1637,22 @@ impl CallFill for Fill { buffer: &Buffer, value: u8, ) -> Result<(), MetalKernelError> { - _call_blit_fill(command_buffer, kernels, elem_count, buffer, value) + let blit = command_buffer.new_blit_command_encoder(); + blit.wait_for_fence(&kernels.fence); + blit.fill_buffer( + &buffer, + metal::NSRange { + location: 0, + length: elem_count as NSUInteger, + }, + value, + ); + blit.update_fence(&kernels.fence); + blit.end_encoding(); + + Ok(()) } } -fn _call_blit_fill( - command_buffer: &CommandBufferRef, - kernels: &Kernels, - elem_count: usize, - buffer: &Buffer, - value: u8, -) -> Result<(), MetalKernelError> { - let blit = command_buffer.new_blit_command_encoder(); - blit.wait_for_fence(&kernels.fence); - blit.fill_buffer( - &buffer, - metal::NSRange { - location: 0, - length: elem_count as NSUInteger, - }, - value, - ); - blit.update_fence(&kernels.fence); - blit.end_encoding(); - - Ok(()) -} - -fn _call_fill( - device: &Device, - command_buffer: &CommandBufferRef, - kernels: &Kernels, - kernel_name: &'static str, - elem_count: usize, - buffer: &Buffer, - value: D, -) -> Result<(), MetalKernelError> { - let pipeline = kernels.load_pipeline(device, Source::Fill, kernel_name)?; - let encoder = command_buffer.new_compute_command_encoder(); - encoder.wait_for_fence(&kernels.fence); - encoder.set_compute_pipeline_state(&pipeline); - encoder.set_threadgroup_memory_length(0, elem_count as NSUInteger); - - set_params!(encoder, (buffer, value, elem_count)); - - let (thread_group_count, thread_group_size) = linear_split(&pipeline, elem_count); - encoder.dispatch_thread_groups(thread_group_count, thread_group_size); - encoder.use_resource(buffer, metal::MTLResourceUsage::Write); - encoder.update_fence(&kernels.fence); - encoder.end_encoding(); - - Ok(()) -} - #[cfg(test)] mod tests; diff --git a/candle-metal-kernels/src/tests.rs b/candle-metal-kernels/src/tests.rs index a4fb726f..b7bff740 100644 --- a/candle-metal-kernels/src/tests.rs +++ b/candle-metal-kernels/src/tests.rs @@ -808,7 +808,7 @@ fn gemm() { fn run_fill(elem_count: usize, value: T) -> Vec where - Fill: CallFill, + Unary: FillOp, { let device = device(); let fence = device.new_fence(); @@ -816,7 +816,7 @@ where let command_queue = device.new_command_queue(); let command_buffer = command_queue.new_command_buffer(); let buffer = new_buffer(&device, &vec![0.0f32; elem_count]); - Fill::::call_fill( + Unary::::fill( &device, command_buffer, &kernels, @@ -835,7 +835,7 @@ where fn fill() { fn assert_fill(value: T) where - Fill: CallFill, + Unary: FillOp, { for i in 0..4 { assert_eq!(run_fill(8 ^ i, value), vec![value; 8 ^ i]); @@ -847,4 +847,4 @@ fn fill() { assert_fill(f16::from_f32(1.23)); assert_fill(bf16::from_f32(4.56)); assert_fill(7.89f32); -} +} \ No newline at end of file