feat: Refactor dynamo export

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

py/torch_tensorrt/_compile.py

@@ -15,8 +15,8 @@ class _IRType(Enum):
 
     ts = 0
     fx = 1
-    fx_ts_compat = 2
-    dynamo_compile = 3
+    dynamo = 2
+    torch_compile = 3
 
 
 class _ModuleType(Enum):
@@ -47,31 +47,29 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
 
     ir_targets_torchscript = any([ir == opt for opt in ["torchscript", "ts"]])
     ir_targets_fx = ir == "fx"
-    ir_targets_dynamo_compile = ir == "dynamo_compile"
-    ir_targets_fx_ts_compat = ir == "fx_ts_compat"
+    ir_targets_dynamo = ir == "dynamo"
+    ir_targets_torch_compile = ir == "torch_compile"
 
     if module_is_tsable and ir_targets_torchscript:
         return _IRType.ts
     elif module_is_fxable and ir_targets_fx:
         return _IRType.fx
-    elif module_is_fxable and ir_targets_fx_ts_compat:
-        return _IRType.fx_ts_compat
-    elif module_is_fxable and ir_targets_dynamo_compile:
-        return _IRType.dynamo_compile
+    elif module_is_fxable and ir_targets_dynamo:
+        return _IRType.dynamo
+    elif module_is_fxable and ir_targets_torch_compile:
+        return _IRType.torch_compile
     else:
         if ir == "default":
             # Options are listed in order of preference
-            if module_is_tsable:
+            if module_is_fxable:
                 logging.log(
-                    logging.Level.Info, "ir was set to default, using TorchScript as ir"
+                    logging.Level.Info, "ir was set to default, using dynamo as ir"
                 )
-                return _IRType.ts
-            elif module_is_fxable:
+                return _IRType.dynamo
+            elif module_is_tsable:
                 raise ValueError(
-                    "Was given a torch.fx.GraphModule, fx is not currently supported by Torch-TensorRT"
+                    "Input graph is a Torchscript module but the ir provided is default (dynamo). Please set ir=torchscript to compile."
                 )
-                # logging.log(logging.Level.Info, "ir was set to default, using TorchScript as fx")
-                # return _IRType.fx
             else:
                 raise ValueError("Module was provided with in an unsupported format")
         else:
@@ -156,12 +154,12 @@ def compile(
             dynamic_batch=False,
             **kwargs,
         )
-    elif target_ir == _IRType.dynamo_compile:
+    elif target_ir == _IRType.dynamo:
         return torch_tensorrt.dynamo.compile(
             module, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs
         )
-    elif target_ir == _IRType.fx_ts_compat:
-        return torch_tensorrt.dynamo.fx_ts_compat.compile(
+    elif target_ir == _IRType.torch_compile:
+        return torch_tensorrt.dynamo.backend.compile(
             module, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs
         )
     else:

py/torch_tensorrt/dynamo/__init__.py

@@ -1,6 +1,2 @@
-import torch
-from packaging import version
-
-if version.parse(torch.__version__) >= version.parse("2.1.dev"):
-    from torch_tensorrt.dynamo import fx_ts_compat
-    from .backend import compile
+from ._settings import *
+from .compile import compile

py/torch_tensorrt/dynamo/common/_settings.py renamed to py/torch_tensorrt/dynamo/_settings.py

@@ -8,7 +8,7 @@
 from torch_tensorrt import EngineCapability, Device
 from torch_tensorrt.fx.utils import LowerPrecision
 
-from torch_tensorrt.dynamo.backend.utils import prepare_inputs, prepare_device
+from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
 from torch_tensorrt.dynamo.backend.backends import torch_tensorrt_backend
 from torch_tensorrt.dynamo._defaults import (
     PRECISION,

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -4,7 +4,7 @@
 from functools import partial
 import torch._dynamo as td
 
-from torch_tensorrt.dynamo.common import CompilationSettings
+from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.backend.lowering._decompositions import (
     get_decompositions,
 )
@@ -15,8 +15,8 @@
     partition,
     get_submod_inputs,
 )
-from torch_tensorrt.dynamo.backend.utils import parse_dynamo_kwargs
-from torch_tensorrt.dynamo.backend.conversion import convert_module
+from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
+from torch_tensorrt.dynamo.conversion import convert_module
 
 from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
 

py/torch_tensorrt/dynamo/backend/backends.py

@@ -7,3 +7,4 @@
 )
 from ._partition import partition, get_submod_inputs, DEFAULT_SINGLE_NODE_PARTITIONS
 from .substitutions import *
