TensorOptions.h

#pragma once

#include <c10/core/DefaultDtype.h>
#include <c10/core/Backend.h>
#include <c10/core/Layout.h>
#include <c10/core/ScalarType.h>
#include <c10/core/Device.h>

#include <c10/util/Optional.h>
#include <c10/util/C++17.h>
#include <c10/macros/Macros.h>

#include <cstddef>
#include <iosfwd>
#include <utility>

namespace c10 {



struct C10_API TensorOptions {
  TensorOptions()
    : requires_grad_(false)
    , is_variable_(false)
    , has_device_(false)
    , has_dtype_(false)
    , has_layout_(false)
    , has_requires_grad_(false)
    , has_is_variable_(false)
    {}

  /* implicit */ TensorOptions(Layout layout) : TensorOptions() {
    this->set_layout(layout);
  }

  template<typename T,
           typename = c10::guts::enable_if_t<std::is_same<c10::guts::decay_t<T>, Device>::value>>
  /* implicit */ TensorOptions(T&& device) : TensorOptions() {
    this->set_device(std::forward<T>(device));
  }

  template <typename... Args,
            typename = c10::guts::enable_if_t<std::is_constructible<Device, Args&&...>::value>>
   /* implicit */ TensorOptions(Args&&... args)
    : TensorOptions(Device(std::forward<Args>(args)...)) {}

  /* implicit */ TensorOptions(caffe2::TypeMeta dtype) : TensorOptions() {
    this->set_dtype(dtype);
  }

  /* implicit */ TensorOptions(ScalarType dtype) : TensorOptions() {
    this->set_dtype(dtype);
  }

  bool operator==(const TensorOptions& other) const noexcept {
    return
        has_dtype_ == other.has_dtype_ &&
        has_layout_ == other.has_layout_ &&
        has_device_ == other.has_device_ &&
        has_requires_grad_ == other.has_requires_grad_ &&
        has_is_variable_ == other.has_is_variable_ &&
        (!has_dtype_ || dtype_ == other.dtype_) &&
        (!has_layout_ || layout_ == other.layout_) &&
        (!has_device_ || device_ == other.device_) &&
        (!requires_grad_ || requires_grad_ == other.requires_grad_) &&
        (!is_variable_ || is_variable_ == other.is_variable_);
  }

  bool operator!=(const TensorOptions& other) const noexcept {
    return !(*this == other);
  }

  C10_NODISCARD TensorOptions device(c10::optional<Device> device) const noexcept {
    TensorOptions r = *this;
    r.set_device(device);
    return r;
  }

  template<typename ... Args>
  C10_NODISCARD TensorOptions device(Args&&... args) const noexcept {
    return device(c10::optional<Device>(c10::in_place, std::forward<Args>(args)...));
  }

  C10_NODISCARD TensorOptions device_index(int16_t device_index) const noexcept {
    return device(Device::Type::CUDA, device_index);
  }

  C10_NODISCARD TensorOptions dtype(c10::optional<caffe2::TypeMeta> dtype) const noexcept {
    TensorOptions r = *this;
    r.set_dtype(dtype);
    return r;
  }

  // legacy function to support ScalarType
  C10_NODISCARD TensorOptions dtype(c10::optional<ScalarType> dtype) const noexcept {
    TensorOptions r = *this;
    r.set_dtype(dtype);
    return r;
  }

  // Since dtype is taken...
  template <typename T>
  TensorOptions& dtype() {
    dtype_ = caffe2::TypeMeta::Make<T>();
    has_dtype_ = true;
    return *this;
  }

  C10_NODISCARD TensorOptions layout(c10::optional<Layout> layout) const noexcept {
    TensorOptions r = *this;
    r.set_layout(layout);
    return r;
  }

  C10_NODISCARD TensorOptions requires_grad(c10::optional<bool> requires_grad) const noexcept {
    TensorOptions r = *this;
    r.set_requires_grad(requires_grad);
    return r;
  }

  C10_NODISCARD TensorOptions is_variable(c10::optional<bool> is_variable) const noexcept {
    TensorOptions r = *this;
    r.set_is_variable(is_variable);
    return r;
  }

  Device device() const noexcept {
    return has_device_ ? device_ : Device(kCPU);
  }

  bool has_device() const noexcept {
    return has_device_;
  }

  c10::optional<Device> device_opt() const noexcept {
    return has_device_ ? c10::make_optional(device_) : c10::nullopt;
  }

