support argmax converter

Signed-off-by: Bo Wang <bowa@nvidia.com>

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1722,3 +1722,22 @@ def aten_ops_reshape(
         input=args[0],
         shape=args[1],
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.argmax.default)  # type: ignore[misc]
+def aten_ops_argmax(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.argmax.argmax(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        dim=args_bounds_check(args, 1),
+        keep_dim=args_bounds_check(args, 2, False),
+    )

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -3,6 +3,7 @@
 from . import (
     activation,
     attention,
+    argmax,
     cast,
     cat,
     condition,

py/torch_tensorrt/dynamo/conversion/impl/argmax.py

@@ -0,0 +1,47 @@
+from typing import Optional
+
+import tensorrt as trt
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    cast_trt_tensor,
+    get_axes_for_reduce_op,
+)
+from torch_tensorrt.fx.converters.converter_utils import (
+    get_positive_dim,
+    set_layer_name,
+)
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+from . import squeeze
+
+
+def argmax(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dim: int = 0,
+    keep_dim: bool = False,
+) -> TRTTensor:
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"argmax received input {input} that is not part " "of the TensorRT region!"
+        )
+    if input.dtype == trt.int32:
+        input = cast_trt_tensor(network, input, trt.float32, name)
+    if dim < 0:
+        dim = len(tuple(input.shape)) + dim
+    reduce_mask = get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape)))
+    topk_layer = network.add_topk(input, trt.TopKOperation.MAX, 1, reduce_mask)
+    set_layer_name(topk_layer, target, name)
+
+    out = topk_layer.get_output(1)
+
+    if not keep_dim:
+        out = squeeze.squeeze(
+            network, target, SourceIR.ATEN, name + "_squeeze", out, dim
+        )
+
+    return out