mkldnn_quantizer.cc 25.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed 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.

#include "paddle/fluid/inference/api/mkldnn_quantizer.h"
16

17
#include <algorithm>
18
#include <limits>
19 20 21 22
#include <map>
#include <numeric>
#include <unordered_map>
#include <utility>
23

24 25 26 27 28 29 30 31
#include "paddle/fluid/framework/eigen.h"
#include "paddle/fluid/framework/ir/fuse_pass_base.h"
#include "paddle/fluid/framework/ir/graph.h"
#include "paddle/fluid/framework/ir/pass.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/type_defs.h"
#include "paddle/fluid/inference/analysis/analyzer.h"
#include "paddle/fluid/inference/api/analysis_predictor.h"
32
#include "paddle/fluid/platform/mkldnn_helper.h"
33 34
#include "paddle/phi/common/place.h"
#include "paddle/utils/string/pretty_log.h"
35 36 37

namespace paddle {

38
using framework::Variable;
39
using framework::ir::Graph;
40
using phi::CPUPlace;
41 42
using ConstEigenVectorArrayMap =
    Eigen::Map<const Eigen::Array<float, Eigen::Dynamic, 1>>;
M
Michał Gallus 已提交
43 44 45 46 47
using EigenMatrixDoubleArray =
    Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
using EigenMatrixArray =
    Eigen::Array<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
using ConstEigenMatrixArrayMap = Eigen::Map<const EigenMatrixArray>;
48
using string::PrettyLogH1;
49
using VariableNameMap = std::map<std::string, std::vector<std::string>>;
50
static phi::DenseTensor CreateScaleTensor(int64_t channels_num = 1);
51

52
static void check_var(const Variable* var, const std::string& var_name) {
53 54 55
  PADDLE_ENFORCE_NOT_NULL(
      var,
      platform::errors::PreconditionNotMet("%s is not in the scope", var_name));
56
  PADDLE_ENFORCE_EQ(
57
      var->IsType<phi::DenseTensor>(),
58
      true,
59 60 61
      platform::errors::PreconditionNotMet("Only support lod tensor now."));
}

62
static void check_tensor(const phi::DenseTensor& tensor) {
63
  PADDLE_ENFORCE_GT(
64 65
      tensor.dims().size(),
      0,
66
      platform::errors::InvalidArgument("Tensor dimension is empty."));
67 68
}

69 70
void AnalysisPredictor::MkldnnQuantizer::CalculateScalesForRNNWeights(
    const paddle::framework::OpDesc* op, bool gru) {
71 72 73 74 75 76 77 78 79
  const auto& wx_names = op->Input("WeightX");
  const auto& wh_names = op->Input("WeightH");
  for (size_t i = 0; i < wx_names.size(); ++i) {
    const auto& wx_name = wx_names[i];
    const auto& wh_name = wh_names[i];
    auto* wx_var = predictor_.sub_scope_->FindVar(wx_name);
    auto* wh_var = predictor_.sub_scope_->FindVar(wh_name);
    check_var(wx_var, wx_name);
    check_var(wh_var, wh_name);
80 81
    phi::DenseTensor* wx_tensor = wx_var->GetMutable<phi::DenseTensor>();
    phi::DenseTensor* wh_tensor = wh_var->GetMutable<phi::DenseTensor>();
82 83 84 85 86
    if (gru) {
      scales_[wx_name] = GetMaxChGRUScalingFactor(*wx_tensor, *wh_tensor);
    } else {
      scales_[wx_name] = GetMaxChLSTMScalingFactor(*wx_tensor, *wh_tensor);
    }
87 88 89 90 91 92
  }
}

void AnalysisPredictor::MkldnnQuantizer::CalculateScalesForOpInputs(
    const paddle::framework::OpDesc* op) {
  if (op->Type() == "fusion_gru" || op->Type() == "multi_gru") {
93 94 95
    CalculateScalesForRNNWeights(op, true);
  } else if (op->Type() == "fusion_lstm") {
    CalculateScalesForRNNWeights(op, false);
96 97 98 99 100 101 102
  }
  for (auto const& input : op->Inputs()) {
    for (const auto& var_name : input.second) {
      // skip if scale already computed
      if (scales_.find(var_name) != scales_.end()) continue;
      auto* var = predictor_.sub_scope_->FindVar(var_name);
      check_var(var, var_name);
103
      phi::DenseTensor* var_tensor = var->GetMutable<phi::DenseTensor>();
104 105
      // force unsigned type if already know it
      bool is_unsigned = false;
106 107
      CalculateSingleScale(
          op->Type(), input.first, var_name, *var_tensor, is_unsigned);
108 109 110 111 112 113 114 115 116 117 118 119
    }
  }
}

void AnalysisPredictor::MkldnnQuantizer::CalculateScalesForOpOutputs(
    const paddle::framework::OpDesc* op) {
  for (auto const& output : op->Outputs()) {
    for (const auto& var_name : output.second) {
      // skip if scale already computed
      if (scales_.find(var_name) != scales_.end()) continue;
      auto* var = predictor_.sub_scope_->FindVar(var_name);
      check_var(var, var_name);
120
      phi::DenseTensor* var_tensor = var->GetMutable<phi::DenseTensor>();
121 122 123
      // force unsigned type if already know it
      bool is_unsigned = false;
      bool compute_scale = true;
124
      if (op->Type() == "conv2d" || op->Type() == "fc") {
125 126 127 128 129 130 131
        // output of conv2d with relu must be unsigned
        std::string fuse_activation =
            op->GetAttrIfExists<std::string>("fuse_activation");
        is_unsigned = (fuse_activation == "relu" || fuse_activation == "relu6");
      } else if (op->Type() == "relu") {
        is_unsigned = true;
      } else if (op->Type() == "transpose2" || op->Type() == "reshape2" ||
132
                 op->Type() == "pool2d" || op->Type() == "nearest_interp" ||
P
Paulina Gacek 已提交
133
                 op->Type() == "nearest_interp_v2" || op->Type() == "split") {
134
        auto input_var_name = op->Input("X")[0];
135 136
        PADDLE_ENFORCE_NE(scales_.find(input_var_name),
                          scales_.end(),
137 138 139
                          platform::errors::PreconditionNotMet(
                              "Input scales must be calculated before the "
                              "output scales to infer if output is unsigned."));
Z
Zuza 已提交
140 141 142 143
        if (scales_.find(input_var_name) != scales_.end()) {
          scales_[var_name] = scales_[input_var_name];
        }
        compute_scale = false;
144
      } else if (op->Type() == "slice" || op->Type() == "shape") {
Z
Zuza 已提交
145
        auto input_var_name = op->Input("Input")[0];
146 147
        PADDLE_ENFORCE_NE(scales_.find(input_var_name),
                          scales_.end(),
Z
Zuza 已提交
148 149 150
                          platform::errors::PreconditionNotMet(
                              "Input scales must be calculated before the "
                              "output scales to infer if output is unsigned."));
151 152 153 154 155 156 157 158 159 160
        if (scales_.find(input_var_name) != scales_.end()) {
          scales_[var_name] = scales_[input_var_name];
        }
        compute_scale = false;
      } else if (op->Type() == "concat") {
        // output of ops with unsigned input must be unsigned
        is_unsigned = true;
        double min_scale = std::numeric_limits<double>::max();
        for (auto input_var_name : op->Input("X")) {
          PADDLE_ENFORCE_NE(
161 162
              scales_.find(input_var_name),
              scales_.end(),
163 164 165 166 167 168 169 170 171 172 173 174 175
              platform::errors::PreconditionNotMet(
                  "Input scales must be calculated before the "
                  "output scales to infer if output is unsigned."));
          is_unsigned = is_unsigned && scales_[input_var_name].first;
          min_scale = std::min(
              min_scale, scales_[input_var_name].second.data<double>()[0]);
        }
        auto scale_tensor = CreateScaleTensor();
        scale_tensor.data<double>()[0] = min_scale;
        scales_[var_name] = {is_unsigned, scale_tensor};
        compute_scale = false;
      }
      if (compute_scale) {
176 177
        CalculateSingleScale(
            op->Type(), output.first, var_name, *var_tensor, is_unsigned);
178 179 180 181 182
      }
    }
  }
}

183 184
bool AnalysisPredictor::MkldnnQuantizer::CalculateScales() {
  PrettyLogH1("--- Calculating scales for quantization");
185
  std::map<std::string, std::map<std::string, phi::DenseTensor>> gathered_data;
186
  for (const auto* op : predictor_.inference_program_->Block(0).AllOps()) {
187
    if (platform::HasOpINT8DataType(op)) {
188
      // handle inputs first to let is_unsigned be inferred for the outputs
189 190
      CalculateScalesForOpInputs(op);
      CalculateScalesForOpOutputs(op);
191 192 193 194 195 196
    }
  }
  return true;
}

void AnalysisPredictor::MkldnnQuantizer::CalculateSingleScale(
197 198 199
    const std::string& op_type_name,
    const std::string& conn_name,
    const std::string& var_name,
200
    const phi::DenseTensor& var_tensor,
201 202 203 204
    bool is_unsigned) {
  auto rule = qconfig_->scale_algo(op_type_name, conn_name);
  if (rule == ScaleAlgo::NONE) return;

205 206 207 208 209 210 211 212 213
  PADDLE_ENFORCE_GT(var_tensor.numel(),
                    0,
                    platform::errors::InvalidArgument(
                        "MkldnnQuantizer: phi::DenseTensor of variable %s for "
                        "quantization of op "
                        "%s of connection %s should not be empty.",
                        var_name,
                        op_type_name,
                        conn_name));
214 215 216 217 218 219

  switch (rule) {
    case ScaleAlgo::MAX:
      scales_[var_name] = GetMaxScalingFactor(var_tensor, is_unsigned);
      break;
    case ScaleAlgo::MAX_CH:
220 221
      scales_[var_name] = GetMaxChScalingFactor(var_tensor,
                                                is_unsigned,
M
Michał Gallus 已提交
222 223 224
                                                /*is_transposed*/ false);
      break;
    case ScaleAlgo::MAX_CH_T:
225 226
      scales_[var_name] = GetMaxChScalingFactor(var_tensor,
                                                is_unsigned,
M
Michał Gallus 已提交
227
                                                /*is_transposed*/ true);
228 229 230 231 232 233 234 235 236 237
      break;
    case ScaleAlgo::KL:
      scales_[var_name] = GetKLScalingFactor(var_tensor, is_unsigned);
      break;
    default:
      throw std::runtime_error(
          "MkldnnQuantizer: Unexpected ScaleAlgo specified.");
  }
}

238 239
static phi::DenseTensor CreateScaleTensor(int64_t channels_num) {
  phi::DenseTensor scale_tensor;
240 241 242 243 244
  scale_tensor.Resize({channels_num});
  scale_tensor.mutable_data<double>(CPUPlace());
  return scale_tensor;
}

245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273
std::vector<int> AnalysisPredictor::MkldnnQuantizer::ExpandQuantizedBins(
    std::vector<int> quantized_bins, std::vector<int> reference_bins) const {
  std::vector<int> expanded_quantized_bins(reference_bins.size(), 0);
  int num_merged_bins = reference_bins.size() / quantized_bins.size();
  int j_start = 0;
  int j_end = num_merged_bins;
  for (size_t idx = 0; idx < quantized_bins.size(); idx++) {
    int zero_count =
        std::count(&reference_bins[j_start], &reference_bins[j_end], 0);
    num_merged_bins = j_end - j_start;
    int avg_bin_ele;
    if (zero_count == num_merged_bins) {
      avg_bin_ele = 0;
    } else {
      avg_bin_ele = quantized_bins[idx] / (num_merged_bins - zero_count + 0.0);
    }
    for (int idx1 = j_start; idx1 < j_end; idx1++) {
      expanded_quantized_bins[idx1] =
          (reference_bins[idx1] == 0) ? 0 : avg_bin_ele;
    }
    j_start += num_merged_bins;
    j_end += num_merged_bins;
    if ((idx + 1) == quantized_bins.size() - 1) {
      j_end = reference_bins.size();
    }
  }
  return expanded_quantized_bins;
}

274
std::pair<bool, phi::DenseTensor>
275
AnalysisPredictor::MkldnnQuantizer::GetKLScalingFactor(
276
    const phi::DenseTensor& var_tensor, bool is_unsigned) const {
277 278
  ConstEigenVectorArrayMap eigen_tensor{
      var_tensor.data<float>(), var_tensor.numel(), 1};
279 280 281 282 283
  int precision_hist_num_bins = 2048;
  float max_val = eigen_tensor.maxCoeff();
  float min_val = eigen_tensor.minCoeff();
  bool is_positive = min_val >= 0.0f;
  if (is_unsigned)
284
    PADDLE_ENFORCE_EQ(
285 286
        is_positive,
        true,
287 288 289
        platform::errors::InvalidArgument(
            "Tensor is claimed to be unsigned, but its min value (%f) is < 0.0",
            min_val));
290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335

  int num_quantized_bins = 255;

  std::vector<int> hist;
  float bin_width;
  int starting_iter;
  int ending_iter = precision_hist_num_bins - 1;
  if (is_positive) {
    std::tie(hist, bin_width) =
        Histogram(var_tensor, min_val, max_val, precision_hist_num_bins);
    starting_iter = static_cast<int>(ending_iter * 0.7);
  } else {
    float th = std::max(std::abs(max_val), std::abs(min_val));
    std::tie(hist, bin_width) =
        Histogram(var_tensor, -th, th, precision_hist_num_bins);
    starting_iter = 0;
    if (std::abs(max_val) > std::abs(min_val)) {
      while (starting_iter < ending_iter) {
        if (hist[starting_iter] == 0) {
          ++starting_iter;
          continue;
        } else {
          break;
        }
      }
      starting_iter += static_cast<int>((ending_iter - starting_iter) * 0.6);
    } else {
      while (ending_iter > 0) {
        if (hist[ending_iter] == 0) {
          --ending_iter;
          continue;
        } else {
          break;
        }
      }
      starting_iter = static_cast<int>(0.6 * ending_iter);
    }
  }
  auto P_sum = eigen_tensor.size();
  int min_kl_divergence = 0;
  int min_kl_index = 0;
  bool kl_inited = false;
  for (int i = starting_iter; i <= ending_iter; i++) {
    std::vector<int> reference_distr_P(&hist[0], &hist[i]);
    auto outliers_count =
        std::accumulate(&hist[i], &hist[precision_hist_num_bins], 0);
336
    if (i <= 0 || reference_distr_P[i - 1] == 0) {
337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382
      continue;
    }
    reference_distr_P[i - 1] += outliers_count;
    auto reference_distr_bins = reference_distr_P;
    std::vector<int> candidate_distr_Q(&hist[0], &hist[i]);
    int num_merged_bins = i / num_quantized_bins;
    std::vector<int> candidate_distr_Q_quantized(num_quantized_bins, 0);
    int j_start = 0;
    int j_end = num_merged_bins;
    for (int idx = 0; idx < num_quantized_bins; idx++) {
      candidate_distr_Q_quantized[idx] = std::accumulate(
          &candidate_distr_Q[j_start], &candidate_distr_Q[j_end], 0);
      j_start += num_merged_bins;
      j_end += num_merged_bins;
      if ((idx + 1) == num_quantized_bins - 1) {
        j_end = i;
      }
    }
    candidate_distr_Q =
        ExpandQuantizedBins(candidate_distr_Q_quantized, reference_distr_bins);
    int Q_sum =
        std::accumulate(candidate_distr_Q.begin(), candidate_distr_Q.end(), 0);
    auto kl_divergence =
        SafeEntropy(reference_distr_P, P_sum, candidate_distr_Q, Q_sum);
    if (!kl_inited) {
      min_kl_divergence = kl_divergence;
      min_kl_index = i;
      kl_inited = true;
    } else if (kl_divergence < min_kl_divergence) {
      min_kl_divergence = kl_divergence;
      min_kl_index = i;
    } else {
    }
  }
  if (min_kl_index == 0) {
    while (starting_iter > 0) {
      if (hist[starting_iter] == 0) {
        starting_iter -= 1;
        continue;
      } else {
        break;
      }
    }
    min_kl_index = starting_iter;
  }

383
  phi::DenseTensor scale_tensor = CreateScaleTensor();
384
  scale_tensor.data<double>()[0] = 1.0 / ((min_kl_index + 0.5) * bin_width);
385 386 387 388

  return std::make_pair(is_unsigned, scale_tensor);
}

389
std::pair<bool, phi::DenseTensor>
390
AnalysisPredictor::MkldnnQuantizer::GetMaxScalingFactor(
391
    const phi::DenseTensor& var_tensor, bool is_unsigned) const {
392 393
  ConstEigenVectorArrayMap eigen_tensor{
      var_tensor.data<float>(), var_tensor.numel(), 1};
394 395 396
  float max_abs = eigen_tensor.abs().maxCoeff();
  float min_val = eigen_tensor.minCoeff();
  if (is_unsigned)
397
    PADDLE_ENFORCE_GE(
398 399
        min_val,
        0.0f,
400 401 402
        platform::errors::InvalidArgument(
            "Tensor is claimed to be unsigned, but its min value (%f) is < 0.0",
            min_val));
403

404
  phi::DenseTensor scale_tensor = CreateScaleTensor();
405
  scale_tensor.data<double>()[0] = 1.0 / max_abs;
406 407 408 409

  return std::make_pair(is_unsigned, scale_tensor);
}

410
std::pair<bool, phi::DenseTensor>
411
AnalysisPredictor::MkldnnQuantizer::GetMaxChScalingFactor(
412 413 414
    const phi::DenseTensor& var_tensor,
    bool is_unsigned,
    bool is_transposed) const {
415
  check_tensor(var_tensor);
416

417 418
  ConstEigenVectorArrayMap eigen_tensor{
      var_tensor.data<float>(), var_tensor.numel(), 1};
419 420
  float min_val = eigen_tensor.minCoeff();
  if (is_unsigned)
421
    PADDLE_ENFORCE_GE(
422 423
        min_val,
        0.0f,
424 425 426
        platform::errors::InvalidArgument(
            "Tensor is claimed to be unsigned, but its min value (%f) is < 0.0",
            min_val));
427

M
Michał Gallus 已提交
428 429
  auto dims = var_tensor.dims();
  constexpr int num_col_dims = 1;
430
  auto flattened_dims = phi::flatten_to_2d(dims, num_col_dims);
M
Michał Gallus 已提交
431 432
  ConstEigenMatrixArrayMap eigen_tensor_mat{
      var_tensor.data<float>(), flattened_dims[0], flattened_dims[1]};
433

M
Michał Gallus 已提交
434 435 436 437 438
  EigenMatrixDoubleArray scales;
  if (is_transposed) {
    scales = 1.0 / eigen_tensor_mat.cast<double>().abs().colwise().maxCoeff();
  } else {
    scales = 1.0 / eigen_tensor_mat.cast<double>().abs().rowwise().maxCoeff();
439
  }
M
Michał Gallus 已提交
440 441
  int output_channel_axis = is_transposed;
  int channels = dims[output_channel_axis];
442
  phi::DenseTensor scale_tensor = CreateScaleTensor(channels);
M
Michał Gallus 已提交
443 444
  auto* scale_ptr = scale_tensor.mutable_data<double>(CPUPlace());
  std::copy(scales.data(), scales.data() + scales.size(), scale_ptr);
445 446 447 448

  return std::make_pair(is_unsigned, scale_tensor);
}

449
std::pair<bool, phi::DenseTensor>
450
AnalysisPredictor::MkldnnQuantizer::GetMaxChGRUScalingFactor(
451 452
    const phi::DenseTensor& wx_tensor,
    const phi::DenseTensor& wh_tensor) const {
453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
  check_tensor(wx_tensor);
  check_tensor(wh_tensor);

  int OC = wh_tensor.dims()[0];
  std::vector<float> scale_ur(2 * OC);
  std::vector<float> scale_o(OC);

  for (int row_id = 0; row_id < wx_tensor.dims()[0]; row_id++) {
    for (int col_id = 0; col_id < 2 * OC; col_id++) {
      int idx = (row_id * wx_tensor.dims()[1]) + col_id;
      auto abs_value = std::abs(wx_tensor.data<float>()[idx]);
      if (row_id == 0) {
        scale_ur[col_id] = abs_value;
      } else {
        if (abs_value > scale_ur[col_id]) scale_ur[col_id] = abs_value;
      }
    }
  }

  for (int i = 0; i < 2 * OC * OC; i++) {
    int col_id = i % (2 * OC);
    auto abs_value = std::abs(wh_tensor.data<float>()[i]);
    if (abs_value > scale_ur[col_id]) scale_ur[col_id] = abs_value;
  }

  for (int row_id = 0; row_id < wx_tensor.dims()[0]; row_id++) {
    for (int col_id = 2 * OC; col_id < wx_tensor.dims()[1]; col_id++) {
      int idx = (row_id * wx_tensor.dims()[1]) + col_id;
      auto abs_value = std::abs(wx_tensor.data<float>()[idx]);
      if (row_id == 0) {
        scale_o[col_id % OC] = abs_value;
      } else {
        if (abs_value > scale_o[col_id]) scale_o[col_id % OC] = abs_value;
      }
    }
  }

  for (int i = 2 * OC * OC; i < OC * wh_tensor.dims()[1]; i++) {
    int col_id = i % OC;
    auto abs_value = std::abs(wh_tensor.data<float>()[i]);
    if (abs_value > scale_o[col_id]) scale_o[col_id] = abs_value;
  }
  scale_ur.insert(scale_ur.end(), scale_o.begin(), scale_o.end());
496 497 498
  transform(scale_ur.begin(), scale_ur.end(), scale_ur.begin(), [](float& c) {
    return 1 / c;
  });
499
  phi::DenseTensor scale_tensor = CreateScaleTensor(scale_ur.size());
500 501 502 503 504 505
  auto* scale_ptr = scale_tensor.mutable_data<double>(CPUPlace());
  std::copy(scale_ur.begin(), scale_ur.end(), scale_ptr);
  bool is_unsigned = false;
  return std::make_pair(is_unsigned, scale_tensor);
}

506
std::pair<bool, phi::DenseTensor>
507
AnalysisPredictor::MkldnnQuantizer::GetMaxChLSTMScalingFactor(
508 509
    const phi::DenseTensor& wx_tensor,
    const phi::DenseTensor& wh_tensor) const {
510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532
  check_tensor(wx_tensor);
  check_tensor(wh_tensor);

  std::vector<float> scale(wx_tensor.dims()[1]);

  for (int row_id = 0; row_id < wx_tensor.dims()[0]; row_id++) {
    for (int col_id = 0; col_id < wx_tensor.dims()[1]; col_id++) {
      int idx = (row_id * wx_tensor.dims()[1]) + col_id;
      auto abs_value = std::abs(wx_tensor.data<float>()[idx]);
      if (row_id == 0) {
        scale[col_id] = abs_value;
      } else {
        if (abs_value > scale[col_id]) scale[col_id] = abs_value;
      }
    }
  }
  for (int row_id = 0; row_id < wh_tensor.dims()[0]; row_id++) {
    for (int col_id = 0; col_id < wh_tensor.dims()[1]; col_id++) {
      int idx = (row_id * wh_tensor.dims()[1]) + col_id;
      auto abs_value = std::abs(wh_tensor.data<float>()[idx]);
      if (abs_value > scale[col_id]) scale[col_id] = abs_value;
    }
  }
533 534 535
  transform(scale.begin(), scale.end(), scale.begin(), [](float& c) {
    return 1 / c;
  });
536
  phi::DenseTensor scale_tensor = CreateScaleTensor(scale.size());
537 538 539 540 541 542
  auto* scale_ptr = scale_tensor.mutable_data<double>(CPUPlace());
  std::copy(scale.begin(), scale.end(), scale_ptr);
  bool is_unsigned = false;
  return std::make_pair(is_unsigned, scale_tensor);
}

543 544
std::pair<std::vector<int>, float>
AnalysisPredictor::MkldnnQuantizer::Histogram(
545
    const phi::DenseTensor& var_tensor,
546 547
    float min_val,
    float max_val,
548
    size_t num_bins) const {
549 550
  PADDLE_ENFORCE_GT(num_bins,
                    0,
551 552 553
                    platform::errors::InvalidArgument(
                        "MkldnnQuantizer: To calculate Histogram, num_bins (" +
                        std::to_string(num_bins) + ") must be positive."));
554 555
  PADDLE_ENFORCE_GT(var_tensor.numel(),
                    0,
556 557 558
                    platform::errors::InvalidArgument(
                        "MkldnnQuantizer: To calculate Histogram, the tensor "
                        "must not be empty."));
559 560
  PADDLE_ENFORCE_GE(max_val,
                    min_val,
561 562
                    platform::errors::InvalidArgument(
                        "MkldnnQuantizer: To calculate Histogram, max_val (" +
563 564 565
                        std::to_string(max_val) +
                        ") must be greater or equal"
                        "to min_val (" +
566
                        std::to_string(min_val) + ")."));
567 568
  ConstEigenVectorArrayMap eigen_tensor{
      var_tensor.data<float>(), var_tensor.numel(), 1};
569 570 571 572 573 574 575 576 577 578 579 580 581
  auto bin_width = std::abs(max_val - min_val) / num_bins;
  std::vector<int> hist(num_bins);

  for (int i = 0; i < eigen_tensor.size(); i++) {
    int bin = std::min(
        num_bins - 1,
        static_cast<size_t>(floor((eigen_tensor[i] - min_val) / bin_width)));
    ++hist[bin];
  }

  return std::make_pair(std::move(hist), std::move(bin_width));
}

582 583
void AnalysisPredictor::MkldnnQuantizer::ClearDeviceContext() const {
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
584 585 586
  phi::OneDNNContext* dev_ctx =
      (phi::OneDNNContext*)pool.Get(predictor_.place_);
  dev_ctx->ResetBlobMap(phi::OneDNNContext::tls().get_curr_exec());
587 588
}

589 590 591 592 593 594
void AnalysisPredictor::MkldnnQuantizer::PrepareArgument() const {
  auto& arg = predictor_.argument_;
  if (!arg.scope_valid()) arg.SetScope(new framework::Scope);
  arg.SetMainProgramNotOwned(predictor_.inference_program_.get());
  auto graph = std::unique_ptr<Graph>(new Graph(arg.main_program()));
  arg.SetMainGraph(graph.release());
595
  auto* scope_ptr = arg.scope_ptr();
596 597 598
  PADDLE_ENFORCE_NOT_NULL(
      scope_ptr,
      platform::errors::PreconditionNotMet("The scope should not be nullptr."));
599
  arg.main_graph().SetNotOwned(framework::ir::kParamScopeAttr, scope_ptr);
600 601

  auto* builder = predictor_.config_.pass_builder();
602 603
  builder->SetPasses({"cpu_quantize_pass",
                      "cpu_quantize_squash_pass",
604 605
                      "int8_scale_calculation_mkldnn_pass",
                      "params_quantization_mkldnn_pass"});
606 607 608
  if (predictor_.config_.ir_debug_) builder->TurnOnDebug();
  auto passes = builder->AllPasses();
  predictor_.argument_.SetIrAnalysisPasses(passes);
609 610
  predictor_.argument_.SetAnalysisPasses(
      {"ir_analysis_pass", "memory_optimize_pass", "ir_graph_to_program_pass"});
611 612 613 614 615 616
  predictor_.argument_.SetQuantVarScales(scales_);
}

bool AnalysisPredictor::MkldnnQuantizer::Quantize() {
  if (!RunWarmup()) return false;
  if (!CalculateScales()) return false;
617
  ClearDeviceContext();
618 619 620 621 622 623 624 625 626
  predictor_.PrepareScope(predictor_.scope_);
  predictor_.CreateExecutor();
  if (!RunQuantizePasses()) return false;
  predictor_.PrepareExecutor();
  predictor_.PrepareFeedFetch();
  return true;
}

bool AnalysisPredictor::MkldnnQuantizer::RunQuantizePasses() const {
627 628
  predictor_.executor_->CreateVariables(
      *predictor_.inference_program_, 0, true, predictor_.sub_scope_);
629 630 631
  PrepareArgument();
  auto& arg = predictor_.argument_;
  Analyzer().Run(&arg);
632
  PADDLE_ENFORCE_EQ(
633 634
      arg.scope_valid(),
      true,
635
      platform::errors::PreconditionNotMet("The scope should be valid."));
636 637 638 639 640
  VLOG(5) << "to prepare executor";
  ARGUMENT_CHECK_FIELD((&arg), ir_analyzed_program);
  predictor_.inference_program_.reset(
      new framework::ProgramDesc(arg.ir_analyzed_program()));
  LOG(INFO) << "== optimize 2 end ==";
641 642
  predictor_.executor_->CreateVariables(
      *predictor_.inference_program_, 0, false, predictor_.sub_scope_);
643 644 645 646 647 648 649
  return true;
}

bool AnalysisPredictor::MkldnnQuantizer::RunWarmup() const {
  VLOG(3) << "Predictor: run a quantization warmup iteration";
  auto warmup_data = qconfig_->warmup_data();
  PADDLE_ENFORCE_NOT_NULL(warmup_data,
650 651
                          platform::errors::PreconditionNotMet(
                              "Warmup data cannot be NULL in the config."));
652 653 654 655 656 657 658 659 660 661
  PrettyLogH1("--- Running warmup iteration for quantization");

  // Run the inference program
  std::vector<PaddleTensor> output_slots;
  predictor_.Run(*warmup_data, &output_slots, qconfig_->warmup_batch_size());

  return true;
}

float AnalysisPredictor::MkldnnQuantizer::SafeEntropy(
662 663 664 665 666 667
    std::vector<int> reference_distr_P,
    int P_sum,
    std::vector<int> candidate_distr_Q,
    int Q_sum) const {
  PADDLE_ENFORCE_EQ(reference_distr_P.size(),
                    candidate_distr_Q.size(),
668 669
                    platform::errors::InvalidArgument(
                        "The P size %d should be equal to Q size %d",
670 671
                        reference_distr_P.size(),
                        candidate_distr_Q.size()));
672 673 674 675 676 677 678 679 680
  float tmp_sum1 = 0;
  float tmp_sum2 = 0;
  for (size_t idx = 0; idx < reference_distr_P.size(); idx++) {
    int p_idx = reference_distr_P[idx];
    int q_idx = candidate_distr_Q[idx];
    if (p_idx == 0) {
      tmp_sum1 += 0;
      tmp_sum2 += 0;
    } else {
681
      PADDLE_ENFORCE_NE(
682 683
          q_idx,
          0,
684 685 686
          platform::errors::PreconditionNotMet(
              "MkldnnQuantizer: Fatal error!, idx = " + std::to_string(idx) +
              " qindex = 0! p_idx = " + std::to_string(p_idx)));
687 688 689 690 691 692 693 694
    }
    tmp_sum1 += p_idx * (log(Q_sum * p_idx));
    tmp_sum2 += p_idx * (log(P_sum * q_idx));
  }
  return (tmp_sum1 - tmp_sum2) / P_sum;
}

}  // namespace paddle