feat: support bmm converter in dynamo #2248

bowang007 · 2023-08-21T01:56:17Z

Description

Support bmm converter for dynamo

New feature (non-breaking change which adds functionality)

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py	2023-08-21 01:56:33.820650+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py	2023-08-21 01:59:17.826690+00:00
@@ -180,21 +180,21 @@
) -> Union[TRTTensor, Sequence[TRTTensor]]:
    return impl.matmul.matrix_multiply(
        network, target, SourceIR.ATEN, name, args[0], args[1]
    )

+
@dynamo_tensorrt_converter(torch.ops.aten.bmm.default)
def aten_ops_bmm(
    network: TRTNetwork,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.matmul.bmm(
-        network, target, SourceIR.ATEN, name, args[0], args[1]
-    )
+    return impl.matmul.bmm(network, target, SourceIR.ATEN, name, args[0], args[1])
+

@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)
def aten_ops_layernorm(
    network: TRTNetwork,
    target: Target,
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/matmul.py	2023-08-21 01:56:33.824650+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/matmul.py	2023-08-21 01:59:17.915372+00:00
@@ -47,10 +47,11 @@
    )
    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
    set_layer_name(layer, target, name)
    return layer.get_output(0)

+
def bmm(
    network: TRTNetwork,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
@@ -73,12 +74,12 @@
    if len(input.shape) != 3:
        raise RuntimeError(f"Expected 3-dimensional tensor, but got ")

    if len(other.shape) != 3:
        raise RuntimeError(f"Expected 3-dimensional tensor, but got")
-    
-    if (input.shape[0] != other.shape[0]):
+
+    if input.shape[0] != other.shape[0]:
        raise RuntimeError("expected input tensors to have same batch size.")
-    
+
    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
    set_layer_name(layer, target, name)
-    return layer.get_output(0)
\ No newline at end of file
+    return layer.get_output(0)
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_bmm.py	2023-08-21 01:56:33.840650+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_bmm.py	2023-08-21 01:59:21.411402+00:00
@@ -2,10 +2,11 @@
import torch.nn as nn
from parameterized import parameterized
from torch.testing._internal.common_utils import run_tests
from torch_tensorrt.dynamo.test_utils import DispatchTestCase
from torch_tensorrt import Input
+

class TestBmmConverter(DispatchTestCase):
    @parameterized.expand(
        [
            ("10_3_5", (10, 3, 4), (9, 4, 5)),
@@ -13,21 +14,20 @@
    )
    def test_bmm(self, _, input_shape, mat2_shape):
        class BMM(nn.Module):
            def __init__(self):
                super().__init__()
-                
+
            def forward(self, input, mat2):
                return torch.bmm(input, mat2)
-            
+
        inputs = [torch.randn(*input_shape), torch.randn(*mat2_shape)]

-
        self.run_test(
-            BMM(), 
+            BMM(),
            inputs,
            expected_ops={},
        )
-        
+

if __name__ == "__main__":
    run_tests()

github-actions

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py	2023-08-21 01:56:46.017488+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py	2023-08-21 01:59:42.153437+00:00
@@ -180,21 +180,21 @@
) -> Union[TRTTensor, Sequence[TRTTensor]]:
    return impl.matmul.matrix_multiply(
        network, target, SourceIR.ATEN, name, args[0], args[1]
    )

+
@dynamo_tensorrt_converter(torch.ops.aten.bmm.default)
def aten_ops_bmm(
    network: TRTNetwork,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.matmul.bmm(
-        network, target, SourceIR.ATEN, name, args[0], args[1]
-    )
+    return impl.matmul.bmm(network, target, SourceIR.ATEN, name, args[0], args[1])
+

@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)
def aten_ops_layernorm(
    network: TRTNetwork,
    target: Target,
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/matmul.py	2023-08-21 01:56:46.017488+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/matmul.py	2023-08-21 01:59:42.159876+00:00
@@ -47,10 +47,11 @@
    )
    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
    set_layer_name(layer, target, name)
    return layer.get_output(0)

+
def bmm(
    network: TRTNetwork,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
@@ -73,12 +74,12 @@
    if len(input.shape) != 3:
        raise RuntimeError(f"Expected 3-dimensional tensor, but got ")

    if len(other.shape) != 3:
        raise RuntimeError(f"Expected 3-dimensional tensor, but got")
-    
-    if (input.shape[0] != other.shape[0]):
+
+    if input.shape[0] != other.shape[0]:
        raise RuntimeError("expected input tensors to have same batch size.")
-    
+
    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
    set_layer_name(layer, target, name)
-    return layer.get_output(0)
\ No newline at end of file
+    return layer.get_output(0)
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_bmm.py	2023-08-21 01:56:46.041488+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_bmm.py	2023-08-21 01:59:47.252864+00:00
@@ -2,10 +2,11 @@
import torch.nn as nn
from parameterized import parameterized
from torch.testing._internal.common_utils import run_tests
from torch_tensorrt.dynamo.test_utils import DispatchTestCase
from torch_tensorrt import Input
+

class TestBmmConverter(DispatchTestCase):
    @parameterized.expand(
        [
            ("10_3_5", (10, 3, 4), (9, 4, 5)),
@@ -13,21 +14,20 @@
    )
    def test_bmm(self, _, input_shape, mat2_shape):
        class BMM(nn.Module):
            def __init__(self):
                super().__init__()
-                
+
            def forward(self, input, mat2):
                return torch.bmm(input, mat2)
-            
+
        inputs = [torch.randn(*input_shape), torch.randn(*mat2_shape)]

-
        self.run_test(
-            BMM(), 
+            BMM(),
            inputs,
            expected_ops={},
        )
-        
+

if __name__ == "__main__":
    run_tests()

narendasan · 2023-08-23T21:23:48Z

@bowang007 lint the PR

tests/py/dynamo/converters/test_bmm.py

gs-olive · 2023-08-25T18:30:34Z

To remove:

TensorRT/py/torch_tensorrt/dynamo/aten_tracer.py

Line 138 in a65c95c

compose_bmm,

gs-olive · 2023-08-25T18:31:06Z

Consider switching to True default:

TensorRT/tests/py/dynamo/converters/harness.py

Line 219 in a65c95c

disable_passes: bool = False,

github-actions

Code conforms to C++ style guidelines

github-actions

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/matmul.py	2023-09-08 18:11:55.418367+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/matmul.py	2023-09-08 18:14:17.787932+00:00
@@ -47,10 +47,11 @@
    )
    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
    set_layer_name(layer, target, name, source_ir)
    return layer.get_output(0)

+
def bmm(
    network: TRTNetwork,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
@@ -73,12 +74,12 @@
    if len(input.shape) != 3:
        raise RuntimeError(f"Expected 3-dimensional tensor, but got ")

    if len(other.shape) != 3:
        raise RuntimeError(f"Expected 3-dimensional tensor, but got")
-    
-    if (input.shape[0] != other.shape[0]):
+
+    if input.shape[0] != other.shape[0]:
        raise RuntimeError("expected input tensors to have same batch size.")
-    
+
    layer = network.add_matrix_multiply(input, input_matrix_op, other, other_matrix_op)
    set_layer_name(layer, target, name)
-    return layer.get_output(0)
\ No newline at end of file
+    return layer.get_output(0)
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py	2023-09-08 18:11:55.418367+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py	2023-09-08 18:14:17.923968+00:00
@@ -313,21 +313,21 @@
        name,
        args[0],
        args[1],
    )

+
@dynamo_tensorrt_converter(torch.ops.aten.bmm.default)
def aten_ops_bmm(
    network: TRTNetwork,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.matmul.bmm(
-        network, target, SourceIR.ATEN, name, args[0], args[1]
-    )
+    return impl.matmul.bmm(network, target, SourceIR.ATEN, name, args[0], args[1])
+

@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
def aten_ops_layernorm(
    network: TRTNetwork,
    target: Target,
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_bmm.py	2023-09-08 18:11:55.434368+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/converters/test_bmm.py	2023-09-08 18:14:21.664385+00:00
@@ -2,10 +2,11 @@
import torch.nn as nn
from parameterized import parameterized
from torch.testing._internal.common_utils import run_tests
from torch_tensorrt.dynamo.test_utils import DispatchTestCase
from torch_tensorrt import Input
+

class TestBmmConverter(DispatchTestCase):
    @parameterized.expand(
        [
            ("10_3_5", (10, 3, 4), (9, 4, 5)),
@@ -13,21 +14,20 @@
    )
    def test_bmm(self, _, input_shape, mat2_shape):
        class BMM(nn.Module):
            def __init__(self):
                super().__init__()
-                
+
            def forward(self, input, mat2):
                return torch.bmm(input, mat2)
-            
+
        inputs = [torch.randn(*input_shape), torch.randn(*mat2_shape)]

-
        self.run_test(
-            BMM(), 
+            BMM(),
            inputs,
            expected_ops={},
        )
-        
+

if __name__ == "__main__":
    run_tests()

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

gs-olive

Looks good to me, pending CI validation for the added test case!

bowang007 · 2023-09-23T00:32:38Z

Looks good to me, pending CI validation for the added test case!

Is CI fixed now?

gs-olive · 2023-09-23T01:12:03Z

It should be working, but I see there is some sort of Torch version mismatch error. I'm not yet sure what's the cause - it seems to be happening across all the recent CI runs, since about 3pm today.

apbose · 2023-09-25T20:03:43Z

tests/py/dynamo/conversion/test_bmm.py

+
+
+class TestBmmConverter(DispatchTestCase):
+    @parameterized.expand(


As such everything looks good. But can more test cases be added with varying dimensions?

gs-olive

CI issues resolved - should just need the additional test cases, and then can merge

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

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

apbose

Looks good!

gs-olive

Looks good to me!

bowang007 requested a review from narendasan August 21, 2023 01:56

github-actions bot approved these changes Aug 21, 2023

View reviewed changes

github-actions bot requested changes Aug 21, 2023

View reviewed changes

facebook-github-bot added the cla signed label Aug 21, 2023

github-actions bot requested a review from apbose August 21, 2023 16:47

gs-olive reviewed Aug 25, 2023

View reviewed changes

tests/py/dynamo/converters/test_bmm.py Outdated Show resolved Hide resolved

bowang007 force-pushed the bmm_converter branch from acf6944 to 1ac8cbd Compare September 8, 2023 18:11

github-actions bot approved these changes Sep 8, 2023

View reviewed changes

github-actions bot requested changes Sep 8, 2023

View reviewed changes

gs-olive reviewed Sep 22, 2023

View reviewed changes

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py Outdated Show resolved Hide resolved

bowang007 force-pushed the bmm_converter branch from 1ac8cbd to 9f579d8 Compare September 22, 2023 22:29

bowang007 changed the title ~~converter: support bmm converter for dynamo~~ feat: support bmm converter in dynamo Sep 22, 2023

bowang007 requested a review from gs-olive September 22, 2023 22:30

github-actions bot approved these changes Sep 22, 2023

View reviewed changes

gs-olive approved these changes Sep 22, 2023

View reviewed changes

apbose reviewed Sep 25, 2023

View reviewed changes

gs-olive requested changes Sep 27, 2023

View reviewed changes

feat: support bmm converter in dynamo

fc12fbd

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

bowang007 force-pushed the bmm_converter branch from 9f579d8 to fc12fbd Compare September 27, 2023 23:40

github-actions bot approved these changes Sep 27, 2023

View reviewed changes

apbose approved these changes Sep 28, 2023

View reviewed changes

gs-olive approved these changes Sep 28, 2023

View reviewed changes

gs-olive merged commit 765933a into main Sep 28, 2023

bowang007 mentioned this pull request Oct 8, 2023

Implement converter for torch.ops.aten.bmm.default #2243

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support bmm converter in dynamo #2248

feat: support bmm converter in dynamo #2248

bowang007 commented Aug 21, 2023

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

narendasan commented Aug 23, 2023

gs-olive commented Aug 25, 2023

gs-olive commented Aug 25, 2023

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

gs-olive left a comment

bowang007 commented Sep 23, 2023

gs-olive commented Sep 23, 2023

apbose Sep 25, 2023

gs-olive left a comment

github-actions bot left a comment

github-actions bot left a comment

apbose left a comment

gs-olive left a comment



		class TestBmmConverter(DispatchTestCase):
		@parameterized.expand(

feat: support bmm converter in dynamo #2248

feat: support bmm converter in dynamo #2248

Conversation

bowang007 commented Aug 21, 2023

Description

Checklist:

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

narendasan commented Aug 23, 2023

gs-olive commented Aug 25, 2023

gs-olive commented Aug 25, 2023

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

gs-olive left a comment

Choose a reason for hiding this comment

bowang007 commented Sep 23, 2023

gs-olive commented Sep 23, 2023

apbose Sep 25, 2023

Choose a reason for hiding this comment

gs-olive left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

apbose left a comment

Choose a reason for hiding this comment

gs-olive left a comment

Choose a reason for hiding this comment