Converter reorg and softmax operation

softmax linting error fix

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -31,6 +31,7 @@
 from torch_tensorrt.fx.converters.impl.elementwise import fmod
 from torch_tensorrt.fx.converters.impl.normalization import batch_norm
 from torch_tensorrt.fx.converters.impl.normalization import layer_norm
+from torch_tensorrt.fx.converters.impl.normalization import softmax
 from torch_tensorrt.fx.converters.impl.unary import sign
 from torch_tensorrt.fx.converters.impl.elementwise.base import (
     convert_binary_elementwise,
@@ -671,37 +672,7 @@ def acc_ops_softmax(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-    input_ranks = len(input_val.shape) + (1 if network.has_implicit_batch_dimension else 0)  # type: ignore[union-attr]
-
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"softmax received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    # Used to get dim when dim is None. Copied from PyTorch softmax implementation.
-    def get_softmax_dim(ndim: int) -> int:
-        if ndim == 0 or ndim == 1 or ndim == 3:
-            ret = 0
-        else:
-            ret = 1
-        return ret
-
-    if kwargs["dim"] is None:
-        dim = get_softmax_dim(input_ranks)
-    else:
-        dim = cast(int, kwargs["dim"])
-
-    dim = get_positive_dim(dim, input_ranks)
-    if network.has_implicit_batch_dimension:
-        assert dim != 0, "Can't apply softmax on batch dimension when it's implicit."
-        dim -= 1
-
-    layer = network.add_softmax(input_val)
-    layer.axes = 1 << dim
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)
+    return softmax(network, target, SourceIR.ACC, name, kwargs["input"], kwargs["dim"])
 
 
 @tensorrt_converter(acc_ops.tile)

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -27,6 +27,7 @@
 from torch_tensorrt.fx.converters.impl.elementwise import rsub
 from torch_tensorrt.fx.converters.impl.normalization import batch_norm
 from torch_tensorrt.fx.converters.impl.normalization import layer_norm
+from torch_tensorrt.fx.converters.impl.normalization import softmax
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -490,6 +491,17 @@ def aten_ops_rsub(
     return rsub(network, target, SourceIR.ATEN, name, args[0], args[1], alpha)
 
 
+@tensorrt_converter(torch.ops.aten._softmax.default)
+def aten_ops_softmax(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return softmax(network, target, SourceIR.ATEN, name, args[0], args[1])
+
+
 @tensorrt_converter(torch.ops.aten.tanh.default)
 def aten_ops_tanh(
     network: TRTNetwork,

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

@@ -1,7 +1,6 @@
 import operator
 import warnings
 from typing import cast, Union, Callable, Any, Optional, Sequence
-import logging
 
 import numpy as np
 
@@ -273,3 +272,43 @@ def layer_norm_no_plugin(
         scale_layer,
         beta_tensor.get_output(0),
     )
+
+
+def softmax(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dim: Optional[Any] = None,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    input_ranks = len(input.shape) + (1 if network.has_implicit_batch_dimension else 0)  # type: ignore[union-attr]
+
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"softmax received input {input} that is not part "
+            "of the TensorRT region!"
+        )
+
+    # Used to get dim when dim is None. Copied from PyTorch softmax implementation.
+    def get_softmax_dim(ndim: int) -> int:
+        if ndim == 0 or ndim == 1 or ndim == 3:
+            ret = 0
+        else:
+            ret = 1
+        return ret
+
+    if dim is None:
+        dim = get_softmax_dim(input_ranks)
+    else:
+        dim = cast(int, dim)
+
+    dim = get_positive_dim(dim, input_ranks)
+    if network.has_implicit_batch_dimension:
+        assert dim != 0, "Can't apply softmax on batch dimension when it's implicit."
+        dim -= 1
+
+    layer = network.add_softmax(input)
+    layer.axes = 1 << dim
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)

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

@@ -0,0 +1,44 @@
+import torch
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestSoftMaxConverter(DispatchTestCase):
+    def test_softmax(self):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.softmax = torch.nn.Softmax(1)
+
+            def forward(self, x):
+                return self.softmax(x)
+
+        inputs = [torch.randn(1, 3, 224, 224)]
+        self.run_test(
+            TestModule(), inputs, expected_ops={torch.ops.aten._softmax.default}
+        )
+
+    def test_softmax_with_dynamic_shape(self):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.softmax = torch.nn.Softmax(2)
+
+            def forward(self, x):
+                return self.softmax(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, 3, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 3, 1, 1), (1, 3, 5, 5), (2, 3, 10, 10))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten._softmax.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()