stabilizes perf test

Author: pranavm-nvidia

tripy/tests/frontend/test_tensor.py

@@ -22,12 +22,12 @@
 import numpy as np
 import pytest
 import torch
+from tests.conftest import DATA_TYPE_TEST_CASES
+from tests.helper import NUMPY_TO_TRIPY
 
 import tripy as tp
-from tests.conftest import DATA_TYPE_TEST_CASES
-from tests.helper import NUMPY_TYPES, np_to_tripy_dtype
-from tripy.utils.stack_info import SourceInfo
 from tripy.common.utils import get_element_type
+from tripy.utils.stack_info import SourceInfo
 
 
 class TestTensor:
@@ -52,12 +52,12 @@ def test_tensor_device(self, kind):
         assert isinstance(a.trace_tensor.producer, tp.frontend.trace.ops.Storage)
         assert a.trace_tensor.producer.device.kind == kind
 
-    @pytest.mark.parametrize("dtype", NUMPY_TYPES)
+    @pytest.mark.parametrize("dtype", NUMPY_TO_TRIPY.keys())
     def test_dtype_from_numpy(self, dtype):
 
         np_array = np.array([1, 2, 3], dtype=dtype)
         tensor = tp.Tensor(np_array)
-        tp_dtype = np_to_tripy_dtype(dtype)
+        tp_dtype = NUMPY_TO_TRIPY[dtype]
         assert tensor.dtype == tp_dtype
 
     def test_bool_tensor(self):

tripy/tests/helper.py

@@ -102,35 +102,22 @@ def check_mlir(mlir, expected):
 
 
 # Supported NumPy data types
-NUMPY_TYPES = [
-    np.int8,
+NUMPY_TO_TRIPY = {
+    bool: tp.bool,
+    np.int8: tp.int8,
+    np.int32: tp.int32,
+    np.int64: tp.int64,
+    np.float16: tp.float16,
+    np.float32: tp.float32,
     # np.int16,  # TODO(#247): Add support for int16
-    np.int32,
-    np.int64,
     # np.uint8,  # TODO(#247): Add support for uint8
     # np.uint16, # TODO(#190): Add support for unsupported MLIR-TensorRT types.
     # np.uint32, # TODO(#190): Add support for unsupported MLIR-TensorRT types.
     # np.uint64, # TODO(#190): Add support for unsupported MLIR-TensorRT types.
-    np.float16,
-    np.float32,
     # np.float64,  # TODO(#247): Add support for float64
-]
-
-
-def np_to_tripy_dtype(dtype):
-    return {
-        bool: tp.bool,
-        np.int8: tp.int8,
-        np.int32: tp.int32,
-        np.int64: tp.int64,
-        np.float16: tp.float16,
-        np.float32: tp.float32,
-    }[dtype]
-
+}
 
-def torch_type_supported(data: np.ndarray):
-    unsupported_dtypes = [np.int16, np.uint16, np.uint32, np.uint64]
-    return data.dtype not in unsupported_dtypes
+TRIPY_TO_NUMPY = {v: k for k, v in NUMPY_TO_TRIPY.items()}
 
 
 TORCH_DTYPES = {

tripy/tests/integration/test_cast.py

@@ -21,7 +21,7 @@
 
 import tripy as tp
 from tests.conftest import skip_if_older_than_sm89
-from tests.helper import np_to_tripy_dtype
+from tests.helper import NUMPY_TO_TRIPY
 
 
 class TestCast:
@@ -50,8 +50,8 @@ class TestCast:
         ],
     )
     def test_cast(self, input_dtype, target_dtype):
-        tp_input_dtype = np_to_tripy_dtype(input_dtype)
-        tp_target_dtype = np_to_tripy_dtype(target_dtype)
+        tp_input_dtype = NUMPY_TO_TRIPY[input_dtype]
+        tp_target_dtype = NUMPY_TO_TRIPY[target_dtype]
 
         # TODO(#222): Integer casts with negative numbers fail in many cases
         input_tensor = tp.Tensor([0, 1, 2], dtype=tp_input_dtype)
@@ -71,7 +71,7 @@ def test_cast_quantized_dtypes_into_bool(self, source_dtype):
 
     @pytest.mark.parametrize("target_dtype", [np.float32, np.int32, np.int64, np.int8])
     def test_cast_from_bool(self, target_dtype):
-        tp_target_dtype = np_to_tripy_dtype(target_dtype)
+        tp_target_dtype = NUMPY_TO_TRIPY[target_dtype]
 
         # in principle, it is not important what *specific* values we convert to,
         # so long as false is mapped to 0 and true to nonzero

tripy/tests/performance/cases/linear_block.py

@@ -12,8 +12,9 @@
 # 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 pytest
+import numpy as np
 import torch
+from tests.helper import TRIPY_TO_NUMPY
 from tests.performance.conftest import PerfParam, perf_fixture
 
 import tripy as tp
@@ -33,10 +34,8 @@ class LinearBlock(tp.Module):
         def __init__(self):
             self.layers = [tp.Linear(256, 256, bias=False, dtype=tripy_dtype) for _ in range(NUM_LAYERS)]
             for layer in self.layers:
-                # Adjust the weights to prevent FP16 overflows.
-                weight = torch.tile(
-                    torch.tensor([[-1, 1], [1, -1]], dtype=torch_dtype, device=torch.device("cuda")), (128, 128)
-                )
+                # Adjust the weights to prevent FP16 overflows:
+                weight = np.tile(np.array([[-1, 1], [1, -1]], dtype=TRIPY_TO_NUMPY[tripy_dtype]), (128, 128))
                 layer.weight = tp.Parameter(weight)
 
         def __call__(self, input):