[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*
This commit is contained in:
Maxim Vafin 2024-05-20 09:05:23 +02:00 committed by GitHub
parent 38fdcbab78
commit ab6e0f8333
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 35 additions and 14 deletions

View File

@ -398,7 +398,9 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::clamp", op::translate_clamp},
{"aten::clamp_", op::inplace_op<op::translate_clamp>},
{"aten::clamp_max", op::translate_1to1_match_2_inputs_align_types<opset10::Minimum>},
{"aten::clamp_max_", op::inplace_op<op::translate_1to1_match_2_inputs_align_types<opset10::Minimum>>},
{"aten::clamp_min", op::translate_1to1_match_2_inputs_align_types<opset10::Maximum>},
{"aten::clamp_min_", op::inplace_op<op::translate_1to1_match_2_inputs_align_types<opset10::Maximum>>},
{"aten::clip", op::translate_clamp},
{"aten::clip_", op::inplace_op<op::translate_clamp>},
{"aten::clone", op::skip_node}, // ignore clone operators that are inserted by PyTorch autograd
@ -617,6 +619,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::roll", op::translate_roll},
{"aten::round", op::translate_round},
{"aten::rsqrt", op::optional_out<op::translate_rsqrt, 1>},
{"aten::rsqrt_", op::inplace_op<op::translate_rsqrt>},
{"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<std::string, CreatorFunction> 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<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Sqrt>, 1>},
{"aten::sqrt_", op::inplace_op<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Sqrt>>},
{"aten::square", op::translate_square},
{"aten::squeeze", op::quantizable_op<op::translate_squeeze>},
// aten::stack - Supported in limited set of patterns

View File

@ -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)

View File

@ -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_",