empty_memory_format evaluator (#2745)

Author: apbose

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -18,6 +18,7 @@
 )
 from torch_tensorrt.dynamo.conversion.impl.elementwise import sub, trunc_div
 from torch_tensorrt.fx.types import TRTTensor
+from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -165,3 +166,43 @@ def aten_ops_randperm(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     return np.random.permutation(args[0])
+
+
+def empty_validator(empty_node: Node) -> bool:
+    device = empty_node.kwargs.get("device", None)
+    if device is not None:
+        _LOGGER.debug(f"Currently we don't support specifying device, got {device}.")
+        return False
+    layout = empty_node.kwargs.get("layout", None)
+    if layout is not None:
+        _LOGGER.debug(f"Currently we don't support specifying layout, got {layout}.")
+        return False
+    memory_format = empty_node.kwargs.get("memory_format", None)
+    if memory_format is not None:
+        _LOGGER.debug(
+            f"Currently we don't support specifying memory_format, got {memory_format}."
+        )
+        return False
+    return True
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.empty.memory_format, capability_validator=empty_validator
+)
+def aten_ops_empty(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    empty_np_tensor = None
+    if kwargs.get("dtype") is not None:
+        empty_np_tensor = np.empty(
+            tuple(args[0]),
+            dtype=unified_dtype_converter(kwargs.get("dtype"), Frameworks.NUMPY),
+        )
+    else:
+        # default returns np.float64. Verify the correctness of this
+        empty_np_tensor = np.empty(tuple(args[0]))
+    return empty_np_tensor

tests/py/dynamo/conversion/test_empty_aten.py

@@ -0,0 +1,90 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_tensorrt
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+empty_ops = [
+    (
+        "empty_one_dimension",
+        [1],
+        None,
+    ),
+    (
+        "empty_two_dimension",
+        [1, 2],
+        None,
+    ),
+    (
+        "empty_three_dimension",
+        [2, 3, 4],
+        None,
+    ),
+    (
+        "empty_one_dimension_dtype",
+        [1],
+        torch.float32,
+    ),
+    (
+        "empty_two_dimension_dtype",
+        [2, 3],
+        torch.float32,
+    ),
+    (
+        "empty_four_dimension_dtype",
+        [1, 2, 2, 1],
+        torch.float32,
+    ),
+    (
+        "empty_five_dimension_dtype",
+        [1, 2, 2, 2, 1],
+        torch.float32,
+    ),
+]
+
+
+class TestEmptyConverter(DispatchTestCase):
+    @parameterized.expand(
+        [(empty_op[0], empty_op[1], empty_op[2]) for empty_op in empty_ops]
+    )
+    def test_empty(self, name, shape_or_input, data_type):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                shape_or_input[0] = x.shape[0]
+                return torch.ops.aten.empty.memory_format(
+                    shape_or_input,
+                    dtype=data_type,
+                )
+
+        empty_model = TestModule()
+
+        inputs = [torch.randint(1, 3, shape_or_input, dtype=torch.int32)]
+        comparator_shape_dtype_device = (
+            lambda x, y, check_dtype: x.shape == y.shape
+            and (x.stride() == y.stride())
+            and (x.dtype == y.dtype if check_dtype else True)
+        )
+        expected_ops = []
+        if "dtype" in name:
+            self.run_test_compare_tensor_attributes_only(
+                empty_model,
+                inputs,
+                expected_ops,
+                [(comparator_shape_dtype_device, [True])],
+                use_dynamo_tracer=True,
+            )
+        else:
+            self.run_test_compare_tensor_attributes_only(
+                empty_model,
+                inputs,
+                expected_ops,
+                [(comparator_shape_dtype_device, [False])],
+                use_dynamo_tracer=True,
+            )
+
+
+if __name__ == "__main__":
+    run_tests()