Qualcomm AI Engine Direct - Requantization Mechanism Implementation

backends/qualcomm/builders/node_visitor.py

@@ -150,19 +150,19 @@ def make_qnn_per_tensor_config(self, quant_attrs: Dict):
             quant_config,
         )
 
-    def get_quant_encoding_conf(self, node: torch.fx.Node) -> Tuple[Any, Dict]:
+    def get_quant_encoding_conf(
+        self, node: torch.fx.Node, is_input_tensor: bool = False
+    ) -> Tuple[Any, Dict]:
         if not node.meta.get("quant_attrs", None):
             return (
                 PyQnnWrapper.Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_UNDEFINED,
                 {},
             )
-
         quant_attrs = (
-            node.meta["requantize"]["dq_attrs"]
-            if "requantize" in node.meta
+            node.meta["requantize"]
+            if "requantize" in node.meta and is_input_tensor
             else node.meta["quant_attrs"]
         )
-
         if quant_attrs["encoding"] in PER_CHANNEL_ENCODING:
             return self.make_qnn_per_channel_config(node, quant_attrs)
 
@@ -275,15 +275,27 @@ def define_custom_tensor_wrapper(
         nodes_to_wrappers[node_name] = tensor_wrapper
         return tensor_wrapper
 
-    def define_value(
+    def define_tensor(
         self,
         node: torch.fx.Node,
         tensor: torch.Tensor,
         tensor_type: PyQnnWrapper.Qnn_TensorType_t,
         nodes_to_wrappers: Dict[str, PyQnnWrapper.TensorWrapper],
-        is_tensor: bool,
+        is_input_tensor: bool,
         node_name: str = None,
+        is_tensor: bool = True,
     ) -> PyQnnWrapper.TensorWrapper:
+        """
+        Covert torch.Tensor to TensorWrapper
+
+        Args:
+            node: EdgeIR Node
+            tensor: EdgeIR Tensor
+            tensor_type: QNN tensor type
+            nodes_to_wrappers: Set contains edge_graph values(node targets)
+            is_input_tensor: Whether tensor is a fake input tensor relatively to
+                             the op builder that is calling this function
+        """
         if node_name is None:
             node_name = node.name
 
@@ -294,7 +306,9 @@ def define_value(
             tensor_name = "output_" + tensor_name
         dims = [1] if len(tensor.size()) == 0 else tensor.size()
         tensor_type = self.get_tensor_type(node, tensor_type)
-        quant_encoding, quant_configs = self.get_quant_encoding_conf(node)
+        quant_encoding, quant_configs = self.get_quant_encoding_conf(
+            node, is_input_tensor
+        )
         dtype = self.get_data_type(tensor, quant_configs, is_tensor)
         if isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor):
             tensor_wrapper = PyQnnWrapper.TensorWrapper(
@@ -330,52 +344,6 @@ def define_value(
         nodes_to_wrappers[node_name] = tensor_wrapper
         return tensor_wrapper
 
-    def define_scalar(
-        self,
-        node: torch.fx.Node,
-        tensor: torch.Tensor,
-        tensor_type: PyQnnWrapper.Qnn_TensorType_t,
-        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper],
-    ) -> PyQnnWrapper.TensorWrapper:
-        """
-        Covert constant scalar to TensorWrapper
-
-        Args:
-            tensor: EdgeIR Tensor
-            nodes_to_wrappers: Set contains edge_graph values(node targets)
-        """
-        return self.define_value(
-            node,
-            tensor,
-            tensor_type,
-            nodes_to_wrappers,
-            is_tensor=False,
-        )
-
-    def define_tensor(
-        self,
-        node: torch.fx.Node,
-        tensor: torch.Tensor,
-        tensor_type: PyQnnWrapper.Qnn_TensorType_t,
-        nodes_to_wrappers: Dict[str, PyQnnWrapper.TensorWrapper],
-        node_name: str = None,
-    ) -> PyQnnWrapper.TensorWrapper:
-        """
-        Covert torch.Tensor to TensorWrapper
-
-        Args:
-            tensor: EdgeIR Tensor
-            nodes_to_wrappers: Set contains edge_graph values(node targets)
-        """
-        return self.define_value(
-            node,
-            tensor,
-            tensor_type,
-            nodes_to_wrappers,
-            is_tensor=True,
-            node_name=node_name,
-        )
-
     def define_node(
         self,
         node: torch.fx.Node,

backends/qualcomm/builders/op_add.py

@@ -31,6 +31,7 @@ def define_node(
             out_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         add_output_tensors = [output_tensor_wrapper]
 
@@ -45,6 +46,7 @@ def define_node(
                 input_tensor,
                 tensor_type,
                 nodes_to_wrappers,
+                is_input_tensor=True,
             )
             add_input_tensors.append(input_tensor_wrapper)
 

backends/qualcomm/builders/op_avg_pool2d.py

@@ -33,6 +33,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         output_tensor = self.get_tensor(node, node)
@@ -41,6 +42,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         # kernel info
         filter_size = cast(List[int], node.args[1])

backends/qualcomm/builders/op_batch_norm.py

@@ -38,6 +38,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         bias_node = node.args[2]
@@ -52,6 +53,7 @@ def define_node(
             bias_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         filter_tensor = filter_tensor / torch.sqrt(var_tensor + eps)
@@ -60,6 +62,7 @@ def define_node(
             filter_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         batch_norm_input_tensors = [
@@ -74,6 +77,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         batch_norm_output_tensors = [output_tensor_wrapper]
 

backends/qualcomm/builders/op_bmm.py

@@ -35,6 +35,7 @@ def define_node(
                 input_tensor,
                 PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
                 nodes_to_wrappers,
+                is_input_tensor=True,
             )
             bmm_input_tensors.append(input_tensor_wrapper)
 
@@ -44,6 +45,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         bmm_output_tensors = [output_tensor_wrapper]
 

backends/qualcomm/builders/op_cast.py

@@ -33,6 +33,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         output_tensor = self.get_tensor(node, node)
@@ -42,6 +43,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         cast_op = PyQnnWrapper.PyQnnOpWrapper(

backends/qualcomm/builders/op_cat.py

@@ -37,6 +37,7 @@ def define_node(
                     input_tensor,
                     PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
                     nodes_to_wrappers,
+                    is_input_tensor=True,
                 )
             )
 
@@ -52,6 +53,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         # node args[1] might not exist

backends/qualcomm/builders/op_ceil.py

@@ -32,6 +32,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         output_tensor = self.get_tensor(node, node)
@@ -40,6 +41,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         ceil_op = PyQnnWrapper.PyQnnOpWrapper(

backends/qualcomm/builders/op_clamp.py

@@ -33,6 +33,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         # default value of output_min and output_max
@@ -53,6 +54,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         clamp_op = PyQnnWrapper.PyQnnOpWrapper(

backends/qualcomm/builders/op_conv2d.py

@@ -96,7 +96,7 @@ def _define_conv1d(
         op_wrapper_list = []  # op_wrapper to return
         unsqueeze_input_node = node.args[0]
         input_quant_encoding, input_quant_configs = self.get_quant_encoding_conf(
-            unsqueeze_input_node
+            unsqueeze_input_node,
         )
 
         unsqueeze_input_tensor = self.get_tensor(unsqueeze_input_node, node)
@@ -105,6 +105,7 @@ def _define_conv1d(
             unsqueeze_input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
         unsqueeze_output_tensor = unsqueeze_input_tensor.unsqueeze(1).contiguous()
         dtype = self.get_data_type(unsqueeze_output_tensor, input_quant_configs, True)
@@ -144,6 +145,7 @@ def _define_conv1d(
             filter_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         conv_input_tensors = [unsqueeze_output_tensor_wrapper, filter_tensor_wrapper]
         if node.args[2] is not None:
@@ -154,6 +156,7 @@ def _define_conv1d(
                 bias_tensor,
                 PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
                 nodes_to_wrappers,
+                is_input_tensor=False,
             )
             conv_input_tensors.append(bias_tensor_wrapper)
 
@@ -221,7 +224,8 @@ def _define_conv1d(
             squeeze_output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
-            node.name,
+            is_input_tensor=False,
+            node_name=node.name,
         )
         squeeze_op.AddInputTensors([conv_output_tensor_wrapper])
         squeeze_op.AddOutputTensors([squeeze_output_tensor_wrapper])
@@ -244,6 +248,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         filter_node = node.args[1]
@@ -256,6 +261,7 @@ def define_node(
             filter_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         conv_input_tensors = [input_tensor_wrapper, filter_tensor_wrapper]
 
@@ -267,6 +273,7 @@ def define_node(
                 bias_tensor,
                 PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
                 nodes_to_wrappers,
+                is_input_tensor=False,
             )
             conv_input_tensors.append(bias_tensor_wrapper)
 
@@ -276,6 +283,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         conv_output_tensors = [output_tensor_wrapper]
 

backends/qualcomm/builders/op_depth_to_space.py

@@ -34,6 +34,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
 
         output_tensor = self.get_tensor(node, node)
@@ -42,6 +43,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
 
         block_size = []

backends/qualcomm/builders/op_dequantize.py

@@ -30,6 +30,7 @@ def define_node(
             input_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=True,
         )
         dequant_input_tensors.append(inp_tensor_wrapper)
 
@@ -39,6 +40,7 @@ def define_node(
             output_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         dequant_output_tensors = [output_tensor_wrapper]
 

backends/qualcomm/builders/op_div.py

@@ -31,6 +31,7 @@ def define_node(
             out_tensor,
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
+            is_input_tensor=False,
         )
         div_output_tensors = [output_tensor_wrapper]
 
@@ -45,6 +46,7 @@ def define_node(
                 input_tensor,
                 tensor_type,
                 nodes_to_wrappers,
+                is_input_tensor=True,
             )
             div_input_tensors.append(input_tensor_wrapper)