quant params from static inputs (#573)

Summary:
Pull Request resolved: #573

Since we allow tensor constants to be quantized inputs, we need to adjust the from_inputs api to search if this input is static or not. If it is static, then we take the first q node in get_attr --> q --> dq. If it is not static, then we just take the dq node to create the QuantParams object.

In the past, we can take in static quant inputs only on weights and biases.

Reviewed By: digantdesai

Differential Revision: D49850149

fbshipit-source-id: 007d594977f144c0fa58b6db01ca63e52e40312d

Author: mcr229

backends/xnnpack/operators/node_visitor.py

@@ -423,7 +423,7 @@ def define_nodes_tensor_inputs_outputs(
                     inp,
                     xnn_graph,
                     vals_to_ids,
-                    quant_params=QuantParams.from_inputs(inp),
+                    quant_params=QuantParams.from_inputs(inp, self._exported_program),
                     convert_to_nhwc=convert_to_nhwc,
                 )
         else:
@@ -434,7 +434,9 @@ def define_nodes_tensor_inputs_outputs(
             )
             # Define Input Node
             input_node = get_input_node(node, input_type_map.node_input)
-            input_quant_params = QuantParams.from_inputs(input_node)
+            input_quant_params = QuantParams.from_inputs(
+                input_node, self._exported_program
+            )
             self.define_tensor(
                 input_node,
                 xnn_graph,

backends/xnnpack/operators/op_add.py

@@ -43,7 +43,7 @@ def define_node(
             input1,
             xnn_graph,
             vals_to_ids,
-            quant_params=QuantParams.from_inputs(input1),
+            quant_params=QuantParams.from_inputs(input1, self._exported_program),
         )
         input1_id = vals_to_ids[input1]
 
@@ -53,7 +53,7 @@ def define_node(
             input2,
             xnn_graph,
             vals_to_ids,
-            quant_params=QuantParams.from_inputs(input2),
+            quant_params=QuantParams.from_inputs(input2, self._exported_program),
         )
         input2_id = vals_to_ids[input2]
 

backends/xnnpack/operators/op_cat.py

@@ -46,7 +46,9 @@ def define_node(
                 tensor_input,
                 xnn_graph,
                 vals_to_ids,
-                quant_params=QuantParams.from_inputs(tensor_input),
+                quant_params=QuantParams.from_inputs(
+                    tensor_input, self._exported_program
+                ),
             )
 
         self.define_tensor(

backends/xnnpack/operators/op_conv2d.py

@@ -46,7 +46,7 @@ def define_node(
         kwargs = {}
         # input
         input_node = get_input_node(node, 0)
-        input_quant_params = QuantParams.from_inputs(input_node)
+        input_quant_params = QuantParams.from_inputs(input_node, self._exported_program)
         self.define_tensor(
             input_node,
             xnn_graph,

backends/xnnpack/operators/op_multiply.py

@@ -43,7 +43,7 @@ def define_node(
             input1,
             xnn_graph,
             vals_to_ids,
-            quant_params=QuantParams.from_inputs(input1),
+            quant_params=QuantParams.from_inputs(input1, self._exported_program),
         )
         input1_id = vals_to_ids[input1]
 
@@ -53,7 +53,7 @@ def define_node(
             input2,
             xnn_graph,
             vals_to_ids,
-            quant_params=QuantParams.from_inputs(input2),
+            quant_params=QuantParams.from_inputs(input2, self._exported_program),
         )
         input2_id = vals_to_ids[input2]
 

backends/xnnpack/operators/op_sub.py

@@ -43,7 +43,7 @@ def define_node(
             input1,
             xnn_graph,
             vals_to_ids,
-            quant_params=QuantParams.from_inputs(input1),
+            quant_params=QuantParams.from_inputs(input1, self._exported_program),
         )
         input1_id = vals_to_ids[input1]
 
@@ -53,7 +53,7 @@ def define_node(
             input2,
             xnn_graph,
             vals_to_ids,
-            quant_params=QuantParams.from_inputs(input2),
+            quant_params=QuantParams.from_inputs(input2, self._exported_program),
         )
         input2_id = vals_to_ids[input2]
 

backends/xnnpack/operators/op_to_copy.py

@@ -51,7 +51,7 @@ def define_node(
         )
 
         input_node = get_input_node(node, 0)
-        input_quant_params = QuantParams.from_inputs(input_node)
+        input_quant_params = QuantParams.from_inputs(input_node, self._exported_program)
         output_quant_params = QuantParams.from_outputs(node)
 
         permute_order = PERM_NCHW_TO_NHWC if to_channels_last else PERM_NHWC_TO_NCHW

backends/xnnpack/operators/quant_params.py

@@ -11,8 +11,9 @@
 import torch
 from executorch.backends.xnnpack.passes.tag_implicit_q_dq_pass import TagImplicitQDqPass
 from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant
-from executorch.backends.xnnpack.utils.utils import check_or_raise
+from executorch.backends.xnnpack.utils.utils import check_or_raise, is_param_node
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import ExportedProgram
 
 
 class QuantParams:
@@ -178,11 +179,18 @@ def from_weights(cls, tensor_node: torch.fx.Node) -> Optional[QuantParams]:
         return cls.from_q_dq_node(q)
 
     @classmethod
-    def from_inputs(cls, tensor_node: torch.fx.Node) -> Optional[QuantParams]:
+    def from_inputs(
+        cls, tensor_node: torch.fx.Node, ep: ExportedProgram
+    ) -> Optional[QuantParams]:
         # tensor_node is quantized if it is produced by a dequant node
         if is_dequant(tensor_node) and TagImplicitQDqPass.is_tagged_as_implicit_q_dq(
             tensor_node
         ):
+            dq_input = cast(torch.fx.Node, tensor_node.args[0])
+            if is_quant(dq_input):
+                q_input = cast(torch.fx.Node, dq_input.args[0])
+                if is_param_node(ep, q_input):
+                    return cls.from_q_dq_node(dq_input)
             return cls.from_q_dq_node(tensor_node)
 
         return None