Parallel reduction 3D
The purpose of this example is to demonstrate how to perform a parallel reduction operation on a 3D tensor using CubeCL. The reduction will sum the elements along the last dimension (depth) of the tensor, resulting in a 2D tensor.
A first try
We will start with a simple implementation of a parallel reduction on a 3D tensor. The goal is to reduce the tensor along the last dimension (depth) by summing the elements. This will result in a 2D tensor where each element is the sum of the corresponding elements in the depth dimension.
use std::marker::PhantomData;
use cubecl::benchmark::{Benchmark, TimingMethod};
use cubecl::{future, prelude::*};
use cubecl_example::gpu_tensor::GpuTensor; // Change to the path of your own module containing the GpuTensor
pub struct ReductionBench<R: Runtime, F: Float + CubeElement> {
input_shape: Vec<usize>,
client: ComputeClient<R::Server, R::Channel>,
_f: PhantomData<F>,
}
const LINE_SIZE: u32 = 4;
impl<R: Runtime, F: Float + CubeElement> Benchmark for ReductionBench<R, F> {
type Input = GpuTensor<R, F>;
type Output = GpuTensor<R, F>;
fn prepare(&self) -> Self::Input {
GpuTensor::<R, F>::arange(self.input_shape.clone(), &self.client)
}
fn name(&self) -> String {
format!("{}-reduction-{:?}", R::name(&self.client), self.input_shape).to_lowercase()
}
fn sync(&self) {
future::block_on(self.client.sync())
}
fn execute(&self, input: Self::Input) -> Self::Output {
let output_shape: Vec<usize> = vec![self.input_shape[0]];
let output = GpuTensor::<R, F>::empty(output_shape, &self.client);
unsafe {
reduce_matrix::launch_unchecked::<F, R>(
&self.client,
CubeCount::Static(1, 1, 1),
CubeDim::new(self.input_shape[0] as u32, self.input_shape[1] as u32, 1),
input.into_tensor_arg(LINE_SIZE as u8),
output.into_tensor_arg(LINE_SIZE as u8),
);
}
output
}
}
// Note the addition of the [Line] struct inside the tensor to guarantee that the data is contiguous and can be parallelized.
#[cube(launch_unchecked)]
fn reduce_matrix<F: Float>(input: &Tensor<Line<F>>, output: &mut Tensor<Line<F>>) {
let mut acc = Line::new(F::new(0.0f32)); // A [Line] is also necessary here
for i in 0..input.shape(2) / LINE_SIZE {
acc = acc + input[UNIT_POS_X * input.stride(0) + UNIT_POS_Y * input.stride(1) + i];
}
output[UNIT_POS_X * input.stride(0) + UNIT_POS_Y] = acc;
}
pub fn launch<R: Runtime, F: Float + CubeElement>(device: &R::Device) {
let client = R::client(&device);
let bench1 = ReductionBench::<R, F> {
input_shape: vec![64, 256, 1024],
client: client.clone(),
_f: PhantomData,
};
let bench2 = ReductionBench::<R, F> {
input_shape: vec![64, 64, 4096],
client: client.clone(),
_f: PhantomData,
};
for bench in [bench1, bench2] {
println!("{}", bench.name());
println!("{}", bench.run(TimingMethod::System));
}
}
fn main() {
launch::<cubecl::wgpu::WgpuRuntime, f32>(&Default::default());
}Let's try to run this code.
wgpu error: Validation Error
Caused by:
In Device::create_compute_pipeline, label = 'reduce_matrix_f32'
Error matching shader requirements against the pipeline
Shader entry point's workgroup size [64, 256, 1] (16384 total invocations) must be less or equal to the per-dimension limit [1024, 1024, 1024] and the total invocation limit 1024
What happened? The error message indicates that the workgroup size exceeds the limits imposed by the WebGPU backend. The total number of invocations (64 * 256 * 1 = 16384) exceeds the maximum allowed invocations per workgroup, which is 1024. In other words, the CubeDim size is too large for the GPU to handle. We needs to find another way to parallelize the reduction operation without exceeding the limits.
A better approach
To address the issue, we will parallelize with the CUBE_COUNT and CUBE_POS variables, which will allow us to launch multiple invocation in parallel without exceeding the limits of the CUBE_DIM. The CUBE_COUNT variable will determine how many invocation we will launch, and the CUBE_POS variable will determine the position of each invocation in the 3D tensor.
use std::marker::PhantomData;
use cubecl::benchmark::{Benchmark, TimingMethod};
use cubecl::{future, prelude::*};
use cubecl_example::gpu_tensor::GpuTensor; // Change to the path of your own module containing the GpuTensor
pub struct ReductionBench<R: Runtime, F: Float + CubeElement> {
input_shape: Vec<usize>,
client: ComputeClient<R::Server, R::Channel>,
_f: PhantomData<F>,
}
const LINE_SIZE: u32 = 4;
impl<R: Runtime, F: Float + CubeElement> Benchmark for ReductionBench<R, F> {
type Input = GpuTensor<R, F>;
type Output = GpuTensor<R, F>;
fn prepare(&self) -> Self::Input {
GpuTensor::<R, F>::arange(self.input_shape.clone(), &self.client)
}
fn name(&self) -> String {
format!("{}-reduction-{:?}", R::name(&self.client), self.input_shape).to_lowercase()
}
fn sync(&self) {
future::block_on(self.client.sync())
}
fn execute(&self, input: Self::Input) -> Self::Output {
let output_shape: Vec<usize> = vec![self.input_shape[0]];
let output = GpuTensor::<R, F>::empty(output_shape, &self.client);
unsafe {
reduce_matrix::launch_unchecked::<F, R>(
&self.client,
CubeCount::Static(self.input_shape[0] as u32, 1, 1),
CubeDim::new(self.input_shape[1] as u32, 1, 1),
input.into_tensor_arg(LINE_SIZE as u8),
output.into_tensor_arg(LINE_SIZE as u8),
);
}
output
}
}
// Note the addition of the [Line] struct inside the tensor to guarantee that the data is contiguous and can be parallelized.
#[cube(launch_unchecked)]
fn reduce_matrix<F: Float>(input: &Tensor<Line<F>>, output: &mut Tensor<Line<F>>) {
let mut acc = Line::new(F::new(0.0f32)); // A [Line] is also necessary here
for i in 0..input.shape(2) / LINE_SIZE {
acc = acc + input[CUBE_POS_X * input.stride(0) + UNIT_POS_X * input.stride(1) + i];
}
output[CUBE_POS_X * output.stride(0) + UNIT_POS_X] = acc;
}
pub fn launch<R: Runtime, F: Float + CubeElement>(device: &R::Device) {
let client = R::client(&device);
let bench1 = ReductionBench::<R, F> {
input_shape: vec![64, 256, 1024],
client: client.clone(),
_f: PhantomData,
};
let bench2 = ReductionBench::<R, F> {
input_shape: vec![64, 64, 4096],
client: client.clone(),
_f: PhantomData,
};
for bench in [bench1, bench2] {
println!("{}", bench.name());
println!("{}", bench.run(TimingMethod::System));
}
}
fn main() {
launch::<cubecl::wgpu::WgpuRuntime, f32>(&Default::default());
}Now, let's run the code again.
wgpu<wgsl>-reduction-[64, 256, 1024]
―――――――― Result ―――――――――
Timing system
Samples 10
Mean 1.483ms
Variance 27.000ns
Median 1.535ms
Min 1.239ms
Max 1.808ms
―――――――――――――――――――――――――
wgpu<wgsl>-reduction-[64, 64, 4096]
―――――――― Result ―――――――――
Timing system
Samples 10
Mean 924.409µs
Variance 189.000ns
Median 945.270µs
Min 600.110µs
Max 2.098ms
―――――――――――――――――――――――――
It runs and it is fast! The reduction operation is now parallelized across multiple invocations, and we can see that the performance is significantly improved compared to the previous implementation. The results show that the reduction operation is efficient and can handle larger tensors without exceeding the GPU limits. It's also almost the same speed as the 2D reduction, even if there's even more elements to reduce. This is because the reduction is now parallelized across multiple cubes and hyper-cubes, allowing the GPU to process the data more efficiently. See the parallel reduction if you need a refresher on the different parallelization level used in CubeCL. It is also worth noting that the performance and optimal CUBE_COUNT and CUBE_DIM values may vary depending on the GPU architecture and the specific workload. You may need to experiment with different values to find the best configuration for your use case.