+from ._fusers import *

py/torch_tensorrt/dynamo/backend/lowering/__init__.py

@@ -0,0 +1,72 @@
+import torch
+from torch_tensorrt.fx.tracer.acc_tracer import acc_ops
+
+
+def check_permute(node: torch.fx.Node):
+    ranks = len(node.meta["tensor_meta"].shape)
+    permutation = list(i % ranks for i in node.kwargs["permutation"])  # type: ignore[union-attr]
+    allowed_permutation = list(i for i in range(ranks))
+    allowed_permutation[-1] = ranks - 2
+    allowed_permutation[-2] = ranks - 1
+    return permutation == allowed_permutation
+
+
+def fuse_permute_matmul(gm: torch.fx.GraphModule):
+    """
+    Fuse pattern like permute + matmul if permute is transposing the last two dimension.
+    """
+    for node in gm.graph.nodes:
+        if node.target == acc_ops.matmul:
+            lhs, rhs = node.kwargs["input"], node.kwargs["other"]
+            lhs_transposed = rhs_tranposed = False
+            skip = False
+
+            if lhs.target == acc_ops.permute and check_permute(lhs):
+                lhs_transposed = True
+                lhs = lhs.kwargs["input"]
+
+            if rhs.target == acc_ops.permute and check_permute(rhs):
+                rhs_tranposed = True
+                rhs = rhs.kwargs["input"]
+
+            if (not skip) and (lhs_transposed or rhs_tranposed):
+                with gm.graph.inserting_before(node):
+                    fused_node = gm.graph.call_function(
+                        trt_transposed_matmul,
+                        args=(lhs, rhs, lhs_transposed, rhs_tranposed),
+                    )
+                node.replace_all_uses_with(fused_node)
+
+    gm.graph.eliminate_dead_code()
+    gm.graph.lint()
+    gm.recompile()
+    return gm
+
+
+def trt_transposed_linear(
+    input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
+):
+    return torch.matmul(input.transpose(-1, -2), weight.t()) + bias
+
+
+def fuse_permute_linear(gm: torch.fx.GraphModule):
+    """
+    Fuse pattern like permute + linear if permute is transposing the last two dimension.
+    """
+    for node in gm.graph.nodes:
+        if node.target == acc_ops.linear:
+            inp = node.kwargs["input"]
+            if inp.target == acc_ops.permute and check_permute(inp):
+                inp = inp.kwargs["input"]
+                weight = node.kwargs["weight"]
+                bias = node.kwargs["bias"]
+                with gm.graph.inserting_before(node):
+                    fused_node = gm.graph.call_function(
+                        trt_transposed_linear, args=(inp, weight, bias)
+                    )
+                    node.replace_all_uses_with(fused_node)
+
+    gm.graph.eliminate_dead_code()
+    gm.graph.lint()
+    gm.recompile()
+    return gm

py/torch_tensorrt/dynamo/backend/lowering/_fusers.py

@@ -2,9 +2,8 @@
 from torch.testing._internal.common_utils import run_tests, TestCase
 import torch
 from copy import deepcopy
-from torch_tensorrt.dynamo import compile
-from utils import lower_graph_testing
-from torch_tensorrt.dynamo.common.test_utils import DECIMALS_OF_AGREEMENT
+from torch_tensorrt.dynamo.backend import compile
+from utils import lower_graph_testing, DECIMALS_OF_AGREEMENT
 
 
 class TestTRTModuleNextCompilation(TestCase):

py/torch_tensorrt/dynamo/backend/test/test_backend_compiler.py

@@ -1,4 +1,4 @@
-from torch_tensorrt.dynamo.backend.utils import prepare_device, prepare_inputs
+from torch_tensorrt.dynamo.utils import prepare_device, prepare_inputs
 from utils import same_output_format
 import torch_tensorrt
 import unittest

py/torch_tensorrt/dynamo/backend/test/test_compiler_utils.py

@@ -1,9 +1,8 @@
 from functools import partial
-from utils import lower_graph_testing
+from utils import lower_graph_testing, DECIMALS_OF_AGREEMENT
 from torch.testing._internal.common_utils import run_tests, TestCase
 import torch
-from torch_tensorrt.dynamo import compile
-from torch_tensorrt.dynamo.common.test_utils import DECIMALS_OF_AGREEMENT
+from torch_tensorrt.dynamo.backend import compile
 
 
 class TestLowering(TestCase):