Converter reorg fmod

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -28,6 +28,7 @@
 from torch_tensorrt.fx.tracer.acc_tracer.acc_ops import contiguous
 from torch_tensorrt.fx.converters.impl import activation
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
+from torch_tensorrt.fx.converters.impl.elementwise import fmod
 from torch_tensorrt.fx.converters.impl.unary import sign
 from torch_tensorrt.fx.converters.impl.elementwise.base import (
     convert_binary_elementwise,
@@ -2091,34 +2092,14 @@ def acc_ops_fmod(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    # NOTE: TRT doesnt currently implement fmod so we need multiple operations to perform it
-    trunc_div_value = trunc_div(
+    return fmod(
         network,
         target,
         SourceIR.ACC,
-        name + "_trunc_div",
+        name,
         kwargs["input"],
         kwargs["other"],
     )
-    prod_value = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        name + "_prod",
-        trt.ElementWiseOperation.PROD,
-        trunc_div_value,
-        kwargs["other"],
-    )
-    sub_value = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        name + "_sub",
-        trt.ElementWiseOperation.SUB,
-        kwargs["input"],
-        prod_value,
-    )
-    return sub_value
 
 
 # T113156424 embedding implemenatation is very limited and shows no usage in hf models due to the indices are int64.

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -23,6 +23,7 @@
 from torch_tensorrt.fx.converters.impl import activation
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
 from torch_tensorrt.fx.converters.impl.elementwise import rsqrt
+from torch_tensorrt.fx.converters.impl.elementwise import fmod
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -219,11 +220,7 @@ def aten_ops_fmod(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-        "other": args[1],
-    }
-    return acc_ops_converters.acc_ops_fmod(network, target, None, kwargs_new, name)
+    return fmod(network, target, SourceIR.ATEN, name, args[0], args[1])
 
 
 @tensorrt_converter(torch.ops.aten.hardtanh.default)

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

@@ -139,3 +139,41 @@ def rsqrt(
     )
 
     return output
+
+
+def fmod(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    other: TRTTensor,
+) -> TRTTensor:
+    # NOTE: TRT doesnt currently implement fmod so we need multiple operations to perform it
+    trunc_div_value = trunc_div(
+        network,
+        target,
+        source_ir,
+        name + "_trunc_div",
+        input,
+        other,
+    )
+    prod_value = convert_binary_elementwise(
+        network,
+        target,
+        source_ir,
+        name + "_prod",
+        trt.ElementWiseOperation.PROD,
+        trunc_div_value,
+        other,
+    )
+    sub_value = convert_binary_elementwise(
+        network,
+        target,
+        SourceIR.ACC,
+        name + "_sub",
+        trt.ElementWiseOperation.SUB,
+        input,
+        prod_value,
+    )
+    return sub_value