Code Review

Author: winskuo-quic

backends/qualcomm/_passes/decompose_binary_alpha.py

@@ -24,12 +24,22 @@ def __init__(self) -> None:
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
         for node in graph.nodes:
-            if node.target in decomp_set and "alpha" in node.kwargs:
+            if (
+                node.target in decomp_set
+                and "alpha" in node.kwargs
+                and node.kwargs["alpha"] != 1
+            ):
                 alpha = node.kwargs["alpha"]
                 # Remove alpha from immutable dict
                 node.kwargs = {k: v for k, v in node.kwargs.items() if k != "alpha"}
+                input2_node = node.args[1]
+                # If input2 is constant, we can just multiply the value for optimization
+                if isinstance(input2_node, (int, float)):
+                    arg_list = list(node.args)
+                    arg_list[1] = input2_node * alpha
+                    node.args = tuple(arg_list)
+                    continue
                 with graph.inserting_before(node):
-                    input2_node = node.args[1]
                     mul_op = torch.ops.aten.mul.Scalar
                     mul_node = graph.create_node(
                         "call_function",
@@ -40,7 +50,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                         ),
                     )
                     mul_node.meta = copy_meta(node.meta)
-                    mul_node.users = {node: None}
+                    node.replace_input_with(input2_node, mul_node)
                     node.args = (
                         node.args[0],
                         mul_node,

backends/qualcomm/builders/op_conv.py

@@ -209,11 +209,11 @@ def define_node(
                     stacklevel=1,
                 )
 
-        if is_depthwise_conv:
+        if is_transpose_conv:
+            op_class = OpTransposeConv2d if is_conv2d else OpTransposeConv3d
+        elif is_depthwise_conv:
             assert is_conv2d, "DepthWise only supports Conv2d"
             op_class = OpDepthWiseConv2d
-        elif is_transpose_conv:
-            op_class = OpTransposeConv2d if is_conv2d else OpTransposeConv3d
         else:
             op_class = OpConv2d if is_conv2d else OpConv3d