add quantized operators to match Caffe2 QNN pack

Author: snoopyisadog

include/caffe/layers/eltwise_layer.hpp

@@ -50,6 +50,7 @@ class EltwiseLayer : public Layer<Dtype> {
   double output_scale_;  //CUSTOMIZATION
   int output_zero_point_; //CUSTOMIZATION
   Dtype saturate_; //CUSTOMIZATION
+  int quantize_method_; //CUSTOMIZATION
   int outer_dim_, inner_dim_, eltwise_dim_, dim_; //CUSTOMIZATION
   int axis_; //CUSTOMIZATION
   Blob<Dtype> eltwise_multiplier_;

src/caffe/layers/conv_layer.cpp

@@ -155,12 +155,20 @@ void ConvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
           }
         }
       }
-      else { // quantize_method_ == PoolingParameter_QuantizeMethod_ONNX
+      else if (quantize_method == ConvolutionParameter_QuantizeMethod_ONNX) {
         float onnx_scale = (float) input_scale * (float) weight_scale / (float) output_scale;
         for (int k = 0; k < count_t; ++k) {
           top_data[k] = std::rint(top_data[k] * onnx_scale);
         }
       }
+      else { // Caffe2
+        float out_scal = (float)input_scale * weight_scale;
+        out_scal /= output_scale;
+        int q_shift;
+        // Caffe2 uses float; cast to double to fit the function-call
+        int q_scal = tfl_QuantizeMultiplier((double)out_scal, &q_shift);
+        MultiplyByQuantizedMultiplierVR(count_t, top_data, q_scal, q_shift, 2);
+      }
     }
 
     if (shift_output) {

src/caffe/layers/eltwise_layer.cpp

@@ -43,6 +43,7 @@ void EltwiseLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
   output_scale_ = this->layer_param_.eltwise_param().output_scale();
   output_zero_point_ = this->layer_param_.eltwise_param().output_zero_point();
   saturate_ = this->layer_param_.eltwise_param().saturate();
+  quantize_method_ = this->layer_param_.eltwise_param().quantize_method();
 
   //<--CUSTOMIZATION, for broadcasting
   const EltwiseParameter& param = this->layer_param_.eltwise_param();
@@ -100,13 +101,6 @@ void EltwiseLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
   //CUSTOMIZATION-->
 }
 
-template <typename Dtype>
-int affine_and_shift(const Dtype x, const int zp_in, const double mul, const int shift) {
-  int r = (int) std::round((x - zp_in) * mul);
-  r = r << shift;
-  return r;
-}
-
 template <typename Dtype>
 void tflite_add_kernel(const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top,
   const vector<double> &input_scale, const vector<int> &input_zero_point,
@@ -151,17 +145,18 @@ void tflite_add_kernel(const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dty
   }
 }
 
-typedef double Stype; // scale type
+typedef float Stype; // scale type; Caffe2 use float as floating-point representation (of tensors and scales)
 template <typename Dtype>
 void caffe2_int8add_kernel(const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top,
   const vector<double> &input_scale, const vector<int> &input_zero_point,
   const double &output_scale, const int &output_zero_point) {
   // refer to https://github.com/pytorch/pytorch/pull/14089#issuecomment-439545562
-  Stype max_scale = std::max(input_scale[0], input_scale[1]) / output_scale;
+  Stype in_s0 = input_scale[0], in_s1 = input_scale[1], out_s = output_scale;
+  Stype max_scale = std::max(in_s0, in_s1) / out_s;
   const int max_22bits = 1 << 21;
   int shift = 0;
-  Stype a_multiplier = input_scale[0] / output_scale;
-  Stype b_multiplier = input_scale[1] / output_scale;
+  Stype a_multiplier = in_s0 / out_s;
+  Stype b_multiplier = in_s1 / out_s;
   while (max_scale < max_22bits) {
     // the result will be 2^22 <= max_scale < 2^23, cast to integer it will occupy 22 bits
     max_scale *= 2;
@@ -252,8 +247,11 @@ void EltwiseLayer<Dtype>::Forward_cpu(
   case EltwiseParameter_EltwiseOp_SUM:
     if (is_quant) {
       // introduce custom computation
-      //caffe2_int8add_kernel(bottom, top, input_scale_, input_zero_point_, output_scale_, output_zero_point_);
-      tflite_add_kernel(bottom, top, input_scale_, input_zero_point_, output_scale_, output_zero_point_);
+      if (quantize_method_ == ConvolutionParameter_QuantizeMethod_tflite) {
+        tflite_add_kernel(bottom, top, input_scale_, input_zero_point_, output_scale_, output_zero_point_);
+      } else {
+        caffe2_int8add_kernel(bottom, top, input_scale_, input_zero_point_, output_scale_, output_zero_point_);
+      }
       break;
     }
     caffe_set(count, Dtype(0), top_data);

src/caffe/layers/pooling_layer.cpp

@@ -461,14 +461,24 @@ void PoolingLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
                   acc += output_zero_point_;
                 }
                 top_data[ph * pooled_width_ + pw] = acc;
-              } else { // quantize_method_ == PoolingParameter_QuantizeMethod_ONNX
+              } else if (quantize_method_ == PoolingParameter_QuantizeMethod_ONNX) {
                 float scale = (float) input_scale_ / ((float)output_scale_ * (float) pool_size);
                 Dtype acc = top_data[ph * pooled_width_ + pw];
                 acc -= input_zero_point_ * pool_size;
                 acc = std::rint(acc * scale);
                 acc += output_zero_point_;
                 top_data[ph * pooled_width_ + pw] = acc;
               }
+              else { // Caffe2
+                // https://github.com/pytorch/QNNPACK/blob/7d2a4e9931a82adc3814275b6219a03e24e36b4c/src/average-pooling.c#L176-L179
+                float scale = (float) input_scale_ / ((float)output_scale_ * (float) pool_size);
+                Dtype acc = top_data[ph * pooled_width_ + pw];
+                acc -= input_zero_point_ * pool_size;
+                //acc = std::round(acc / pool_size);
+                acc = std::round(acc * scale);
+                acc += output_zero_point_;
+                top_data[ph * pooled_width_ + pw] = acc;
+              }
             }
             else {
               top_data[ph * pooled_width_ + pw] /= pool_size;

src/caffe/layers/power_layer.cpp

@@ -41,8 +41,12 @@ void PowerLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
       for (int i = 0; i < count_t; ++i) {
         top_data[i] = std::round(top_data[i]);
       }
-    } else { // QuantizeMethod_ONNX
+    } else if (quantize_method_ == ConvolutionParameter_QuantizeMethod_ONNX) { // QuantizeMethod_ONNX
       caffe_cpu_round(count, top_data);
+    } else {
+      for (int i = 0; i < count_t; ++i) {
+        top_data[i] = std::nearbyint(top_data[i]);
+      }
     }
   }
   if (shift_ != Dtype(0)) {