Caffe2 - C++ API
A deep learning, cross platform ML framework
TensorOptions.h
1 #pragma once
2 
3 #include <c10/core/DefaultDtype.h>
4 #include <c10/core/Backend.h>
5 #include <c10/core/Layout.h>
6 #include <c10/core/ScalarType.h>
7 #include <c10/core/Device.h>
8 
9 #include <c10/util/Optional.h>
10 #include <c10/util/C++17.h>
11 #include <c10/macros/Macros.h>
12 
13 #include <cstddef>
14 #include <iosfwd>
15 #include <utility>
16 
17 namespace c10 {
54 
97 
98 
99 struct C10_API TensorOptions {
100  TensorOptions()
101  : requires_grad_(false)
102  , is_variable_(false)
103  , has_device_(false)
104  , has_dtype_(false)
105  , has_layout_(false)
106  , has_requires_grad_(false)
107  , has_is_variable_(false)
108  {}
109 
111  /* implicit */ TensorOptions(Layout layout) : TensorOptions() {
112  this->set_layout(layout);
113  }
114 
117  template<typename T,
118  typename = c10::guts::enable_if_t<std::is_same<c10::guts::decay_t<T>, Device>::value>>
119  /* implicit */ TensorOptions(T&& device) : TensorOptions() {
120  this->set_device(std::forward<T>(device));
121  }
122 
131  template <typename... Args,
132  typename = c10::guts::enable_if_t<std::is_constructible<Device, Args&&...>::value>>
133  /* implicit */ TensorOptions(Args&&... args)
134  : TensorOptions(Device(std::forward<Args>(args)...)) {}
135 
137  /* implicit */ TensorOptions(caffe2::TypeMeta dtype) : TensorOptions() {
138  this->set_dtype(dtype);
139  }
140 
142  /* implicit */ TensorOptions(ScalarType dtype) : TensorOptions() {
143  this->set_dtype(dtype);
144  }
145 
147  bool operator==(const TensorOptions& other) const noexcept {
148  return
149  has_dtype_ == other.has_dtype_ &&
150  has_layout_ == other.has_layout_ &&
151  has_device_ == other.has_device_ &&
152  has_requires_grad_ == other.has_requires_grad_ &&
153  has_is_variable_ == other.has_is_variable_ &&
154  (!has_dtype_ || dtype_ == other.dtype_) &&
155  (!has_layout_ || layout_ == other.layout_) &&
156  (!has_device_ || device_ == other.device_) &&
157  (!requires_grad_ || requires_grad_ == other.requires_grad_) &&
158  (!is_variable_ || is_variable_ == other.is_variable_);
159  }
160 
163  bool operator!=(const TensorOptions& other) const noexcept {
164  return !(*this == other);
165  }
166 
169  C10_NODISCARD TensorOptions device(c10::optional<Device> device) const noexcept {
170  TensorOptions r = *this;
171  r.set_device(device);
172  return r;
173  }
174 
178  template<typename ... Args>
179  C10_NODISCARD TensorOptions device(Args&&... args) const noexcept {
180  return device(c10::optional<Device>(c10::in_place, std::forward<Args>(args)...));
181  }
182 
188  C10_NODISCARD TensorOptions device_index(int16_t device_index) const noexcept {
189  return device(Device::Type::CUDA, device_index);
190  }
191 
193  C10_NODISCARD TensorOptions dtype(c10::optional<caffe2::TypeMeta> dtype) const noexcept {
194  TensorOptions r = *this;
195  r.set_dtype(dtype);
196  return r;
197  }
198 
199  // legacy function to support ScalarType
200  C10_NODISCARD TensorOptions dtype(c10::optional<ScalarType> dtype) const noexcept {
201  TensorOptions r = *this;
202  r.set_dtype(dtype);
203  return r;
204  }
205 
206  // Since dtype is taken...
207  template <typename T>
208  TensorOptions& dtype() {
209  dtype_ = caffe2::TypeMeta::Make<T>();
210  has_dtype_ = true;
211  return *this;
212  }
213 
215  C10_NODISCARD TensorOptions layout(c10::optional<Layout> layout) const noexcept {
216  TensorOptions r = *this;
217  r.set_layout(layout);
218  return r;
219  }
220 
222  C10_NODISCARD TensorOptions requires_grad(c10::optional<bool> requires_grad) const noexcept {
223  TensorOptions r = *this;
224  r.set_requires_grad(requires_grad);
225  return r;
226  }
227 
229  C10_NODISCARD TensorOptions is_variable(c10::optional<bool> is_variable) const noexcept {
230  TensorOptions r = *this;
231  r.set_is_variable(is_variable);
232  return r;
233  }
234 
236  Device device() const noexcept {
237  return has_device_ ? device_ : Device(kCPU);
238  }
239 
241  bool has_device() const noexcept {
242  return has_device_;
243  }
244 
248  return has_device_ ? c10::make_optional(device_) : c10::nullopt;
249  }
250 
252  int32_t device_index() const noexcept {
253  return device().index();
254  }
255 
257  caffe2::TypeMeta dtype() const noexcept {
258  return has_dtype_ ? dtype_ : get_default_dtype();
259  }
260 
262  bool has_dtype() const noexcept {
263  return has_dtype_;
264  }
265 
269  return has_dtype_ ? c10::make_optional(dtype_) : c10::nullopt;
270  }
271 
273  Layout layout() const noexcept {
274  return has_layout_ ? layout_ : kStrided;
275  }
276 
278  bool has_layout() const noexcept {
279  return has_layout_;
280  }
281 
285  return has_layout_ ? c10::make_optional(layout_) : c10::nullopt;
286  }
287 
289  bool requires_grad() const noexcept {
290  return has_requires_grad_ ? requires_grad_ : false;
291  }
292 
294  bool has_requires_grad() const noexcept {
295  return has_requires_grad_;
296  }
297 
301  return has_requires_grad_ ? c10::make_optional(requires_grad_)
302  : c10::nullopt;
303  }
304 
306  bool is_variable() const noexcept {
307  return has_is_variable_ ? is_variable_ : false;
308  }
309 
311  bool has_is_variable() const noexcept {
312  return has_is_variable_;
313  }
314 
318  return has_is_variable_ ? c10::make_optional(is_variable_) : c10::nullopt;
319  }
320 
321  // Resolves the ATen backend specified by the current construction axes.
322  Backend backend() const noexcept {
323  return at::tensorTypeIdToBackend(computeTensorTypeId());
324  }
325 
326  inline TensorTypeId computeTensorTypeId() const {
327  switch (layout()) {
328  case Layout::Strided:
329  switch (device().type()) {
330  case DeviceType::CPU:
331  return CPUTensorId();
332  case DeviceType::CUDA:
333  return CUDATensorId();
334  case DeviceType::MKLDNN:
335  return MKLDNNTensorId();
336  case DeviceType::OPENGL:
337  return OpenGLTensorId();
338  case DeviceType::OPENCL:
339  return OpenCLTensorId();
340  case DeviceType::IDEEP:
341  return IDEEPTensorId();
342  case DeviceType::HIP:
343  return HIPTensorId();
344  case DeviceType::MSNPU:
345  return MSNPUTensorId();
346  case DeviceType::XLA:
347  return XLATensorId();
348  default:
349  AT_ERROR("Unsupported device type for dense layout: ", device().type());
350  }
351  case Layout::Sparse:
352  switch (device().type()) {
353  case DeviceType::CPU:
354  return SparseCPUTensorId();
355  case DeviceType::CUDA:
356  return SparseCUDATensorId();
357  case DeviceType::HIP:
358  return SparseHIPTensorId();
359  default:
360  AT_ERROR("Unsupported device type for sparse layout: ", device().type());
361  }
362  default:
363  AT_ERROR("Unsupported layout: ", layout());
364  }
365  }
366 
367  private:
368 
369  // These methods are currently private because I'm not sure if it's wise
370  // to actually publish them. They are methods because I need them in
371  // the constructor and the functional API implementation.
372  //
373  // If you really, really need it, you can make these public, but check if you
374  // couldn't just do what you need with the functional API. Similarly, these
375  // methods are not chainable, because if you wanted chaining, you probably
376  // want to use the functional API instead. (It's probably OK to make
377  // these chainable, because these functions are all explicitly annotated
378  // with a ref-qualifier, the trailing &, that makes them illegal to call
379  // on temporaries.)
380 
382  void set_device(c10::optional<Device> device) & noexcept {
383  if (device) {
384  device_ = *device;
385  has_device_ = true;
386  } else {
387  has_device_ = false;
388  }
389  }
390 
392  void set_dtype(c10::optional<caffe2::TypeMeta> dtype) & noexcept {
393  if (dtype) {
394  dtype_ = *dtype;
395  has_dtype_ = true;
396  } else {
397  has_dtype_ = false;
398  }
399  }
400 
401  // legacy function to support ScalarType
402  void set_dtype(c10::optional<ScalarType> dtype) & noexcept {
403  if (dtype) {
404  dtype_ = scalarTypeToTypeMeta(*dtype);
405  has_dtype_ = true;
406  } else {
407  has_dtype_ = false;
408  }
409  }
410 
412  void set_layout(c10::optional<Layout> layout) & noexcept {
413  if (layout) {
414  layout_ = *layout;
415  has_layout_ = true;
416  } else {
417  has_layout_ = false;
418  }
419  }
420 
422  void set_requires_grad(c10::optional<bool> requires_grad) & noexcept {
423  if (requires_grad) {
424  requires_grad_ = *requires_grad;
425  has_requires_grad_ = true;
426  } else {
427  has_requires_grad_ = false;
428  }
429  }
430 
432  void set_is_variable(c10::optional<bool> is_variable) & noexcept {
433  if (is_variable) {
434  is_variable_ = *is_variable;
435  has_is_variable_ = true;
436  } else {
437  has_is_variable_ = false;
438  }
439  }
440 
441  // WARNING: If you edit TensorOptions to add more options, you
442  // must adjust the implementation of Tensor::options
443 
444  // NB: We didn't use c10::optional here, because then we can't pack
445  // the has_***_ boolean fields.
446 
447  caffe2::TypeMeta dtype_ = caffe2::TypeMeta::Make<float>(); // 64-bit
448  Device device_ = at::kCPU; // 32-bit
449  Layout layout_ = at::kStrided; // 8-bit
450 
451  // Bitmask required here to get this to fit inside 32 bits (or even 64 bits,
452  // for that matter)
453 
454  bool requires_grad_ : 1;
455  bool is_variable_ : 1;
456 
457  bool has_device_ : 1;
458  bool has_dtype_ : 1;
459  bool has_layout_ : 1;
460  bool has_requires_grad_ : 1;
461  bool has_is_variable_ : 1;
462 };
463 
464 // We should aspire to fit in one machine-size word; but a size greater than two
465 // words is too much. (We are doing terribly on 32-bit archs, where we require
466 // three machine size words to store tensor options. Eek!)
467 static_assert( sizeof(TensorOptions) <= sizeof(int64_t) * 2,
468  "TensorOptions must fit in 128-bits" );
469 
473  return TensorOptions().dtype(dtype);
474 }
475 
476 // legacy function to support ScalarType
477 inline TensorOptions dtype(ScalarType dtype) {
478  return TensorOptions().dtype(scalarTypeToTypeMeta(dtype));
479 }
480 
483 inline TensorOptions layout(Layout layout) {
484  return TensorOptions().layout(layout);
485 }
486 
490  return TensorOptions().device(std::move(device));
491 }
492 
496  return TensorOptions().device_index(device_index);
497 }
498 
502  return TensorOptions().requires_grad(requires_grad);
503 }
504 
505 C10_API std::ostream& operator<<(
506  std::ostream& stream,
507  const TensorOptions& options);
508 
509 template <typename T>
510 inline TensorOptions dtype() {
511  return dtype(caffe2::TypeMeta::Make<T>());
512 }
513 
514 // This is intended to be a centralized location by which we can determine
515 // what an appropriate TensorTypeId for a tensor is.
516 //
517 // This takes a TensorOptions, rather than just a DeviceType and Layout, because
518 // we reserve the right to change dispatch based on *any* aspect of
519 // TensorOptions. WARNING: If you do this, you need to fix the calls
520 // to computeTensorTypeId in caffe2/tensor.h
521 inline TensorTypeId computeTensorTypeId(TensorOptions options) {
522  return options.computeTensorTypeId();
523 }
524 
525 inline DeviceType computeDeviceType(TensorTypeId tid) {
526  if (tid == CPUTensorId()) {
527  return DeviceType::CPU;
528  } else if (tid == CUDATensorId()) {
529  return DeviceType::CUDA;
530  } else if (tid == HIPTensorId()) {
531  return DeviceType::HIP;
532  } else if (tid == MKLDNNTensorId()) {
533  return DeviceType::MKLDNN;
534  } else if (tid == OpenGLTensorId()) {
535  return DeviceType::IDEEP;
536  } else if (tid == OpenCLTensorId()) {
537  return DeviceType::OPENCL;
538  } else if (tid == IDEEPTensorId()) {
539  return DeviceType::IDEEP;
540  } else if (tid == HIPTensorId()) {
541  return DeviceType::HIP;
542  } else if (tid == MSNPUTensorId()) {
543  return DeviceType::MSNPU;
544  } else if (tid == XLATensorId()) {
545  return DeviceType::XLA;
546  } else if (tid == SparseCPUTensorId()) {
547  return DeviceType::CPU;
548  } else if (tid == SparseCUDATensorId()) {
549  return DeviceType::CUDA;
550  } else if (tid == SparseHIPTensorId()) {
551  return DeviceType::HIP;
552  } else {
553  AT_ASSERTM(false, "Unknown TensorTypeId: ", tid);
554  }
555 }
556 
557 } // namespace c10
C10_NODISCARD TensorOptions device(c10::optional< Device > device) const noexcept
Return a copy of TensorOptions with device set to the given one, or cleared if device is nullopt...
C10_NODISCARD TensorOptions device(Args &&...args) const noexcept
Return a copy of TensorOptions with device set to the given one.
TensorOptions device_index(int16_t device_index)
Convenience function that returns a TensorOptions object with the device set to CUDA and the device_i...
TensorOptions(Layout layout)
Constructs a TensorOptions object with the given layout.
c10::optional< bool > requires_grad_opt() const noexcept
Returns the requires_grad property of the TensorOptions, or c10::nullopt if requires_grad is not spec...
bool has_device() const noexcept
Returns whether the device is specified.
bool has_requires_grad() const noexcept
Returns whether the requires_grad is specified.
bool has_layout() const noexcept
Returns whether the layout is specified.
c10::optional< Device > device_opt() const noexcept
Returns the device of the TensorOptions, or c10::nullopt if device is not specified.
C10_NODISCARD TensorOptions device_index(int16_t device_index) const noexcept
Return a copy of TensorOptions, but with device set to CUDA, and the device index set to the given on...
TensorOptions(Args &&...args)
Constructs a TensorOptions object from arguments allowed in Device constructors.
c10::optional< caffe2::TypeMeta > dtype_opt() const noexcept
Returns the dtype of the TensorOptions, or c10::nullopt if device is not specified.
int32_t device_index() const noexcept
Returns the device index of the TensorOptions.
TensorOptions device(Device device)
Convenience function that returns a TensorOptions object with the device set to the given one...
bool has_dtype() const noexcept
Returns whether the dtype is specified.
Represents a a compute device on which a tensor is located.
Definition: Device.h:30
bool operator!=(const TensorOptions &other) const noexcept
True if any of the elements of this TensorOptions do not match that of the other. ...
TensorOptions(caffe2::TypeMeta dtype)
Constructs a TensorOptions object with the given dtype.
Backend
This legacy enum class defines the set of backends supported by old school, code generated Type-based...
Definition: Backend.h:23
Device device() const noexcept
Returns the device of the TensorOptions.
c10::optional< bool > is_variable_opt() const noexcept
Returns the is_variable property of the TensorOptions, or c10::nullopt if is_variable is not specifie...
Dynamic type ID of a Tensor argument.
Definition: TensorTypeId.h:19
TensorOptions(T &&device)
A class to encapsulate construction axes of an Tensor.
Definition: TensorOptions.h:80
Layout layout() const noexcept
Returns the layout of the TensorOptions.
TensorOptions(ScalarType dtype)
legacy constructor to support ScalarType
TensorOptions layout(Layout layout)
Convenience function that returns a TensorOptions object with the layout set to the given one...
C10_NODISCARD TensorOptions dtype(c10::optional< caffe2::TypeMeta > dtype) const noexcept
Return a copy of TensorOptions with dtype set to the given one.
bool is_variable() const noexcept
Returns the is_variable property of the TensorOptions.
C10_NODISCARD TensorOptions is_variable(c10::optional< bool > is_variable) const noexcept
Sets the is_variable property on the TensorOptions.
bool operator==(const TensorOptions &other) const noexcept
True if all elements of the TensorOptions match that of the other.
To register your own kernel for an operator, do in one (!) cpp file: C10_REGISTER_KERNEL(OperatorHand...
Definition: alias_info.h:7
c10::optional< Layout > layout_opt() const noexcept
Returns the layout of the TensorOptions, or c10::nullopt if layout is not specified.
TensorOptions requires_grad(bool requires_grad=true)
Convenience function that returns a TensorOptions object with the requires_grad set to the given one...
bool requires_grad() const noexcept
Returns the requires_grad property of the TensorOptions.
TypeMeta is a thin class that allows us to store the type of a container such as a blob...
Definition: typeid.h:324
bool has_is_variable() const noexcept
Returns whether the is_variable is specified.
caffe2::TypeMeta dtype() const noexcept
Returns the dtype of the TensorOptions.
C10_NODISCARD TensorOptions layout(c10::optional< Layout > layout) const noexcept
Sets the layout of the TensorOptions.
TensorOptions(T &&device)
Constructs a TensorOptions object with the given device.
C10_NODISCARD TensorOptions requires_grad(c10::optional< bool > requires_grad) const noexcept
Sets the requires_grad property of the TensorOptions.
TensorOptions dtype(caffe2::TypeMeta dtype)
Convenience function that returns a TensorOptions object with the dtype set to the given one...