fix: Improve input handling in aten bmm lowering pass

py/torch_tensorrt/fx/passes/lower_basic_pass_aten.py

@@ -412,33 +412,50 @@ def compose_bmm(
     modified = False
     for n in module.graph.nodes:
         if n.op == "call_function" and n.target in (torch.ops.aten.bmm.default,):
-            modified = True
+            modified = False
             node = n
             input_n = node.all_input_nodes[0]
             other_n = node.all_input_nodes[1]
+
+            # If no input nodes are available, the bmm argument itself could be an input
+            # Alternatively, if the node has no users, it can be eliminated
+            if len(input_n.all_input_nodes) == 0 or len(node.users) == 0:
+                return PassResult(module, modified)
+
             output = next(iter(node.users))
             input_input_n = input_n.all_input_nodes[0]
             if (
                 input_input_n.target != torch.ops.aten.expand.default
                 and input_n.target != torch.ops.aten.view.default
             ):
-                raise RuntimeError(
-                    "Bmm is addressed in fixed pattern. A new pattern is met!"
+                _LOGGER.warn(
+                    "Bmm is addressed in fixed pattern. "
+                    + f"A new pattern {input_input_n.target}, {input_n.target} is met! "
+                    + "Skipping bmm lowering on this operation"
                 )
+                return PassResult(module, modified)
+
             real_input = input_input_n.all_input_nodes[0]
             input_other_n = other_n.all_input_nodes[0]
             if (
                 input_other_n.target != torch.ops.aten.expand.default
                 and other_n.target != torch.ops.aten.view.default
             ):
-                raise RuntimeError(
-                    "Bmm is addressed in fixed pattern. A new pattern is met!"
+                _LOGGER.warn(
+                    "Bmm is addressed in fixed pattern. "
+                    + f"A new pattern {input_other_n.target}, {other_n.target} is met! "
+                    + "Skipping bmm lowering on this operation"
                 )
+                return PassResult(module, modified)
+
             real_other = input_other_n.all_input_nodes[0]
             if len(real_other.meta["val"].size()) == 2:
                 new_func = aten_compose_bmm_2d
-            if len(real_other.meta["val"].size()) == 3:
+            elif len(real_other.meta["val"].size()) == 3:
                 new_func = aten_compose_bmm_3d
+            else:
+                # No valid bmm replacement exists for the specified dimensions
+                return PassResult(module, modified)
 
             with module.graph.inserting_after(node):
                 new_args = (real_input, real_other)
@@ -449,6 +466,7 @@ def compose_bmm(
                     kwargs=None,
                 )
             output.replace_all_uses_with(new_node)
+            modified = True
 
     module.graph.eliminate_dead_code()
     module.recompile()

py/torch_tensorrt/fx/test/passes/test_compose_bmm.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests, TestCase
+from torch_tensorrt.fx.tracer.dispatch_tracer.aten_tracer import trace
+from torch_tensorrt.fx.passes.lower_basic_pass_aten import compose_bmm
+
+
+class TestComposeBMM(TestCase):
+    @parameterized.expand(
+        [
+            ("3_dim", (2, 3, 4), (2, 4, 3)),
+            ("3_dim_same_shape", (4, 4, 4), (4, 4, 4)),
+        ]
+    )
+    def test_compose_bmm(self, test_name, x_shape, y_shape):
+        class BMM(nn.Module):
+            def forward(self, x, y):
+                return torch.bmm(x, y)
+
+        inputs = [torch.randn(x_shape), torch.randn(y_shape)]
+        fx_model, _ = trace(BMM(), inputs)
+        composed_module = compose_bmm(fx_model)
+        out = composed_module.graph_module(*inputs)
+
+
+if __name__ == "__main__":
+    run_tests()