qtensor.h

#ifndef CAFFE2_CORE_QTENSOR_H_
#define CAFFE2_CORE_QTENSOR_H_

#include <algorithm>
#include <climits>
#include <cstddef>
#include <vector>

#include "caffe2/core/common.h"
#include "caffe2/core/context.h"
#include "caffe2/core/tensor.h"
#include <c10/util/typeid.h>

namespace caffe2 {

template <class Context>
class C10_EXPORT QTensor {
 public:
  QTensor() {}
  virtual ~QTensor() {}
  // TODO: changing at::ArrayRef<int> to at::ArrayRef<int64_t>?
  explicit QTensor(
      at::ArrayRef<int> dims,
      const unsigned char precision,
      const bool signbit = false)
      : precision_(precision), signed_(signbit) {
    Resize(dims);
  }

  void Resize(at::ArrayRef<int> dim_source) {
    if (dims_ != dim_source) {
      size_t source_size = std::accumulate(
          dim_source.begin(), dim_source.end(), 1, std::multiplies<int>());
      if ((source_size * (precision_ + signed_)) > capacity_) {
        data_ptr_.clear();
        capacity_ = 0;
      }
      dims_ = dim_source.vec();
      size_ = source_size;
    }
  }

  void
  SetBitAtIndex(const unsigned char bit, const size_t index, const bool value) {
    // Get the mutable data at bit depth `bit`.
    unsigned char* d = mutable_data();

    CAFFE_ENFORCE(
        bit < precision_ + signed_,
        "Attempted to a set a bit that is not allocated.");
    CAFFE_ENFORCE(bit * aligned_size() < capacity_);

    auto idx = (aligned_size() * bit) / CHAR_BIT;
    d = &d[idx];

    idx = index / CHAR_BIT;
    auto shift = CHAR_BIT - (index % CHAR_BIT) - 1;

    if (value) {
      d[idx] |= 1 << shift;
    } else {
      d[idx] &= ~(1 << shift);
    }
  }

  bool GetBitAtIndex(const unsigned char bit, const size_t index) const {
    // Get the data at bit depth `bit`
    const unsigned char* d = data();
    auto idx = (aligned_size() * bit) / CHAR_BIT;
    d = &d[idx];

    idx = index / CHAR_BIT;
    auto shift = CHAR_BIT - (index % CHAR_BIT) - 1;

    return d[idx] & (1 << shift);
  }

  void SetPrecision(const unsigned char precision) {
    precision_ = precision;
    data_ptr_.clear();
  }

  void SetSigned(const bool make_signed = true) {
    signed_ = make_signed;
    data_ptr_.clear();
  }

  void SetScale(const double scale) {
    scale_ = scale;
  }

  void SetBias(const double bias) {
    bias_ = bias;
  }

  unsigned char* mutable_data() {
    if (!data_ptr_) {
      data_ptr_ = Context::New(nbytes());
      capacity_ = nbytes() * CHAR_BIT;
    }
    CAFFE_ENFORCE(capacity_ == nbytes() * CHAR_BIT);
    return static_cast<unsigned char*>(data_ptr_.get());
  }

  inline const unsigned char* data() const {
    return static_cast<unsigned char*>(data_ptr_.get());
  }

  inline size_t size() const {
    return size_;
  }

  inline unsigned char alignment() const {
    return alignment_;
  }

  inline unsigned char precision() const {
    return precision_;
  }

  inline at::ArrayRef<int> sizes() const {
    return dims_;
  }

  // TODO: deprecate?
  inline at::ArrayRef<int> dims() const {
    return dims_;
  }

  inline bool is_signed() const {
    return signed_;
  }

  inline int ndim() const {
    return dims_.size();
  }

  inline size_t aligned_size() const {
    return alignment_ * ((size_ + alignment_ - 1) / alignment_);
  }

  inline size_t nbytes() const {
    return (aligned_size() * (precision_ + signed_)) / CHAR_BIT;
  }

  inline double scale() const {
    return scale_;
  }

  inline double bias() const {
    return bias_;
  }

  inline int dim32(const int i) const {
    DCHECK_LT(i, dims_.size()) << "Exceeding ndim limit " << dims_.size();
    DCHECK_GE(i, 0) << "Cannot have negative index";
    CAFFE_ENFORCE_LT(dims_[i], std::numeric_limits<int>::max());
    return static_cast<int>(dims_[i]);
  }

  inline int canonical_axis_index(int axis_index) const {
    CAFFE_ENFORCE_GE(axis_index, -ndim());
    CAFFE_ENFORCE_LT(axis_index, ndim());
    if (axis_index < 0) {
      return axis_index + ndim();
    }
    return axis_index;
  }

  inline int64_t size_from_dim(int k) const {
    int64_t r = 1;
    for (int i = k; i < dims_.size(); ++i) {
      r *= dims_[i];
    }
    return r;
  }

  inline int64_t size_to_dim(int k) const {
    CAFFE_ENFORCE(k < dims_.size());
    int64_t r = 1;
    for (int i = 0; i < k; ++i) {
      r *= dims_[i];
    }
    return r;
  }

 protected:
  std::vector<int> dims_;
  size_t size_ = 0;

  // Precision in bits.
  unsigned char precision_ = CHAR_BIT;
  // Bit alignment.
  unsigned char alignment_ = CHAR_BIT;

  // Allocated data.
  at::DataPtr data_ptr_;

  // value = scale_ * (x + bias_)
  double scale_;
  double bias_;
  bool signed_ = false;

  // Capacity in bits.
  size_t capacity_ = 0;
};

} // namespace caffe2
#endif // CAFFE2_CORE_QTENSOR_H_