PackBMatrix Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/qnnpack/include/qnnpack_func.h lines 41–99

class PackBMatrix final {
 public:
  PackBMatrix(
      size_t input_channels,
      size_t output_channels,
      const uint8_t* kernel_zero_points,
      const float* requantization_scale,
      const uint8_t* kernel,
      const int32_t* bias);

  // This constructor is to be used for dynamic mode
  // quantization. In dynamic mode, we dont yet support
  // per channel quantization, and paying the cost of
  // memory allocation for per channel zero point and
  // requant scale will hurt performance.
  PackBMatrix(
      size_t input_channels,
      size_t output_channels,
      const uint8_t kernel_zero_point,
      const float requantization_scale,
      const uint8_t* kernel,
      const int32_t* bias);

  void* getPackedWeights() const
  {
    return packed_weights_;
  }

  void unpackWeights(
      const uint8_t* kernel_zero_points,
      int8_t* kernel
    ) const;

  size_t getInputChannels() const
  {
    return input_channels_;
  }

  size_t getOutputChannels() const
  {
    return output_channels_;
  }

  ~PackBMatrix()
  {
    if (packed_weights_ != nullptr) {
      free(packed_weights_);
    }
  }

  PackBMatrix() = delete;
  PackBMatrix(const PackBMatrix&) = delete;
  PackBMatrix& operator=(const PackBMatrix&) = delete;

 private:
  void* packed_weights_ = nullptr;
  size_t input_channels_;
  size_t output_channels_;
};