Converter reorg and where operator

adding where aten op

aten::where correction and linting error changes

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -32,6 +32,7 @@
 from torch_tensorrt.fx.converters.impl.select import select
 from torch_tensorrt.fx.converters.impl.slice import slice_op
 from torch_tensorrt.fx.converters.impl.matmul import matrix_multiply
+from torch_tensorrt.fx.converters.impl.condition import where
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -563,6 +564,25 @@ def aten_ops_tanh(
     )
 
 
+@tensorrt_converter(torch.ops.aten.where.self)
+def aten_ops_where(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return where(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[1],
+        args[2],
+        args[0],
+    )
+
+
 @tensorrt_converter(torch.ops.aten.cat.default)
 def aten_ops_cat(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -409,6 +409,37 @@ def broadcast(
     return a, b
 
 
+def broadcastable(
+    a: TRTTensor,
+    b: TRTTensor,
+) -> bool:
+    "Check if two tensors are broadcastable according to torch rules"
+    a_shape = tuple(a.shape)
+    b_shape = tuple(b.shape)
+    # check from the trailing
+    diff = len(a_shape) - len(b_shape)
+    if diff == 0:
+        return True
+    if diff > 0:
+        max = len(a_shape)
+        min = len(b_shape)
+        greater_tensor = a_shape
+        lesser_tensor = b_shape
+    elif diff < 0:
+        max = len(b_shape)
+        min = len(a_shape)
+        greater_tensor = b_shape
+        lesser_tensor = a_shape
+    j = min - 1
+    for i in range(max - 1, diff - 1, -1):
+        if not (
+            greater_tensor[i] != lesser_tensor[j]
+            and (greater_tensor[i] == 1 or lesser_tensor[i] == 1)
+        ):
+            return False
+    return True
+
+
 def squeeze_left(const: Union[torch.Tensor, np.ndarray]):
     """
     Squeeze the size-1 dimensions on the left side of the shape tuple.

py/torch_tensorrt/fx/converters/impl/condition/__init__.py

@@ -0,0 +1 @@
+from .ops import *

py/torch_tensorrt/fx/converters/impl/condition/ops.py

@@ -0,0 +1,111 @@
+import operator
+import warnings
+from typing import Optional, cast
+
+import numpy as np
+
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, Shape
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    broadcast,
+    broadcastable,
+    get_trt_tensor,
+    set_layer_name,
+)
+from torch_tensorrt.fx.converters.impl.slice import expand
+
+
+def where(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    other: TRTTensor,
+    condition: TRTTensor,
+) -> TRTTensor:
+    input_dim = len(tuple(input.shape))
+    other_dim = len(tuple(other.shape))
+    condition_dim = len(tuple(condition.shape))
+
+    if type(input) != TRTTensor:
+        assert type(input) is torch.Tensor, f"value {input} is not torch.Tensor!"
+
+    if type(other) != TRTTensor:
+        assert type(other) is torch.Tensor, f"value {other} is not torch.Tensor!"
+
+    if not (broadcastable(input, other)):
+        assert f"The two torch tensors should be broadcastable"
+
+    # get output shape
+    # purpose of this is to bring input and other rank same as
+    # output_shape to input it to the add_expand operation
+    # condition will have dimension of either input or other
+    input, other = broadcast(network, input, other, f"{name}_x", f"{name}_y")
+    if len(tuple(condition.shape)) != len(tuple(input.shape)):
+        condition, input = broadcast(
+            network, condition, input, f"{name}_condition", f"{name}_x"
+        )
+
+    x_shape = list(input.shape)
+    y_shape = list(other.shape)
+    condition_shape = list(condition.shape)
+    output_shape = list(torch.broadcast_shapes(condition_shape, x_shape, y_shape))
+
+    # expand shape
+    if type(condition) != TRTTensor:
+        assert condition.dtype == torch.bool, "condition dtype is not bool"
+        if condition_shape != output_shape:
+            condition.expand(output_shape)
+        condition = condition.to(torch.int32)
+        condition_const = get_trt_tensor(network, condition, f"{name}_condition")
+        condition_layer = network.add_identity(condition_const)
+        condition_layer.set_output_type(0, trt.bool)
+        set_layer_name(condition_layer, target, f"{name}_condition")
+        condition_val = condition_layer.get_output(0)
+    else:
+        assert condition.dtype == trt.bool, "mask dtype is not bool!"
+        if condition_shape != condition_dim:
+            condition_val = expand(
+                network, target, source_ir, f"{name}_expand", condition, output_shape
+            )
+        else:
+            condition_val = condition
+
+    if type(input) != TRTTensor:
+        if x_shape != input_dim:
+            # special case where 1 element in input
+            if len(input.shape) == 0:
+                input = input.unsqueeze(0)
+            input = input.expand(output_shape)
+        x_val = get_trt_tensor(network, input, f"{name}_x")
+    else:
+        x_val = input
+        if x_shape != output_shape:
+            x_val = expand(
+                network, target, source_ir, f"{name}_x_expand", input, output_shape
+            )
+
+    if type(other) != TRTTensor:
+        if y_shape != output_shape:
+            # special case where 1 element in other
+            if len(other.shape) == 0:
+                other = other.unsqueeze(0)
+            other = other.expand(output_shape)
+        y_val = get_trt_tensor(network, other, f"{name}_y")
+    else:
+        y_val = other
+        if y_shape != other_dim:
+            y_val = expand(
+                network, target, source_ir, f"{name}_y_expand", y_val, output_shape
+            )
+
+    select_layer = network.add_select(condition_val, x_val, y_val)
+
+    set_layer_name(select_layer, target, f"{name}_select")
+
+    return select_layer.get_output(0)

py/torch_tensorrt/fx/converters/impl/slice/ops.py

@@ -10,12 +10,15 @@
 import torch
 from torch.fx.node import Target
 
-from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, Shape
+
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     set_layer_name,
     get_positive_dim,
     has_dynamic_shape,
+    broadcast,
+    get_trt_tensor,
 )
 from torch_tensorrt.fx.converters.impl.shape import get_shape_with_dynamic_shape
 from torch_tensorrt.fx.converters.impl.slice.base import slice
@@ -62,3 +65,40 @@ def slice_op(
     output_shape[dim] = math.ceil((stop_int - start_int) / step_int)
 
     return slice(network, target, source_ir, name, input, start, output_shape, stride)
+
+
+def expand(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    sizes: Shape,
+) -> TRTTensor:
+    shape = list(sizes)
+
+    input_val = get_trt_tensor(network, input, f"{name}_input")
+
+    if network.has_implicit_batch_dimension:
+        shape = shape[1:]
+
+    ranks = len(input_val.shape)
+    # TRT does not support different dimension size
+    # though this condition is not seen in the case of bmm
+    # where input_t and shape dimensions are not equal
+    assert len(shape) >= ranks
+    if len(shape) != ranks:
+        shape_tuple = tuple([0] * len(shape))
+        shape_tensor = get_trt_tensor(network, input, f"{name}_shape")
+        input_val, shape_tensor = broadcast(
+            network, input_val, shape_tensor, f"{name}_input_val", f"{name}_shape_val"
+        )
+        ranks = len(shape)
+
+    inshape = tuple(input_val.shape)
+    shape = tuple(shape)
+    start = tuple([0] * ranks)
+    stride = tuple(
+        [int(i == o) for i, o in zip(inshape, shape)]
+    )  # stride == 1 if dimensions match, 0 otherwise
+    return slice(network, target, source_ir, name, input_val, start, shape, stride)

py/torch_tensorrt/fx/test/converters/aten_op/test_where_aten.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestWhereConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d_condition_xshape_yshape", (2, 2), (2, 2)),
+            ("2d_broadcast_condition_xshape_yshape", (2, 2), (2, 1)),
+            ("3d_condition_xshape_yshape", (2, 2, 1), (2, 2, 1)),
+            ("2d_3d_condition_xshape_yshape", (2, 2), (1, 2, 2)),
+        ]
+    )
+    def test_(self, _, x_size, y_size):
+        class Where(nn.Module):
+            def forward(self, condition, x, y):
+                return torch.where(condition, x, y)
+
+        inputX = torch.randn(*x_size)
+        inputOther = torch.randn(*y_size)
+        condition = inputX < 0
+        self.run_test(
+            Where(),
+            (condition, inputX, inputOther),
+            expected_ops={torch.ops.aten.where.self},
+        )
+
+
+if __name__ == "__main__":
+    run_tests()