Scalar.h

#pragma once

#include <cstdint>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <c10/core/OptionalRef.h>
#include <c10/core/ScalarType.h>
#include <c10/core/SymBool.h>
#include <c10/core/SymFloat.h>
#include <c10/core/SymInt.h>
#include <c10/core/SymNodeImpl.h>
#include <c10/macros/Export.h>
#include <c10/macros/Macros.h>
#include <c10/util/Deprecated.h>
#include <c10/util/Exception.h>
#include <c10/util/Half.h>
#include <c10/util/TypeCast.h>
#include <c10/util/complex.h>
#include <c10/util/intrusive_ptr.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)) {}

  void destroy() {
    if (Tag::HAS_si == tag || Tag::HAS_sd == tag || Tag::HAS_sb == tag) {
      raw::intrusive_ptr::decref(v.p);
      v.p = nullptr;
    }
  }

  ~Scalar() {
    destroy();
  }

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

  AT_FORALL_SCALAR_TYPES_AND7(
      Half,
      BFloat16,
      Float8_e5m2,
      Float8_e4m3fn,
      Float8_e5m2fnuz,
      Float8_e4m3fnuz,
      ComplexHalf,
      DEFINE_IMPLICIT_CTOR)
  AT_FORALL_COMPLEX_TYPES(DEFINE_IMPLICIT_CTOR)

  // Helper constructors to allow Scalar creation from long and long long types
  // As std::is_same_v<long, long long> is false(except Android), one needs to
  // provide a constructor from either long or long long in addition to one from
  // int64_t
#if defined(__APPLE__) || defined(__MACOSX)
  static_assert(
      std::is_same_v<long long, int64_t>,
      "int64_t is the same as long long on MacOS");
  Scalar(long vv) : Scalar(vv, true) {}
#endif
#if defined(__linux__) && !defined(__ANDROID__)
  static_assert(
      std::is_same_v<long, int64_t>,
      "int64_t is the same as long on Linux");
  Scalar(long long vv) : Scalar(vv, true) {}
#endif

  Scalar(uint16_t vv) : Scalar(vv, true) {}
  Scalar(uint32_t vv) : Scalar(vv, true) {}
  Scalar(uint64_t vv) {
    if (vv > static_cast<uint64_t>(INT64_MAX)) {
      tag = Tag::HAS_u;
      v.u = vv;
    } else {
      tag = Tag::HAS_i;
      // NB: no need to use convert, we've already tested convertibility
      v.i = static_cast<int64_t>(vv);
    }
  }

#undef DEFINE_IMPLICIT_CTOR

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

  template <
      typename T,
      typename std::enable_if_t<std::is_same_v<T, c10::SymBool>, bool>* =
          nullptr>
  Scalar(T vv) : tag(Tag::HAS_sb) {
    v.i = convert<int64_t, c10::SymBool>(vv);
  }

