c10/core/Scalar.h

#pragma once

#include <assert.h>
#include <stdint.h>
#include <stdexcept>
#include <string>
#include <utility>

#include <c10/core/ScalarType.h>
#include <c10/macros/Macros.h>
#include <c10/util/Half.h>

namespace c10 {

/**
 * Scalar represents a 0-dimensional tensor which contains a single element.
 * Unlike a tensor, numeric literals (in C++) are implicitly convertible to
 * Scalar (which is why, for example, we provide both add(Tensor) and
 * add(Scalar) overloads for many operations). It may also be used in
 * circumstances where you statically know a tensor is 0-dim and single size,
 * but don't know its type.
 */
class C10_API Scalar {
 public:
  Scalar() : Scalar(int64_t(0)) {}

#define DEFINE_IMPLICIT_CTOR(type, name)      \
  Scalar(type vv) : Scalar(vv, true) { }

  AT_FORALL_SCALAR_TYPES_AND2(Half, BFloat16, DEFINE_IMPLICIT_CTOR)

#undef DEFINE_IMPLICIT_CTOR

  // Value* is both implicitly convertible to SymbolicVariable and bool which
  // causes ambiguosity error. Specialized constructor for bool resolves this
  // problem.
  template <
      typename T,
      typename std::enable_if<std::is_same<T, bool>::value, bool>::type* =
          nullptr>
  Scalar(T vv) : tag(Tag::HAS_b) {
    v.i = convert<int64_t, bool>(vv);
  }

#define DEFINE_IMPLICIT_COMPLEX_CTOR(type, name, member) \
  Scalar(type vv) : tag(Tag::HAS_##member) {             \
    v.member[0] = c10::convert<double>(vv.real());       \
    v.member[1] = c10::convert<double>(vv.imag());       \
  }

  DEFINE_IMPLICIT_COMPLEX_CTOR(at::ComplexHalf, ComplexHalf, z)
  DEFINE_IMPLICIT_COMPLEX_CTOR(std::complex<float>, ComplexFloat, z)
  DEFINE_IMPLICIT_COMPLEX_CTOR(std::complex<double>, ComplexDouble, z)

#undef DEFINE_IMPLICIT_COMPLEX_CTOR

#define DEFINE_ACCESSOR(type, name)                       \
  type to##name() const {                                 \
    if (Tag::HAS_d == tag) {                              \
      return checked_convert<type, double>(v.d, #type);   \
    } else if (Tag::HAS_z == tag) {                       \
      return checked_convert<type, std::complex<double>>( \
          {v.z[0], v.z[1]}, #type);                       \
    } if (Tag::HAS_b == tag) {                            \
      return checked_convert<type, bool>(v.i, #type);     \
    } else {                                              \
      return checked_convert<type, int64_t>(v.i, #type);  \
    }                                                     \
  }

  // TODO: Support ComplexHalf accessor
  AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_ACCESSOR)

  // also support scalar.to<int64_t>();
  template <typename T>
  T to() const;

#undef DEFINE_ACCESSOR
  bool isFloatingPoint() const {
    return Tag::HAS_d == tag;
  }

  C10_DEPRECATED_MESSAGE("isIntegral is deprecated. Please use the overload with 'includeBool' parameter instead.")
  bool isIntegral() const {
    return Tag::HAS_i == tag;
  }
  bool isIntegral(bool includeBool) const {
    return Tag::HAS_i == tag || (includeBool && isBoolean());
  }

  bool isComplex() const {
    return Tag::HAS_z == tag;
  }
  bool isBoolean() const {
    return Tag::HAS_b == tag;
  }

  Scalar operator-() const;

 private:
    template<typename T,
             typename std::enable_if<std::numeric_limits<T>::is_integer && ! std::is_same<T, bool>::value, bool>::type* =
                 nullptr>
    Scalar(T vv, bool) : tag(Tag::HAS_i) {
      v.i = convert<decltype(v.i), T>(vv);
    }

    template<typename T,
             typename std::enable_if<!std::numeric_limits<T>::is_integer, bool>::type* =
                 nullptr>
    Scalar(T vv, bool) : tag(Tag::HAS_d) {
      v.d = convert<decltype(v.d), T>(vv);
    }

  // We can't set v in the initializer list using the
  // syntax v{ .member = ... } because it doesn't work on MSVC

  enum class Tag { HAS_d, HAS_i, HAS_z, HAS_b };
  Tag tag;
  union {
    double d;
    int64_t i;
    // Can't do put std::complex in the union, because it triggers
    // an nvcc bug:
    //    error: designator may not specify a non-POD subobject
    double z[2];
  } v;
};

// define the scalar.to<int64_t>() specializations
template <typename T>
inline T Scalar::to() const {
  throw std::runtime_error("to() cast to unexpected type.");
}

#define DEFINE_TO(T, name)    \
  template <>                 \
  inline T Scalar::to<T>() const {  \
    return to##name();        \
  }
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_TO)
#undef DEFINE_TO
} // namespace c10