From ab6e0f8333501c447680fc247f71806aa04b521d Mon Sep 17 00:00:00 2001 From: Maxim Vafin Date: Mon, 20 May 2024 09:05:23 +0200 Subject: [PATCH] [PT FE] Support inplace aten::sqrt, aten::rsqrt, aten::clamp_{min, max} (#24354) ### Details: - *Support inplace `aten::sqrt_`, `aten::rsqrt_`, `aten::clamp_min_`, `aten::clamp_max_`* ### Tickets: - *ticket-id* --- src/frontends/pytorch/src/op_table.cpp | 4 ++ tests/layer_tests/pytorch_tests/test_clamp.py | 41 ++++++++++++------- .../pytorch_tests/test_unary_ops.py | 4 ++ 3 files changed, 35 insertions(+), 14 deletions(-) diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index f0749cbb82b..231197de7c4 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -398,7 +398,9 @@ const std::map get_supported_ops_ts() { {"aten::clamp", op::translate_clamp}, {"aten::clamp_", op::inplace_op}, {"aten::clamp_max", op::translate_1to1_match_2_inputs_align_types}, + {"aten::clamp_max_", op::inplace_op>}, {"aten::clamp_min", op::translate_1to1_match_2_inputs_align_types}, + {"aten::clamp_min_", op::inplace_op>}, {"aten::clip", op::translate_clamp}, {"aten::clip_", op::inplace_op}, {"aten::clone", op::skip_node}, // ignore clone operators that are inserted by PyTorch autograd @@ -617,6 +619,7 @@ const std::map get_supported_ops_ts() { {"aten::roll", op::translate_roll}, {"aten::round", op::translate_round}, {"aten::rsqrt", op::optional_out}, + {"aten::rsqrt_", op::inplace_op}, {"aten::rsub", op::translate_rsub}, {"aten::ScalarImplicit", op::skip_node}, {"aten::scaled_dot_product_attention", op::translate_scaled_dot_product_attention}, @@ -647,6 +650,7 @@ const std::map get_supported_ops_ts() { // aten::split - Supported in limited set of patterns // aten::split_with_sizes - Supported in limited set of patterns {"aten::sqrt", op::optional_out, 1>}, + {"aten::sqrt_", op::inplace_op>}, {"aten::square", op::translate_square}, {"aten::squeeze", op::quantizable_op}, // aten::stack - Supported in limited set of patterns diff --git a/tests/layer_tests/pytorch_tests/test_clamp.py b/tests/layer_tests/pytorch_tests/test_clamp.py index b8a977a8970..4550146dfe8 100644 --- a/tests/layer_tests/pytorch_tests/test_clamp.py +++ b/tests/layer_tests/pytorch_tests/test_clamp.py @@ -59,27 +59,34 @@ class TestClampMin(PytorchLayerTest): import numpy as np return (np.random.randn(1, 3, 224, 224).astype(np.float32),) - def create_model(self, minimum, as_tensor=False): + def create_model(self, minimum, as_tensor=False, inplace=False): import torch class aten_clamp_min(torch.nn.Module): - def __init__(self, minimum, as_tensor): + def __init__(self, minimum, as_tensor, inplace): super(aten_clamp_min, self).__init__() self.min = torch.tensor(minimum) if as_tensor else minimum + if inplace: + self.forward = self.forward_inplace def forward(self, x): return torch.clamp_min(x, self.min) - ref_net = None - op_name = "aten::clamp_min" - return aten_clamp_min(minimum, as_tensor), ref_net, op_name + def forward_inplace(self, x): + torch.clamp_min_(x, self.min) + return x + + op_name = "aten::clamp_min_" if inplace else "aten::clamp_min" + return aten_clamp_min(minimum, as_tensor, inplace), None, op_name @pytest.mark.parametrize("minimum", [0., 1., -1., 0.5, 2]) @pytest.mark.parametrize("as_tensor", [True, False]) + @pytest.mark.parametrize("inplace", [True, False]) @pytest.mark.nightly + @pytest.mark.precommit @pytest.mark.precommit_fx_backend - def test_clamp_min(self, minimum, as_tensor, ie_device, precision, ir_version): - self._test(*self.create_model(minimum, as_tensor), ie_device, + def test_clamp_min(self, minimum, as_tensor, inplace, ie_device, precision, ir_version): + self._test(*self.create_model(minimum, as_tensor, inplace), ie_device, precision, ir_version, use_convert_model=True, trace_model=True) @@ -88,27 +95,33 @@ class TestClampMax(PytorchLayerTest): import numpy as np return (np.random.randn(1, 3, 224, 224).astype(np.float32),) - def create_model(self, maximum, as_tensor=False): + def create_model(self, maximum, as_tensor=False, inplace=False): import torch class aten_clamp_max(torch.nn.Module): - def __init__(self, maximum, as_tensor): + def __init__(self, maximum, as_tensor, inplace): super(aten_clamp_max, self).__init__() self.max = torch.tensor(maximum) if as_tensor else maximum + if inplace: + self.forward = self.forward_inplace def forward(self, x): return torch.clamp_max(x, self.max) - ref_net = None - op_name = "aten::clamp_max" - return aten_clamp_max(maximum, as_tensor), ref_net, op_name + def forward_inplace(self, x): + torch.clamp_max_(x, self.max) + return x + + op_name = "aten::clamp_max_" if inplace else "aten::clamp_max" + return aten_clamp_max(maximum, as_tensor, inplace), None, op_name @pytest.mark.parametrize("maximum", [0., 1., -1., 0.5, 2]) @pytest.mark.parametrize("as_tensor", [True, False]) + @pytest.mark.parametrize("inplace", [skip_if_export(True), False]) @pytest.mark.nightly @pytest.mark.precommit @pytest.mark.precommit_torch_export @pytest.mark.precommit_fx_backend - def test_clamp(self, maximum, as_tensor, ie_device, precision, ir_version): - self._test(*self.create_model(maximum, as_tensor), ie_device, + def test_clamp_max(self, maximum, as_tensor, inplace, ie_device, precision, ir_version): + self._test(*self.create_model(maximum, as_tensor, inplace), ie_device, precision, ir_version, use_convert_model=True, trace_model=True) diff --git a/tests/layer_tests/pytorch_tests/test_unary_ops.py b/tests/layer_tests/pytorch_tests/test_unary_ops.py index d54d1102737..689ad67e64d 100644 --- a/tests/layer_tests/pytorch_tests/test_unary_ops.py +++ b/tests/layer_tests/pytorch_tests/test_unary_ops.py @@ -11,7 +11,9 @@ OPS = { "aten::abs": torch.abs, "aten::abs_": torch.abs_, "aten::rsqrt": torch.rsqrt, + "aten::rsqrt_": torch.rsqrt_, "aten::sqrt": torch.sqrt, + "aten::sqrt_": torch.sqrt_, "aten::erf": torch.erf, "aten::erf_": torch.erf_, "aten::erfc": torch.erfc, @@ -171,6 +173,8 @@ class TestUnaryOp(PytorchLayerTest): "aten::erf_", "aten::erfc_", "aten::exp_", + "aten::rsqrt_", + "aten::sqrt_", "aten::expm1_", "aten::sigmoid_", "aten::reciprocal_",