empty_memory_format evaluator

Author: apbose

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -11,6 +11,7 @@
     dynamo_tensorrt_converter,
 )
 from torch_tensorrt.fx.types import TRTTensor
+from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -118,3 +119,50 @@ def aten_ops_randperm(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     return np.random.permutation(args[0])
+
+
+def empty_validator(empty_node: Node) -> bool:
+    layout = empty_node.kwargs.get("layout", None)
+    pin_memory = empty_node.kwargs.get("pin_memory", None)
+    memory_format = empty_node.kwargs.get("memory_format", None)
+    if layout is not None:
+        _LOGGER.debug(f"Currently we don't support specifying layout, got {layout}.")
+        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
+    memory_format = kwargs.get("memory_format")
+    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]))
+
+    empty_tensor = torch.Tensor(empty_np_tensor)
+    # device
+    if kwargs.get("device") is not None:
+        empty_tensor = empty_tensor.to(device=kwargs.get("device"))
+
+    # memory_format. default is torch.contiguous_format
+    if memory_format == torch.channels_last:
+        # shape of args[0] must be 4
+        empty_tensor = empty_tensor.to(memory_format=torch.channels_last)
+    elif memory_format == torch.channels_last_3d:
+        # shape of args[0] must be 5
+        empty_tensor = empty_tensor.to(memory_format=torch.channels_last_3d)
+
+    return empty_tensor

tests/py/dynamo/conversion/test_empty_aten.py

@@ -0,0 +1,126 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_tensorrt
+from harness import DispatchTestCase
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+empty_ops = [
+    (
+        "empty_one_dimension",
+        [1],
+        None,
+        None,
+        None,
+    ),
+    (
+        "empty_two_dimension",
+        [1, 2],
+        None,
+        None,
+        None,
+    ),
+    (
+        "empty_three_dimension",
+        [2, 3, 4],
+        None,
+        None,
+        None,
+    ),
+    (
+        "empty_one_dimension_dtype",
+        [1],
+        torch.float32,
+        None,
+        None,
+    ),
+    (
+        "empty_two_dimension_dtype",
+        [2, 3],
+        torch.float32,
+        None,
+        None,
+    ),
+    (
+        "empty_one_dimension_dtype_device",
+        [1],
+        torch.float32,
+        "cuda",
+        None,
+    ),
+    (
+        "empty_two_dimension_dtype_device",
+        [2, 3],
+        torch.float32,
+        "cuda",
+        None,
+    ),
+    (
+        "empty_four_dimension_memformat",
+        [1, 2, 2, 1],
+        torch.float32,
+        "cuda",
+        torch.channels_last,
+    ),
+    (
+        "empty_five_dimension_memformat",
+        [1, 2, 2, 2, 1],
+        torch.float32,
+        "cuda",
+        torch.channels_last_3d,
+    ),
+]
+
+
+class TestRandConverter(DispatchTestCase):
+    @parameterized.expand(
+        [(empty_op[0], empty_op[1], empty_op[2], empty_op[3]) for empty_op in empty_ops]
+    )
+    def test_empty(self, name, shape_or_input, data_type, device):
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                shape_or_input[0] = x.shape[0]
+                return torch.empty(shape_or_input)
+
+        empty_model = TestModule()
+
+        inputs = [torch.randint(1, 3, shape_or_input, dtype=torch.int32)]
+        comparator_shape_dtype_device = (
+            lambda x, y, check_dtype, check_device: x.shape == y.shape
+            and (x.stride() == y.stride())
+            and (x.dtype == y.dtype if check_dtype else True)
+            and (x.get_device() == y.get_device() if check_device else True)
+        )
+        expected_ops = []
+        if "device" in name:
+            self.run_test_compare_tensor_attributes_only(
+                empty_model,
+                inputs,
+                expected_ops,
+                [(comparator_shape_dtype_device, [True, True])],
+                use_dynamo_tracer=True,
+            )
+        elif "dtype" in name:
+            self.run_test_compare_tensor_attributes_only(
+                empty_model,
+                inputs,
+                expected_ops,
+                [(comparator_shape_dtype_device, [True, False])],
+                use_dynamo_tracer=True,
+            )
+        else:
+            self.run_test_compare_tensor_attributes_only(
+                empty_model,
+                inputs,
+                expected_ops,
+                [(comparator_shape_dtype_device, [False, False])],
+                use_dynamo_tracer=True,
+            )
+
+
+if __name__ == "__main__":
+    run_tests()