Caffe2 - C++ API
A deep learning, cross platform ML framework
qtensor.h
1 
17 #ifndef CAFFE2_CORE_QTENSOR_H_
18 #define CAFFE2_CORE_QTENSOR_H_
19 
20 #include <algorithm>
21 #include <climits>
22 #include <cstddef>
23 #include <vector>
24 
25 #include "caffe2/core/common.h"
26 #include "caffe2/core/context.h"
27 #include "caffe2/core/tensor.h"
28 #include "caffe2/core/typeid.h"
29 
30 namespace caffe2 {
31 
32 template <class Context>
33 class QTensor {
34  public:
35  QTensor() {}
36  virtual ~QTensor() {}
65  explicit QTensor(
66  const std::vector<int>& dims,
67  const unsigned char precision,
68  const bool signbit = false)
69  : precision_(precision), signed_(signbit) {
70  Resize(dims);
71  }
72 
73  void Resize(std::vector<int> dim_source) {
74  if (dims_ != dim_source) {
75  size_t source_size = std::accumulate(
76  dim_source.begin(), dim_source.end(), 1, std::multiplies<int>());
77  if ((source_size * (precision_ + signed_)) > capacity_) {
78  data_.reset();
79  capacity_ = 0;
80  }
81  dims_ = dim_source;
82  size_ = source_size;
83  }
84  }
85 
86  void
87  SetBitAtIndex(const unsigned char bit, const size_t index, const bool value) {
88  // Get the mutable data at bit depth `bit`.
89  unsigned char* d = mutable_data();
90 
91  CAFFE_ENFORCE(
92  bit < precision_ + signed_,
93  "Attempted to a set a bit that is not allocated.");
94  CAFFE_ENFORCE(bit * aligned_size() < capacity_);
95 
96  auto idx = (aligned_size() * bit) / CHAR_BIT;
97  d = &d[idx];
98 
99  idx = index / CHAR_BIT;
100  auto shift = CHAR_BIT - (index % CHAR_BIT) - 1;
101 
102  if (value) {
103  d[idx] |= 1 << shift;
104  } else {
105  d[idx] &= ~(1 << shift);
106  }
107  }
108 
109  bool GetBitAtIndex(const unsigned char bit, const size_t index) const {
110  // Get the data at bit depth `bit`
111  const unsigned char* d = data();
112  auto idx = (aligned_size() * bit) / CHAR_BIT;
113  d = &d[idx];
114 
115  idx = index / CHAR_BIT;
116  auto shift = CHAR_BIT - (index % CHAR_BIT) - 1;
117 
118  return d[idx] & (1 << shift);
119  }
120 
121  void SetPrecision(const unsigned char precision) {
122  precision_ = precision;
123  data_.reset();
124  }
125 
126  void SetSigned(const bool make_signed = true) {
127  signed_ = make_signed;
128  data_.reset();
129  }
130 
131  void SetScale(const double scale) {
132  scale_ = scale;
133  }
134 
135  void SetBias(const double bias) {
136  bias_ = bias;
137  }
138 
139  unsigned char* mutable_data() {
140  if (!data_) {
141  auto ptr_and_deleter = Context::New(nbytes());
142  data_.reset(
143  static_cast<unsigned char*>(ptr_and_deleter.first),
144  ptr_and_deleter.second);
145  capacity_ = nbytes() * CHAR_BIT;
146  }
147  CAFFE_ENFORCE(capacity_ == nbytes() * CHAR_BIT);
148  return data_.get();
149  }
150 
151  inline const unsigned char* data() const {
152  return data_.get();
153  }
154 
155  inline size_t size() const {
156  return size_;
157  }
158 
159  inline unsigned char alignment() const {
160  return alignment_;
161  }
162 
163  inline unsigned char precision() const {
164  return precision_;
165  }
166 
167  inline const vector<int>& dims() const {
168  return dims_;
169  }
170 
171  inline bool is_signed() const {
172  return signed_;
173  }
174 
178  inline int ndim() const {
179  return dims_.size();
180  }
181 
182  inline size_t aligned_size() const {
183  return alignment_ * ((size_ + alignment_ - 1) / alignment_);
184  }
185 
186  inline size_t nbytes() const {
187  return (aligned_size() * (precision_ + signed_)) / CHAR_BIT;
188  }
189 
190  inline double scale() const {
191  return scale_;
192  }
193 
194  inline double bias() const {
195  return bias_;
196  }
197 
201  inline int dim32(const int i) const {
202  DCHECK_LT(i, dims_.size()) << "Exceeding ndim limit " << dims_.size();
203  DCHECK_GE(i, 0) << "Cannot have negative index";
204  CAFFE_ENFORCE_LT(dims_[i], std::numeric_limits<int>::max());
205  return static_cast<int>(dims_[i]);
206  }
207 
219  inline int canonical_axis_index(int axis_index) const {
220  CAFFE_ENFORCE_GE(axis_index, -ndim());
221  CAFFE_ENFORCE_LT(axis_index, ndim());
222  if (axis_index < 0) {
223  return axis_index + ndim();
224  }
225  return axis_index;
226  }
227 
231  inline TIndex size_from_dim(int k) const {
232  TIndex r = 1;
233  for (int i = k; i < dims_.size(); ++i) {
234  r *= dims_[i];
235  }
236  return r;
237  }
238 
242  inline TIndex size_to_dim(int k) const {
243  CAFFE_ENFORCE(k < dims_.size());
244  TIndex r = 1;
245  for (int i = 0; i < k; ++i) {
246  r *= dims_[i];
247  }
248  return r;
249  }
250 
251  protected:
252  std::vector<int> dims_;
253  size_t size_ = 0;
254 
255  // Precision in bits.
256  unsigned char precision_ = CHAR_BIT;
257  // Bit alignment.
258  unsigned char alignment_ = CHAR_BIT;
259 
260  // Allocated data.
261  std::shared_ptr<unsigned char> data_;
262 
263  // value = scale_ * (x + bias_)
264  double scale_;
265  double bias_;
266  bool signed_ = false;
267 
268  // Capacity in bits.
269  size_t capacity_ = 0;
270 };
271 
272 } // namespace caffe2
273 #endif // CAFFE2_CORE_QTENSOR_H_
int canonical_axis_index(int axis_index) const
Returns the &#39;canonical&#39; version of a (usually) user-specified axis, allowing for negative indexing (e...
Definition: qtensor.h:219
int ndim() const
Returns the number of dimensions of the data.
Definition: qtensor.h:178
QTensor(const std::vector< int > &dims, const unsigned char precision, const bool signbit=false)
Creates a quantized tensor of the given dimension.
Definition: qtensor.h:65
TIndex size_from_dim(int k) const
Return product of all dimensions starting from K.
Definition: qtensor.h:231
Copyright (c) 2016-present, Facebook, Inc.
int dim32(const int i) const
Returns the i-th dimension of the qtensor in int.
Definition: qtensor.h:201
TIndex size_to_dim(int k) const
Product of all dims up to.
Definition: qtensor.h:242