From 103fb811e9f241f3d45f59b8b06995055fbc258b Mon Sep 17 00:00:00 2001 From: Karan Jakhar Date: Fri, 19 Jan 2024 00:25:25 +0530 Subject: [PATCH] [PT FE] Add aten::gcd (#21740) * Add gcd and test * constant with static * exec_condition set to true * boolean scalar, output condition, zero shape [] * removing unused headers * more tests added * code format * Update src/frontends/pytorch/src/op/gcd.cpp Co-authored-by: Maxim Vafin * Update src/frontends/pytorch/src/op/gcd.cpp Co-authored-by: Maxim Vafin * Update src/frontends/pytorch/src/op/gcd.cpp Co-authored-by: Maxim Vafin * Update src/frontends/pytorch/src/op/gcd.cpp Co-authored-by: Maxim Vafin * ngraph header removed * int64 * Add int64 support * Update tests/layer_tests/pytorch_tests/test_gcd.py Co-authored-by: Maxim Vafin --------- Co-authored-by: Maxim Vafin --- src/frontends/pytorch/src/op/gcd.cpp | 68 +++++++++++ src/frontends/pytorch/src/op_table.cpp | 2 + tests/layer_tests/pytorch_tests/test_gcd.py | 124 ++++++++++++++++++++ 3 files changed, 194 insertions(+) create mode 100644 src/frontends/pytorch/src/op/gcd.cpp create mode 100644 tests/layer_tests/pytorch_tests/test_gcd.py diff --git a/src/frontends/pytorch/src/op/gcd.cpp b/src/frontends/pytorch/src/op/gcd.cpp new file mode 100644 index 00000000000..70301185afc --- /dev/null +++ b/src/frontends/pytorch/src/op/gcd.cpp @@ -0,0 +1,68 @@ +// Copyright (C) 2018-2023 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/loop.hpp" +#include "openvino/op/mod.hpp" +#include "openvino/op/not_equal.hpp" +#include "openvino/op/parameter.hpp" +#include "openvino/op/reduce_logical_or.hpp" +#include "openvino/op/select.hpp" +#include "openvino/openvino.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +using namespace ov::op; + +OutputVector translate_gcd(const NodeContext& context) { + num_inputs_check(context, 2, 2); + auto x = context.get_input(0); + auto y = context.get_input(1); + align_eltwise_input_types(context, x, y, true); + auto zero_i32 = ov::op::v0::Constant::create(element::i32, Shape{}, {0}); + + auto trip_count = std::make_shared(element::i32, Shape{}, 1000); + auto exec_condition = std::make_shared(element::boolean, Shape{}, true); + + auto loop = std::make_shared(trip_count, exec_condition); + + auto x_input = std::make_shared(x.get_element_type(), x.get_partial_shape()); + auto y_input = std::make_shared(y.get_element_type(), y.get_partial_shape()); + + x_input->set_element_type(x.get_element_type()); + y_input->set_element_type(y.get_element_type()); + auto zero = std::make_shared(zero_i32, x_input); + auto condition = std::make_shared(y_input, zero); + auto mod = std::make_shared(x_input, y_input); + auto new_x = std::make_shared(condition, y_input, x_input); + auto new_y = std::make_shared(condition, mod, zero); + + auto reduced_condition = std::make_shared(condition, zero); + + auto body = + std::make_shared(OutputVector{new_x, new_y, reduced_condition}, ParameterVector{x_input, y_input}); + loop->set_function(body); + + loop->set_special_body_ports({-1, 2}); + + loop->set_merged_input(x_input, x, new_x); + loop->set_merged_input(y_input, y, new_y); + + auto gcd_output = loop->get_iter_value(new_x, -1); + auto gcd_node = gcd_output.get_node_shared_ptr(); + + auto marked_gcd_node = context.mark_node(gcd_node); + return {marked_gcd_node}; +} + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index eb0338f52bb..82674ea59db 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -84,6 +84,7 @@ OP_CONVERTER(translate_frobenius_norm); OP_CONVERTER(translate_full); OP_CONVERTER(translate_full_like); OP_CONVERTER(translate_gather); +OP_CONVERTER(translate_gcd); OP_CONVERTER(translate_gelu); OP_CONVERTER(translate_get_attr); OP_CONVERTER(translate_getitem); @@ -386,6 +387,7 @@ const std::map get_supported_ops_ts() { {"aten::full", op::translate_full}, {"aten::full_like", op::translate_full_like}, {"aten::gather", op::translate_gather}, + {"aten::gcd", op::translate_gcd}, {"aten::ge", op::translate_1to1_match_2_inputs_align_types}, {"aten::gelu", op::translate_gelu}, {"aten::glu", op::translate_glu}, diff --git a/tests/layer_tests/pytorch_tests/test_gcd.py b/tests/layer_tests/pytorch_tests/test_gcd.py new file mode 100644 index 00000000000..a1b816b853e --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_gcd.py @@ -0,0 +1,124 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +import torch + +from pytorch_layer_test_class import PytorchLayerTest + + +class TestGcd(PytorchLayerTest): + def _prepare_input(self): + return self.input_data + + def create_model_tensor_input(self): + class aten_gcd_tensor(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, tensor_a, tensor_b): + return torch.gcd(tensor_a, tensor_b) + + ref_net = None + + return aten_gcd_tensor(), ref_net, "aten::gcd" + + def create_model_int_input(self): + class aten_gcd_int(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, int_a: int, int_b: int): + return torch.tensor(torch.gcd(int_a, int_b)) + + ref_net = None + + return aten_gcd_int(), ref_net, "aten::gcd" + + @pytest.mark.nightly + @pytest.mark.precommit + def test_gcd_int(self, ie_device, precision, ir_version): + self.input_data = (np.array(11, dtype=np.int32), np.array(17, dtype=np.int32)) + self._test( + *self.create_model_int_input(), + ie_device, + precision, + ir_version, + use_convert_model=True, + trace_model=True + ) + + @pytest.mark.nightly + @pytest.mark.precommit + def test_gcd_tensor(self, ie_device, precision, ir_version): + self.input_data = ( + np.array([14, 4, 12, 10, 3, 0], dtype=np.int32), + np.array([121, 2, 16, 0, 1, 8], dtype=np.int32), + ) + self._test( + *self.create_model_tensor_input(), + ie_device, + precision, + ir_version, + use_convert_model=True, + trace_model=True + ) + + @pytest.mark.nightly + @pytest.mark.precommit + def test_gcd_int64(self, ie_device, precision, ir_version): + self.input_data = (np.array(11, dtype=np.int64), np.array(17, dtype=np.int64)) + self._test( + *self.create_model_int_input(), + ie_device, + precision, + ir_version, + use_convert_model=True, + trace_model=True + ) + + @pytest.mark.nightly + @pytest.mark.precommit + def test_gcd_tensor64(self, ie_device, precision, ir_version): + self.input_data = ( + np.array([14, 4, 12, 10, 3, 0], dtype=np.int64), + np.array([121, 2, 16, 0, 1, 8], dtype=np.int64), + ) + self._test( + *self.create_model_tensor_input(), + ie_device, + precision, + ir_version, + use_convert_model=True, + trace_model=True + ) + + @pytest.mark.nightly + @pytest.mark.precommit + def test_gcd_int_diff_dtypes(self, ie_device, precision, ir_version): + self.input_data = (np.array(11, dtype=np.int64), np.array(17, dtype=np.int32)) + self._test( + *self.create_model_int_input(), + ie_device, + precision, + ir_version, + use_convert_model=True, + trace_model=True + ) + + @pytest.mark.nightly + @pytest.mark.precommit + def test_gcd_tensor_diff_dtypes(self, ie_device, precision, ir_version): + self.input_data = ( + np.array([14, 4, 12, 10, 3, 0], dtype=np.int64), + np.array([121, 2, 16, 0, 1, 8], dtype=np.int32), + ) + self._test( + *self.create_model_tensor_input(), + ie_device, + precision, + ir_version, + use_convert_model=True, + trace_model=True + )