quant params from static inputs

Differential Revision: D49850149

fbshipit-source-id: fb8910116d1b2ac5eecdccdd2f3a64eb275fd122

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