Precompute T1 offset for quantized conv2d NHWC in TIE kernel (#18960)

Differential Revision: D100690813

Pull Request resolved: #18960

Author: abeakkas

backends/cadence/aot/functions.yaml

@@ -384,7 +384,7 @@
     - arg_meta: null
       kernel_name: impl::generic::quantized_conv2d_nchw_per_tensor_out
 
-- func: cadence::quantized_conv2d_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)
+- func: cadence::quantized_conv2d_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, Tensor? offset=None, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null
       kernel_name: impl::generic::quantized_conv2d_nhwc_per_tensor_out

backends/cadence/aot/functions_hifi.yaml

@@ -395,7 +395,7 @@
     - arg_meta: null
       kernel_name: impl::HiFi::quantized_conv2d_nchw_per_tensor_out
 
-- func: cadence::quantized_conv2d_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)
+- func: cadence::quantized_conv2d_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, Tensor? offset=None, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null
       kernel_name: impl::HiFi::quantized_conv2d_nhwc_per_tensor_out

backends/cadence/aot/ops_registrations.py

@@ -233,10 +233,10 @@ def register_fake(
     "quantized_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)"
 )
 lib.define(
-    "quantized_conv2d_nhwc.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)"
+    "quantized_conv2d_nhwc.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, Tensor? offset=None) -> (Tensor Z)"
 )
 lib.define(
-    "quantized_conv2d_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)"
+    "quantized_conv2d_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, Tensor? offset=None, *, Tensor(a!) out) -> Tensor(a!)"
 )
 lib.define(
     "quantized_conv1d_ncl(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)"
@@ -1444,6 +1444,7 @@ def quantized_conv2d_nhwc_per_tensor_meta(
     output_zero_point: int,
     out_multiplier: int,
     out_shift: int,
+    offset: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     in_size = input.shape
     # Assert that the input tensor has at least 3 dimensions, and at most 6

backends/cadence/aot/ref_implementations.py

@@ -1449,6 +1449,7 @@ def quantized_conv2d_nhwc_per_tensor(
     output_zero_point: int,
     out_multiplier: int,
     out_shift: int,
+    offset: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """
     Quantized convolution operation.

backends/cadence/generic/operators/op_quantized_conv2d.cpp

@@ -16,6 +16,7 @@ namespace impl {
 namespace generic {
 namespace native {
 
+using ::executorch::aten::optional;
 using ::executorch::aten::ScalarType;
 using ::executorch::aten::Tensor;
 using ::executorch::runtime::KernelRuntimeContext;
@@ -935,6 +936,7 @@ Tensor& quantized_conv2d_nhwc_per_tensor_out(
     int64_t output_zero_point,
     ET_UNUSED int64_t out_multiplier,
     ET_UNUSED int64_t out_shift,
+    ET_UNUSED const ::executorch::aten::optional<Tensor>& offset,
     Tensor& out) {
   quantized_conv2d_nhwc(
       input,

backends/cadence/generic/operators/op_quantized_conv2d.h

@@ -205,6 +205,7 @@ ::executorch::aten::Tensor& quantized_conv2d_nhwc_per_tensor_out(
     int64_t output_zero_point,
     int64_t out_multiplier,
     int64_t out_shift,
+    const ::executorch::aten::optional<Tensor>& offset,
     Tensor& out);
 
 ::executorch::aten::Tensor&

backends/cadence/hifi/operators/op_quantized_conv2d_nhwc_out.cpp

@@ -17,6 +17,7 @@ using Tensor = executorch::aten::Tensor;
 using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
 using ScalarType = executorch::aten::ScalarType;
 using ::executorch::aten::IntArrayRef;
+using ::executorch::aten::optional;
 
 namespace impl {
 namespace HiFi {
@@ -378,6 +379,7 @@ void xa_opt_quantized_conv2d_nhwc(
       output_zero_point,
       0, // out_multiplier (unused)
       0, // out_shift (unused)
+      optional<Tensor>(), // offset (unused)
       out);
 }
 
@@ -568,6 +570,7 @@ void quantized_conv2d_nhwc_per_tensor_out(
     int64_t output_zero_point,
     __ET_UNUSED int64_t out_multiplier,
     __ET_UNUSED int64_t out_shift,
+    const optional<Tensor>& offset,
     Tensor& out) {
   // Handle W8A16 heterogeneous type (int16_t activations, int8_t weights)
   if (out.scalar_type() == ::executorch::aten::ScalarType::Short &&
@@ -589,6 +592,7 @@ void quantized_conv2d_nhwc_per_tensor_out(
         output_zero_point,
         out_multiplier,
         out_shift,
+        offset,
         out);
     return;
   }

backends/cadence/hifi/operators/operators.h

@@ -158,6 +158,7 @@ void quantized_conv2d_nhwc_per_tensor_out(
     int64_t output_zero_point,
     int64_t out_multiplier,
     int64_t out_shift,
+    const ::executorch::aten::optional<::executorch::aten::Tensor>& offset,
     ::executorch::aten::Tensor& out);
 
 ::executorch::aten::Tensor& cat_out(

backends/cadence/vision/operators/op_quantized_conv_out.cpp

@@ -582,6 +582,7 @@ void quantized_conv2d_nhwc_per_tensor_out(
     int64_t output_zero_point,
     int64_t out_multiplier,
     int64_t out_shift,
+    ET_UNUSED const ::executorch::aten::optional<Tensor>& offset,
     Tensor& out) {
   quantized_conv_per_tensor_out(
       ctx,