Fused quant convolution kernel (#19491)

backends/cadence/fused_quant/op_convolution.cpp

@@ -0,0 +1,284 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/cadence/fused_quant/op_convolution.h>
+#include <executorch/backends/cadence/fused_quant/quant_utils.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace cadence {
+namespace fused_quant {
+namespace native {
+
+using executorch::aten::IntArrayRef;
+using executorch::aten::optional;
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+using executorch::runtime::KernelRuntimeContext;
+
+namespace {
+
+void conv2d_kernel(
+    const float* inp,
+    const float* weight,
+    const float* bias,
+    float* out,
+    int64_t N,
+    int64_t C_in,
+    int64_t H_in,
+    int64_t W_in,
+    int64_t C_out,
+    int64_t kH,
+    int64_t kW,
+    int64_t stride_h,
+    int64_t stride_w,
+    int64_t pad_h,
+    int64_t pad_w,
+    int64_t dil_h,
+    int64_t dil_w,
+    int64_t groups,
+    int64_t H_out,
+    int64_t W_out) {
+  int64_t C_in_per_group = C_in / groups;
+  int64_t C_out_per_group = C_out / groups;
+
+  for (int64_t n = 0; n < N; ++n) {
+    for (int64_t g = 0; g < groups; ++g) {
+      for (int64_t oc = 0; oc < C_out_per_group; ++oc) {
+        int64_t oc_global = g * C_out_per_group + oc;
+        for (int64_t oh = 0; oh < H_out; ++oh) {
+          for (int64_t ow = 0; ow < W_out; ++ow) {
+            float sum = bias ? bias[oc_global] : 0.0f;
+            for (int64_t ic = 0; ic < C_in_per_group; ++ic) {
+              int64_t ic_global = g * C_in_per_group + ic;
+              for (int64_t kh = 0; kh < kH; ++kh) {
+                for (int64_t kw = 0; kw < kW; ++kw) {
+                  int64_t ih = oh * stride_h - pad_h + kh * dil_h;
+                  int64_t iw = ow * stride_w - pad_w + kw * dil_w;
+                  if (ih >= 0 && ih < H_in && iw >= 0 && iw < W_in) {
+                    float inp_val =
+                        inp[((n * C_in + ic_global) * H_in + ih) * W_in + iw];
+                    float w_val = weight
+                        [((oc_global * C_in_per_group + ic) * kH + kh) * kW +
+                         kw];
+                    sum += inp_val * w_val;
+                  }
+                }
+              }
+            }
+            out[((n * C_out + oc_global) * H_out + oh) * W_out + ow] = sum;
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace
+
+Tensor& convolution_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& inp,
+    const Tensor& weight,
+    const optional<Tensor>& bias,
+    // inp qparams
+    const optional<Tensor>& inp_scale,
+    const optional<Tensor>& inp_zero_point,
+    ScalarType inp_dtype,
+    int64_t inp_quant_min,
+    int64_t inp_quant_max,
+    optional<int64_t> inp_axis,
+    // weight qparams
+    const optional<Tensor>& weight_scale,
+    const optional<Tensor>& weight_zero_point,
+    ScalarType weight_dtype,
+    int64_t weight_quant_min,
+    int64_t weight_quant_max,
+    optional<int64_t> weight_axis,
+    // bias qparams
+    const optional<Tensor>& bias_scale,
+    const optional<Tensor>& bias_zero_point,
+    ScalarType bias_dtype,
+    int64_t bias_quant_min,
+    int64_t bias_quant_max,
+    optional<int64_t> bias_axis,
+    // out qparams
+    const optional<Tensor>& out_scale,
+    const optional<Tensor>& out_zero_point,
+    ScalarType out_dtype,
+    int64_t out_quant_min,
+    int64_t out_quant_max,
+    optional<int64_t> out_axis,
+    // conv params
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    int64_t groups,
+    Tensor& out) {
+  // Extract dimensions from input tensor [N, C_in, H_in, W_in]
+  int64_t N = inp.size(0);
+  int64_t C_in = inp.size(1);
+  int64_t H_in = inp.size(2);
+  int64_t W_in = inp.size(3);
+
+  // Extract dimensions from weight tensor [C_out, C_in/groups, kH, kW]
+  int64_t C_out = weight.size(0);
+  int64_t kH = weight.size(2);
+  int64_t kW = weight.size(3);
+
+  int64_t stride_h = stride[0];
+  int64_t stride_w = stride[1];
+  int64_t pad_h = padding[0];
+  int64_t pad_w = padding[1];
+  int64_t dil_h = dilation[0];
+  int64_t dil_w = dilation[1];
+
+  int64_t H_out = (H_in + 2 * pad_h - dil_h * (kH - 1) - 1) / stride_h + 1;
+  int64_t W_out = (W_in + 2 * pad_w - dil_w * (kW - 1) - 1) / stride_w + 1;
+
+  int64_t inp_numel = inp.numel();
+  int64_t weight_numel = weight.numel();
+  int64_t out_numel = N * C_out * H_out * W_out;
+
+  bool inp_quantized = inp_scale.has_value();
+  bool weight_quantized = weight_scale.has_value();
+  bool bias_quantized = bias_scale.has_value();
+  bool out_quantized = out_scale.has_value();
+
+  // Dequantize input if quantized
+  std::vector<float> inp_buf;
+  const float* const inp_float = [&]() -> const float* {
+    if (!inp_quantized) {
+      return inp.const_data_ptr<float>();
+    }
+    inp_buf.resize(inp_numel);
+    QParams qp = extract_qparams(
+        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp_axis, inp);
+    FUSED_QUANT_DTYPE_SWITCH(
+        inp.scalar_type(),
+        scalar_t,
+        dequantize_buffer(
+            inp.const_data_ptr<scalar_t>(), inp_buf.data(), inp_numel, qp);)
+    return inp_buf.data();
+  }();
+
+  // Dequantize weight if quantized
+  std::vector<float> weight_buf;
+  const float* const weight_float = [&]() -> const float* {
+    if (!weight_quantized) {
+      return weight.const_data_ptr<float>();
+    }
+    weight_buf.resize(weight_numel);
+    QParams qp = extract_qparams(
+        weight_scale,
+        weight_zero_point,
+        weight_quant_min,
+        weight_quant_max,
+        weight_axis,
+        weight);
+    FUSED_QUANT_DTYPE_SWITCH(weight.scalar_type(),
+                             scalar_t,
+                             dequantize_buffer(
+                                 weight.const_data_ptr<scalar_t>(),
+                                 weight_buf.data(),
+                                 weight_numel,
+                                 qp);)
+    return weight_buf.data();
+  }();
+
+  // Dequantize bias if present and quantized
+  std::vector<float> bias_buf;
+  const float* bias_float = nullptr;
+  if (bias.has_value()) {
+    const Tensor& bias_tensor = bias.value();
+    if (bias_quantized) {
+      int64_t bias_numel = bias_tensor.numel();
+      bias_buf.resize(bias_numel);
+      QParams qp = extract_qparams(
+          bias_scale,
+          bias_zero_point,
+          bias_quant_min,
+          bias_quant_max,
+          bias_axis,
+          bias_tensor);
+      FUSED_QUANT_DTYPE_SWITCH(bias_tensor.scalar_type(),
+                               scalar_t,
+                               dequantize_buffer(
+                                   bias_tensor.const_data_ptr<scalar_t>(),
+                                   bias_buf.data(),
+                                   bias_numel,
+                                   qp);)
+      bias_float = bias_buf.data();
+    } else {
+      bias_float = bias_tensor.const_data_ptr<float>();
+    }
+  }
+
+  // Run convolution
+  if (out_quantized) {
+    std::vector<float> result_float(out_numel);
+    conv2d_kernel(
+        inp_float,
+        weight_float,
+        bias_float,
+        result_float.data(),
+        N,
+        C_in,
+        H_in,
+        W_in,
+        C_out,
+        kH,
+        kW,
+        stride_h,
+        stride_w,
+        pad_h,
+        pad_w,
+        dil_h,
+        dil_w,
+        groups,
+        H_out,
+        W_out);
+
+    QParams qp = extract_qparams(
+        out_scale, out_zero_point, out_quant_min, out_quant_max, out_axis, out);
+    FUSED_QUANT_DTYPE_SWITCH(out.scalar_type(),
+                             scalar_t,
+                             quantize_buffer(
+                                 result_float.data(),
+                                 out.mutable_data_ptr<scalar_t>(),
+                                 out_numel,
+                                 qp);)
+  } else {
+    conv2d_kernel(
+        inp_float,
+        weight_float,
+        bias_float,
+        out.mutable_data_ptr<float>(),
+        N,
+        C_in,
+        H_in,
+        W_in,
+        C_out,
+        kH,
+        kW,
+        stride_h,
+        stride_w,
+        pad_h,
+        pad_w,
+        dil_h,
+        dil_w,
+        groups,
+        H_out,
+        W_out);
+  }
+
+  return out;
+}
+
+} // namespace native
+} // namespace fused_quant
+} // namespace cadence

backends/cadence/fused_quant/op_convolution.h

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace cadence {
+namespace fused_quant {
+namespace native {
+
+executorch::aten::Tensor& convolution_out(
+    executorch::runtime::KernelRuntimeContext& ctx,
+    const executorch::aten::Tensor& inp,
+    const executorch::aten::Tensor& weight,
+    const executorch::aten::optional<executorch::aten::Tensor>& bias,
+    // inp qparams (6)
+    const executorch::aten::optional<executorch::aten::Tensor>& inp_scale,
+    const executorch::aten::optional<executorch::aten::Tensor>& inp_zero_point,
+    executorch::aten::ScalarType inp_dtype,
+    int64_t inp_quant_min,
+    int64_t inp_quant_max,
+    executorch::aten::optional<int64_t> inp_axis,
+    // weight qparams (6)
+    const executorch::aten::optional<executorch::aten::Tensor>& weight_scale,
+    const executorch::aten::optional<executorch::aten::Tensor>&
+        weight_zero_point,
+    executorch::aten::ScalarType weight_dtype,
+    int64_t weight_quant_min,
+    int64_t weight_quant_max,
+    executorch::aten::optional<int64_t> weight_axis,
+    // bias qparams (6)
+    const executorch::aten::optional<executorch::aten::Tensor>& bias_scale,
+    const executorch::aten::optional<executorch::aten::Tensor>& bias_zero_point,
+    executorch::aten::ScalarType bias_dtype,
+    int64_t bias_quant_min,
+    int64_t bias_quant_max,
+    executorch::aten::optional<int64_t> bias_axis,
+    // out qparams (6)
+    const executorch::aten::optional<executorch::aten::Tensor>& out_scale,
+    const executorch::aten::optional<executorch::aten::Tensor>& out_zero_point,
+    executorch::aten::ScalarType out_dtype,
+    int64_t out_quant_min,
+    int64_t out_quant_max,
+    executorch::aten::optional<int64_t> out_axis,
+    // conv params
+    executorch::aten::IntArrayRef stride,
+    executorch::aten::IntArrayRef padding,
+    executorch::aten::IntArrayRef dilation,
+    int64_t groups,
+    executorch::aten::Tensor& out);
+
+} // namespace native
+} // namespace fused_quant
+} // namespace cadence