use cubecl::{calculate_cube_count_elemwise, prelude::*};
use burn_tensor::{
ops::{conv::calculate_conv_output_size, ConvOptions},
Shape,
};
use crate::{
kernel::into_contiguous,
ops::{
numeric::{empty_device, zeros_device},
reshape,
},
tensor::JitTensor,
FloatElement, JitRuntime,
};
#[derive(CubeLaunch)]
struct Conv3dArgs {
conv_stride_0: u32,
conv_stride_1: u32,
conv_stride_2: u32,
dilation_0: u32,
dilation_1: u32,
dilation_2: u32,
padding_0: u32,
padding_1: u32,
padding_2: u32,
groups: u32,
}
#[cube(launch)]
fn conv3d_kernel<F: Float>(
input: Tensor<F>,
weight: Tensor<F>,
bias: Tensor<F>,
mut output: Tensor<F>,
args: Conv3dArgs,
#[comptime] kernel_size_0_unroll: Option<u32>,
#[comptime] kernel_size_1_unroll: Option<u32>,
#[comptime] kernel_size_2_unroll: Option<u32>,
) {
if ABSOLUTE_POS >= output.len() {
return;
}
let in_channels = weight.shape(1);
let kernel_size_0 = kernel_size_0_unroll.unwrap_or_else(|| weight.shape(2));
let unroll_0 = kernel_size_0_unroll.is_some();
let kernel_size_1 = kernel_size_1_unroll.unwrap_or_else(|| weight.shape(3));
let unroll_1 = kernel_size_1_unroll.is_some();
let kernel_size_2 = kernel_size_2_unroll.unwrap_or_else(|| weight.shape(4));
let unroll_2 = kernel_size_2_unroll.is_some();
let b = ABSOLUTE_POS / output.stride(0) % output.shape(0);
let oc = ABSOLUTE_POS / output.stride(1) % output.shape(1);
let od = ABSOLUTE_POS / output.stride(2) % output.shape(2);
let oh = ABSOLUTE_POS / output.stride(3) % output.shape(3);
let ow = ABSOLUTE_POS / output.stride(4) % output.shape(4);
let g = (weight.shape(0) + oc) % args.groups;
let ic_start = in_channels * g;
let ic_end = ic_start + in_channels;
let mut sum = bias[oc];
let id_base = od * args.conv_stride_0;
let ih_base = oh * args.conv_stride_1;
let iw_base = ow * args.conv_stride_2;
let weight_stride_1 = weight.stride(1);
let weight_stride_2 = weight.stride(2);
let weight_stride_3 = weight.stride(3);
let weight_stride_4 = weight.stride(4);
let input_stride_1 = input.stride(1);
let input_stride_2 = input.stride(2);
let input_stride_3 = input.stride(3);
let input_stride_4 = input.stride(4);
let input_shape_2 = input.shape(2);
let input_shape_3 = input.shape(3);
let input_shape_4 = input.shape(4);
let border_front = args.padding_0;
let border_top = args.padding_1;
let border_left = args.padding_2;
let border_back = input_shape_2 + args.padding_0;
let border_bottom = input_shape_3 + args.padding_1;
let border_right = input_shape_4 + args.padding_2;
let index_input_0 = b * input.stride(0);
let index_weight_0 = oc * weight.stride(0);
for ic in ic_start..ic_end {
let index_input_1 = ic * input_stride_1;
let index_weight_1 = (ic - ic_start) * weight_stride_1;
#[unroll(unroll_0)]
for kd in 0..kernel_size_0 {
#[unroll(unroll_1)]
for kh in 0..kernel_size_1 {
#[unroll(unroll_2)]
for kw in 0..kernel_size_2 {
let id = kd * args.dilation_0 + id_base;
let ih = kh * args.dilation_1 + ih_base;
let iw = kw * args.dilation_2 + iw_base;
let within_padding = id >= border_front
&& id < border_back
&& ih >= border_top
&& ih < border_bottom
&& iw >= border_left
&& iw < border_right;
if within_padding {
let id_pad = id - args.padding_0;
let ih_pad = ih - args.padding_1;
let iw_pad = iw - args.padding_2;
let index_input = index_input_0
+ index_input_1
+ id_pad * input_stride_2
+ ih_pad * input_stride_3
+ iw_pad * input_stride_4;
let index_weight = index_weight_0
+ index_weight_1
+ kd * weight_stride_2
+ kh * weight_stride_3
+ kw * weight_stride_4;
sum += input[index_input] * weight[index_weight];
}
}
}
}
}
output[ABSOLUTE_POS] = sum;
}
pub(crate) fn conv3d<R: JitRuntime, E: FloatElement>(
input: JitTensor<R>,
weight: JitTensor<R>,
bias: Option<JitTensor<R>>,
options: ConvOptions<3>,
) -> JitTensor<R> {
let input = into_contiguous(input);
let weight = into_contiguous(weight);
let [batch_size, _, in_depth, in_height, in_width] = input.shape.dims();
let [out_channels, _, kernel_0, kernel_1, kernel_2] = weight.shape.dims();
let out_0 = calculate_conv_output_size(
kernel_0,
options.stride[0],
options.padding[0],
options.dilation[0],
in_depth,
);
let out_1 = calculate_conv_output_size(
kernel_1,
options.stride[1],
options.padding[1],
options.dilation[1],
in_height,
);
let out_2 = calculate_conv_output_size(
kernel_2,
options.stride[2],
options.padding[2],
options.dilation[2],
in_width,
);
let shape_out = Shape::new([batch_size, out_channels, out_0, out_1, out_2]);
let output = empty_device::<R, E>(
input.client.clone(),
input.device.clone(),
shape_out.clone(),
);
let bias = match bias {
Some(bias) => {
let shape = Shape::from([bias.shape.dims[0], 1, 1, 1, 1]);
reshape(bias, shape)
}
None => {
let shape = Shape::from([output.shape.dims[0], 1, 1, 1, 1]);
zeros_device::<R, E>(input.client.clone(), input.device.clone(), shape)
}
};
let cube_dim = CubeDim::default();
let cube_count = calculate_cube_count_elemwise(output.shape.num_elements(), cube_dim);
conv3d_kernel::launch::<E, R>(
&input.client,
cube_count,
cube_dim,
input.as_tensor_arg::<E>(1),
weight.as_tensor_arg::<E>(1),
bias.as_tensor_arg::<E>(1),
output.as_tensor_arg::<E>(1),
Conv3dArgsLaunch::new(
ScalarArg::new(options.stride[0] as u32),
ScalarArg::new(options.stride[1] as u32),
ScalarArg::new(options.stride[2] as u32),
ScalarArg::new(options.dilation[0] as u32),
ScalarArg::new(options.dilation[1] as u32),
ScalarArg::new(options.dilation[2] as u32),
ScalarArg::new(options.padding[0] as u32),
ScalarArg::new(options.padding[1] as u32),
ScalarArg::new(options.padding[2] as u32),
ScalarArg::new(options.groups as u32),
),
Some(kernel_0 as u32),
Some(kernel_1 as u32),
Some(kernel_2 as u32),
);
output
}