  int32_t device_index() const noexcept {
    return device().index();
  }

  caffe2::TypeMeta dtype() const noexcept {
    return has_dtype_ ? dtype_ : get_default_dtype();
  }

  bool has_dtype() const noexcept {
    return has_dtype_;
  }

  c10::optional<caffe2::TypeMeta> dtype_opt() const noexcept {
    return has_dtype_ ? c10::make_optional(dtype_) : c10::nullopt;
  }

  Layout layout() const noexcept {
    return has_layout_ ? layout_ : kStrided;
  }

  bool has_layout() const noexcept {
    return has_layout_;
  }

  c10::optional<Layout> layout_opt() const noexcept {
    return has_layout_ ? c10::make_optional(layout_) : c10::nullopt;
  }

  bool requires_grad() const noexcept {
    return has_requires_grad_ ? requires_grad_ : false;
  }

  bool has_requires_grad() const noexcept {
    return has_requires_grad_;
  }

  c10::optional<bool> requires_grad_opt() const noexcept {
    return has_requires_grad_ ? c10::make_optional(requires_grad_)
                              : c10::nullopt;
  }

  bool is_variable() const noexcept {
    return has_is_variable_ ? is_variable_ : false;
  }

  bool has_is_variable() const noexcept {
    return has_is_variable_;
  }

  c10::optional<bool> is_variable_opt() const noexcept {
    return has_is_variable_ ? c10::make_optional(is_variable_) : c10::nullopt;
  }

  // Resolves the ATen backend specified by the current construction axes.
  Backend backend() const noexcept {
    return at::tensorTypeIdToBackend(computeTensorTypeId());
  }

  inline TensorTypeId computeTensorTypeId() const {
    switch (layout()) {
      case Layout::Strided:
        switch (device().type()) {
          case DeviceType::CPU:
            return CPUTensorId();
          case DeviceType::CUDA:
            return CUDATensorId();
          case DeviceType::MKLDNN:
            return MKLDNNTensorId();
          case DeviceType::OPENGL:
            return OpenGLTensorId();
          case DeviceType::OPENCL:
            return OpenCLTensorId();
          case DeviceType::IDEEP:
            return IDEEPTensorId();
          case DeviceType::HIP:
            return HIPTensorId();
          case DeviceType::MSNPU:
            return MSNPUTensorId();
          case DeviceType::XLA:
            return XLATensorId();
          default:
            AT_ERROR("Unsupported device type for dense layout: ", device().type());
        }
      case Layout::Sparse:
        switch (device().type()) {
          case DeviceType::CPU:
            return SparseCPUTensorId();
          case DeviceType::CUDA:
            return SparseCUDATensorId();
          case DeviceType::HIP:
            return SparseHIPTensorId();
          default:
            AT_ERROR("Unsupported device type for sparse layout: ", device().type());
        }
      default:
        AT_ERROR("Unsupported layout: ", layout());
    }
  }

 private:

  // These methods are currently private because I'm not sure if it's wise
  // to actually publish them.  They are methods because I need them in
  // the constructor and the functional API implementation.
  //
  // If you really, really need it, you can make these public, but check if you
  // couldn't just do what you need with the functional API.  Similarly, these
  // methods are not chainable, because if you wanted chaining, you probably
  // want to use the functional API instead.  (It's probably OK to make
  // these chainable, because these functions are all explicitly annotated
  // with a ref-qualifier, the trailing &, that makes them illegal to call
  // on temporaries.)

  void set_device(c10::optional<Device> device) & noexcept {
    if (device) {
      device_ = *device;
      has_device_ = true;
    } else {
      has_device_ = false;
    }
  }

  void set_dtype(c10::optional<caffe2::TypeMeta> dtype) & noexcept {
    if (dtype) {
      dtype_ = *dtype;
      has_dtype_ = true;
    } else {
      has_dtype_ = false;
    }
  }

  // legacy function to support ScalarType
  void set_dtype(c10::optional<ScalarType> dtype) & noexcept {
    if (dtype) {
      dtype_ = scalarTypeToTypeMeta(*dtype);
      has_dtype_ = true;
    } else {
      has_dtype_ = false;
    }
  }

  void set_layout(c10::optional<Layout> layout) & noexcept {
    if (layout) {
      layout_ = *layout;
      has_layout_ = true;
    } else {
      has_layout_ = false;
    }
  }