#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, c10::complex<double>>(v.z, #type); \
    }                                                                 \
    if (Tag::HAS_b == tag) {                                          \
      return checked_convert<type, bool>(v.i, #type);                 \
    } else if (Tag::HAS_i == tag) {                                   \
      return checked_convert<type, int64_t>(v.i, #type);              \
    } else if (Tag::HAS_u == tag) {                                   \
      return checked_convert<type, uint64_t>(v.u, #type);             \
    } else if (Tag::HAS_si == tag) {                                  \
      return checked_convert<type, int64_t>(                          \
          toSymInt().guard_int(__FILE__, __LINE__), #type);           \
    } else if (Tag::HAS_sd == tag) {                                  \
      return checked_convert<type, int64_t>(                          \
          toSymFloat().guard_float(__FILE__, __LINE__), #type);       \
    } else if (Tag::HAS_sb == tag) {                                  \
      return checked_convert<type, int64_t>(                          \
          toSymBool().guard_bool(__FILE__, __LINE__), #type);         \
    }                                                                 \
    TORCH_CHECK(false)                                                \
  }

  // TODO: Support ComplexHalf accessor
  AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_ACCESSOR)
  DEFINE_ACCESSOR(uint16_t, UInt16)
  DEFINE_ACCESSOR(uint32_t, UInt32)
  DEFINE_ACCESSOR(uint64_t, UInt64)

#undef DEFINE_ACCESSOR

  SymInt toSymInt() const {
    if (Tag::HAS_si == tag) {
      return c10::SymInt(intrusive_ptr<SymNodeImpl>::reclaim_copy(
          static_cast<SymNodeImpl*>(v.p)));
    } else {
      return toLong();
    }
  }

  SymFloat toSymFloat() const {
    if (Tag::HAS_sd == tag) {
      return c10::SymFloat(intrusive_ptr<SymNodeImpl>::reclaim_copy(
          static_cast<SymNodeImpl*>(v.p)));
    } else {
      return toDouble();
    }
  }

  SymBool toSymBool() const {
    if (Tag::HAS_sb == tag) {
      return c10::SymBool(intrusive_ptr<SymNodeImpl>::reclaim_copy(
          static_cast<SymNodeImpl*>(v.p)));
    } else {
      return toBool();
    }
  }

  // also support scalar.to<int64_t>();
  // Deleted for unsupported types, but specialized below for supported types
  template <typename T>
  T to() const = delete;

  // audit uses of data_ptr
  const void* data_ptr() const {
    TORCH_INTERNAL_ASSERT(!isSymbolic());
    return static_cast<const void*>(&v);
  }

  bool isFloatingPoint() const {
    return Tag::HAS_d == tag || Tag::HAS_sd == tag;
  }

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

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

  // you probably don't actually want these; they're mostly for testing
  bool isSymInt() const {
    return Tag::HAS_si == tag;
  }
  bool isSymFloat() const {
    return Tag::HAS_sd == tag;
  }
  bool isSymBool() const {
    return Tag::HAS_sb == tag;
  }

  bool isSymbolic() const {
    return Tag::HAS_si == tag || Tag::HAS_sd == tag || Tag::HAS_sb == tag;
  }

  C10_ALWAYS_INLINE Scalar& operator=(Scalar&& other) noexcept {
    if (&other == this) {
      return *this;
    }

    destroy();
    moveFrom(std::move(other));
    return *this;
  }

  C10_ALWAYS_INLINE Scalar& operator=(const Scalar& other) {
    if (&other == this) {
      return *this;
    }

    *this = Scalar(other);
    return *this;
  }

  Scalar operator-() const;
  Scalar conj() const;
  Scalar log() const;

  template <
      typename T,
      typename std::enable_if_t<!c10::is_complex<T>::value, int> = 0>
  bool equal(T num) const {
    if (isComplex()) {
      TORCH_INTERNAL_ASSERT(!isSymbolic());
      auto val = v.z;
      return (val.real() == num) && (val.imag() == T());
    } else if (isFloatingPoint()) {
      TORCH_CHECK(!isSymbolic(), "NYI SymFloat equality");
      return v.d == num;
    } else if (tag == Tag::HAS_i) {
      if (overflows<T>(v.i, /* strict_unsigned */ true)) {
        return false;
      } else {
        return static_cast<T>(v.i) == num;
      }
    } else if (tag == Tag::HAS_u) {
      if (overflows<T>(v.u, /* strict_unsigned */ true)) {
        return false;
      } else {
        return static_cast<T>(v.u) == num;
      }
    } else if (tag == Tag::HAS_si) {
      TORCH_INTERNAL_ASSERT(false, "NYI SymInt equality");
    } else if (isBoolean()) {
      // boolean scalar does not equal to a non boolean value
      TORCH_INTERNAL_ASSERT(!isSymbolic());
      return false;
    } else {
      TORCH_INTERNAL_ASSERT(false);
    }
  }

  template <
      typename T,
      typename std::enable_if_t<c10::is_complex<T>::value, int> = 0>
  bool equal(T num) const {
    if (isComplex()) {
      TORCH_INTERNAL_ASSERT(!isSymbolic());
      return v.z == num;
    } else if (isFloatingPoint()) {
      TORCH_CHECK(!isSymbolic(), "NYI SymFloat equality");
      return (v.d == num.real()) && (num.imag() == T());
    } else if (tag == Tag::HAS_i) {
      if (overflows<T>(v.i, /* strict_unsigned */ true)) {
        return false;
      } else {
        return static_cast<T>(v.i) == num.real() && num.imag() == T();
      }
    } else if (tag == Tag::HAS_u) {
      if (overflows<T>(v.u, /* strict_unsigned */ true)) {
        return false;
      } else {
        return static_cast<T>(v.u) == num.real() && num.imag() == T();
      }
    } else if (tag == Tag::HAS_si) {
      TORCH_INTERNAL_ASSERT(false, "NYI SymInt equality");
    } else if (isBoolean()) {
      // boolean scalar does not equal to a non boolean value
      TORCH_INTERNAL_ASSERT(!isSymbolic());
      return false;
    } else {
      TORCH_INTERNAL_ASSERT(false);
    }
  }

  bool equal(bool num) const {
    if (isBoolean()) {
      TORCH_INTERNAL_ASSERT(!isSymbolic());
      return static_cast<bool>(v.i) == num;
    } else {
      return false;
    }
  }

  ScalarType type() const {
    if (isComplex()) {
      return ScalarType::ComplexDouble;
    } else if (isFloatingPoint()) {
      return ScalarType::Double;
    } else if (isIntegral(/*includeBool=*/false)) {
      // Represent all integers as long, UNLESS it is unsigned and therefore
      // unrepresentable as long
      if (Tag::HAS_u == tag) {
        return ScalarType::UInt64;
      }
      return ScalarType::Long;
    } else if (isBoolean()) {
      return ScalarType::Bool;
    } else {
      throw std::runtime_error("Unknown scalar type.");
    }
  }

  Scalar(Scalar&& rhs) noexcept : tag(rhs.tag) {
    moveFrom(std::move(rhs));
  }

  Scalar(const Scalar& rhs) : tag(rhs.tag), v(rhs.v) {
    if (isSymbolic()) {
      c10::raw::intrusive_ptr::incref(v.p);
    }
  }

  Scalar(c10::SymInt si) {
    if (auto m = si.maybe_as_int()) {
      tag = Tag::HAS_i;
      v.i = *m;
    } else {
      tag = Tag::HAS_si;
      v.p = std::move(si).release();
    }
  }

  Scalar(c10::SymFloat sd) {
    if (sd.is_symbolic()) {
      tag = Tag::HAS_sd;
      v.p = std::move(sd).release();
    } else {
      tag = Tag::HAS_d;
      v.d = sd.as_float_unchecked();
    }
  }

  Scalar(c10::SymBool sb) {
    if (auto m = sb.maybe_as_bool()) {
      tag = Tag::HAS_b;
      v.i = *m;
    } else {
      tag = Tag::HAS_sb;
      v.p = std::move(sb).release();
    }
  }

  // We can't set v in the initializer list using the
  // syntax v{ .member = ... } because it doesn't work on MSVC
 private:
  enum class Tag { HAS_d, HAS_i, HAS_u, HAS_z, HAS_b, HAS_sd, HAS_si, HAS_sb };

  // Note [Meaning of HAS_u]
  // ~~~~~~~~~~~~~~~~~~~~~~~
  // HAS_u is a bit special.  On its face, it just means that we
  // are holding an unsigned integer.  However, we generally don't
  // distinguish between different bit sizes in Scalar (e.g., we represent
  // float as double), instead, it represents a mathematical notion
  // of some quantity (integral versus floating point).  So actually,
  // HAS_u is used solely to represent unsigned integers that could
  // not be represented as a signed integer.  That means only uint64_t
  // potentially can get this tag; smaller types like uint8_t fits into a
  // regular int and so for BC reasons we keep as an int.

  // NB: assumes that self has already been cleared
  // NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved)
  C10_ALWAYS_INLINE void moveFrom(Scalar&& rhs) noexcept {
    v = rhs.v;
    tag = rhs.tag;
    if (rhs.tag == Tag::HAS_si || rhs.tag == Tag::HAS_sd ||
        rhs.tag == Tag::HAS_sb) {
      // Move out of scalar
      rhs.tag = Tag::HAS_i;
      rhs.v.i = 0;
    }
  }

  Tag tag;

  union v_t {
    double d{};
    int64_t i;
    // See Note [Meaning of HAS_u]
    uint64_t u;
    c10::complex<double> z;
    c10::intrusive_ptr_target* p;
    // NOLINTNEXTLINE(modernize-use-equals-default)
    v_t() {} // default constructor
  } v;

  template <
      typename T,
      typename std::enable_if_t<
          std::is_integral_v<T> && !std::is_same_v<T, bool>,
          bool>* = nullptr>
  Scalar(T vv, bool) : tag(Tag::HAS_i) {
    v.i = convert<decltype(v.i), T>(vv);
  }

  template <
      typename T,
      typename std::enable_if_t<
          !std::is_integral_v<T> && !c10::is_complex<T>::value,
          bool>* = nullptr>
  Scalar(T vv, bool) : tag(Tag::HAS_d) {
    v.d = convert<decltype(v.d), T>(vv);
  }

  template <
      typename T,
      typename std::enable_if_t<c10::is_complex<T>::value, bool>* = nullptr>
  Scalar(T vv, bool) : tag(Tag::HAS_z) {
    v.z = convert<decltype(v.z), T>(vv);
  }
};

using OptionalScalarRef = c10::OptionalRef<Scalar>;

// define the scalar.to<int64_t>() specializations
#define DEFINE_TO(T, name)         \
  template <>                      \
  inline T Scalar::to<T>() const { \
    return to##name();             \
  }
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_TO)
DEFINE_TO(uint16_t, UInt16)
DEFINE_TO(uint32_t, UInt32)
DEFINE_TO(uint64_t, UInt64)
#undef DEFINE_TO

} // namespace c10