From edfa797c15bf9175e9c403be8fd136b13c636f78 Mon Sep 17 00:00:00 2001 From: Wei Wei Date: Wed, 9 Aug 2023 13:40:40 -0700 Subject: [PATCH] Remove aten2ait test cases (#891) Summary: Pull Request resolved: https://github.com/facebookincubator/AITemplate/pull/891 This project is decided to not move forward. Delete those test cases. Reviewed By: aakhundov Differential Revision: D48197622 fbshipit-source-id: 470d21ccaeabe28154be6cf35424a13aa9e8b2e0 --- .../test_ait_adaptive_avg_pool2d_aten.py | 98 -------- .../test_ait_avg_pool2d_aten.py | 83 ------- .../test_ait_batch_norm_aten.py | 94 -------- .../test_ait_binary_op_aten.py | 129 ---------- .../test/converters_aten/test_ait_cat_aten.py | 82 ------- .../converters_aten/test_ait_chunk_aten.py | 63 ----- .../converters_aten/test_ait_clamp_aten.py | 80 ------- .../converters_aten/test_ait_conv2d_aten.py | 127 ---------- .../test_ait_convtranspose2d_aten.py | 107 --------- .../converters_aten/test_ait_flatten_aten.py | 78 ------ .../converters_aten/test_ait_hardtanh_aten.py | 78 ------ .../test_ait_layer_norm_aten.py | 116 --------- .../converters_aten/test_ait_linear_aten.py | 88 ------- .../converters_aten/test_ait_matmul_aten.py | 97 -------- .../test_ait_max_pool2d_aten.py | 83 ------- .../converters_aten/test_ait_model_aten.py | 107 --------- .../converters_aten/test_ait_nan2num_aten.py | 89 ------- .../converters_aten/test_ait_permute_aten.py | 78 ------ .../test/converters_aten/test_ait_pow_aten.py | 73 ------ .../converters_aten/test_ait_reduce_aten.py | 131 ----------- .../converters_aten/test_ait_relu_aten.py | 61 ----- .../converters_aten/test_ait_reshape_aten.py | 124 ---------- .../converters_aten/test_ait_size_aten.py | 69 ------ .../test_ait_slice_tensor_aten.py | 222 ------------------ .../converters_aten/test_ait_split_aten.py | 82 ------- .../converters_aten/test_ait_squeeze_aten.py | 165 ------------- .../test_ait_unary_ops_aten.py | 74 ------ 27 files changed, 2678 deletions(-) delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_adaptive_avg_pool2d_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_avg_pool2d_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_batch_norm_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_binary_op_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_cat_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_chunk_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_clamp_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_conv2d_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_convtranspose2d_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_flatten_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_hardtanh_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_layer_norm_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_linear_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_matmul_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_max_pool2d_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_model_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_nan2num_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_permute_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_pow_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_reduce_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_relu_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_reshape_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_size_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_slice_tensor_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_split_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_squeeze_aten.py delete mode 100644 fx2ait/fx2ait/test/converters_aten/test_ait_unary_ops_aten.py diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_adaptive_avg_pool2d_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_adaptive_avg_pool2d_aten.py deleted file mode 100644 index 2d1e8a635..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_adaptive_avg_pool2d_aten.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -#!/usr/bin/env fbpython -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. - -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestAdaptiveAvgPool2dConverter(DispatchTestCase): - @parameterized.expand( - [ - ((64, 64), torch.ops.aten._adaptive_avg_pool2d.default), - ((128, 128), torch.ops.aten._adaptive_avg_pool2d.default), - (64, torch.ops.aten._adaptive_avg_pool2d.default), - ( - (1, 1), - torch.ops.aten.mean.dim, - ), - ] - ) - def test_adaptive_avgpool2d(self, output_size, op_check): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.pool = torch.nn.AdaptiveAvgPool2d(output_size) - - def forward(self, x): - return self.pool(x) - - model = TestModule().half().cuda() - inputs = [torch.randn(1, 32, 256, 256).cuda().half()] - if op_check == torch.ops.aten.mean.dim: - permute_inputs = None - permute_outputs = None - else: - permute_inputs = [0, 2, 3, 1] - permute_outputs = [0, 3, 1, 2] - self.run_test( - model, - inputs, - expected_ops={op_check}, - permute_inputs=permute_inputs, - permute_outputs=permute_outputs, - ) - - @parameterized.expand( - [ - ((64, 64),), - ((128, 128),), - (64,), - ] - ) - def test_dynamic_adaptive_avgpool2d( - self, - output_size, - ): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.pool = torch.nn.AdaptiveAvgPool2d(output_size) - - def forward(self, x): - return self.pool(x) - - model = TestModule().half().cuda() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 32, 256, 256], - ], - inputs_max=[ - [10, 32, 256, 256], - ], - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten._adaptive_avg_pool2d.default}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_avg_pool2d_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_avg_pool2d_aten.py deleted file mode 100644 index 681f8b407..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_avg_pool2d_aten.py +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestAvgPool2dConverter(DispatchTestCase): - @parameterized.expand( - [ - (1, 1, 0), - ((2, 2), 2, 1), - ((4, 4), (4, 4), 0), - ] - ) - def test_avgpool2d(self, kernel_size, stride, padding): - class TestModule(torch.nn.Module): - def __init__(self, kernel_size, stride, padding): - super().__init__() - self.pool = torch.nn.AvgPool2d(kernel_size, stride, padding) - - def forward(self, x): - return self.pool(x) - - model = TestModule(kernel_size, stride, padding).half().cuda() - inputs = [torch.randn(1, 4, 256, 256).cuda().half()] - self.run_test( - model, - inputs, - expected_ops={torch.ops.aten.avg_pool2d.default}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) - - @parameterized.expand( - [ - (1, 1, 0), - ((2, 2), 2, 1), - ((4, 4), (4, 4), 0), - ] - ) - def test_dynamic_avgpool2d(self, kernel_size, stride, padding): - class TestModule(torch.nn.Module): - def __init__(self, kernel_size, stride, padding): - super().__init__() - self.pool = torch.nn.AvgPool2d(kernel_size, stride, padding) - - def forward(self, x): - return self.pool(x) - - model = TestModule(kernel_size, stride, padding).half().cuda() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [1, 4, 256, 256], - ], - inputs_max=[ - [10, 4, 256, 256], - ], - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.avg_pool2d.default}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_batch_norm_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_batch_norm_aten.py deleted file mode 100644 index 6450643fb..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_batch_norm_aten.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. - -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase - - -class TestAdaptiveAvgPool2dConverter(DispatchTestCase): - def test_batch_norm(self): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.bn = torch.nn.BatchNorm2d(3) - - def forward(self, x): - y = self.bn(x) - y = y.mul(1) - return y - - model = TestModule().half().cuda() - inputs = [torch.randn(1, 3, 244, 244).cuda().half()] - self.run_test( - model, - inputs, - expected_ops={}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) - - def test_batch_norm_2layers(self): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.bn = torch.nn.BatchNorm2d(3) - self.bn2 = torch.nn.BatchNorm2d(3) - - def forward(self, x): - y = self.bn(x) - y = y.mul(1) - y = self.bn2(y) - return y - - model = TestModule().half().cuda() - inputs = [torch.randn(1, 3, 244, 244).cuda().half()] - self.run_test( - model, - inputs, - expected_ops={}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) - - def test_dynamic_batch_norm(self): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.bn = torch.nn.BatchNorm2d(3) - - def forward(self, x): - return self.bn(x) - - model = TestModule().half().cuda() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [1, 3, 244, 244], - ], - inputs_max=[ - [10, 3, 256, 256], - ], - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.batch_norm}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_binary_op_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_binary_op_aten.py deleted file mode 100644 index d2a9f7611..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_binary_op_aten.py +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import operator -from typing import Callable, List, Tuple - -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -TWO_TENSOR_INPUTS = [ - (torch.randn(2, 3, 4), torch.randn(2, 3, 4)), - (torch.randn(3, 4), torch.randn(2, 3, 4)), - (torch.randn(2, 3, 4), torch.randn(3, 4)), - (torch.randn(1, 1, 1), torch.randn(2, 3, 4)), - (torch.randn(1), torch.randn(2, 3, 4)), - (torch.randn(2, 3, 4), torch.randn(1)), - (torch.randn(2, 3, 4), torch.randn(1, 1, 1)), - (torch.randn(1, 3, 4), torch.randn(5, 1, 4)), -] - - -class TestATenBinaryOpConverter(DispatchTestCase): - @parameterized.expand( - [ - [ - "add", - operator.add, - torch.ops.aten.add.Tensor, - TWO_TENSOR_INPUTS, - ], - [ - "sub", - operator.sub, - torch.ops.aten.sub.Tensor, - TWO_TENSOR_INPUTS, - ], - [ - "mul", - operator.mul, - torch.ops.aten.mul.Tensor, - TWO_TENSOR_INPUTS, - ], - [ - "div", - operator.truediv, - torch.ops.aten.div.Tensor, - [(lhs, rhs.clamp(min=0.01)) for lhs, rhs in TWO_TENSOR_INPUTS], - ], - ] - ) - def test_two_tensors( - self, - name: str, - op: Callable, - aten_op: Callable, - inputs: List[Tuple[torch.Tensor, torch.Tensor]], - ) -> None: - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return op(x, y) - - for lhs, rhs in inputs: - model = TestModule().cuda() - lhs = lhs.half().cuda() - rhs = rhs.half().cuda() - self.run_test(model, [lhs, rhs], expected_ops={aten_op}) - - @parameterized.expand( - [ - [ - "dynamic_add", - operator.add, - torch.ops.aten.add.Tensor, - ], - [ - "dynamic_sub", - operator.sub, - torch.ops.aten.sub.Tensor, - ], - [ - "dynamic_sub", - operator.mul, - torch.ops.aten.mul.Tensor, - ], - [ - "dynamic_div", - operator.truediv, - torch.ops.aten.div.Tensor, - ], - ] - ) - def test_dynamic_two_tensors( - self, - name: str, - op: Callable, - aten_op: Callable, - ) -> None: - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return op(x, y) - - m = TensorSpec.gen_int_var_min_max(1, 32, "dynamic_m") - n = TensorSpec.gen_int_var_min_max(3, 1024, "dynamic_n") - k = TensorSpec.gen_int_var_min_max(4, 2048, "dynamic_k") - model = TestModule().cuda().half() - # AIT can automatically calculate broadcast - input_spec = TensorSpec.create_spec_from_int_vars( - [[m, n, k], [n, k]], dtype_list=[torch.float16] * 2 - ) - - self.run_test_with_dynamic_shape( - model, - input_spec, - expected_ops={aten_op}, - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_cat_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_cat_aten.py deleted file mode 100644 index b971f735b..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_cat_aten.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. - -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestCatConverter(DispatchTestCase): - @parameterized.expand( - [ - ["default", 0], - ["positive_dim", 1], - ["negative_dim", -1], - ] - ) - def test_cat(self, name: str, dim: int): - class TestModule(torch.nn.Module): - def forward( - self, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor - ) -> torch.Tensor: - return torch.cat([x, y, z], dim=dim) - - model = TestModule().cuda() - inputs = [ - torch.randn(2, 3, 4).half().cuda(), - torch.randn(2, 3, 4).half().cuda(), - torch.randn(2, 3, 4).half().cuda(), - ] - self.run_test(model, inputs, expected_ops={torch.ops.aten.cat.default}) - - @parameterized.expand( - [ - ["default", 0], - ["positive_dim", 1], - ["negative_dim", -1], - ] - ) - def test_cat_dynamic_shape(self, name: str, dim: int): - class TestModule(torch.nn.Module): - def forward( - self, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor - ) -> torch.Tensor: - return torch.cat([x, y, z], dim=dim) - - model = TestModule().cuda() - - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - [2, 3, 4], - [2, 3, 4], - ], - inputs_max=[ - [20, 3, 4], - [20, 3, 4], - [20, 3, 4], - ], - dtype_list=[ - torch.float16, - torch.float16, - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.cat.default} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_chunk_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_chunk_aten.py deleted file mode 100644 index 3acdb6ad6..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_chunk_aten.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.fx2ait import TensorSpec -from fx2ait.passes.lower_basic_pass_aten import aten_compose_chunk -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestChunkConverter(DispatchTestCase): - @parameterized.expand( - [ - param("default", 2, [3, 10, 2], 1), - param("no_dim", 2, [3, 10, 2]), - param("neg_dim", 1, [3, 10, 2], -2), - param("chunk_bigger_than_dim", 4, [2, 10, 2], 2), - ] - ) - def test_chunk(self, name, chunks, shape, dim=None): - class TestModule(torch.nn.Module): - def forward(self, x): - x = ( - torch.chunk(x, chunks=chunks, dim=dim) - if dim is not None - else torch.chunk(x, chunks=chunks) - ) - # For AIT, all chunk results must be used - return x[0] - - model = TestModule().cuda().half() - inputs = [torch.randn(shape).half().cuda()] - self.run_test(model, inputs, expected_ops={aten_compose_chunk}) - - def test_chunk_dynamic(self): - class TestModule(torch.nn.Module): - def forward(self, x): - x = torch.chunk(x, chunks=2, dim=1) - return x[0] - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[[20, 10, 8]], - inputs_max=[[50, 10, 8]], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={aten_compose_chunk} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_clamp_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_clamp_aten.py deleted file mode 100644 index 31d74189e..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_clamp_aten.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestClampConverter(DispatchTestCase): - @parameterized.expand( - [ - param("default", min=-1, max=0, use_clamp=True), - param("min", min=0.5, use_clamp=False), - param("max", max=0.5, use_clamp=True), - param("minBiggerThanMax", min=1, max=0, use_clamp=False), - ] - ) - def test_clamp(self, name, min=None, max=None, use_clamp=True): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.op = torch.clamp if use_clamp else torch.clip - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x, min=min, max=max) - - model = TestModule().cuda().half() - inputs = [ - torch.randn(1, 2, 3).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.clamp.default}) - - @parameterized.expand( - [ - param("default", min=-1, max=0, use_clamp=True), - param("min", min=0.5, use_clamp=False), - param("max", max=0.5, use_clamp=True), - param("minBiggerThanMax", min=1, max=0, use_clamp=False), - ] - ) - def test_dynamic_clamp(self, name, min=None, max=None, use_clamp=True): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.op = torch.clamp if use_clamp else torch.clip - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x, min=min, max=max) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 8, 10], - ], - inputs_max=[ - [20, 12, 32], - ], - dtype_list=[ - torch.float16, - torch.float16, - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.clamp.default} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_conv2d_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_conv2d_aten.py deleted file mode 100644 index b15dae226..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_conv2d_aten.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. - -import torch -from aitemplate.compiler.public import DynamicProfileStrategy -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestConv2dConverter(DispatchTestCase): - @parameterized.expand( - [ - param("default", 1), - param("no_bias", 1, bias=False), - param("tuple_parameters", 1, (1, 1), (1, 1)), - param("non_zero_padding", 1, padding=1), - param("non_unary_params", 3, 2, padding=1, bias=False), - param("dilation", 1, dilation=2), - param("multi_group", 1, 1, 1, 1, 3, bias=True), - ] - ) - def test_conv2d( - self, - name, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - bias=True, - ): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.conv = torch.nn.Conv2d( - 3, 36, kernel_size, stride, padding, dilation, groups, bias - ) - self.relu = torch.nn.ReLU() - - def forward(self, x): - return self.relu(self.conv(x)) - - model = TestModule().cuda().half() - inputs = [torch.randn(1, 3, 24, 24).cuda().half()] - self.run_test( - model, - inputs, - expected_ops={ - torch.ops.aten.convolution.default, - torch.ops.aten.relu.default, - }, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) - - @parameterized.expand( - [ - param("default", 1), - param("no_bias", 1, bias=False), - param("tuple_parameters", 1, (1, 1), (1, 1)), - param("non_zero_padding", 1, padding=1), - param("non_unary_params", 3, 2, padding=1, bias=False), - param("dilation", 1, dilation=2), - param("multi_group", 1, 1, 1, 1, 3, bias=True), - ] - ) - def test_dynamic_conv2d( - self, - name, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - bias=True, - ): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.conv = torch.nn.Conv2d( - 3, 36, kernel_size, stride, padding, dilation, groups, bias - ) - self.relu = torch.nn.ReLU() - - def forward(self, x): - return self.relu(self.conv(x)) - - model = TestModule().cuda().half() - - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 24, 24], - ], - inputs_max=[ - [32, 3, 24, 24], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={ - torch.ops.aten.convolution.default, - torch.ops.aten.relu.default, - }, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - # AIT conv2d only support HINTS as dyanmic profiliing strategy. - dynamic_profile_strategy=DynamicProfileStrategy.HINTS, - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_convtranspose2d_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_convtranspose2d_aten.py deleted file mode 100644 index 50c13008e..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_convtranspose2d_aten.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestConvtTranspose2dConverter(DispatchTestCase): - @parameterized.expand( - [ - param("default", 1), - param("no_bias", 2, bias=False), - param("tuple_parameters", 1, (1, 1), (1, 1)), - param("non_zero_padding", 1, padding=1), - param("non_unary_params", 3, 2, padding=1, bias=False), - ] - ) - def test_convtranspose( - self, - name, - kernel_size, - stride=2, - padding=0, - dilation=1, # only support dilation = 1 - groups=1, - bias=True, - ): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.convtranspose = torch.nn.ConvTranspose2d( - 192, - 256, - kernel_size=kernel_size, - stride=stride, - padding=padding, - output_padding=0, - groups=groups, - bias=bias, - dilation=dilation, - ) - - def forward(self, x): - return self.convtranspose(x) - - model = TestModule().cuda().half().eval() - inputs = [torch.randn(1, 192, 28, 28).cuda().half()] - _ = model(*inputs) - self.run_test( - model, - inputs, - expected_ops={torch.ops.aten.convolution.default}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) - - # # only works when in_ch == out_ch - def test_convtranspose_multi_group( - self, - name="multi_group", - kernel_size=2, - stride=2, - padding=0, - dilation=1, # only support dilation = 1 - groups=2, - bias=True, - ): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.convtranspose = torch.nn.ConvTranspose2d( - 192, - 192, # must to divisblce by 8 - kernel_size=kernel_size, - stride=stride, - padding=padding, - output_padding=0, - groups=groups, - bias=bias, - dilation=dilation, - ) - - def forward(self, x): - return self.convtranspose(x) - - model = TestModule().cuda().half().eval() - inputs = [torch.randn(1, 192, 28, 28).cuda().half()] - _ = model(*inputs) - self.run_test( - model, - inputs, - expected_ops={torch.ops.aten.convolution.default}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_flatten_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_flatten_aten.py deleted file mode 100644 index fc215d96f..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_flatten_aten.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -import torch.nn as nn -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestFlattenConverter(DispatchTestCase): - @parameterized.expand( - [ - ("flatten_middle_dims", 1, 2), - ("flatten_last_3_dims", 1, 3), - ("flatten_all", 0, 3), - ] - ) - def test_flatten(self, _, start_dim, end_dim): - class TestModule(nn.Module): - def __init__(self, start, end): - super().__init__() - self.start = start - self.end = end - - def forward(self, x): - return torch.flatten(x, self.start, self.end) - - model = TestModule(start_dim, end_dim).cuda().half() - inputs = (torch.randn(2, 3, 4, 5).half().cuda(),) - - self.run_test(model, inputs, expected_ops={torch.ops.aten.view.default}) - - @parameterized.expand( - [ - ("flatten_middle_dims", 1, 2), - ("flatten_last_3_dims", 1, 3), - ] - ) - def test_flatten_with_dynamic_shape(self, _, start_dim, end_dim): - class TestModule(nn.Module): - def __init__(self, start, end): - super().__init__() - self.start = start - self.end = end - - def forward(self, x): - return torch.flatten(x, self.start, self.end) - - model = TestModule(start_dim, end_dim).cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4, 5], - ], - inputs_max=[ - [10, 20, 4, 5], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.view.default}, - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_hardtanh_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_hardtanh_aten.py deleted file mode 100644 index f974ebbb3..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_hardtanh_aten.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestHardTanhConverter(DispatchTestCase): - @parameterized.expand( - [ - param("default", min=-1.5, max=3), - param("min", min=-1.5), - param("max", max=3), - ] - ) - def test_hardtanh(self, name, min=-1, max=1): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.op = torch.nn.Hardtanh(min, max) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x) - - model = TestModule().cuda().half() - inputs = [ - torch.randn(1, 2, 3).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.hardtanh.default}) - - @parameterized.expand( - [ - param("default", min=-1.2, max=2), - param("min", min=-1.2), - param("max", max=2), - ] - ) - def test_dynamic_hardtanh(self, name, min=-1, max=1): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.op = torch.nn.Hardtanh(min, max) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 8, 10], - ], - inputs_max=[ - [20, 12, 32], - ], - dtype_list=[ - torch.float16, - torch.float16, - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.hardtanh.default} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_layer_norm_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_layer_norm_aten.py deleted file mode 100644 index 7a6036dd3..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_layer_norm_aten.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized -from torch import nn - - -class TestLayernormConverter(DispatchTestCase): - @parameterized.expand( - [ - param("1d_normalized_shape", [10], [2, 10]), - param("1d_normalized_shape_3d_input", [10], [2, 6, 10]), - param("2d_normalized_shape", [6, 10], [2, 6, 10]), - # FIXME: Enable test case once layernorm support expand - # param("2d_normalized_shape", [5, 10], [5, 10]), - ] - ) - def test_layer_norm(self, name, normalized_shape, input_shape): - class TestModule(torch.nn.Module): - def __init__(self, normalized_shape): - super().__init__() - # TODO remove hard code eps once layernorm api expose eps setting - self.mod = nn.LayerNorm(normalized_shape, eps=1e-5) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.mod(x) - - model = TestModule(normalized_shape).cuda().half() - inputs = [ - torch.randn(input_shape).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.layer_norm.default}) - - def test_layer_norm_IntImm_shape(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - shape = x.shape - normalized_shape = shape[1:] - return torch.nn.functional.layer_norm(x, normalized_shape, eps=1e-5) - - model = TestModule().cuda().half() - inputs = [ - torch.randn([10, 10]).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.layer_norm.default}) - - @parameterized.expand( - [ - param("1d_normalized_shape", [10], [[2, 10], [12, 10]]), - param("1d_normalized_shape_3d_input", [10], [[2, 6, 10], [12, 20, 10]]), - param("2d_normalized_shape", [6, 10], [[2, 6, 10], [12, 6, 10]]), - ] - ) - def test_dynamic_layer_norm(self, name, normalized_shape, input_shape): - class TestModule(torch.nn.Module): - def __init__(self, normalized_shape): - super().__init__() - # TODO remove hard code eps once layernorm api expose eps setting - self.mod = nn.LayerNorm(normalized_shape, eps=1e-5) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.mod(x) - - model = TestModule(normalized_shape).cuda().half() - - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - input_shape[0], - ], - inputs_max=[ - input_shape[1], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.layer_norm.default} - ) - - def test_dynamic_layer_norm_IntImm_shape(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - shape = x.shape - normalized_shape = shape[1:] - return torch.nn.functional.layer_norm(x, normalized_shape, eps=1e-5) - - model = TestModule().cuda().half() - inputs = TensorSpec.create_spec_from_shapes( - inputs_min=[[10, 30]], - inputs_max=[[20, 30]], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs, expected_ops={torch.ops.aten.layer_norm.default} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_linear_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_linear_aten.py deleted file mode 100644 index 618669f06..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_linear_aten.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.passes.lower_basic_pass_aten import aten_compose_mm_2d -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestLinearConverter(DispatchTestCase): - @parameterized.expand( - [ - ("default", [1, 512], True, torch.ops.aten.linear), - ("matrix", [5, 512], True, torch.ops.aten.linear), - ("no_bias", [1, 512], False, torch.ops.aten.linear), - ("multi_dim_matrix", [4, 5, 512], True, torch.ops.aten.linear), - ( - "multi_dim_matrix", - [4, 5, 512], - False, - aten_compose_mm_2d, - ), - ] - ) - def test_linear(self, test_name, shape, bias, expected): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(512, 256, bias) - - def forward(self, x): - return self.linear(x) - - model = TestModule().cuda().half() - inputs = [torch.randn(shape).half().cuda()] - self.run_test(model, inputs, expected_ops={expected}) - - @parameterized.expand( - [ - ("default", [[1, 5], [512, 512]], True), - ("no_bias", [[1, 4], [512, 512]], False), - ( - "multi_dim_matrix", - [[2, 4], [512, 512]], - True, - ), - ( - "multi_dim_matrix_no_bias", - [[2, 4], [512, 512]], - False, - ), - ] - ) - def test_dynamic_linear(self, test_name, shape, bias): - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(512, 256, bias) - - def forward(self, x): - return self.linear(x) - - model = TestModule().cuda().half() - - input_shape = [] - for i, s in enumerate(shape): - input_shape.append( - TensorSpec.gen_int_var_min_max(s[0], s[1], "dynamic" + str(i)) - ) - input_spec = TensorSpec.create_spec_from_int_vars( - [input_shape], dtype_list=[torch.float16] - ) - - self.run_test_with_dynamic_shape( - model, input_spec, expected_ops={torch.ops.aten.linear} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_matmul_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_matmul_aten.py deleted file mode 100644 index 4819280d7..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_matmul_aten.py +++ /dev/null @@ -1,97 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.passes.lower_basic_pass_aten import aten_compose_bmm_3d, aten_compose_mm_2d -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase - -from parameterized import parameterized - - -class TestMatMulConverter(DispatchTestCase): - @parameterized.expand( - [ - [[2, 3], [3, 4], torch.ops.aten.mm.default], - # TODO check again in future since there is a diff about not decompose https://fburl.com/nysuuf7q - [ - [2, 3, 4], - [4, 6], - aten_compose_mm_2d, - ], - [[2, 3, 4], [2, 4, 6], aten_compose_bmm_3d], - [[2, 2, 2, 3, 4], [4, 6], torch.ops.aten.mm.default], - ] - ) - def test_simple(self, lhs_shape, rhs_shape, op): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - x = x.mul(1) - y = y.mul(1) - return torch.matmul(x, y) - - model = TestModule().cuda() - inputs = [ - torch.randn(*lhs_shape).half().cuda(), - torch.randn(*rhs_shape).half().cuda(), - ] - self.run_test(model, inputs, expected_ops={op}) - - def test_mm(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return torch.mm(x, y) - - model = TestModule().cuda().half() - inputs = [ - torch.randn(2, 3).half().cuda(), - torch.randn(3, 4).half().cuda(), - ] - self.run_test(model, inputs, expected_ops={torch.ops.aten.mm.default}) - - @parameterized.expand( - [ - # Only M can be dynamic: https://github.com/fairinternal/AITemplate/blob/main/tests/unittest/ops/test_gemm.py - [[[2, 3], [3, 3], [6, 6]], torch.ops.aten.mm.default], - [[[2, 3], [2, 3], [3, 3], [6, 6]], aten_compose_mm_2d], - # Cannot test with size=1, we will one specialize - # [[[1, 3], [2, 3], [6, 8], [3, 3], [6, 6]], torch.ops.aten.mm.default], - # FIXME: batch_size cannot be dynamic because the permutation of shape change the names: P544607056 - # b, m, k, n - [[[2, 2], [6, 8], [3, 3], [6, 6]], aten_compose_bmm_3d, True], - ] - ) - def test_dynamic(self, shape, op, bmm=False): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return torch.matmul(x, y) - - model = TestModule().cuda() - - input0_shape = [] - for i, s in enumerate(shape): - if i == len(shape) - 1: - break - input0_shape.append( - TensorSpec.gen_int_var_min_max(s[0], s[1], "dynamic" + str(i)) - ) - input1_shape = [input0_shape[-1]] + [ - TensorSpec.gen_int_var_min_max(shape[-1][0], shape[-1][1], "dynamic_last") - ] - if bmm: - input1_shape = [input0_shape[0]] + input1_shape - input_spec = TensorSpec.create_spec_from_int_vars( - [input0_shape, input1_shape], dtype_list=[torch.float16] * 2 - ) - self.run_test_with_dynamic_shape(model, input_spec, expected_ops={op}) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_max_pool2d_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_max_pool2d_aten.py deleted file mode 100644 index a391617b1..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_max_pool2d_aten.py +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestMaxPool2dConverter(DispatchTestCase): - @parameterized.expand( - [ - (1, 1, 0), - ((2, 2), 2, 1), - ((4, 4), (4, 4), 0), - ] - ) - def test_maxpool2d(self, kernel_size, stride, padding): - class TestModule(torch.nn.Module): - def __init__(self, kernel_size, stride, padding): - super().__init__() - self.pool = torch.nn.MaxPool2d(kernel_size, stride, padding) - - def forward(self, x): - return self.pool(x) - - model = TestModule(kernel_size, stride, padding).half().cuda() - inputs = [torch.randn(1, 4, 256, 256).cuda().half()] - self.run_test( - model, - inputs, - expected_ops={torch.ops.aten.max_pool2d}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) - - @parameterized.expand( - [ - (1, 1, 0), - ((2, 2), 2, 1), - ((4, 4), (4, 4), 0), - ] - ) - def test_dynamic_maxpool2d(self, kernel_size, stride, padding): - class TestModule(torch.nn.Module): - def __init__(self, kernel_size, stride, padding): - super().__init__() - self.pool = torch.nn.MaxPool2d(kernel_size, stride, padding) - - def forward(self, x): - return self.pool(x) - - model = TestModule(kernel_size, stride, padding).half().cuda() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [1, 4, 224, 224], - ], - inputs_max=[ - [10, 4, 256, 256], - ], - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.max_pool2d}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=[0, 3, 1, 2], - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_model_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_model_aten.py deleted file mode 100644 index 383d96e4e..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_model_aten.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -import torchvision -from fx2ait.passes.lower_basic_pass_aten import nchw2nhwc_pass, replace_inplace_ops -from fx2ait.tools.common_aten2ait import DispatchTestCase - - -class TestModelConverter(DispatchTestCase): - def test_resnet50(self): - torch.manual_seed(0) - - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.mod = torchvision.models.resnet18() - - def forward(self, x): - return self.mod(x) - - model = TestModule().cuda().half() - inputs = [torch.randn(32, 3, 224, 224).half().cuda()] - customized_passes = [ - replace_inplace_ops, - nchw2nhwc_pass, - ] - - self.run_test( - model, - inputs, - expected_ops={}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=None, - customized_passes=customized_passes, - ) - - def test_densenet(self): - torch.manual_seed(0) - - class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.mod = torchvision.models.densenet121(pretrained=True) - - def forward(self, x): - return self.mod(x) - - inputs = [torch.randn(1, 3, 224, 224).cuda().half()] - model = TestModule().cuda().half() - self.run_test( - model, - inputs, - atol=0.18, - expected_ops={}, - permute_inputs=[0, 2, 3, 1], - permute_outputs=None, - customized_passes=[ - replace_inplace_ops, - nchw2nhwc_pass, - ], - ) - - # def test_hf_albert_base(self): - # # config = AutoConfig.from_pretrained("albert-base-v2") - # # config = AutoConfig.from_pretrained("gpt2") - # # config = BertConfig() - # config = AutoConfig.from_pretrained("allenai/longformer-base-4096") - # max_length = 128 - # batch_size = 32 - # device = "cuda" - - # class TestModule(torch.nn.Module): - # def __init__(self): - # super().__init__() - # self.mod = AutoModelForMaskedLM.from_config(config).to(device) - - # def forward(self, x): - # return self.mod(x).logits - - # model = TestModule().cuda().half() - # input_ids = torch.randint(0, config.vocab_size, (batch_size, max_length)).to( - # device - # ) - # inputs = [input_ids] - # self.run_test( - # model, - # inputs, - # expected_ops={}, - # # permute_inputs=[0, 2, 3, 1], - # # permute_outputs=None, - # customized_passes=[ - # replace_inplace_ops, - # nchw2nhwc_pass, - # ], - # ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_nan2num_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_nan2num_aten.py deleted file mode 100644 index 5b84e3f2d..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_nan2num_aten.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestCatConverter(DispatchTestCase): - @parameterized.expand( - [ - param("default"), - param("nan", nan=1.0), - param("posinf", posinf=1.0), - param("neginf", neginf=-1.0), - ] - ) - def test_nan_to_num(self, name, nan=None, posinf=None, neginf=None): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) - - model = TestModule().cuda().half() - inputs = [ - torch.tensor([float("nan"), float("inf"), -float("inf"), 3.14]) - .half() - .cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.nan_to_num.default}) - - @parameterized.expand( - [ - param("default"), - param("nan", nan=1.0), - param("posinf", posinf=1.0), - param("neginf", neginf=-1.0), - ] - ) - def test_dynamic_nan_to_num(self, name, nan=None, posinf=None, neginf=None): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - ], - inputs_max=[ - [3, 8, 10], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.nan_to_num.default}, - specify_num=float("nan"), - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.nan_to_num.default}, - specify_num=float("inf"), - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.nan_to_num.default}, - specify_num=-float("inf"), - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_permute_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_permute_aten.py deleted file mode 100644 index 25215b557..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_permute_aten.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestPermuteConverter(DispatchTestCase): - @parameterized.expand( - [ - param((128, 512), (1, 0)), - param((80, 300, 2), (0, 2, 1)), - param((80, 300, 2), (1, 0, 2)), - param((80, 300, 2), (2, 1, 0)), - param((5, 113, 15, 31), (0, 2, 1, 3)), - param((2, 3, 4, 5), (3, 2, 1, 0)), - param((3, 5, 128, 514), (2, 3, 0, 1)), - param((32, 12, 4096, 64), (0, 2, 1, 3)), - param((3, 1, 113, 15, 64), (2, 0, 3, 1, 4)), - ] - ) - def test_permute(self, shape, dims): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.permute(x, dims) - - model = TestModule().cuda() - inputs = [ - torch.randn(shape).half().cuda(), - ] - self.run_test(model, inputs, expected_ops={torch.ops.aten.permute.default}) - - @parameterized.expand( - [ - param((128, 500), (256, 512), (1, 0)), - param((80, 300, 2), (98, 512, 20), (0, 2, 1)), - param((80, 300, 2), (98, 512, 20), (1, 0, 2)), - param((80, 300, 2), (98, 512, 20), (2, 1, 0)), - param((3, 5, 128, 512), (6, 10, 256, 520), (0, 2, 1, 3)), - param((3, 5, 128, 512), (6, 10, 256, 520), (3, 2, 1, 0)), - param((3, 5, 128, 512), (6, 10, 256, 520), (2, 3, 0, 1)), - param((3, 1, 113, 15, 64), (6, 10, 128, 16, 128), (2, 0, 3, 1, 4)), - ] - ) - def test_permute_dynamic_shape(self, input_min, input_max, dims): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.permute(x, dims) - - model = TestModule().cuda() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - input_min, - ], - inputs_max=[ - input_max, - ], - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.permute.default} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_pow_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_pow_aten.py deleted file mode 100644 index 695fe94db..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_pow_aten.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestPowConverter(DispatchTestCase): - @parameterized.expand( - [ - param("exp", size=[10], exp=5), - param("3d_exp", size=[2, 5, 32], exp=5), - param("4d_float_exp", size=[2, 5, 32, 128], exp=2.2), - ] - ) - def test_pow(self, name, size, exp): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.pow(x, exponent=exp) - - model = TestModule().cuda().half() - inputs = [ - torch.randn([2, 5, 32, 128]).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.pow.Tensor_Scalar}) - - @parameterized.expand( - [ - param("exp", inputs_min=[10], inputs_max=[15], exp=5), - param("3d_exp", inputs_min=[2, 5, 32], inputs_max=[3, 7, 64], exp=5), - param( - "4d_float_exp", - inputs_min=[2, 5, 32, 128], - inputs_max=[20, 7, 35, 140], - exp=2.2, - ), - ] - ) - def test_dynamic_pow(self, name, inputs_min, inputs_max, exp): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.pow(x, exponent=exp) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - ], - inputs_max=[ - [3, 8, 10], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.pow.Tensor_Scalar} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_reduce_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_reduce_aten.py deleted file mode 100644 index 2e627480c..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_reduce_aten.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -class TestSumConverter(DispatchTestCase): - @parameterized.expand( - [ - ["default", (1), False], - ["keepdim", (1), True], - ["negative_dim", (-1), False], - ["keepdim_2d", (0, 1), True], - ["nokeepdim_2d", (0, 1), False], - ["negative_2d", (-1, -2), False], - ["keepdim_3d", (0, 1, 2), True], - ] - ) - def test_sum(self, test_name, dim, keepdim): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.sum(x, dim=dim, keepdim=keepdim) - - model = TestModule().cuda() - inputs = [torch.randn(1, 2, 3).half().cuda()] - self.run_test(model, inputs, expected_ops={torch.ops.aten.sum.dim_IntList}) - - @parameterized.expand( - [ - ["default", (1), False], - ["keepdim", (1), True], - ["negative_dim", (-1), False], - ["keepdim_2d", (0, 1), True], - ["nokeepdim_2d", (0, 1), False], - ["negative_2d", (-1, -2), False], - ["keepdim_3d", (0, 1, 2), True], - ] - ) - def test_dynamic_sum(self, test_name, dim, keepdim): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.sum(x, dim=dim, keepdim=keepdim) - - model = TestModule().cuda() - # The last dim has to be static to pre-compute vector_length: - # https://fburl.com/code/1x07doen - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - ], - inputs_max=[ - [20, 6, 4], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.sum.dim_IntList} - ) - - -class TestMeanConverter(DispatchTestCase): - @parameterized.expand( - [ - ["default", (1), False], - ["keepdim", (1), True], - ["negative_dim", (-1), False], - ["keepdim_2d", (0, 1), True], - ["nokeepdim_2d", (0, 1), False], - ["negative_2d", (-1, -2), False], - ["keepdim_3d", (0, 1, 2), True], - ] - ) - def test_mean(self, test_name, dim, keepdim): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.mean(x, dim=dim, keepdim=keepdim) - - model = TestModule().cuda() - inputs = [torch.randn(1, 2, 3).half().cuda()] - self.run_test(model, inputs, expected_ops={torch.ops.aten.mean.dim}) - - @parameterized.expand( - [ - ["default", (1), False], - ["keepdim", (1), True], - ["negative_dim", (-1), False], - ["keepdim_2d", (0, 1), True], - ["nokeepdim_2d", (0, 1), False], - ["negative_2d", (-1, -2), False], - ["keepdim_3d", (0, 1, 2), True], - ] - ) - def test_dynamic_mean(self, test_name, dim, keepdim): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.mean(x, dim=dim, keepdim=keepdim) - - model = TestModule().cuda() - - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - ], - inputs_max=[ - [20, 6, 8], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.mean.dim} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_relu_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_relu_aten.py deleted file mode 100644 index b7f1ce91b..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_relu_aten.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.op = torch.relu - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x) - - -class TestATenReluConverter(DispatchTestCase): - @parameterized.expand( - [ - param("small", size=(2, 3)), - param("large", size=(1024, 4096, 8)), - ] - ) - def test_relu(self, name, size): - model = TestModule().cuda().half() - inputs = (torch.randn(size).half().cuda(),) - - self.run_test(model, inputs, expected_ops={torch.ops.aten.relu.default}) - - def test_relu_with_dynamic_shape(self): - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [1, 3, 4], - ], - inputs_max=[ - [32, 1024, 2048], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.relu.default}, - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_reshape_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_reshape_aten.py deleted file mode 100644 index 5a1febf44..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_reshape_aten.py +++ /dev/null @@ -1,124 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import unittest - -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase - - -class TestATenReshapeConverter(DispatchTestCase): - def test_reshape(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.reshape(x, (2, 12)) - - size = (2, 3, 4) - model = TestModule().cuda().half() - inputs = (torch.randn(size).half().cuda(),) - - self.run_test(model, inputs, expected_ops={torch.ops.aten.view.default}) - - def test_reshape_size(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - dim1_y = y.shape[0] - return torch.reshape(x, (dim1_y, -1, 128)) - - model = TestModule().cuda().half() - inputs = ( - torch.randn(2, 10, 128).half().cuda(), - torch.randn(2, 10, 128).half().cuda(), - ) - - self.run_test(model, inputs, expected_ops={torch.ops.aten.view.default}) - - def test_reshape_with_dynamic_shape(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.reshape(x, (x.size(0), x.size(1) * x.size(2))) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - ], - inputs_max=[ - [10, 30, 4], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.view.default}, - ) - - def test_reshape_neg_with_dynamic_shape(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.reshape(x, (x.size(0), -1)) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 3, 4], - ], - inputs_max=[ - [10, 30, 4], - ], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.view.default}, - ) - - # TODO: trigger assertion in AIT: AssertionError: When there is no unknown index, we expect dim products to be equal, got current shape numel=2560 != new shape prod=256 - @unittest.skip - def test_reshape_size_with_dynamic_shape(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - dim1_y = y.shape[0] - return torch.reshape(x, (dim1_y, -1, 128)) - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 10, 128], - [2, 10, 128], - ], - inputs_max=[ - [20, 10, 128], - [20, 10, 128], - ], - dtype_list=[ - torch.float16, - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, - inputs_spec, - expected_ops={torch.ops.aten.view.default}, - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_size_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_size_aten.py deleted file mode 100644 index bd076b77a..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_size_aten.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch - -# from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase - - -class TestATenSizeConverter(DispatchTestCase): - def test_size(self): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - t = x.size() - return y.reshape(t) - - xsize = (2, 3, 4) - ysize = (2, 12) - model = TestModule().cuda().half() - inputs = (torch.randn(xsize).half().cuda(), torch.randn(ysize).half().cuda()) - - self.run_test(model, inputs, expected_ops={torch.ops.aten.sym_size}) - - ## AIT not support now - # def test_size_dim(self): - # class TestModule(torch.nn.Module): - # def forward(self, x: torch.Tensor) -> torch.Tensor: - # return x.size(1) - - # size = (2, 3, 4) - # model = TestModule().cuda().half() - # inputs = (torch.randn(size).half().cuda(),) - - # self.run_test(model, inputs, expected_ops={torch.ops.aten.sym_size}) - - # def test_size_dim_with_dynamic_shape(self): - # class TestModule(torch.nn.Module): - # def forward(self, x: torch.Tensor) -> torch.Tensor: - # return x.size(1) - - # model = TestModule().cuda().half() - # inputs_spec = TensorSpec.create_spec_from_shapes( - # inputs_min=[ - # [2, 3, 4], - # ], - # inputs_max=[ - # [10, 30, 4], - # ], - # dtype_list=[ - # torch.float16, - # ], - # ) - - # self.run_test_with_dynamic_shape( - # model, - # inputs_spec, - # expected_ops={torch.ops.aten.sym_size}, - # ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_slice_tensor_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_slice_tensor_aten.py deleted file mode 100644 index 724d7d2c1..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_slice_tensor_aten.py +++ /dev/null @@ -1,222 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.passes.lower_basic_pass_aten import ( - aten_compose_getitem_slice, - compose_getitem_slice, -) -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized -from torch import nn - - -class TestSliceTensor(DispatchTestCase): - @parameterized.expand( - [ - ( - "integer_slice", - 1, - { - torch.ops.aten.select.int, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_batch_dim", - slice(None, None, None), - { - torch.ops.aten.slice.Tensor, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_basic", - (slice(None, None, None), slice(0, 3, 1)), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_full", - (slice(None, None, None), slice(0, 10, 1)), - { - torch.ops.aten.slice.Tensor, - torch.ops.aten.add.Tensor, - }, - None, - ), - ## Trace problem in support of ellipsis - # ( - # "ellipsis", # It seems there is some problem in tracing ellipsis: P539875442 - # (slice(None, None, None), ..., slice(0, 3, 1)), - # { - # torch.ops.aten.add.Tensor, - # }, - # ), - ( - "slice_all_none", - ( - slice(None, None, None), - slice(None, None, None), - ), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_start_none", - ( - slice(None, None, None), - slice(None, 2, 1), - ), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_end_none", - (slice(None, None, None), slice(1, None, 1)), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_step_none", - ( - slice(None, None, None), - slice(0, 3, None), - ), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_neg_idx", - (slice(None, None, None), -1), - { - torch.ops.aten.slice.Tensor, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_neg_slice", - (slice(None, None, None), slice(-8, -2, 1)), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "multi_dim", - (slice(None, None, None), 0, 1), - { - torch.ops.aten.slice.Tensor, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_multi_dim", - (slice(None, None, None), slice(0, 3, 1), slice(1, -1, 1)), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "none", - (slice(None, None, None), None, slice(1, -1, 1), 1), - { - torch.ops.aten.slice.Tensor, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "with_squeeze", - (slice(None, None, None), 1, slice(1, -1, 1), None), - { - torch.ops.aten.slice.Tensor, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_zero_slice", - (slice(None, None, None), slice(None, None, None), slice(0, 0, None)), - { - aten_compose_getitem_slice, - torch.ops.aten.add.Tensor, - }, - None, - ), - ( - "slice_basic_compose", - (slice(None, None, None), slice(None, None, None), slice(0, 3, 1)), - { - torch.ops.aten.add.Tensor, - aten_compose_getitem_slice, - }, - [ - compose_getitem_slice, - ], - ), - ( - "slice_zero_slice_compose", - (slice(None, None, None), slice(None, None, None), slice(0, 0, None)), - { - torch.ops.aten.add.Tensor, - aten_compose_getitem_slice, - }, - [ - compose_getitem_slice, - ], - ), - ] - ) - def test_slice_tensor(self, name, idx, expected_ops, customized_passes): - class SliceTensor(nn.Module): - def __init__(self, idx): - super().__init__() - self.idx = idx - - def forward(self, x): - y = x + x - return y[self.idx] - - mod = SliceTensor(idx).half().cuda() - - inputs = [torch.randn(2, 10, 10, 10).half().cuda()] - self.run_test( - mod, - inputs, - expected_ops=expected_ops, - customized_passes=customized_passes, - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_split_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_split_aten.py deleted file mode 100644 index ad4897dec..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_split_aten.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestSplitConverter(DispatchTestCase): - @parameterized.expand( - [ - param( - "dim1", - dim=1, - split_size=3, - expected_ops={torch.ops.aten.split.Tensor}, - ), - param( - "dim0", - dim=0, - split_size=3, - expected_ops={torch.ops.aten.split.Tensor}, - ), - ] - ) - def test_split(self, name, dim, split_size, expected_ops): - class TestModule(torch.nn.Module): - def forward(self, y: torch.Tensor) -> torch.Tensor: - res = torch.split(y, split_size, dim) - return res[0] - - model = TestModule().cuda().half() - inputs = [ - torch.randn(20, 10, 8).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops=expected_ops) - - def test_split_dynamic(self): - class TestModule(torch.nn.Module): - def forward(self, y: torch.Tensor) -> torch.Tensor: - res = torch.split(y, 4, 1) - return res[0] - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[[20, 10, 8]], - inputs_max=[[50, 10, 8]], - dtype_list=[ - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.split.Tensor} - ) - - # TODO low priority. May need to support it in future. - # def test_split_imm(self): - # class TestModule(torch.nn.Module): - # def forward(self, y: torch.Tensor) -> torch.Tensor: - # dim1 = y.size(1) - # split_size = dim1 // 2 - # return torch.split(y, split_size, 1) - - # model = TestModule().cuda().half() - # inputs = [ - # torch.randn(2, 10, 20).half().cuda(), - # ] - # self.run_test(model, inputs, expected_ops={torch.ops.aten.split.Tensor}) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_squeeze_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_squeeze_aten.py deleted file mode 100644 index c205f3cf6..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_squeeze_aten.py +++ /dev/null @@ -1,165 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import param, parameterized - - -class TestSqueezeConverter(DispatchTestCase): - @parameterized.expand( - [ - param( - "default", - dim=None, - shape=[2, 1, 1, 3], - expected_ops={torch.ops.aten.squeeze.default}, - ), - param( - "1", - dim=1, - shape=[2, 1, 1, 3], - expected_ops={torch.ops.aten.squeeze.dim}, - ), - param( - "-1", - dim=-1, - shape=[2, 1, 3, 1], - expected_ops={torch.ops.aten.squeeze.dim}, - ), - ] - ) - def test_squeeze(self, name, dim, shape, expected_ops): - class TestModule(torch.nn.Module): - def forward(self, y: torch.Tensor) -> torch.Tensor: - squeeze = ( - torch.squeeze(y, dim=dim) if dim is not None else torch.squeeze(y) - ) - return squeeze - - model = TestModule().cuda().half() - inputs = [ - torch.randn(shape).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops=expected_ops) - - @parameterized.expand( - [ - param( - "default", - dim=None, - shape1=[[2, 1, 3, 1]], - shape2=[[4, 1, 10, 1]], - expected_ops={torch.ops.aten.squeeze.default}, - ), - param( - "1", - dim=1, - shape1=[[2, 1, 3, 1]], - shape2=[[4, 1, 10, 1]], - expected_ops={torch.ops.aten.squeeze.dim}, - ), - param( - "-1", - dim=-1, - shape1=[[2, 1, 3, 1]], - shape2=[[4, 1, 10, 1]], - expected_ops={torch.ops.aten.squeeze.dim}, - ), - ] - ) - def test_dynamic_squeeze(self, name, dim, shape1, shape2, expected_ops): - class TestModule(torch.nn.Module): - def forward(self, y: torch.Tensor) -> torch.Tensor: - squeeze = ( - torch.squeeze(y, dim=dim) if dim is not None else torch.squeeze(y) - ) - return squeeze - - model = TestModule().cuda().half() - - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=shape1, - inputs_max=shape2, - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape(model, inputs_spec, expected_ops=expected_ops) - - -class TestUnSqueezeConverter(DispatchTestCase): - @parameterized.expand( - [ - param("1", dim=1, shape=[2, 1, 1, 3]), - param("-1", dim=-1, shape=[2, 1, 3, 1]), - ] - ) - def test_unsqueeze(self, name, dim, shape): - class TestModule(torch.nn.Module): - def forward(self, y: torch.Tensor) -> torch.Tensor: - unsqueeze = ( - torch.unsqueeze(y, dim=dim) - if dim is not None - else torch.unsqueeze(y) - ) - return unsqueeze - - model = TestModule().cuda().half() - inputs = [ - torch.randn(shape).half().cuda(), - ] - - self.run_test(model, inputs, expected_ops={torch.ops.aten.unsqueeze.default}) - - @parameterized.expand( - [ - param( - "1", - dim=1, - shape1=[[2, 1, 3, 1]], - shape2=[[4, 1, 10, 1]], - ), - param( - "-1", - dim=-1, - shape1=[[2, 1, 3, 1]], - shape2=[[4, 1, 10, 1]], - ), - ] - ) - def test_dynamic_squeeze(self, name, dim, shape1, shape2): - class TestModule(torch.nn.Module): - def forward(self, y: torch.Tensor) -> torch.Tensor: - unsqueeze = ( - torch.unsqueeze(y, dim=dim) - if dim is not None - else torch.unsqueeze(y) - ) - return unsqueeze - - model = TestModule().cuda().half() - - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=shape1, - inputs_max=shape2, - dtype_list=[ - torch.float16, - ], - ) - self.run_test_with_dynamic_shape( - model, inputs_spec, expected_ops={torch.ops.aten.unsqueeze.default} - ) diff --git a/fx2ait/fx2ait/test/converters_aten/test_ait_unary_ops_aten.py b/fx2ait/fx2ait/test/converters_aten/test_ait_unary_ops_aten.py deleted file mode 100644 index 884e85312..000000000 --- a/fx2ait/fx2ait/test/converters_aten/test_ait_unary_ops_aten.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -from typing import Callable - -import torch -from fx2ait.tensor_spec import TensorSpec -from fx2ait.tools.common_aten2ait import DispatchTestCase -from parameterized import parameterized - - -unary_ops = [ - (torch.abs, torch.ops.aten.abs.default), - (torch.log, torch.ops.aten.log.default), - (torch.sigmoid, torch.ops.aten.sigmoid.default), - (torch.sign, torch.ops.aten.sign.default), - (torch.tanh, torch.ops.aten.tanh.default), - (torch.sin, torch.ops.aten.sin.default), - (torch.cos, torch.ops.aten.cos.default), - (torch.sqrt, torch.ops.aten.sqrt.default), - ( - torch.clone, - torch.ops.aten.mul.Tensor, - ), # clone op can not be the output directly so expected is the op after it(aten.mul) -] - - -class TestUnaryOpsConverter(DispatchTestCase): - @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in unary_ops]) - def test_unary_ops(self, name, orig_op: Callable, expected_op): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return orig_op(x) * 2 - - model = TestModule().cuda().half() - inputs = [ - torch.randn(1, 2, 3).half().cuda(), - ] - _ = model(*inputs) - self.run_test(model, inputs, expected_ops={expected_op}) - - @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in unary_ops]) - def test_dynamic_unary_ops(self, name, orig_op: Callable, expected_op): - class TestModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return orig_op(x) * 2 - - model = TestModule().cuda().half() - inputs_spec = TensorSpec.create_spec_from_shapes( - inputs_min=[ - [2, 8, 10], - ], - inputs_max=[ - [20, 12, 32], - ], - dtype_list=[ - torch.float16, - torch.float16, - torch.float16, - ], - ) - - self.run_test_with_dynamic_shape(model, inputs_spec, expected_ops={expected_op})