  void set_requires_grad(c10::optional<bool> requires_grad) & noexcept {
    if (requires_grad) {
      requires_grad_ = *requires_grad;
      has_requires_grad_ = true;
    } else {
      has_requires_grad_ = false;
    }
  }

  void set_is_variable(c10::optional<bool> is_variable) & noexcept {
    if (is_variable) {
      is_variable_ = *is_variable;
      has_is_variable_ = true;
    } else {
      has_is_variable_ = false;
    }
  }

  // WARNING: If you edit TensorOptions to add more options, you
  // must adjust the implementation of Tensor::options

  // NB: We didn't use c10::optional here, because then we can't pack
  // the has_***_ boolean fields.

  caffe2::TypeMeta dtype_ = caffe2::TypeMeta::Make<float>(); // 64-bit
  Device device_ = at::kCPU; // 32-bit
  Layout layout_ = at::kStrided; // 8-bit

  // Bitmask required here to get this to fit inside 32 bits (or even 64 bits,
  // for that matter)

  bool requires_grad_     : 1;
  bool is_variable_       : 1;

  bool has_device_        : 1;
  bool has_dtype_         : 1;
  bool has_layout_        : 1;
  bool has_requires_grad_ : 1;
  bool has_is_variable_   : 1;
};

// We should aspire to fit in one machine-size word; but a size greater than two
// words is too much.  (We are doing terribly on 32-bit archs, where we require
// three machine size words to store tensor options.  Eek!)
static_assert( sizeof(TensorOptions) <= sizeof(int64_t) * 2,
               "TensorOptions must fit in 128-bits" );

inline TensorOptions dtype(caffe2::TypeMeta dtype) {
  return TensorOptions().dtype(dtype);
}

// legacy function to support ScalarType
inline TensorOptions dtype(ScalarType dtype) {
  return TensorOptions().dtype(scalarTypeToTypeMeta(dtype));
}

inline TensorOptions layout(Layout layout) {
  return TensorOptions().layout(layout);
}

inline TensorOptions device(Device device) {
  return TensorOptions().device(std::move(device));
}

inline TensorOptions device_index(int16_t device_index) {
  return TensorOptions().device_index(device_index);
}

inline TensorOptions requires_grad(bool requires_grad = true) {
  return TensorOptions().requires_grad(requires_grad);
}

C10_API std::ostream& operator<<(
    std::ostream& stream,
    const TensorOptions& options);

template <typename T>
inline TensorOptions dtype() {
  return dtype(caffe2::TypeMeta::Make<T>());
}

// This is intended to be a centralized location by which we can determine
// what an appropriate TensorTypeId for a tensor is.
//
// This takes a TensorOptions, rather than just a DeviceType and Layout, because
// we reserve the right to change dispatch based on *any* aspect of
// TensorOptions.  WARNING: If you do this, you need to fix the calls
// to computeTensorTypeId in caffe2/tensor.h
inline TensorTypeId computeTensorTypeId(TensorOptions options) {
  return options.computeTensorTypeId();
}

inline DeviceType computeDeviceType(TensorTypeId tid) {
  if (tid == CPUTensorId()) {
    return DeviceType::CPU;
  } else if (tid == CUDATensorId()) {
    return DeviceType::CUDA;
  } else if (tid == HIPTensorId()) {
    return DeviceType::HIP;
  } else if (tid == MKLDNNTensorId()) {
    return DeviceType::MKLDNN;
  } else if (tid == OpenGLTensorId()) {
    return DeviceType::IDEEP;
  } else if (tid == OpenCLTensorId()) {
    return DeviceType::OPENCL;
  } else if (tid == IDEEPTensorId()) {
    return DeviceType::IDEEP;
  } else if (tid == HIPTensorId()) {
    return DeviceType::HIP;
  } else if (tid == MSNPUTensorId()) {
    return DeviceType::MSNPU;
  } else if (tid == XLATensorId()) {
    return DeviceType::XLA;
  } else if (tid == SparseCPUTensorId()) {
    return DeviceType::CPU;
  } else if (tid == SparseCUDATensorId()) {
    return DeviceType::CUDA;
  } else if (tid == SparseHIPTensorId()) {
    return DeviceType::HIP;
  } else {
    AT_ASSERTM(false, "Unknown TensorTypeId: ", tid);
  }
}

} // namespace c10