265 lines
10 KiB
C++
265 lines
10 KiB
C++
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
|
#include <ATen/PadNd.h>
|
|
#include <ATen/core/Tensor.h>
|
|
|
|
#include <c10/util/irange.h>
|
|
|
|
#ifndef AT_PER_OPERATOR_HEADERS
|
|
#include <ATen/Functions.h>
|
|
#include <ATen/NativeFunctions.h>
|
|
#else
|
|
#include <ATen/ops/_empty_affine_quantized.h>
|
|
#include <ATen/ops/_pad_circular.h>
|
|
#include <ATen/ops/_pad_circular_native.h>
|
|
#include <ATen/ops/_pad_enum_native.h>
|
|
#include <ATen/ops/constant_pad_nd.h>
|
|
#include <ATen/ops/constant_pad_nd_native.h>
|
|
#include <ATen/ops/empty.h>
|
|
#include <ATen/ops/pad_native.h>
|
|
#include <ATen/ops/reflection_pad1d.h>
|
|
#include <ATen/ops/reflection_pad2d.h>
|
|
#include <ATen/ops/reflection_pad3d.h>
|
|
#include <ATen/ops/replication_pad1d.h>
|
|
#include <ATen/ops/replication_pad2d.h>
|
|
#include <ATen/ops/replication_pad3d.h>
|
|
#endif
|
|
|
|
namespace at::native {
|
|
|
|
Tensor constant_pad_nd(const Tensor& self, IntArrayRef pad, const Scalar& value) {
|
|
TORCH_CHECK(pad.size() % 2 == 0, "Length of pad must be even but instead it equals ",
|
|
pad.size());
|
|
|
|
auto input_sizes = self.sizes();
|
|
auto l_inp = self.dim();
|
|
|
|
auto l_pad = pad.size() / 2;
|
|
auto l_diff = l_inp - l_pad;
|
|
TORCH_CHECK(l_inp >= (int64_t)l_pad, "Length of pad should be no more than twice the number of "
|
|
"dimensions of the input. Pad length is ", pad.size(), "while the input has ",
|
|
l_inp, "dimensions.");
|
|
|
|
|
|
bool all_pads_non_positive = true;
|
|
|
|
auto c_input = self;
|
|
for (const auto i : c10::irange(l_diff, l_inp)) {
|
|
auto pad_idx = 2 * (l_inp - i - 1);
|
|
if (pad[pad_idx] < 0) {
|
|
c_input = c_input.narrow(i, -pad[pad_idx], c_input.size(i) + pad[pad_idx]);
|
|
} else if (pad[pad_idx] != 0) {
|
|
all_pads_non_positive = false;
|
|
}
|
|
if (pad[pad_idx + 1] < 0) {
|
|
c_input = c_input.narrow(i, 0, c_input.size(i) + pad[pad_idx + 1]);
|
|
} else if (pad[pad_idx + 1] != 0) {
|
|
all_pads_non_positive = false;
|
|
}
|
|
}
|
|
|
|
// if none of the pads are positive we can optimize and just return the result
|
|
// of calling .narrow() on the input
|
|
if (all_pads_non_positive) {
|
|
return c_input.clone();
|
|
}
|
|
|
|
|
|
std::vector<int64_t> new_shape;
|
|
new_shape.reserve(l_diff);
|
|
for (size_t i = 0; i < l_diff; i ++) {
|
|
new_shape.emplace_back(input_sizes[i]);
|
|
}
|
|
|
|
for (const auto i : c10::irange((size_t)l_pad)) {
|
|
auto pad_idx = pad.size() - ((i + 1) * 2);
|
|
auto new_dim = input_sizes[l_diff + i] + pad[pad_idx] + pad[pad_idx + 1];
|
|
TORCH_CHECK(new_dim > 0, "The input size ", input_sizes[l_diff + i], ", plus negative padding ",
|
|
pad[pad_idx], " and ", pad[pad_idx + 1], " resulted in a negative output size, "
|
|
"which is invalid. Check dimension ", l_diff + i, " of your input.");
|
|
new_shape.emplace_back(new_dim);
|
|
}
|
|
|
|
at::Tensor output;
|
|
const auto memory_format = self.suggest_memory_format();
|
|
if (self.is_quantized()) {
|
|
const auto qscheme = self.qscheme();
|
|
TORCH_CHECK(qscheme == kPerTensorAffine || qscheme == kPerTensorSymmetric,
|
|
"Only per-tensor padding is supported.");
|
|
output = at::_empty_affine_quantized(
|
|
new_shape, self.options().memory_format(memory_format),
|
|
self.q_scale(), self.q_zero_point(), std::nullopt);
|
|
} else {
|
|
output = at::empty(new_shape, self.options().memory_format(memory_format));
|
|
}
|
|
output.fill_(value);
|
|
|
|
auto c_output = output;
|
|
for (const auto i : c10::irange(l_diff, l_inp)) {
|
|
auto pad_idx = 2 * (l_inp - i - 1);
|
|
if (pad[pad_idx] > 0) {
|
|
c_output = c_output.narrow(i, pad[pad_idx], c_output.size(i) - pad[pad_idx]);
|
|
}
|
|
if (pad[pad_idx + 1] > 0) {
|
|
c_output = c_output.narrow(i, 0, c_output.size(i) - pad[pad_idx + 1]);
|
|
}
|
|
}
|
|
c_output.copy_(c_input);
|
|
return output;
|
|
}
|
|
|
|
Tensor _pad_circular_symint(const Tensor &self, c10::SymIntArrayRef padding) {
|
|
const auto in_shape = self.sym_sizes();
|
|
const auto self_ndim = static_cast<int64_t>(in_shape.size());
|
|
|
|
// number of dimensions that are padded
|
|
const auto ndim_padded = padding.size() / 2;
|
|
// number of preceding non_padded dimensions (1 for no_batch_dim case or 2)
|
|
const auto ndim_nonpadded = self_ndim - ndim_padded;
|
|
|
|
TORCH_CHECK(ndim_nonpadded == 1 || ndim_nonpadded == 2,
|
|
"Invalid padding size, expected 1 or 2 non-padded dimensions, ",
|
|
"which would be equivalent to padding of length ",
|
|
(self_ndim - 1) * 2,
|
|
" or ",
|
|
(self_ndim - 2) * 2,
|
|
" respectively but got ",
|
|
padding.size());
|
|
|
|
c10::SymDimVector out_shape(in_shape.size());
|
|
for (const auto i: c10::irange(ndim_nonpadded)) {
|
|
out_shape[i] = in_shape[i];
|
|
}
|
|
|
|
// Get shape of padded tensor
|
|
for (const auto i : c10::irange(ndim_padded)) {
|
|
const auto& pad_l = padding[2 * (ndim_padded - i - 1) + 0];
|
|
const auto& pad_r = padding[2 * (ndim_padded - i - 1) + 1];
|
|
const auto& size = in_shape[ndim_nonpadded + i];
|
|
out_shape[ndim_nonpadded + i] = size + pad_l + pad_r;
|
|
|
|
TORCH_CHECK(
|
|
pad_l <= size && pad_r <= size,
|
|
"Padding value causes wrapping around more than once.");
|
|
TORCH_CHECK(
|
|
out_shape[ndim_nonpadded + i] >= 0,
|
|
"Negative padding value is resulting in an empty dimension");
|
|
}
|
|
|
|
auto out = self.new_empty_symint(out_shape, self.options());
|
|
|
|
// Put original array into the padded array
|
|
Tensor out_slice = out;
|
|
Tensor in_slice = self;
|
|
const SymInt zero = 0;
|
|
for (const auto i : c10::irange(ndim_padded)) {
|
|
const auto dim = ndim_padded - i + ndim_nonpadded - 1;
|
|
const auto& pad_l = padding[2*i + 0];
|
|
const auto& pad_r = padding[2*i + 1];
|
|
out_slice = out_slice.slice_symint(dim, std::max(pad_l, zero), out_shape[dim] - std::max(pad_r, zero));
|
|
in_slice = in_slice.slice_symint(dim, std::max(-pad_l, zero), in_shape[dim] - std::max(-pad_r, zero));
|
|
}
|
|
out_slice.copy_(in_slice);
|
|
|
|
// The following steps first pad the beginning of the tensor (left side),
|
|
// and then pad the end of the tensor (right side).
|
|
// Note: Corners will be written more than once when ndim_padded > 1.
|
|
//
|
|
// Only in cases where padding values are > 0 are when additional copying
|
|
// is required.
|
|
for (const auto i : c10::irange(ndim_padded)) {
|
|
const auto dim = ndim_padded - i + ndim_nonpadded - 1;
|
|
const auto& pad_l = padding[2*i + 0];
|
|
const auto& pad_r = padding[2*i + 1];
|
|
|
|
if (pad_l > 0) {
|
|
out_slice = out.slice_symint(dim, 0, pad_l);
|
|
in_slice = out.slice_symint(dim,
|
|
out_shape[dim] - pad_l - std::max(pad_r, zero),
|
|
out_shape[dim] - std::max(pad_r, zero));
|
|
out_slice.copy_(in_slice);
|
|
}
|
|
|
|
if (pad_r > 0) {
|
|
out_slice = out.slice_symint(dim, out_shape[dim] - pad_r, out_shape[dim]);
|
|
in_slice = out.slice_symint(dim, std::max(pad_l, zero), std::max(pad_l, zero) + pad_r);
|
|
out_slice.copy_(in_slice);
|
|
}
|
|
}
|
|
|
|
return out;
|
|
}
|
|
|
|
static std::string_view padding_mode_string(padding_mode m) {
|
|
switch (m) {
|
|
case padding_mode::reflect:
|
|
return "reflect";
|
|
case padding_mode::replicate:
|
|
return "replicate";
|
|
case padding_mode::circular:
|
|
return "circular";
|
|
case padding_mode::constant:
|
|
return "constant";
|
|
}
|
|
TORCH_CHECK(false, "Invalid padding mode (", static_cast<int64_t>(m), ")");
|
|
}
|
|
|
|
|
|
Tensor _pad_enum_symint(const Tensor &self, c10::SymIntArrayRef pad, int64_t mode_int, std::optional<double> value) {
|
|
const auto input_dim = self.dim();
|
|
TORCH_CHECK(pad.size() % 2 == 0, "Padding length must be divisible by 2");
|
|
TORCH_CHECK(static_cast<int64_t>(pad.size()) <= input_dim * 2,
|
|
"Padding length should be less than or equal to two times the input dimension but got padding length ", pad.size(), " and input of dimension ", input_dim);
|
|
auto mode = static_cast<at::padding_mode>(mode_int);
|
|
|
|
if (mode == at::padding_mode::constant) {
|
|
return at::constant_pad_nd_symint(self, pad, value.value_or(0.0));
|
|
}
|
|
TORCH_CHECK(!value.has_value() || *value == 0,
|
|
"Padding mode \"", padding_mode_string(mode),
|
|
"\" doesn't take in value argument");
|
|
|
|
if (pad.size() == 2 && (input_dim == 2 || input_dim == 3)) {
|
|
switch (mode) {
|
|
case at::padding_mode::reflect: return at::reflection_pad1d_symint(self, pad);
|
|
case at::padding_mode::replicate: return at::replication_pad1d_symint(self, pad);
|
|
case at::padding_mode::circular: return at::_pad_circular_symint(self, pad);
|
|
default: {}
|
|
}
|
|
} else if(pad.size() == 4 && (input_dim == 3 || input_dim == 4)) {
|
|
switch (mode) {
|
|
case at::padding_mode::reflect: return at::reflection_pad2d_symint(self, pad);
|
|
case at::padding_mode::replicate: return at::replication_pad2d_symint(self, pad);
|
|
case at::padding_mode::circular: return at::_pad_circular_symint(self, pad);
|
|
default: {}
|
|
}
|
|
} else if (pad.size() == 6 && (input_dim == 4 || input_dim == 5)) {
|
|
switch (mode) {
|
|
case at::padding_mode::reflect: return at::reflection_pad3d_symint(self, pad);
|
|
case at::padding_mode::replicate: return at::replication_pad3d_symint(self, pad);
|
|
case at::padding_mode::circular: return at::_pad_circular_symint(self, pad);
|
|
default: {}
|
|
}
|
|
}
|
|
C10_THROW_ERROR(NotImplementedError,
|
|
"Only 2D, 3D, 4D, 5D padding with non-constant padding are supported for now");
|
|
}
|
|
|
|
Tensor pad_symint(const Tensor &self, c10::SymIntArrayRef pad, std::string_view mode, std::optional<double> value) {
|
|
const auto mode_enum = [&] {
|
|
if (mode == "reflect") {
|
|
return at::padding_mode::reflect;
|
|
} else if (mode == "constant") {
|
|
return at::padding_mode::constant;
|
|
} else if (mode == "replicate") {
|
|
return at::padding_mode::replicate;
|
|
} else if (mode == "circular") {
|
|
return at::padding_mode::circular;
|
|
}
|
|
C10_THROW_ERROR(NotImplementedError,
|
|
c10::str("Unrecognised padding mode ", mode));
|
|
}();
|
|
return at::native::_pad_enum_symint(self, pad, static_cast<int64_t>(mode_enum), value);
|
|
}
|
|
|
|
} // namespace at::native
|