text: // Implements the math functions for CPU.
  // The implementation in this file allows us to route the underlying numerical
  // computation library to different compiler options (-mno-avx2 or -mavx2).
  
  #include <algorithm>
  #include <array>
  #include <atomic>
  #include <chrono>
  #include <cmath>
  #include <cstring>
  #include <functional>
  #include <limits>
  #include <numeric>
  #include <random>
  #include <tuple>
  #include <unordered_set>
  #include <vector>
  
  #include "caffe2/core/context.h"
  #include "caffe2/perfkernels/math.h"
  #include "caffe2/utils/cpu_neon.h"
  #include "caffe2/utils/cpuid.h"
  #include "caffe2/utils/eigen_utils.h"
  #include "caffe2/utils/math.h"
  
  #include "Eigen/Core"
  #include "Eigen/Dense"
  
  #ifdef CAFFE2_USE_MKL
  #include <mkl.h>
  #endif // CAFFE2_USE_MKL
  
  #ifdef CAFFE2_USE_HPTT
  #include <hptt.h>
  #endif // CAFFE2_USE_HPTT
  
  #if defined(_MSC_VER)
  #include <process.h>
  #endif
  
  namespace caffe2 {
  
  namespace math {
  #define QEPSILON 1e-8
  
  void quantize_and_compress__avx2(
      const float* input_data,
      uint8_t* output_data,
      size_t input_size,
      size_t bitwidth,
      bool random,
  #ifdef FUSED_ROWWISE_RANDOM_QUANTIZATION_USE_MKL
      VSLStreamStatePtr& vslStream,
      std::vector<float>& random_buffer
  #else
      std::unique_ptr<std::uniform_real_distribution<float>>& dis,
      std::minstd_rand& gen
  #endif
  ) {
    CAFFE_ENFORCE(
        bitwidth == 1 || bitwidth == 2 || bitwidth == 4 || bitwidth == 8,
        "Unsupported bitwidth");
  
  #ifdef __AVX2__
    __m256i shuffle_mask_v = _mm256_set_epi8(
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0x0c,
        0x08,
        0x04,
        0x00,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0xff,
        0x0c,
        0x08,
        0x04,
        0x00);
    __m256i permute_mask_v =
        _mm256_set_epi32(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00);
  #endif // __AVX2__
  
    // memory pointers
    ConstEigenVectorArrayMap<float> input_row(input_data, input_size);
    uint8_t* output_row = output_data;
    EigenVectorArrayMap<uint8_t> output_bitwidth_tail(output_row, 2);
    EigenVectorArrayMap<float> output_row_min_max(
        reinterpret_cast<float*>(output_row + 2), 2);
  
    size_t data_per_byte = 8 / bitwidth;
    size_t tail = input_size % data_per_byte;
    tail = tail ? data_per_byte - tail : 0;
    size_t segment_size = (input_size + data_per_byte - 1) / data_per_byte;
  
    // basic info
    const float minimum_element = input_row.minCoeff();
    const float maximum_element = input_row.maxCoeff();
    output_bitwidth_tail(0) = bitwidth;
    output_bitwidth_tail(1) = tail;
    output_row_min_max(0) = minimum_element;
    output_row_min_max(1) = maximum_element;
  
    float gap = (maximum_element - minimum_element) / ((1 << bitwidth) - 1.0f);
    float gap_inverse = 1. / (gap + QEPSILON);
    uint8_t max_q = (1 << bitwidth) - 1;
    size_t bit_start = 0;
    if (random) {
  #ifdef FUSED_ROWWISE_RANDOM_QUANTIZATION_USE_MKL
      int status = vsRngUniform(
          VSL_RNG_METHOD_UNIFORM_STD,
          vslStream,
          input_size,
          random_buffer.data(),
          0.0f,
          1.0f);
      if (status != VSL_ERROR_OK) {
        LOG(WARNING) << "vsRngUniform returns " << status;
      }
  #endif
      for (int start = 0; start < input_size; start += segment_size) {
        size_t stride = start + segment_size <= input_size ? segment_size
                                                           : input_size - start;
        int i = 0;
  #ifdef __AVX2__
        constexpr int VLEN = 8;
        for (; i < stride / VLEN * VLEN; i += VLEN) {
  #ifdef FUSED_ROWWISE_RANDOM_QUANTIZATION_USE_MKL
          __m256 r_v = _mm256_loadu_ps(&random_buffer[start + i]);
  #else
          float random_buffer[VLEN];
          for (int j = 0; j < VLEN; ++j) {
            random_buffer[j] = (*dis)(gen);
          }
          __m256 r_v = _mm256_loadu_ps(random_buffer);
  #endif
          __m256 fval_v = _mm256_loadu_ps(input_data + start + i);
          __m256 thetimes_v = _mm256_mul_ps(
              _mm256_sub_ps(fval_v, _mm256_set1_ps(minimum_element)),
              _mm256_set1_ps(gap_inverse));
          __m256 rounded_v = _mm256_floor_ps(_mm256_add_ps(thetimes_v, r_v));
          rounded_v = _mm256_max_ps(
              _mm256_setzero_ps(),
              _mm256_min_ps(_mm256_set1_ps(max_q), rounded_v));
          __m256i qval_v = _mm256_cvtps_epi32(rounded_v);
          __m256i orval_v = _mm256_cvtepu8_epi32(_mm_lddqu_si128(
              reinterpret_cast<const __m128i*>(output_row + 10 + i)));
          orval_v =
              _mm256_or_si256(orval_v, _mm256_slli_epi32(qval_v, bit_start));
          orval_v = _mm256_shuffle_epi8(orval_v, shuffle_mask_v);
          orval_v = _mm256_permutevar8x32_epi32(orval_v, permute_mask_v);
          *reinterpret_cast<int64_t*>(output_row + 10 + i) =
              _mm256_extract_epi64(orval_v, 0);
        }
  #endif // __AVX2__
        for (; i < stride; ++i) {
          float fval = input_data[start + i];
          float thetimes = (fval - minimum_element) * gap_inverse;
  #ifdef FUSED_ROWWISE_RANDOM_QUANTIZATION_USE_MKL
          float rounded = floor(thetimes + random_buffer[start + i]);
  #else
          float rounded = floor(thetimes + (*dis)(gen));
  #endif
          rounded = std::max(0.0f, std::min(static_cast<float>(max_q), rounded));
          uint8_t qval = rounded;
  
          uint8_t orval = output_row[10 + i];
          output_row[10 + i] = orval | static_cast<uint8_t>(qval << bit_start);
        }
        bit_start += bitwidth;
      }
    } else {
      for (int start = 0; start < input_size; start += segment_size) {
        size_t stride = start + segment_size <= input_size ? segment_size
                                                           : input_size - start;
        int i = 0;
  #ifdef __AVX2__
        constexpr int VLEN = 8;
        for (; i < stride / VLEN * VLEN; i += VLEN) {
          __m256 fval_v = _mm256_loadu_ps(input_data + start + i);
          __m256 thetimes_v = _mm256_mul_ps(
              _mm256_sub_ps(fval_v, _mm256_set1_ps(minimum_element)),
              _mm256_set1_ps(gap_inverse));
          thetimes_v = _mm256_max_ps(
              _mm256_setzero_ps(),
              _mm256_min_ps(_mm256_set1_ps(max_q), thetimes_v));
          __m256i qval_v = _mm256_cvtps_epi32(_mm256_round_ps(
              thetimes_v, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
          __m256i orval_v = _mm256_cvtepu8_epi32(_mm_lddqu_si128(
              reinterpret_cast<const __m128i*>(output_row + 10 + i)));
          orval_v =
              _mm256_or_si256(orval_v, _mm256_slli_epi32(qval_v, bit_start));
          orval_v = _mm256_shuffle_epi8(orval_v, shuffle_mask_v);
          orval_v = _mm256_permutevar8x32_epi32(orval_v, permute_mask_v);
          *reinterpret_cast<int64_t*>(output_row + 10 + i) =
              _mm256_extract_epi64(orval_v, 0);
        }
  #endif // __AVX2__
        for (; i < stride; ++i) {
          float fval = input_data[start + i];
          float thetimes = (fval - minimum_element) * gap_inverse;
          thetimes =
              std::max(0.0f, std::min(static_cast<float>(max_q), thetimes));
          uint8_t qval = nearbyint(thetimes);
  
          uint8_t orval = output_row[10 + i];
          output_row[10 + i] = orval | static_cast<uint8_t>(qval << bit_start);
        }
        bit_start += bitwidth;
      }
    }
  }
  
  void decompress_and_dequantize__avx2(
      const uint8_t* input_data,
      float* output_data,
      size_t input_size) {
    ConstEigenVectorArrayMap<float> input_row_min_max(
        reinterpret_cast<const float*>(input_data + 2), 2);
  
    // basic info
    const float minimum_element = input_row_min_max(0);
    const float maximum_element = input_row_min_max(1);
    const size_t bitwidth = input_data[0];
    const float gap =
        (maximum_element - minimum_element) / ((1 << bitwidth) - 1.f) +
        QEPSILON; // for exact recovering
  
    CAFFE_ENFORCE(
        bitwidth == 1 || bitwidth == 2 || bitwidth == 4 || bitwidth == 8,
        "Unsupported bitwidth");
    const size_t tail = input_data[1];
  
    const size_t output_size = (input_size - 10) * (8 / bitwidth) - tail;
    EigenVectorArrayMap<float> output_row(output_data, output_size);
    // decoding
    size_t bit_start = 0;
    const size_t segment_size = input_size - 10;
    for (int start = 0; start < output_size; start += segment_size) {
      size_t stride = start + segment_size <= output_size ? segment_size
                                                          : output_size - start;
      uint8_t mask = (1 << bitwidth) - 1;
      int i = 0;
  #ifdef __AVX2__
      // Can process 8 elements at a time because we need to expand uint8_t
      // to int32_t to use epi32 vector instructions.
      constexpr int VLEN = 8;
      for (; i < stride / VLEN * VLEN; i += VLEN) {
        __m128i in_v = _mm_lddqu_si128(
            reinterpret_cast<const __m128i*>(input_data + 10 + i));
        __m256i out_epi32_v = _mm256_and_si256(
            _mm256_srli_epi32(_mm256_cvtepu8_epi32(in_v), bit_start),
            _mm256_set1_epi32(mask));
        _mm256_storeu_ps(
            output_data + start + i, _mm256_cvtepi32_ps(out_epi32_v));
      }
  #endif
      for (; i < stride; ++i) {
        output_data[start + i] = ((input_data[10 + i] >> bit_start) & mask);
      }
      bit_start += bitwidth;
    }
    // scaling and biasing
    output_row = output_row * gap + minimum_element;
  }
  
  #undef QEPSILON
  } // namespace math
  } // namespace caffe2
  