binary operator changes in aten

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -144,11 +144,11 @@ def aten_ops_div(
     }
     rounding_mode = kwargs.get("rounding_mode")
     if rounding_mode is None:
-        return add_div(network, target, None, kwargs_new, name)
+        return add_div(network, target, kwargs_new, name)
     elif rounding_mode == "floor":
-        return add_floor_div(network, target, None, kwargs_new, name)
+        return add_floor_div(network, target, kwargs_new, name)
     elif rounding_mode == "trunc":
-        return add_trunc_div(network, target, None, kwargs_new, name)
+        return add_trunc_div(network, target, kwargs_new, name)
     else:
         raise RuntimeError(
             f"Target {target} does not support rounding mode {rounding_mode}"
@@ -167,7 +167,7 @@ def aten_ops_floor_div(
         "input": args[0],
         "other": args[1],
     }
-    return add_floor_div(network, target, None, kwargs_new, name)
+    return add_floor_div(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.fmod.Scalar)
@@ -183,7 +183,7 @@ def aten_ops_fmod(
         "input": args[0],
         "other": args[1],
     }
-    return add_fmod(network, target, None, kwargs_new, name)
+    return add_fmod(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.linear)
@@ -200,7 +200,7 @@ def aten_ops_linear(
         "bias": args[2],
     }
 
-    return add_linear(network, target, None, kwargs_new, name)
+    return add_linear(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.max_pool3d)
@@ -249,7 +249,7 @@ def aten_ops_mul(
         "input": args[0],
         "other": args[1],
     }
-    return add_mul(network, target, None, kwargs_new, name)
+    return add_mul(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.matmul)
@@ -265,7 +265,7 @@ def aten_ops_matmul(
         "input": args[0],
         "other": args[1],
     }
-    return add_matmul(network, target, None, kwargs_new, name)
+    return add_matmul(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.pow.Tensor_Scalar)
@@ -310,7 +310,7 @@ def aten_ops_sub(
         "input": args[0],
         "other": args[1],
     }
-    return add_sub(network, target, None, kwargs_new, name)
+    return add_sub(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.view.default)
@@ -392,7 +392,7 @@ def aten_ops_operator_floordiv(
         "input": args[0],
         "other": args[1],
     }
-    return add_floor_div(network, target, None, kwargs_new, name)
+    return add_floor_div(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(operator.mul)
@@ -407,7 +407,7 @@ def aten_ops_operator_mul(
         "input": args[0],
         "other": args[1],
     }
-    return add_mul(network, target, None, kwargs_new, name)
+    return add_mul(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(operator.add)
@@ -422,7 +422,7 @@ def aten_ops_operator_add(
         "input": args[0],
         "other": args[1],
     }
-    return add_add(network, target, None, kwargs_new, name)
+    return add_add(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(operator.sub)
@@ -437,7 +437,7 @@ def aten_ops_operator_sub(
         "input": args[0],
         "other": args[1],
     }
-    return add_sub(network, target, None, kwargs_new, name)
+    return add_sub(network, target, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.sym_numel)
@@ -499,22 +499,6 @@ def aten_ops_slice(
     return add_slice(network, target.kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.select)
-def aten_ops_select(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-        "dim": args[1],
-        "index": args[2],
-    }
-    return add_select(network, target.kwargs_new, name)
-
-
 @tensorrt_converter(torch.ops.aten.leaky_relu.default)
 def aten_ops_leaky_relu(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -524,122 +524,6 @@ def get_inputs_from_args_and_kwargs(args, kwargs, input_names):
     return inputs
 
 
-def sign(
-    network: TRTNetwork, input_val: TRTTensor, target: Target, name: str
-) -> TRTTensor:
-    """
-    Sign is calculated as below:
-       x = input
-       sign = (exp(x) // exp(abs(x))) * 2 - 1
-       For positive number and 0, (exp(x) // exp(abs(x))) yield 1; for negative number, (exp(x) // exp(abs(x))) yield 0.
-       With multiply 2, the value become 2(for pos and 0) and 0(for neg).
-       Finally minus 1, the value become 1(for pos and 0) and -1(for neg).
-
-    Args:
-        network (TRTNetwork): TensorRT network object.
-        input_val (TRTTensor): The input tensor.
-        target (Target): fx node target.
-        name (str): Name of the fx node with optional suffix.
-
-    Returns:
-        A TensorRT tensor represent the result of sign operator.
-    """
-    input_exp_output = add_unary_layer(
-        network, input_val, trt.UnaryOperation.EXP, target, f"{name}_prod_exp"
-    )
-    input_abs_output = add_unary_layer(
-        network, input_val, trt.UnaryOperation.ABS, target, f"{name}_prod_abs"
-    )
-    input_abs_exp_output = add_unary_layer(
-        network,
-        input_abs_output,
-        trt.UnaryOperation.EXP,
-        target,
-        f"{name}_prod_abs_exp",
-    )
-    floor_div_output = add_binary_elementwise_layer(
-        network,
-        input_exp_output,
-        input_abs_exp_output,
-        trt.ElementWiseOperation.FLOOR_DIV,
-        target,
-        f"{name}_exp_floor_div",
-    )
-    double_floor_div_output = add_binary_elementwise_layer(
-        network,
-        floor_div_output,
-        2,
-        trt.ElementWiseOperation.PROD,
-        target,
-        f"{name}_floor_div*2",
-    )
-    return add_binary_elementwise_layer(
-        network,
-        double_floor_div_output,
-        1,
-        trt.ElementWiseOperation.SUB,
-        target,
-        f"{name}_sign",
-    )
-
-
-def trunc_div(
-    input: TRTTensor, other: TRTTensor, network: TRTNetwork, target: Target, name: str
-) -> TRTTensor:
-    """
-    Perform trunc divide on Tensor, result of divide will be round toward zero.
-    This means for positive number, it will be floor round; for negative number,
-    it will be ceil round. Example: [2.1, 0.8, -3.2] -> [2, 0, -3].
-
-    Args:
-        input: divisor.
-        other: dividend.
-        network: INetworkDefinition.
-        target: node target.
-        name: namespace for the op
-
-    Returns:
-        A TensorRT tensor represent the result of trunc divide.
-    """
-    prod_output = add_binary_elementwise_layer(
-        network, input, other, trt.ElementWiseOperation.PROD, target, f"{name}_prod"
-    )
-    sign_output = sign(network, prod_output, target, name)
-
-    # Convert constant input into ITensor for UnaryOperation
-    if not isinstance(input, trt.tensorrt.ITensor):
-        input = get_trt_tensor(network, input, f"{name}_input")
-    if not isinstance(other, trt.tensorrt.ITensor):
-        other = get_trt_tensor(
-            network, other, f"{name}_other", dtype=torch_dtype_from_trt(input.dtype)
-        )
-
-    abs_input_output = add_unary_layer(
-        network, input, trt.UnaryOperation.ABS, target, f"{name}_abs_input"
-    )
-    abs_other_output = add_unary_layer(
-        network, other, trt.UnaryOperation.ABS, target, f"{name}_abs_other"
-    )
-    abs_floor_output = add_binary_elementwise_layer(
-        network,
-        abs_input_output,
-        abs_other_output,
-        trt.ElementWiseOperation.FLOOR_DIV,
-        target,
-        f"{name}_floor_div",
-    )
-    output = add_binary_elementwise_layer(
-        network,
-        abs_floor_output,
-        sign_output,
-        trt.ElementWiseOperation.PROD,
-        target,
-        f"{name}_output",
-    )
-
-    return output
-
-
 def dtype_uniform(
     network: TRTNetwork, target: Target, name: str, input: TRTTensor, other: TRTTensor
 ):

py/torch_tensorrt/fx/converters/operator.py

@@ -162,6 +162,65 @@ def add_binary_elementwise_layer(
     return output
 
 
+def sign(
+    network: TRTNetwork, input_val: TRTTensor, target: Target, name: str
+) -> TRTTensor:
+    """
+    Sign is calculated as below:
+       x = input
+       sign = (exp(x) // exp(abs(x))) * 2 - 1
+       For positive number and 0, (exp(x) // exp(abs(x))) yield 1; for negative number, (exp(x) // exp(abs(x))) yield 0.
+       With multiply 2, the value become 2(for pos and 0) and 0(for neg).
+       Finally minus 1, the value become 1(for pos and 0) and -1(for neg).
+
+    Args:
+        network (TRTNetwork): TensorRT network object.
+        input_val (TRTTensor): The input tensor.
+        target (Target): fx node target.
+        name (str): Name of the fx node with optional suffix.
+
+    Returns:
+        A TensorRT tensor represent the result of sign operator.
+    """
+    input_exp_output = add_unary_layer(
+        network, input_val, trt.UnaryOperation.EXP, target, f"{name}_prod_exp"
+    )
+    input_abs_output = add_unary_layer(
+        network, input_val, trt.UnaryOperation.ABS, target, f"{name}_prod_abs"
+    )
+    input_abs_exp_output = add_unary_layer(
+        network,
+        input_abs_output,
+        trt.UnaryOperation.EXP,
+        target,
+        f"{name}_prod_abs_exp",
+    )
+    floor_div_output = add_binary_elementwise_layer(
+        network,
+        input_exp_output,
+        input_abs_exp_output,
+        trt.ElementWiseOperation.FLOOR_DIV,
+        target,
+        f"{name}_exp_floor_div",
+    )
+    double_floor_div_output = add_binary_elementwise_layer(
+        network,
+        floor_div_output,
+        2,
+        trt.ElementWiseOperation.PROD,
+        target,
+        f"{name}_floor_div*2",
+    )
+    return add_binary_elementwise_layer(
+        network,
+        double_floor_div_output,
+        1,
+        trt.ElementWiseOperation.SUB,
+        target,
+        f"{name}_sign",
+    )
+
+
 def trunc_div(
     input: TRTTensor, other: TRTTensor, network: TRTNetwork, target: Target, name: str
 ) -> TRTTensor:
@@ -701,7 +760,7 @@ def add_pow(network, target, kwargs, name):
     return add_binary_elementwise_layer(
         network,
         kwargs["input"],
-        kwargs["other"],
+        kwargs["exponent"],
         trt.ElementWiseOperation.POW,
         target,
         name,