[executorch][nvidia][tensorrt][14/n] Add additional converters (mul, sub, div, mm, relu)

backends/nvidia/tensorrt/converters/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -8,3 +8,8 @@
 
 # Import converters to trigger registration via @converter decorator
 from executorch.backends.nvidia.tensorrt.converters import add  # noqa: F401
+from executorch.backends.nvidia.tensorrt.converters import div  # noqa: F401
+from executorch.backends.nvidia.tensorrt.converters import mm  # noqa: F401
+from executorch.backends.nvidia.tensorrt.converters import mul  # noqa: F401
+from executorch.backends.nvidia.tensorrt.converters import relu  # noqa: F401
+from executorch.backends.nvidia.tensorrt.converters import sub  # noqa: F401

backends/nvidia/tensorrt/converters/targets.bzl

@@ -12,6 +12,11 @@ def define_common_targets():
         srcs = [
             "__init__.py",
             "add.py",
+            "div.py",
+            "mm.py",
+            "mul.py",
+            "relu.py",
+            "sub.py",
         ],
         visibility = ["PUBLIC"],
         deps = [

backends/nvidia/tensorrt/partitioner/operator_support.py

@@ -22,10 +22,19 @@ class TensorRTOperatorSupport(OperatorSupportBase):
     3. Its output dtype is in SUPPORTED_DTYPES
     """
 
-    # Operations that have TensorRT converters.
-    # Format: "op_name.overload" (e.g., "add.Tensor")
+    # Operations that have TensorRT converters (sorted alphabetically).
     SUPPORTED_OPS: Set[str] = {
         "add.Tensor",
+        "add_.Tensor",
+        "div.Tensor",
+        "div.Tensor_mode",
+        "mm.default",
+        "mul.Scalar",
+        "mul.Tensor",
+        "mul_.Tensor",
+        "relu.default",
+        "relu_.default",
+        "sub.Tensor",
     }
 
     # Glue operations that don't compute but are needed to keep partitions connected.
@@ -86,8 +95,15 @@ def _get_op_name(self, node: torch.fx.Node) -> str:
         if hasattr(target, "_schema"):
             schema = target._schema
             base_name = schema.name.replace("::", ".")
-            if hasattr(schema, "overload_name") and schema.overload_name:
-                return f"{base_name}.{schema.overload_name}"
+            # Note: For the "default" overload, overload_name is an empty string "",
+            # so we need to check for that and use "default" as the overload name.
+            if hasattr(schema, "overload_name"):
+                overload_name = schema.overload_name
+                if overload_name:
+                    return f"{base_name}.{overload_name}"
+                else:
+                    # Empty overload_name means "default" overload
+                    return f"{base_name}.default"
             return base_name
 
         # Callable with module info (e.g., operator.getitem)

examples/nvidia/tensorrt/export.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -36,6 +36,7 @@
     "add",
     "add_mul",
     "mul",
+    "softmax",
 }
 
 

examples/nvidia/tensorrt/tests/test_export.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -16,7 +16,7 @@
 import shutil
 import unittest
 
 import torch
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -85,7 +85,7 @@
         from executorch.examples.models import MODEL_NAME_TO_MODEL
         from executorch.examples.models.model_factory import EagerModelFactory
         from executorch.exir import to_edge_transform_and_lower
         import torch
         from torch.export import export
         model, example_inputs, _, _ = EagerModelFactory.create_model(*MODEL_NAME_TO_MODEL["add"])
         model = model.eval()
@@ -95,3 +95,6 @@
         exec_prog = edge.to_executorch()
         self.assertIsNotNone(exec_prog)
         logger.info("PASS: add model exported with BF16 precision")
+
+    def test_softmax(self) -> None:
+        _export_and_verify("softmax")