[MKLDNN] add quantized sum

src/operator/quantization/mkldnn/mkldnn_quantized_sum-inl.h

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_QUANTIZED_SUM_INL_H_
+#define MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_QUANTIZED_SUM_INL_H_
+#if MXNET_USE_MKLDNN == 1
+
+#include <utility>
+#include <vector>
+#include <string>
+#include "../../tensor/elemwise_unary_op.h"
+
+namespace mxnet {
+namespace op {
+
+struct RequantizeSumParam : public dmlc::Parameter<RequantizeSumParam> {
+  dmlc::optional<float> min_calib_range;  // min float value calculated from calibration dataset
+  dmlc::optional<float> max_calib_range;  // max float value calculated from calibration dataset
+  DMLC_DECLARE_PARAMETER(RequantizeSumParam) {
+    DMLC_DECLARE_FIELD(min_calib_range)
+    .set_default(dmlc::optional<float>())
+    .describe("The minimum scalar value in the form of float32 obtained "
+              "through calibration. If present, it will be used to requantize the "
+              "int8 output data.");
+    DMLC_DECLARE_FIELD(max_calib_range)
+    .set_default(dmlc::optional<float>())
+    .describe("The maximum scalar value in the form of float32 obtained "
+              "through calibration. If present, it will be used to requantize the "
+              "int8 output data.");
+  }
+};
+
+namespace quantized_sum_enum {
+enum QuantizedSumOutputs { kOut, kMin, kMax };
+enum QuantizedSumInputs { kDataA, kDataB, kAMin, kAMax, kBMin, kBMax};
+}
+
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // MXNET_USE_MKLDNN == 1
+#endif  // MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_QUANTIZED_SUM_INL_H_

src/operator/quantization/mkldnn/mkldnn_quantized_sum.cc

@@ -0,0 +1,206 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2019 by Contributors
+ * \file mkldnn_quantized_sum.cc
+ * \brief
+ */
+
+#if MXNET_USE_MKLDNN == 1
+#include "./mkldnn_quantized_sum-inl.h"
+#include "../../nn/mkldnn/mkldnn_ops-inl.h"
+#include "../../nn/mkldnn/mkldnn_base-inl.h"
+#include "../quantization_utils.h"
+
+namespace mxnet {
+namespace op {
+
+DMLC_REGISTER_PARAMETER(RequantizeSumParam);
+
+static float GetScale(const NDArray& data, float min, float max) {
+  auto data_range = (data.dtype() == mshadow::kInt8) ? kInt8Range : kUint8Range;
+  return data_range / MaxAbs(min, max);
+}
+
+static void MKLDNNQuantizedSumForward(const nnvm::NodeAttrs& attrs, const OpContext& ctx,
+                                         const std::vector<NDArray>& in_data,
+                                         const std::vector<OpReqType>& req,
+                                         const std::vector<NDArray>& out_data) {
+  const RequantizeSumParam& params = nnvm::get<RequantizeSumParam>(attrs.parsed);
+  // A, B, A_min, A_max, B_min, B_max
+  CHECK_EQ(in_data.size(), 6U);
+  // C, C_min, C_max
+  CHECK_EQ(out_data.size(), 3U);
+
+  // Collect data min,max,absmax
+  const float dataA_min = in_data[quantized_sum_enum::kAMin].data().dptr<float>()[0];
+  const float dataB_min = in_data[quantized_sum_enum::kBMin].data().dptr<float>()[0];
+  const float dataA_max = in_data[quantized_sum_enum::kAMax].data().dptr<float>()[0];
+  const float dataB_max = in_data[quantized_sum_enum::kBMax].data().dptr<float>()[0];
+  const float dataA_absmax = MaxAbs(dataA_min, dataA_max);
+  const float dataB_absmax = MaxAbs(dataB_min, dataB_max);
+
+  auto dataA_mem  = in_data[quantized_sum_enum::kDataA].GetMKLDNNData();
+  auto dataB_mem  = in_data[quantized_sum_enum::kDataB].GetMKLDNNData();
+  const bool dataA_int8 = (in_data[quantized_sum_enum::kDataA].dtype() == mshadow::kInt8)
+                          ? true : false;
+  const size_t dataA_range = dataA_int8 ? kInt8Range : kUint8Range;
+
+  const float A_scale = GetScale(in_data[quantized_sum_enum::kDataA], dataA_min, dataA_max);
+  const float B_scale = GetScale(in_data[quantized_sum_enum::kDataB], dataB_min, dataB_max);
+  // rescaled_mem is for reorder mkldnn memory
+  std::shared_ptr<mkldnn::memory> rescaled_mem;
+  // output default set as int32
+  size_t output_data_range = kInt32Range;
+  auto output_data_type = mkldnn::memory::s32;
+  // dataA && dataB are uint8
+  if (out_data[quantized_sum_enum::kDataA].dtype() == mshadow::kInt8) {
+    output_data_range = kInt8Range;
+    output_data_type = mkldnn::memory::s8;
+  } else if (out_data[quantized_sum_enum::kDataA].dtype() == mshadow::kUint8) {
+    output_data_range = kUint8Range;
+    output_data_type = mkldnn::memory::u8;
+  }
+
+  float output_min = 0;
+  float output_max = 0;
+  float out_data_scale = 0;
+  if (params.max_calib_range.has_value() && params.min_calib_range.has_value()) {
+    output_min = params.min_calib_range.value();
+    output_max = params.max_calib_range.value();
+    out_data_scale = output_data_range/MaxAbs(output_min, output_max);
+  } else {
+    output_max = dataA_absmax + dataB_absmax;
+    output_min = 0 - output_max;
+  }
+  // 2: scale 0 for dataA, scale 1 for data B
+  const int scales_num = 2;
+  std::vector<float> scales;
+  scales.reserve(scales_num);
+  if (in_data[quantized_sum_enum::kDataA].dtype() != in_data[quantized_sum_enum::kDataB].dtype()) {
+    auto s8_pd = (dataA_int8 == true)
+                 ? dataA_mem->get_primitive_desc()
+                 : dataB_mem->get_primitive_desc();
+    rescaled_mem = std::make_shared<mkldnn::memory>(s8_pd);
+    float u8_reorder_scale = 0;
+    if (params.max_calib_range.has_value() && params.min_calib_range.has_value()) {
+      if (dataA_int8 == true) {
+        u8_reorder_scale = out_data_scale/B_scale;
+        scales.push_back(out_data_scale/A_scale);
+        scales.push_back(1);
+      } else {
+        u8_reorder_scale = out_data_scale/A_scale;
+        scales.push_back(1);
+        scales.push_back(out_data_scale/B_scale);
+      }
+    } else {
+      // x*dataA_absmax/dataA_range = y*(dataA_absmax+dataB_absmax)/output_range
+      if (dataA_int8 == true) {
+        u8_reorder_scale = dataB_absmax*output_data_range
+                           /((dataA_absmax + dataB_absmax)*kUint8Range);
+        scales.push_back(dataA_absmax*output_data_range
+                         /((dataA_absmax + dataB_absmax)*dataA_range));
+        scales.push_back(1);
+      } else {
+        u8_reorder_scale = dataA_absmax*output_data_range
+                           /((dataA_absmax + dataB_absmax)*dataA_range);
+        scales.push_back(1);
+        scales.push_back(dataB_absmax*output_data_range
+                         /((dataA_absmax + dataB_absmax)*kInt8Range));
+      }
+    }
+    std::vector<float> reorder_scale = {u8_reorder_scale};
+    primitive_attr reorder_attr;
+    reorder_attr.set_int_output_round_mode(round_mode::round_nearest);
+    reorder_attr.set_output_scales(0, reorder_scale);
+    auto u8_mem = (dataA_int8 == true) ? dataB_mem : dataA_mem;
+    const auto reorder_pd = mkldnn::reorder::primitive_desc(u8_mem->get_primitive_desc(),
+                                                            s8_pd,
+                                                            reorder_attr);
+    MKLDNNStream::Get()->RegisterPrim(mkldnn::reorder(reorder_pd, *u8_mem, *rescaled_mem));
+
+    if (dataA_int8 == true) {
+      dataB_mem = rescaled_mem.get();
+    } else {
+      dataA_mem = rescaled_mem.get();
+    }
+  } else {
+    // same data type and has same data range
+    if (params.max_calib_range.has_value() && params.min_calib_range.has_value()) {
+      scales.push_back(out_data_scale/A_scale);
+      scales.push_back(out_data_scale/B_scale);
+    } else {
+      scales.push_back(dataA_absmax*output_data_range/((dataA_absmax + dataB_absmax)*dataA_range));
+      scales.push_back(dataB_absmax*output_data_range/((dataA_absmax + dataB_absmax)*dataA_range));
+    }
+  }
+
+  std::vector<mkldnn::primitive::at> in_prims;
+  std::vector<mkldnn::memory::primitive_desc> in_pds;
+  in_prims.push_back(*dataA_mem);
+  in_prims.push_back(*dataB_mem);
+  in_pds.push_back(dataA_mem->get_primitive_desc());
+  in_pds.push_back(dataB_mem->get_primitive_desc());
+  size_t i_ndim = in_data[quantized_sum_enum::kDataA].shape().ndim();
+  mkldnn::memory::dims i_dims = mkldnn::memory::dims(i_ndim);
+  for (size_t i = 0; i < i_ndim; i++) {
+    i_dims[i] = static_cast<int>(in_data[quantized_sum_enum::kDataA].shape()[i]);
+  }
+  mkldnn::memory::format i_fmt = static_cast<mkldnn::memory::format>(
+                                   in_pds[quantized_sum_enum::kDataA].desc().data.format);
+  auto output_desc = memory::desc(i_dims, output_data_type, i_fmt);
+  mkldnn::sum::primitive_desc pdesc(output_desc, scales, in_pds);
+  auto mem = CreateMKLDNNMem(out_data[quantized_sum_enum::kOut],
+                             pdesc.dst_primitive_desc(),
+                             req[0],
+                             &in_data[0]);
+  MKLDNNStream *stream = MKLDNNStream::Get();
+  stream->RegisterPrim(mkldnn::sum(pdesc, in_prims, *mem.second));
+  CommitOutput(out_data[quantized_sum_enum::kOut], mem);
+  stream->Submit();
+
+  out_data[quantized_sum_enum::kMin].data().dptr<float>()[0] = output_min;
+  out_data[quantized_sum_enum::kMax].data().dptr<float>()[0] = output_max;
+}
+
+inline static bool SumStorageType(const nnvm::NodeAttrs& attrs, const int dev_mask,
+                                     DispatchMode* dispatch_mode, std::vector<int>* in_attrs,
+                                     std::vector<int>* out_attrs) {
+  // A, B, A_min, A_max, B_min, B_max
+  CHECK_EQ(in_attrs->size(), 6U);
+  // C, C_min, C_max
+  CHECK_EQ(out_attrs->size(), 3U);
+
+  return MKLDNNStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs, out_attrs);
+}
+
+NNVM_REGISTER_OP(_contrib_quantized_sum)
+.set_attr<FInferStorageType>("FInferStorageType", SumStorageType)
+.set_attr<FComputeEx>("FComputeEx<cpu>", MKLDNNQuantizedSumForward)
+.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& n) {
+  return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+})
+.set_attr<bool>("TIsMKLDNN", true)
+.set_attr_parser(ParamParser<RequantizeSumParam>)
+.add_arguments(RequantizeSumParam::__FIELDS__());
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // MXNET_USE_MKLDNN == 1

src/operator/quantization/quantization_utils.h

@@ -34,6 +34,7 @@ namespace op {
 
 static const size_t kUint8Range = 255;
 static const size_t kInt8Range = 127;
+static const size_t kInt32Range = 0x7fffffff;
 
 template<typename T>
 MSHADOW_XINLINE int Sign(T val) {