From a4f14d3ae9cc505557ea52eb483c843fe3e441f7 Mon Sep 17 00:00:00 2001 From: Maxim Vafin Date: Wed, 21 Jan 2026 18:47:13 +0100 Subject: [PATCH 1/2] [PT FE] Refactor PixelShuffle operations Signed-off-by: Maxim Vafin --- src/frontends/pytorch/src/op/neg.cpp | 29 --- .../pytorch/src/op/pixel_shuffle.cpp | 198 +++++++++--------- src/frontends/pytorch/src/op_table.cpp | 7 +- 3 files changed, 97 insertions(+), 137 deletions(-) delete mode 100644 src/frontends/pytorch/src/op/neg.cpp diff --git a/src/frontends/pytorch/src/op/neg.cpp b/src/frontends/pytorch/src/op/neg.cpp deleted file mode 100644 index 0de89d300c133c..00000000000000 --- a/src/frontends/pytorch/src/op/neg.cpp +++ /dev/null @@ -1,29 +0,0 @@ -// Copyright (C) 2018-2026 Intel Corporation -// SPDX-License-Identifier: Apache-2.0 -// - -#include "openvino/frontend/pytorch/node_context.hpp" -#include "openvino/op/constant.hpp" -#include "openvino/op/convert_like.hpp" -#include "openvino/op/multiply.hpp" -#include "utils.hpp" - -namespace ov { -namespace frontend { -namespace pytorch { -namespace op { - -using namespace ov::op; - -OutputVector translate_neg(const NodeContext& context) { - num_inputs_check(context, 1, 1); - auto x = context.get_input(0); - auto const_neg_1 = context.mark_node(v0::Constant::create(element::i32, Shape{}, {-1})); - auto cast = context.mark_node(std::make_shared(const_neg_1, x)); - return {context.mark_node(std::make_shared(x, cast))}; -}; - -} // namespace op -} // namespace pytorch -} // namespace frontend -} // namespace ov \ No newline at end of file diff --git a/src/frontends/pytorch/src/op/pixel_shuffle.cpp b/src/frontends/pytorch/src/op/pixel_shuffle.cpp index 72531e40f16a0b..1ebf3ae14b971a 100644 --- a/src/frontends/pytorch/src/op/pixel_shuffle.cpp +++ b/src/frontends/pytorch/src/op/pixel_shuffle.cpp @@ -2,21 +2,21 @@ // SPDX-License-Identifier: Apache-2.0 // +#include + #include "openvino/frontend/pytorch/node_context.hpp" -#include "openvino/op/add.hpp" #include "openvino/op/concat.hpp" #include "openvino/op/constant.hpp" +#include "openvino/op/depth_to_space.hpp" #include "openvino/op/divide.hpp" #include "openvino/op/gather.hpp" #include "openvino/op/multiply.hpp" -#include "openvino/op/range.hpp" #include "openvino/op/reshape.hpp" #include "openvino/op/shape_of.hpp" +#include "openvino/op/shuffle_channels.hpp" #include "openvino/op/slice.hpp" +#include "openvino/op/space_to_depth.hpp" #include "openvino/op/split.hpp" -#include "openvino/op/squeeze.hpp" -#include "openvino/op/transpose.hpp" -#include "openvino/op/unsqueeze.hpp" #include "utils.hpp" namespace ov { @@ -26,116 +26,106 @@ namespace op { using namespace ov::op; +namespace { + +struct PixelSpatialInfo { + Output dims_before; + Output channels; + Output height; + Output width; +}; + +PixelSpatialInfo get_pixel_spatial_info(const NodeContext& context, const Output& x) { + auto zero_vec = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); + auto zero_scalar = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + auto neg_three = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-3})); + auto one_vec = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1})); + auto shape = context.mark_node(std::make_shared(x, element::i32)); + auto dims_before = context.mark_node(std::make_shared(shape, zero_vec, neg_three, one_vec)); + auto indices = context.mark_node(v0::Constant::create(element::i32, Shape{3}, {-3, -2, -1})); + auto dims = context.mark_node(std::make_shared(shape, indices, zero_scalar)); + auto split = context.mark_node(std::make_shared(dims, zero_scalar, 3)); + return {dims_before, split->output(0), split->output(1), split->output(2)}; +} + +Output make_flatten_shape(const NodeContext& context, + const Output& channels, + const Output& height, + const Output& width) { + auto neg_one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); + auto chw = context.mark_node(std::make_shared(OutputVector{channels, height, width}, 0)); + return context.mark_node(std::make_shared(OutputVector{neg_one, chw}, 0)); +} + +Output make_final_shape(const NodeContext& context, + const Output& dims_before, + const Output& new_c, + const Output& new_h, + const Output& new_w) { + auto tail = context.mark_node(std::make_shared(OutputVector{new_c, new_h, new_w}, 0)); + return context.mark_node(std::make_shared(OutputVector{dims_before, tail}, 0)); +} + +OutputVector translate_pixel_transform(const NodeContext& context, bool is_shuffle) { + num_inputs_check(context, 2, 2); + const auto x = context.get_input(0); + const auto block = context.const_input(1); + PYTORCH_OP_CONVERSION_CHECK(block > 0, "Upscale factor for pixel shuffle ops must be positive"); + + const auto block_size = static_cast(block); + const auto block_scalar = + context.mark_node(v0::Constant::create(element::i32, Shape{}, {static_cast(block)})); + const auto block_sq_scalar = + context.mark_node(v0::Constant::create(element::i32, Shape{}, {static_cast(block * block)})); + + const auto [dims_before, channels, height, width] = get_pixel_spatial_info(context, x); + const auto flat_shape = make_flatten_shape(context, channels, height, width); + const auto flattened = context.mark_node(std::make_shared(x, flat_shape, false)); + + Output transformed; + Output new_c; + Output new_h; + Output new_w; + + if (is_shuffle) { + transformed = context.mark_node( + std::make_shared(flattened, v0::DepthToSpace::DepthToSpaceMode::DEPTH_FIRST, block_size)); + new_c = context.mark_node(std::make_shared(channels, block_sq_scalar)); + new_h = context.mark_node(std::make_shared(height, block_scalar)); + new_w = context.mark_node(std::make_shared(width, block_scalar)); + } else { + transformed = context.mark_node( + std::make_shared(flattened, v0::SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST, block_size)); + new_c = context.mark_node(std::make_shared(channels, block_sq_scalar)); + new_h = context.mark_node(std::make_shared(height, block_scalar, true)); + new_w = context.mark_node(std::make_shared(width, block_scalar, true)); + } + + const auto final_shape = make_final_shape(context, dims_before, new_c, new_h, new_w); + auto reshaped = context.mark_node(std::make_shared(transformed, final_shape, false)); + return {std::move(reshaped)}; +} + +} // namespace + OutputVector translate_pixel_shuffle(const NodeContext& context) { // aten::pixel_shuffle(Tensor self, int upscale_factor) -> Tensor - num_inputs_check(context, 2, 2); - auto x = context.get_input(0); - auto upscale_factor = get_input_as_i32(context, 1); - auto neg_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); - auto neg_3 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-3})); - auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); - auto zero_s = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); - auto one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1})); - auto one_s = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); - Output shape; - Output rank; - std::tie(shape, rank) = get_shape_rank(context, x, true); - // 1. Reshape input to [*, -1, r, r, H, W], where r is upscale factor - auto indices = context.mark_node(v0::Constant::create(element::i32, Shape{3}, {-3, -2, -1})); - auto dims = context.mark_node(std::make_shared(shape, indices, zero_s)); - auto dims_splitted = context.mark_node(std::make_shared(dims, zero_s, 3)); - auto c = dims_splitted->output(0); - auto h = dims_splitted->output(1); - auto w = dims_splitted->output(2); - auto dims_before = context.mark_node(std::make_shared(shape, zero, neg_3, one)); - auto upscale_factor_1d = context.mark_node(std::make_shared(upscale_factor, neg_1, false)); - auto intermediate_shape = context.mark_node( - std::make_shared(OutputVector{dims_before, neg_1, upscale_factor_1d, upscale_factor_1d, h, w}, 0)); - auto reshape = context.mark_node(std::make_shared(x, intermediate_shape, false)); - // 2. Transpose tensor to [*, C, r, H, r, W] - auto dims_before_len = context.mark_node(std::make_shared(dims_before, element::i32)); - auto dims_before_len_s = context.mark_node(std::make_shared(dims_before_len, zero)); - auto order_begin = context.mark_node(std::make_shared(zero_s, dims_before_len_s, one_s, element::i32)); - auto order_end_neg = context.mark_node( - v0::Constant::create(element::i32, Shape{5}, {-3, 0, -2, 1, -1})); // +2 because rank is expanded - auto order_end = context.mark_node(std::make_shared(order_end_neg, rank)); - auto order = context.mark_node(std::make_shared(OutputVector{order_begin, order_end}, 0)); - auto transpose = context.mark_node(std::make_shared(reshape, order)); - // 3. Reshape to [*, -1, r * H, r * W] - auto new_h = context.mark_node(std::make_shared(h, upscale_factor)); - auto new_w = context.mark_node(std::make_shared(w, upscale_factor)); - auto shape_after = - context.mark_node(std::make_shared(OutputVector{dims_before, neg_1, new_h, new_w}, 0)); - return {context.mark_node(std::make_shared(transpose, shape_after, false))}; + return translate_pixel_transform(context, true); }; OutputVector translate_pixel_unshuffle(const NodeContext& context) { // aten::pixel_unshuffle(Tensor self, int upscale_factor) -> Tensor - num_inputs_check(context, 2, 2); - auto x = context.get_input(0); - auto upscale_factor = get_input_as_i32(context, 1); - auto neg_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); - auto neg_3 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-3})); - auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); - auto zero_s = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); - auto one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1})); - auto one_s = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); - Output shape; - Output rank; - std::tie(shape, rank) = get_shape_rank(context, x, true); - // 1. Reshape input to [-1, C, H / r, r, W / r, r], where r is upscale factor - auto indices = context.mark_node(v0::Constant::create(element::i32, Shape{3}, {-3, -2, -1})); - auto dims = context.mark_node(std::make_shared(shape, indices, zero_s)); - auto dims_splitted = context.mark_node(std::make_shared(dims, zero_s, 3)); - auto c = dims_splitted->output(0); - auto h = dims_splitted->output(1); - auto w = dims_splitted->output(2); - auto dims_before = context.mark_node(std::make_shared(shape, zero, neg_3, one)); - auto r = context.mark_node(std::make_shared(upscale_factor, zero)); - auto new_h = context.mark_node(std::make_shared(h, upscale_factor, true)); - auto new_w = context.mark_node(std::make_shared(w, upscale_factor, true)); - auto intermediate_shape = - context.mark_node(std::make_shared(OutputVector{neg_1, c, new_h, r, new_w, r}, 0)); - auto x_reshaped = context.mark_node(std::make_shared(x, intermediate_shape, false)); - // 2. Transpose to [-1, C, r, r, H / r, W / r] - auto transpose_order = context.mark_node(v0::Constant::create(element::i32, Shape{6}, {0, 1, 3, 5, 2, 4})); - auto x_transposed = context.mark_node(std::make_shared(x_reshaped, transpose_order)); - // 3. Reshape to [*, C*r*r, H / r, W / r] - auto r_sqr = context.mark_node(std::make_shared(r, r)); - auto new_c = context.mark_node(std::make_shared(c, r_sqr)); - auto final_shape = - context.mark_node(std::make_shared(OutputVector{dims_before, new_c, new_h, new_w}, 0)); - return {context.mark_node(std::make_shared(x_transposed, final_shape, false))}; + return translate_pixel_transform(context, false); }; OutputVector translate_channel_shuffle(const NodeContext& context) { // aten::channel_shuffle(Tensor self, int groups) -> Tensor num_inputs_check(context, 2, 2); - auto x = context.get_input(0); - auto groups = context.get_input(1); - auto neg_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); - auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); - auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); - auto shape = context.mark_node(std::make_shared(x, element::i32)); - // PyTorch realization uses assumption that channels dim is always 1 - auto indices = context.mark_node(v0::Constant::create(element::i32, Shape{2}, {0, 1})); - auto dims = context.mark_node(std::make_shared(shape, indices, zero)); - auto dims_splitted = context.mark_node(std::make_shared(dims, zero, 2)); - auto c = dims_splitted->output(1); - auto n = dims_splitted->output(0); - groups = context.mark_node(std::make_shared(groups, element::i32)); - auto k = context.mark_node(std::make_shared(c, groups, true)); - auto g = context.mark_node(std::make_shared(groups, zero)); - // 1. Reshape input [N, G, K=C/G, -1] - auto reshape_indices = context.mark_node( - std::make_shared(OutputVector{std::move(n), std::move(g), std::move(k), std::move(neg_1)}, 0)); - x = context.mark_node(std::make_shared(x, reshape_indices, false)); - // 2. Transpose to [N, K, G, -1] - auto permute_indices = context.mark_node(v0::Constant::create(element::i32, Shape{4}, {0, 2, 1, 3})); - auto y = context.mark_node(std::make_shared(x, permute_indices)); - // 3. Reshape back to original shape - auto result = context.mark_node(std::make_shared(y, shape, false)); - return {result}; + const auto x = context.get_input(0); + const auto groups = context.const_input(1); + PYTORCH_OP_CONVERSION_CHECK(groups > 0, "groups argument for channel_shuffle must be positive"); + auto shuffled = context.mark_node(std::make_shared(x, 1, static_cast(groups))); + return {std::move(shuffled)}; }; } // namespace op diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index 89bb929581a20e..2a14b0e601f3de 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -180,7 +180,6 @@ OP_CONVERTER(translate_movedim); OP_CONVERTER(translate_multinomial); OP_CONVERTER(translate_narrow); OP_CONVERTER(translate_native_multi_head_attention); -OP_CONVERTER(translate_neg); OP_CONVERTER(translate_new_full); OP_CONVERTER(translate_new_ones); OP_CONVERTER(translate_new_zeros); @@ -638,7 +637,7 @@ const std::unordered_map get_supported_ops_ts() { {"aten::multinomial", op::translate_multinomial}, {"aten::narrow", op::translate_narrow}, {"aten::ne", op::translate_1to1_match_2_inputs_align_types}, - {"aten::neg", op::translate_neg}, + {"aten::neg", op::translate_1to1_match_1_inputs}, {"aten::new_empty", op::translate_new_zeros}, {"aten::new_full", op::translate_new_full}, {"aten::new_ones", op::translate_new_ones}, @@ -827,7 +826,7 @@ const std::unordered_map get_supported_ops_fx() { {"", op::translate_floor_divide}, {"", op::translate_getitem}, // TODO: Check if there is any other way to handle this {"", op::translate_mul}, - {"", op::translate_neg}, + {"", op::translate_1to1_match_1_inputs}, {"", op::translate_sub}, {"aten._adaptive_avg_pool1d.default", op::translate_adaptive_avg_pool1d}, {"aten._adaptive_avg_pool2d.default", op::translate_adaptive_avg_pool2d}, @@ -1014,7 +1013,7 @@ const std::unordered_map get_supported_ops_fx() { {"aten.native_layer_norm.default", op::translate_layer_norm_fx}, {"aten.ne.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.ne.Tensor", op::translate_1to1_match_2_inputs_align_types}, - {"aten.neg.default", op::translate_neg}, + {"aten.neg.default", op::translate_1to1_match_1_inputs}, {"aten.new_full.default", op::translate_new_full_fx}, {"aten.new_ones.default", op::translate_new_ones_fx}, {"aten.new_zeros.default", op::translate_new_zeros_fx}, From b6083fc5f7f7f4e329f9650e17967d0ba0ff4a5d Mon Sep 17 00:00:00 2001 From: Maxim Vafin Date: Wed, 21 Jan 2026 19:06:48 +0100 Subject: [PATCH 2/2] Add comment Signed-off-by: Maxim Vafin --- src/frontends/pytorch/src/op/pixel_shuffle.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/frontends/pytorch/src/op/pixel_shuffle.cpp b/src/frontends/pytorch/src/op/pixel_shuffle.cpp index 1ebf3ae14b971a..781c74d58c8869 100644 --- a/src/frontends/pytorch/src/op/pixel_shuffle.cpp +++ b/src/frontends/pytorch/src/op/pixel_shuffle.cpp @@ -28,6 +28,7 @@ using namespace ov::op; namespace { +// Holds the split shape information; dims_before captures all pre-spatial dims struct PixelSpatialInfo { Output dims_before; Output channels;