Files
paddlepaddle--paddle/paddle/phi/kernels/impl/slow_conv_kernel_impl.cuh
T
2026-07-13 12:40:42 +08:00

954 lines
35 KiB
Plaintext

// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <numeric>
#include <type_traits>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include "paddle/common/flags.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/concat_kernel.h"
#include "paddle/phi/kernels/contiguous_kernel.h"
#include "paddle/phi/kernels/conv_kernel.h"
#include "paddle/phi/kernels/cpu/conv_util.h"
#include "paddle/phi/kernels/elementwise_add_kernel.h"
#include "paddle/phi/kernels/fill_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/batch_norm_utils.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#include "paddle/phi/kernels/funcs/im2col_slow.cuh"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/funcs/vol2col_slow.cuh"
#include "paddle/phi/kernels/reduce_sum_kernel.h"
#include "paddle/phi/kernels/slice_kernel.h"
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
namespace phi {
template <typename T>
static inline T div_rtn(T x, T y) {
int q = x / y;
int r = x % y;
if ((r != 0) && ((r < 0) != (y < 0))) --q;
return q;
}
template <typename C,
std::enable_if_t<std::is_integral_v<typename C::value_type>, int> = 0>
inline int64_t multiply_integers(const C& container) {
return std::accumulate(container.begin(),
container.end(),
static_cast<int64_t>(1),
std::multiplies<>());
}
template <typename Iter,
std::enable_if_t<std::is_integral_v<
typename std::iterator_traits<Iter>::value_type>,
int> = 0>
inline int64_t multiply_integers(Iter begin, Iter end) {
return std::accumulate(
begin, end, static_cast<int64_t>(1), std::multiplies<>());
}
template <int64_t dim>
std::vector<int64_t> GetOutputSpatialSize(
const DenseTensor& input,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& stride_size,
const std::vector<int64_t>& pad_size,
const std::vector<int64_t>& dilation_size) {
std::vector<int64_t> sizes;
auto input_dim = input.dims().size();
for (int64_t index = 0; index < dim; ++index) {
int64_t input_size = input.dims()[index + input_dim - dim];
int64_t kernel = kernel_size[index];
int64_t stride = stride_size[index];
int64_t pad = pad_size[index];
int64_t dilation = dilation_size[index];
int64_t numerator = input_size + 2 * pad - (dilation * (kernel - 1) + 1);
int64_t size = div_rtn<int64_t>(numerator, stride) + 1;
sizes.push_back(size);
}
return sizes;
}
template <int64_t dim>
std::vector<int64_t> GetOutputSize(const DenseTensor& input,
const DenseTensor& weight,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& stride_size,
const std::vector<int64_t>& pad_size,
const std::vector<int64_t>& dilation_size) {
auto output_size = GetOutputSpatialSize<dim>(
input, kernel_size, stride_size, pad_size, dilation_size);
output_size.insert(output_size.begin(), weight.dims()[0]);
if (input.dims().size() == dim + 2) {
output_size.insert(output_size.begin(), input.dims()[0]);
}
return output_size;
}
template <typename T, typename Context, int64_t dim>
void hvol2col(const Context& dev_ctx,
const T* data_hvol,
int channels,
const std::vector<int64_t>& input_size,
const std::vector<int64_t>& output_size,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& stride_size,
const std::vector<int64_t>& pad_size,
const std::vector<int64_t>& dilation_size,
T* data_col) {
if (dim == 3) {
funcs::vol2col_slow<T, Context>(dev_ctx,
data_hvol,
channels,
input_size[0],
input_size[1],
input_size[2],
output_size[0],
output_size[1],
output_size[2],
kernel_size[0],
kernel_size[1],
kernel_size[2],
pad_size[0],
pad_size[1],
pad_size[2],
stride_size[0],
stride_size[1],
stride_size[2],
dilation_size[0],
dilation_size[1],
dilation_size[2],
data_col);
} else if (dim == 2) {
funcs::im2col_slow<T, Context>(dev_ctx,
data_hvol,
channels,
input_size[0],
input_size[1],
output_size[0],
output_size[1],
kernel_size[0],
kernel_size[1],
pad_size[0],
pad_size[1],
stride_size[0],
stride_size[1],
dilation_size[0],
dilation_size[1],
data_col);
}
}
template <typename T, typename Context, int64_t dim>
void col2hvol(const Context& dev_ctx,
const T* data_col,
const int channels,
const std::vector<int64_t>& input_size,
const std::vector<int64_t>& output_size,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& stride_size,
const std::vector<int64_t>& pad_size,
const std::vector<int64_t>& dilation_size,
T* data_hvol) {
if (dim == 3) {
funcs::col2vol_slow<T, T, Context>(dev_ctx,
data_col,
channels,
input_size[0],
input_size[1],
input_size[2],
output_size[0],
output_size[1],
output_size[2],
kernel_size[0],
kernel_size[1],
kernel_size[2],
pad_size[0],
pad_size[1],
pad_size[2],
stride_size[0],
stride_size[1],
stride_size[2],
dilation_size[0],
dilation_size[1],
dilation_size[2],
data_hvol);
}
if (dim == 2) {
funcs::col2im_slow<T, T, Context>(dev_ctx,
data_col,
channels,
input_size[0],
input_size[1],
output_size[0],
output_size[1],
kernel_size[0],
kernel_size[1],
pad_size[0],
pad_size[1],
stride_size[0],
stride_size[1],
dilation_size[0],
dilation_size[1],
data_hvol);
}
}
// Select View function
template <typename T>
DenseTensor Select(const DenseTensor& src, int64_t index) {
DenseTensor out;
out.ShareDataWith(src);
auto dims = src.dims();
std::vector<int64_t> new_dims;
for (int i = 1; i < dims.size(); ++i) {
new_dims.push_back(dims[i]);
}
out.Resize(new_dims);
int64_t stride_0 = src.numel() / dims[0];
size_t offset_bytes = index * stride_0 * sizeof(T);
out.set_offset(src.offset() + offset_bytes);
return out;
}
template <typename T, typename Context, int Dims>
void SlowConvDilatedAllCUDAImpl(const Context& dev_ctx,
DenseTensor* output,
const DenseTensor* input,
const DenseTensor* weight,
const DenseTensor* bias,
const DenseTensor* grad_output,
DenseTensor* grad_input,
DenseTensor* grad_weight,
DenseTensor* grad_bias,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& strides,
const std::vector<int64_t>& paddings,
const std::vector<int64_t>& dilations) {
const int64_t batch_size = input->dims()[0];
const int64_t input_channels = weight->dims()[1];
const int64_t output_channels = weight->dims()[0];
std::vector<int64_t> input_spatial_size;
for (int i = 2; i < input->dims().size(); ++i) {
input_spatial_size.push_back(input->dims()[i]);
}
std::vector<int64_t> output_spatial_size = GetOutputSpatialSize<Dims>(
*input, kernel_size, strides, paddings, dilations);
int64_t kernel_volume = multiply_integers(kernel_size);
int64_t output_volume = multiply_integers(output_spatial_size);
// Buffer
int64_t col_dim0 = input_channels * kernel_volume;
int64_t col_dim1 = output_volume;
DenseTensor columns;
if (output || grad_weight || grad_input) {
columns.Resize({col_dim0, col_dim1});
dev_ctx.template Alloc<T>(&columns);
}
// Initialize
funcs::SetConstant<Context, T> set_zero;
if (grad_weight) set_zero(dev_ctx, grad_weight, static_cast<T>(0));
if (grad_bias) set_zero(dev_ctx, grad_bias, static_cast<T>(0));
if (output && !bias) set_zero(dev_ctx, output, static_cast<T>(0));
// Bias CPU Mirror
DenseTensor bias_cpu;
const T* bias_cpu_data = nullptr;
if (output && bias) {
Copy(dev_ctx, *bias, CPUPlace(), true, &bias_cpu);
bias_cpu_data = bias_cpu.data<T>();
}
DenseTensor grad_output_n;
std::vector<int64_t> sum_axes;
for (int i = 0; i < Dims; ++i) sum_axes.push_back(i + 1);
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
for (int elt = 0; elt < batch_size; ++elt) {
T* columns_ptr = columns.data<T>();
// Prepare Input Slice View
DenseTensor input_n = Select<T>(*input, elt);
const T* input_ptr_raw = input_n.data<T>();
// Forward
if (output) {
DenseTensor output_n = Select<T>(*output, elt);
T* output_ptr_raw = output_n.data<T>();
if (bias) {
for (int n = 0; n < output_channels; ++n) {
DenseTensor out_slice = Select<T>(output_n, n);
FillKernel<T, Context>(
dev_ctx, out_slice, Scalar(bias_cpu_data[n]), &out_slice);
}
}
hvol2col<T, Context, Dims>(dev_ctx,
input_ptr_raw,
input_channels,
input_spatial_size,
output_spatial_size,
kernel_size,
strides,
paddings,
dilations,
columns_ptr);
blas.GEMM(false, // TransA
false, // TransB
static_cast<int>(output_channels), // M
static_cast<int>(col_dim1), // N
static_cast<int>(col_dim0), // K
static_cast<T>(1), // alpha
weight->data<T>(), // A
static_cast<int>(col_dim0), // lda
columns_ptr, // B
static_cast<int>(col_dim1), // ldb
static_cast<T>(1), // beta = 1 (Accumulate)
output_ptr_raw, // C
static_cast<int>(col_dim1) // ldc
);
} else {
grad_output_n = Select<T>(*grad_output, elt);
}
// Backward Grad Input
if (grad_input) {
DenseTensor grad_input_n = Select<T>(*grad_input, elt);
T* grad_input_ptr_raw = grad_input_n.data<T>();
const T* grad_output_ptr_raw = grad_output_n.data<T>();
blas.GEMM(true, // TransA
false, // TransB
static_cast<int>(col_dim0), // M
static_cast<int>(col_dim1), // N
static_cast<int>(output_channels), // K
static_cast<T>(1), // alpha
weight->data<T>(), // A
static_cast<int>(col_dim0), // lda
grad_output_ptr_raw, // B
static_cast<int>(col_dim1), // ldb
static_cast<T>(0), // beta
columns_ptr, // C
static_cast<int>(col_dim1) // ldc
);
col2hvol<T, Context, Dims>(dev_ctx,
columns_ptr,
input_channels,
input_spatial_size,
output_spatial_size,
kernel_size,
strides,
paddings,
dilations,
grad_input_ptr_raw);
}
// Backward Grad Weight
if (grad_weight) {
const T* grad_output_ptr_raw = grad_output_n.data<T>();
hvol2col<T, Context, Dims>(dev_ctx,
input_ptr_raw,
input_channels,
input_spatial_size,
output_spatial_size,
kernel_size,
strides,
paddings,
dilations,
columns_ptr);
blas.GEMM(false, // TransA
true, // TransB
static_cast<int>(output_channels), // M
static_cast<int>(col_dim0), // N
static_cast<int>(col_dim1), // K
static_cast<T>(1), // alpha
grad_output_ptr_raw, // A
static_cast<int>(col_dim1), // lda
columns_ptr, // B
static_cast<int>(col_dim1), // ldb
static_cast<T>(1), // beta
grad_weight->data<T>(), // C
static_cast<int>(col_dim0) // ldc
);
}
// Backward Grad Bias
if (grad_bias) {
DenseTensor sum_result = Sum<T, Context>(dev_ctx,
grad_output_n,
IntArray(sum_axes),
CppTypeToDataType<T>::Type(),
false);
Add<T, Context>(dev_ctx, *grad_bias, sum_result, grad_bias);
}
}
}
template <typename T, typename Context, int64_t dim>
void SlowConvBackwardNoGroup(const Context& dev_ctx,
const DenseTensor& grad_output,
const DenseTensor& input,
const DenseTensor& weight,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& strides,
const std::vector<int64_t>& paddings,
const std::vector<int64_t>& dilations,
DenseTensor* grad_input,
DenseTensor* grad_weight,
DenseTensor* grad_bias) {
int64_t rank = input.dims().size();
bool is_batch = (rank == (dim + 2));
// tensor.unsqueeze(0)
auto make_batch_view = [&](const DenseTensor& src, DenseTensor& dst) {
if (!is_batch) {
dst.ShareDataWith(src);
std::vector<int64_t> new_shape = {1};
for (int i = 0; i < src.dims().size(); ++i)
new_shape.push_back(src.dims()[i]);
dst.Resize(new_shape);
} else {
dst.ShareDataWith(src);
}
};
DenseTensor grad_output_;
make_batch_view(grad_output, grad_output_);
DenseTensor input_;
make_batch_view(input, input_);
const DenseTensor& weight_ = weight;
DenseTensor grad_input_view;
DenseTensor* grad_input_ptr = nullptr;
if (grad_input) {
dev_ctx.template Alloc<T>(grad_input);
if (!is_batch) {
grad_input_view.ShareDataWith(*grad_input);
std::vector<int64_t> new_shape = {1};
for (int i = 0; i < grad_input->dims().size(); ++i)
new_shape.push_back(grad_input->dims()[i]);
grad_input_view.Resize(new_shape);
grad_input_ptr = &grad_input_view;
} else {
grad_input_ptr = grad_input;
}
}
DenseTensor* grad_weight_ptr = nullptr;
if (grad_weight) {
dev_ctx.template Alloc<T>(grad_weight);
grad_weight_ptr = grad_weight;
}
DenseTensor* grad_bias_ptr = nullptr;
if (grad_bias) {
dev_ctx.template Alloc<T>(grad_bias);
grad_bias_ptr = grad_bias;
}
SlowConvDilatedAllCUDAImpl<T, Context, dim>(
dev_ctx,
nullptr, // [Output]
&input_, // [Input]
&weight_, // [Weight]
nullptr, // [Bias]
&grad_output_, // [GradOutput]
grad_input_ptr, // [GradInput]
grad_weight_ptr, // [GradWeight]
grad_bias_ptr, // [GradBias] (New)
kernel_size,
strides,
paddings,
dilations);
}
template <typename T, typename Context, int64_t dim>
void SlowConvNoGroup(const Context& dev_ctx,
const DenseTensor& input,
const DenseTensor& weight,
const DenseTensor* bias,
const std::vector<int64_t>& kernel_size,
const std::vector<int64_t>& strides,
const std::vector<int64_t>& paddings,
const std::vector<int64_t>& dilations,
DenseTensor* output) {
int64_t rank = input.dims().size();
bool is_batch = (rank == (dim + 2));
// (is_batch ? input.contiguous() : input.contiguous().unsqueeze(0));
DenseTensor input_;
if (!is_batch) {
input_.ShareDataWith(input);
std::vector<int64_t> new_shape = {1};
for (int i = 0; i < rank; ++i) new_shape.push_back(input.dims()[i]);
input_.Resize(new_shape);
} else {
input_.ShareDataWith(input);
}
const DenseTensor& weight_ = weight;
// (is_batch ? output : output.unsqueeze(0));
if (output) dev_ctx.template Alloc<T>(output);
DenseTensor output_;
if (!is_batch) {
output_.ShareDataWith(*output);
std::vector<int64_t> out_shape = {1};
for (int i = 0; i < output->dims().size(); ++i) {
out_shape.push_back(output->dims()[i]);
}
output_.Resize(out_shape);
} else {
output_.ShareDataWith(*output);
}
SlowConvDilatedAllCUDAImpl<T, Context, dim>(dev_ctx,
&output_, // [Output]
&input_, // [Input]
&weight_, // [Weight]
bias, // [Bias]
nullptr, // [GradOutput]
nullptr, // [GradInput]
nullptr, // [GradWeight]
nullptr, // [GradBias]
kernel_size,
strides,
paddings,
dilations);
}
template <typename T, typename Context, int64_t dim>
void SlowConvForward(const Context& dev_ctx,
const DenseTensor& input,
const DenseTensor& filter_t,
const paddle::optional<DenseTensor>& bias,
const std::vector<int>& strides,
const std::vector<int>& paddings_t,
const std::string& padding_algorithm,
int groups,
const std::vector<int>& dilations_t,
const std::string& data_format,
DenseTensor* output) {
std::vector<int> paddings = paddings_t;
std::vector<int> dilations = dilations_t;
DenseTensor filter = filter_t;
if (input.numel() == 0 || filter.numel() == 0) {
Full<T, Context>(dev_ctx, output->dims(), 0, output);
return;
}
dev_ctx.template Alloc<T>(output);
const bool channel_last = (data_format == "NHWC" || data_format == "NDHWC");
DenseTensor transformed_input(input.type());
DenseTensor transformed_output(output->type());
if (channel_last) {
ResizeToChannelFirst<Context, T>(dev_ctx, &input, &transformed_input);
TransToChannelFirst<Context, T>(dev_ctx, &input, &transformed_input);
ResizeToChannelFirst<Context, T>(dev_ctx, output, &transformed_output);
} else {
transformed_input = input;
transformed_output = *output;
}
// update padding and dilation
auto trans_in_dims = transformed_input.dims();
auto filter_dims = filter.dims();
DDim in_data_dims = slice_ddim(trans_in_dims, 2, trans_in_dims.size());
DDim filter_data_dims = slice_ddim(filter_dims, 2, filter_dims.size());
std::vector<int> ksize = vectorize<int>(filter_data_dims);
UpdatePaddingAndDilation(
&paddings, &dilations, padding_algorithm, in_data_dims, strides, ksize);
// =================================================================
// Contiguous & Grouping
// =================================================================
DenseTensor input_contiguous;
ContiguousKernel<T, Context>(dev_ctx, transformed_input, &input_contiguous);
DenseTensor weight_contiguous;
ContiguousKernel<T, Context>(dev_ctx, filter_t, &weight_contiguous);
auto to_int64_vec = [](const std::vector<int>& in) {
return std::vector<int64_t>(in.begin(), in.end());
};
const DenseTensor* bias_ptr = bias.get_ptr();
DenseTensor bias_contiguous;
if (bias_ptr) {
ContiguousKernel<T, Context>(dev_ctx, *bias_ptr, &bias_contiguous);
bias_ptr = &bias_contiguous;
}
if (groups == 1) {
SlowConvNoGroup<T, Context, dim>(dev_ctx,
input_contiguous,
weight_contiguous,
bias_ptr,
to_int64_vec(ksize),
to_int64_vec(strides),
to_int64_vec(paddings),
to_int64_vec(dilations),
&transformed_output);
} else {
int64_t in_rank = input_contiguous.dims().size();
bool has_batch = (in_rank == dim + 2);
int channel_dim = has_batch ? 1 : 0;
int64_t in_channels = input_contiguous.dims()[channel_dim];
int64_t out_channels = weight_contiguous.dims()[0];
int64_t in_g_sz = in_channels / groups;
int64_t out_g_sz = out_channels / groups;
std::vector<DenseTensor> outputs(groups);
for (int g = 0; g < groups; ++g) {
// Slice Input (Channel)
DenseTensor input_g;
SliceKernel<T, Context>(dev_ctx,
input_contiguous,
{channel_dim},
{g * in_g_sz},
{(g + 1) * in_g_sz},
{1},
{},
&input_g);
// Slice Weight (OutChannel dim 0)
DenseTensor weight_g;
SliceKernel<T, Context>(dev_ctx,
weight_contiguous,
{0},
{g * out_g_sz},
{(g + 1) * out_g_sz},
{1},
{},
&weight_g);
// Slice Bias (OutChannel dim 0)
DenseTensor bias_g;
const DenseTensor* bias_g_ptr = nullptr;
if (bias_ptr) {
SliceKernel<T, Context>(dev_ctx,
*bias_ptr,
{0},
{g * out_g_sz},
{(g + 1) * out_g_sz},
{1},
{},
&bias_g);
bias_g_ptr = &bias_g;
}
DenseTensor output_g;
auto out_shape = transformed_output.dims();
out_shape[channel_dim] = out_g_sz;
output_g.Resize(out_shape);
dev_ctx.template Alloc<T>(&output_g);
SlowConvNoGroup<T, Context, dim>(dev_ctx,
input_g,
weight_g,
bias_g_ptr,
to_int64_vec(ksize),
to_int64_vec(strides),
to_int64_vec(paddings),
to_int64_vec(dilations),
&output_g);
outputs[g] = output_g;
}
// Concat
std::vector<const DenseTensor*> outputs_ptr;
for (auto& t : outputs) outputs_ptr.push_back(&t);
ConcatKernel<T, Context>(
dev_ctx, outputs_ptr, channel_dim, &transformed_output);
}
if (channel_last) {
TransToChannelLast<Context, T>(dev_ctx, &transformed_output, output);
}
}
template <typename T, typename Context, int64_t dim>
void SlowConvBackward(const Context& dev_ctx,
const DenseTensor& input,
const DenseTensor& filter_t,
const DenseTensor& output_grad,
const std::vector<int>& strides,
const std::vector<int>& paddings_t,
const std::string& padding_algorithm,
const std::vector<int>& dilations_t,
int groups,
const std::string& data_format,
DenseTensor* input_grad,
DenseTensor* filter_grad,
DenseTensor* bias_grad) {
if (!input_grad && !filter_grad && !bias_grad) return;
std::vector<int> paddings = paddings_t;
std::vector<int> dilations = dilations_t;
DenseTensor filter = filter_t;
// 0-size
if (input.numel() == 0 || filter_t.numel() == 0) {
if (input_grad) dev_ctx.template Alloc<T>(input_grad);
if (filter_grad) {
Full<T, Context>(dev_ctx, filter_grad->dims(), 0, filter_grad);
}
if (bias_grad) {
dev_ctx.template Alloc<T>(bias_grad);
Full<T, Context>(dev_ctx, bias_grad->dims(), 0, bias_grad);
}
return;
}
if (input_grad) dev_ctx.template Alloc<T>(input_grad);
if (filter_grad) dev_ctx.template Alloc<T>(filter_grad);
if (bias_grad) dev_ctx.template Alloc<T>(bias_grad);
const bool channel_last = (data_format == "NHWC" || data_format == "NDHWC");
DenseTensor transformed_input(input.type());
DenseTensor transformed_output_grad(output_grad.type());
if (channel_last) {
ResizeToChannelFirst<Context, T>(dev_ctx, &input, &transformed_input);
TransToChannelFirst<Context, T>(dev_ctx, &input, &transformed_input);
ResizeToChannelFirst<Context, T>(
dev_ctx, &output_grad, &transformed_output_grad);
TransToChannelFirst<Context, T>(
dev_ctx, &output_grad, &transformed_output_grad);
} else {
transformed_input = input;
transformed_output_grad = output_grad;
}
// update padding and dilation
auto in_dims = transformed_input.dims();
auto filter_dims = filter.dims();
DDim in_data_dims = slice_ddim(in_dims, 2, in_dims.size());
DDim filter_data_dims = slice_ddim(filter_dims, 2, filter_dims.size());
std::vector<int> ksize = vectorize<int>(filter_data_dims);
UpdatePaddingAndDilation<int>(
&paddings, &dilations, padding_algorithm, in_data_dims, strides, ksize);
// =================================================================
// Contiguous & Grouping
// =================================================================
DenseTensor tmp_input_grad;
DenseTensor* t_input_grad_ptr = nullptr;
DenseTensor* t_filter_grad_ptr = filter_grad;
DenseTensor* t_bias_grad_ptr = bias_grad;
if (input_grad) {
if (channel_last) {
tmp_input_grad.Resize(transformed_input.dims());
t_input_grad_ptr = &tmp_input_grad;
} else {
t_input_grad_ptr = input_grad;
}
}
// Contiguous
DenseTensor grad_output_cont;
ContiguousKernel<T, Context>(
dev_ctx, transformed_output_grad, &grad_output_cont);
DenseTensor input_cont;
ContiguousKernel<T, Context>(dev_ctx, transformed_input, &input_cont);
DenseTensor weight_cont;
ContiguousKernel<T, Context>(dev_ctx, filter, &weight_cont);
auto to_int64_vec = [](const std::vector<int>& in) {
return std::vector<int64_t>(in.begin(), in.end());
};
// Group
if (groups == 1) {
SlowConvBackwardNoGroup<T, Context, dim>(dev_ctx,
grad_output_cont,
input_cont,
weight_cont,
to_int64_vec(ksize),
to_int64_vec(strides),
to_int64_vec(paddings),
to_int64_vec(dilations),
t_input_grad_ptr,
t_filter_grad_ptr,
t_bias_grad_ptr);
} else {
int64_t in_rank = input_cont.dims().size();
bool has_batch = (in_rank == dim + 2);
int channel_dim = has_batch ? 1 : 0;
int64_t in_channels = input_cont.dims()[channel_dim];
int64_t out_channels = grad_output_cont.dims()[channel_dim];
int64_t in_g_sz = in_channels / groups;
int64_t out_g_sz = out_channels / groups;
std::vector<DenseTensor> grad_inputs_g(groups);
std::vector<DenseTensor> grad_weights_g(groups);
std::vector<DenseTensor> grad_biases_g(groups);
for (int g = 0; g < groups; ++g) {
// Slice GradOutput (Channel)
DenseTensor grad_output_g;
SliceKernel<T, Context>(dev_ctx,
grad_output_cont,
{channel_dim},
{g * out_g_sz},
{(g + 1) * out_g_sz},
{1},
{},
&grad_output_g);
// Slice Input (Channel)
DenseTensor input_g;
SliceKernel<T, Context>(dev_ctx,
input_cont,
{channel_dim},
{g * in_g_sz},
{(g + 1) * in_g_sz},
{1},
{},
&input_g);
// Slice Weight (Output Channel / dim 0)
DenseTensor weight_g;
SliceKernel<T, Context>(dev_ctx,
weight_cont,
{0},
{g * out_g_sz},
{(g + 1) * out_g_sz},
{1},
{},
&weight_g);
DenseTensor grad_input_g_tensor;
DenseTensor grad_weight_g_tensor;
DenseTensor grad_bias_g_tensor;
if (t_input_grad_ptr) {
auto g_shape = t_input_grad_ptr->dims();
g_shape[channel_dim] = in_g_sz;
grad_input_g_tensor.Resize(g_shape);
}
if (t_filter_grad_ptr) {
auto w_shape = t_filter_grad_ptr->dims();
w_shape[0] = out_g_sz;
grad_weight_g_tensor.Resize(w_shape);
}
if (t_bias_grad_ptr) {
auto b_shape = t_bias_grad_ptr->dims();
b_shape[0] = out_g_sz;
grad_bias_g_tensor.Resize(b_shape);
}
SlowConvBackwardNoGroup<T, Context, dim>(
dev_ctx,
grad_output_g,
input_g,
weight_g,
to_int64_vec(ksize),
to_int64_vec(strides),
to_int64_vec(paddings),
to_int64_vec(dilations),
(t_input_grad_ptr ? &grad_input_g_tensor : nullptr),
(t_filter_grad_ptr ? &grad_weight_g_tensor : nullptr),
(t_bias_grad_ptr ? &grad_bias_g_tensor : nullptr));
if (t_input_grad_ptr) grad_inputs_g[g] = grad_input_g_tensor;
if (t_filter_grad_ptr) grad_weights_g[g] = grad_weight_g_tensor;
if (t_bias_grad_ptr) grad_biases_g[g] = grad_bias_g_tensor;
}
// Concat Input Grad
if (t_input_grad_ptr) {
std::vector<const DenseTensor*> ptrs;
for (auto& t : grad_inputs_g) ptrs.push_back(&t);
ConcatKernel<T, Context>(dev_ctx, ptrs, channel_dim, t_input_grad_ptr);
}
// Concat Weight Grad
if (t_filter_grad_ptr) {
std::vector<const DenseTensor*> ptrs;
for (auto& t : grad_weights_g) ptrs.push_back(&t);
ConcatKernel<T, Context>(dev_ctx, ptrs, 0, t_filter_grad_ptr);
}
// Concat Bias Grad
if (t_bias_grad_ptr) {
std::vector<const DenseTensor*> ptrs;
for (auto& t : grad_biases_g) ptrs.push_back(&t);
ConcatKernel<T, Context>(dev_ctx, ptrs, 0, t_bias_grad_ptr);
}
}
if (channel_last && input_grad) {
TransToChannelLast<Context, T>(dev_ctx, t_input_grad_ptr, input_grad);
}
}
} // namespace phi