linear fused relu in xnnpack delegate (#1551)

Summary:
Pull Request resolved: #1551

Add support for linear relu fusion in XNNPACK delegate

Reviewed By: balakv504, digantdesai

Differential Revision: D52574085

fbshipit-source-id: 88222fc8fdd7525ccb1f56afd45433ec31696669

Author: mcr229

backends/xnnpack/operators/op_linear.py

@@ -9,17 +9,20 @@
 import torch
 from executorch.backends.xnnpack.operators.node_visitor import (
     get_input_node,
-    InputTypeToIndex,
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.xnnpack.operators.quant_params import QuantParams
 from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    OutputMinMax,
     XNNFullyConnected,
     XNNGraph,
     XNode,
 )
+from executorch.backends.xnnpack.utils.utils import get_relu_fused_node
 
 from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_INVALID_VALUE_ID
+from executorch.exir.dialects._ops import ops as exir_ops
 
 
 @register_node_visitor
@@ -36,30 +39,62 @@ def define_node(
         vals_to_ids: Dict[torch.fx.Node, int],
         debug_handle: int,
     ) -> None:
-        input_type_map = (
-            InputTypeToIndex(node_input=0, node_weight=1, node_bias=2)
-            if len(node.args) == 3
-            else InputTypeToIndex(node_input=0, node_weight=1)
-        )
-        self.define_nodes_tensor_inputs_outputs(
-            node, xnn_graph, vals_to_ids, input_type_map=input_type_map
-        )
-
-        # bias
-        bias_id = (
-            XNN_INVALID_VALUE_ID
-            if len(node.args) == 2
-            else vals_to_ids[get_input_node(node, input_type_map.node_bias)]
-        )
 
         # input
-        input_id = vals_to_ids[get_input_node(node, input_type_map.node_input)]
+        input_node = get_input_node(node, 0)
+        input_quant_params = QuantParams.from_inputs(input_node, self._exported_program)
+        self.define_tensor(
+            input_node,
+            xnn_graph,
+            vals_to_ids,
+            quant_params=input_quant_params,
+        )
+        input_id = vals_to_ids[input_node]
 
         # filter
-        filter_id = vals_to_ids[get_input_node(node, input_type_map.node_weight)]
+        weight_node = get_input_node(node, 1)
+        weight_quant_params = QuantParams.from_weights(
+            weight_node, self._exported_program
+        )
+        self.define_tensor(
+            weight_node,
+            xnn_graph,
+            vals_to_ids,
+            quant_params=weight_quant_params,
+        )
+        filter_id = vals_to_ids[weight_node]
+
+        # bias
+        if len(node.args) > 2:
+            bias_node = get_input_node(node, 2)
+            bias_quant_params = QuantParams.from_bias(
+                bias_node, weight_quant_params, input_quant_params
+            )
+            self.define_tensor(
+                get_input_node(node, 2),
+                xnn_graph,
+                vals_to_ids,
+                quant_params=bias_quant_params,
+            )
+            bias_id = vals_to_ids[bias_node]
+        else:
+            bias_id = XNN_INVALID_VALUE_ID
 
         # output
-        output_id = vals_to_ids[node]
+        output_node = get_relu_fused_node(node) or node
+        output_min_max = None
+        if output_node.target == exir_ops.edge.aten.relu.default:
+            output_node.meta["XNNPACK_FUSED"] = True
+            output_min_max = OutputMinMax(output_min=0, output_max="+inf")
+
+        output_quant_params = QuantParams.from_outputs(output_node)
+        self.define_tensor(
+            output_node,
+            xnn_graph,
+            vals_to_ids,
+            quant_params=output_quant_params,
+        )
+        output_id = vals_to_ids[output_node]
 
         ser_node = XNode(
             xnode_union=XNNFullyConnected(
@@ -70,5 +105,6 @@ def define_node(
                 flags=0,
             ),
             debug_handle=debug_handle,
+            output_min_max=output_min_max,
         )
         xnn_graph.xnodes.append(ser_node)

backends/xnnpack/passes/tag_implicit_q_dq_pass.py

@@ -64,6 +64,11 @@ class TagImplicitQDqPass(XNNPACKPass):
                 _END_OF_CHAIN_MARKER: True,
             }
         },
+        exir_ops.edge.aten.linear.default.name(): {
+            exir_ops.edge.aten.relu.default.name(): {
+                _END_OF_CHAIN_MARKER: True,
+            }
+        },
     }
     IS_IMPLICIT_Q_DQ_TAG = "IS_IMPLICIT_Q_DQ_TAG"
 

backends/xnnpack/test/ops/linear.py

@@ -204,6 +204,45 @@ def forward(self, x, y):
             LinearModule(), inputs, linear_count=3, is_per_channel=True, uses_bias=True
         )
 
+    def test_fp32_linear_fused_relu(self):
+        class LinearReluModule(torch.nn.Module):
+            def __init__(self, in_size, out_size, use_bias):
+                super().__init__()
+                self.linear = torch.nn.Linear(in_size, out_size, bias=use_bias)
+
+            def forward(self, x):
+                return torch.nn.functional.relu(self.linear(x))
+
+        for use_bias in (True, False):
+            self._test_linear(
+                lambda in_size, out_size: LinearReluModule(
+                    in_size,
+                    out_size,
+                    use_bias,  # noqa
+                ),
+                uses_bias=use_bias,
+            )
+
+    def test_qs8_linear_fused_relu(self):
+        class LinearReluModule(torch.nn.Module):
+            def __init__(self, in_size, out_size, use_bias):
+                super().__init__()
+                self.linear = torch.nn.Linear(in_size, out_size, bias=use_bias)
+
+            def forward(self, x):
+                return torch.nn.functional.relu(self.linear(x))
+
+        for use_bias in (True, False):
+            self._test_linear(
+                lambda in_size, out_size: LinearReluModule(
+                    in_size,
+                    out_size,
+                    use_bias,  # noqa
+                ),
+                uses_bias=use_bias,
+                quant=True,
+            )
+
     def _test_linear(self, make_module, uses_bias, quant=False):
         aten_op, edge_op = (
             (
@@ -256,7 +295,7 @@ def _test_linear(self, make_module, uses_bias, quant=False):
             tester.to_executorch()
             tester.serialize()
             tester.run_method()
-            tester.compare_outputs()
+            tester.compare_outputs(qtol=quant)
 
     def _test_dqlinear(
         self, module, inputs, linear_count=1, is_per_channel=False, uses_bias=False