Type.hpp

Go to the documentation of this file.

#ifndef CYTNX_TYPE_H_
#define CYTNX_TYPE_H_
 
#include <complex>
#include <cstdint>
#include <string>
#include <type_traits>
#include <tuple>
#include <array>
#include <utility>
#include <vector>
#include <variant>
 
#include "cytnx_error.hpp"  // also brings in cuComplex.h
 
#ifdef UNI_GPU
  #include <cuda/std/complex>
#endif
 
#define MKL_Complex8 std::complex<float>
#define MKL_Complex16 std::complex<double>
 
#ifdef BACKEND_TORCH
typedef int32_t blas_int;
#else
 
  #ifdef UNI_MKL
    #include <mkl.h>
typedef MKL_INT blas_int;
  #else
typedef int32_t blas_int;
  #endif
 
#endif
 
// @cond
namespace cytnx {
 
  template <class T>
  using vec3d = std::vector<std::vector<std::vector<T>>>;
 
  template <class T>
  using vec2d = std::vector<std::vector<T>>;
 
  typedef double cytnx_double;
  typedef float cytnx_float;
  typedef uint64_t cytnx_uint64;
  typedef uint32_t cytnx_uint32;
  typedef uint16_t cytnx_uint16;
  typedef int64_t cytnx_int64;
  typedef int32_t cytnx_int32;
  typedef int16_t cytnx_int16;
  typedef std::size_t cytnx_size_t;
  typedef std::complex<float> cytnx_complex64;
  typedef std::complex<double> cytnx_complex128;
  typedef bool cytnx_bool;
 
#ifdef UNI_GPU
  using cytnx_cuda_complex64 = cuda::std::complex<float>;
  using cytnx_cuda_complex128 = cuda::std::complex<double>;
#endif
 
  namespace internal {
    template <class>
    struct is_complex_impl : std::false_type {};
 
    template <class T>
    struct is_complex_impl<std::complex<T>> : std::true_type {};
 
#ifdef UNI_GPU
    template <class T>
    struct is_complex_impl<cuda::std::complex<T>> : std::true_type {};
#endif
 
    template <typename>
    struct is_complex_floating_point_impl : std::false_type {};
 
    template <typename T>
    struct is_complex_floating_point_impl<std::complex<T>> : std::is_floating_point<T> {};
 
#ifdef UNI_GPU
    template <typename T>
    struct is_complex_floating_point_impl<cuda::std::complex<T>> : std::is_floating_point<T> {};
#endif
 
    template <std::size_t Idx, typename T, typename Tuple>
    constexpr std::size_t index_in_tuple_helper() {
      static_assert(Idx < std::tuple_size_v<Tuple>, "Type not found!");
      if constexpr (std::is_same_v<T, std::tuple_element_t<Idx, Tuple>>) {
        return Idx;
      } else {
        return index_in_tuple_helper<Idx + 1, T, Tuple>();
      }
    }
 
  }  // namespace internal
 
  // helper metafunction to transform a variant into another variant via a
  // transform template alias
  template <typename V, template <typename> class Transform>
  struct make_variant_from_transform;
 
  template <template <typename> class Transform, typename... Args>
  struct make_variant_from_transform<std::variant<Args...>, Transform> {
    using type = std::variant<typename Transform<Args>::type...>;
  };
 
  // helper type alias for make_variant_from_transform
  template <typename V, template <typename> class Transform>
  using make_variant_from_transform_t = typename make_variant_from_transform<V, Transform>::type;
 
  template <typename T>
  using is_complex = internal::is_complex_impl<std::remove_cv_t<T>>;
 
  template <typename T>
  using is_complex_floating_point = internal::is_complex_floating_point_impl<std::remove_cv_t<T>>;
 
  // is_complex_v checks if a data type is of type std::complex
  // usage: is_complex_v<T> returns true or false for a data type T
  template <typename T>
  constexpr bool is_complex_v = is_complex<T>::value;
 
  // is_complex_floating_point_v<T> is a template constant that is true if T is of type
  // std::complex<U> where U is a floating point type, and false otherwise.
  template <typename T>
  constexpr bool is_complex_floating_point_v = is_complex_floating_point<T>::value;
 
  template <typename>
  inline constexpr bool always_false_v = false;
 
  template <typename T, typename Variant>
  inline constexpr bool variant_contains_v = false;
 
  template <typename T, typename... Types>
  inline constexpr bool variant_contains_v<T, std::variant<Types...>> = (std::is_same_v<T, Types> ||
                                                                         ...);
 
  // variant_index<T, Variant> returns the index of type T in the Variant.
  template <typename T, typename Variant>
  struct variant_index;
 
  template <typename T>
  struct variant_index<T, std::variant<>> {
    static_assert(always_false_v<T>, "variant_index<T, Variant>: T is not in Variant");
    static constexpr std::size_t value = 0;
  };
 
  template <typename T, typename... Types>
  struct variant_index<T, std::variant<T, Types...>> {
    static constexpr std::size_t value = 0;
  };
 
  template <typename T, typename U, typename... Types>
  struct variant_index<T, std::variant<U, Types...>> {
    static constexpr std::size_t value = 1 + variant_index<T, std::variant<Types...>>::value;
  };
 
  // helper template variable
  template <typename T, typename Variant>
  static constexpr std::size_t variant_index_v = variant_index<T, Variant>::value;
 
  namespace internal {
    // type_size returns the sizeof(T) for the supported types. This is the same as
    // sizeof(T), except that size_type<void> is 0.
    template <typename T>
    inline constexpr int type_size = sizeof(T);
    template <>
    inline constexpr int type_size<void> = 0;
  }  // namespace internal
 
  // the list of supported types. The dtype() of an object is an index into this list.
  // std::variant works better than std::tuple here since a variant is constrained to only
  // hold each type once, and we have std::variant_alternative_t<n> to get the n'th type,
  // as well as the variant_index_v helper to get the index of a given type
  using Type_list =
    std::variant<void, cytnx_complex128, cytnx_complex64, cytnx_double, cytnx_float, cytnx_int64,
                 cytnx_uint64, cytnx_int32, cytnx_uint32, cytnx_int16, cytnx_uint16, cytnx_bool>;
 
  // For GPU kernels, use cuda::std::complex for complex arithmetic. Low-level CUDA library calls
  // that require cuComplex ABI pointers should cast explicitly at those call boundaries.
#ifdef UNI_GPU
  using Type_list_gpu =
    std::variant<void, cytnx_cuda_complex128, cytnx_cuda_complex64, cytnx_double, cytnx_float,
                 cytnx_int64, cytnx_uint64, cytnx_int32, cytnx_uint32, cytnx_int16, cytnx_uint16,
                 cytnx_bool>;
#endif
 
  // CytnxType<T> is satisfied by the element types that have a cytnx dtype (the members of
  // Type_list, excluding the Void placeholder). Storage_base::data/at/back are constrained to it
  // so that requesting an unsupported T is a compile-time error at the call site. The GPU cuComplex
  // / cuda::std::complex pointer views are non-cytnx-dtype types and are provided separately as
  // explicit specializations.
  template <typename T>
  concept CytnxType = variant_contains_v<T, Type_list> && !std::is_void_v<T>;
 
#ifdef UNI_GPU
  // The GPU complex pointer-view types that Storage_base::data<T>() specializes for. They are not
  // cytnx dtypes: cuDoubleComplex/cuFloatComplex are the cuComplex ABI types for CUDA library
  // calls, and cuda::std::complex<...> is the representation GPU kernels use internally.
  template <typename T>
  concept GpuComplexView =
    std::is_same_v<T, cuDoubleComplex> || std::is_same_v<T, cuFloatComplex> ||
    std::is_same_v<T, cytnx_cuda_complex128> || std::is_same_v<T, cytnx_cuda_complex64>;
 
  // The element types data<T>() accepts: cytnx dtypes plus the GPU complex pointer views.
  template <typename T>
  concept StorageDataType = CytnxType<T> || GpuComplexView<T>;
#else
  template <typename T>
  concept StorageDataType = CytnxType<T>;
#endif
 
  // The number of supported types
  constexpr int N_Type = std::variant_size_v<Type_list>;
  constexpr int N_fType = 5;
 
  // The friendly name of each type
  template <typename T>
  inline constexpr char* Type_names = nullptr;
  template <>
  inline constexpr const char* Type_names<void> = "Void";
  template <>
  inline constexpr const char* Type_names<cytnx_complex128> = "Complex Double (Complex Float64)";
  template <>
  inline constexpr const char* Type_names<cytnx_complex64> = "Complex Float (Complex Float32)";
  template <>
  inline constexpr const char* Type_names<cytnx_double> = "Double (Float64)";
  template <>
  inline constexpr const char* Type_names<cytnx_float> = "Float (Float32)";
  template <>
  inline constexpr const char* Type_names<cytnx_int64> = "Int64";
  template <>
  inline constexpr const char* Type_names<cytnx_uint64> = "Uint64";
  template <>
  inline constexpr const char* Type_names<cytnx_int32> = "Int32";
  template <>
  inline constexpr const char* Type_names<cytnx_uint32> = "Uint32";
  template <>
  inline constexpr const char* Type_names<cytnx_int16> = "Int16";
  template <>
  inline constexpr const char* Type_names<cytnx_uint16> = "Uint16";
  template <>
  inline constexpr const char* Type_names<cytnx_bool> = "Bool";
 
  // The corresponding Python enumeration name
  template <typename T>
  inline constexpr char* Type_enum_name = nullptr;
  template <>
  inline constexpr const char* Type_enum_name<void> = "Void";
  // std::monostate is Void's stand-in inside value variants (e.g.
  // Scalar::ScalarVariant, where a real alternative is needed at index 0);
  // give it the same display name so visit-based error messages can name it.
  template <>
  inline constexpr const char* Type_enum_name<std::monostate> = "Void";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_complex128> = "ComplexDouble";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_complex64> = "ComplexFloat";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_double> = "Double";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_float> = "Float";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_int64> = "Int64";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_uint64> = "Uint64";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_int32> = "Int32";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_uint32> = "Uint32";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_int16> = "Int16";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_uint16> = "Uint16";
  template <>
  inline constexpr const char* Type_enum_name<cytnx_bool> = "Bool";
 
  struct Type_struct {
    const char* name;  // char* is OK here, it is only ever initialized from a string literal
    const char* enum_name;
    bool is_unsigned;
    bool is_complex;
    bool is_float;
    bool is_int;
    unsigned int typeSize;
  };
 
  template <typename T>
  struct Type_struct_t {
    static constexpr unsigned int cy_typeid = variant_index_v<T, Type_list>;
#ifdef UNI_GPU
    static constexpr unsigned int cy_typeid_gpu = variant_index_v<T, Type_list_gpu>;
#endif
    static constexpr const char* name = Type_names<T>;
    static constexpr const char* enum_name = Type_enum_name<T>;
    static constexpr bool is_complex = is_complex_v<T>;
    static constexpr bool is_unsigned = std::is_unsigned_v<T>;
    static constexpr bool is_float = std::is_floating_point_v<T> || is_complex_floating_point_v<T>;
    static constexpr bool is_int = std::is_integral_v<T> && !std::is_same_v<T, bool>;
    static constexpr std::size_t typeSize = internal::type_size<T>;
 
    static constexpr Type_struct construct() {
      return {name, enum_name, is_unsigned, is_complex, is_float, is_int, typeSize};
    }
  };
 
  namespace internal {
    template <typename Variant, std::size_t... Indices>
    constexpr auto make_type_array_helper(std::index_sequence<Indices...>) {
      return std::array<Type_struct, sizeof...(Indices)>{
        Type_struct_t<std::variant_alternative_t<Indices, Variant>>::construct()...};
    }
    template <typename Variant>
    constexpr auto make_type_array() {
      return make_type_array_helper<Variant>(
        std::make_index_sequence<std::variant_size_v<Variant>>());
    }
  }  // namespace internal
 
  class Type_class {
   private:
   public:
    // Typeinfos is a std::array<Type_struct> for each type in Type_list
    static constexpr auto Typeinfos = internal::make_type_array<Type_list>();
 
    template <typename T>
    static constexpr unsigned int cy_typeid_v = variant_index_v<T, Type_list>;
 
#ifdef UNI_GPU
    template <typename T>
    static constexpr unsigned int cy_typeid_gpu_v = variant_index_v<T, Type_list_gpu>;
#endif
 
    enum Type : unsigned int {
      Void = cy_typeid_v<void>,
      ComplexDouble = cy_typeid_v<cytnx_complex128>,
      ComplexFloat = cy_typeid_v<cytnx_complex64>,
      Double = cy_typeid_v<cytnx_double>,
      Float = cy_typeid_v<cytnx_float>,
      Int64 = cy_typeid_v<cytnx_int64>,
      Uint64 = cy_typeid_v<cytnx_uint64>,
      Int32 = cy_typeid_v<cytnx_int32>,
      Uint32 = cy_typeid_v<cytnx_uint32>,
      Int16 = cy_typeid_v<cytnx_int16>,
      Uint16 = cy_typeid_v<cytnx_uint16>,
      Bool = cy_typeid_v<cytnx_bool>
    };
 
    static constexpr void check_type(unsigned int type_id) {
      cytnx_error_msg(type_id >= N_Type, "[ERROR] invalid type_id: %s", type_id);
    }
 
    // This could be constexpr returning constexpr char*, but there is lots of code that
    // assumes that it returns a std::string and calls getname(n).c_str()
    static std::string getname(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].name;
    }
    // This cannot be constexpr as we define it in a .cpp file,
    // and typeid(T).name() is not constexpr
    static unsigned int c_typename_to_id(const std::string& c_name);
    static char const* enum_name(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].enum_name;
    }
    static constexpr unsigned int typeSize(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].typeSize;
    }
    static constexpr bool is_unsigned(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].is_unsigned;
    }
    static constexpr bool is_complex(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].is_complex;
    }
    static constexpr bool is_float(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].is_float;
    }
    static constexpr bool is_int(unsigned int type_id) {
      check_type(type_id);
      return Typeinfos[type_id].is_int;
    }
 
    template <class T>
    static constexpr unsigned int cy_typeid(const T& rc) {
      return cy_typeid_v<T>;
    }
 
    // Real counterpart of a dtype: ComplexDouble -> Double, ComplexFloat -> Float,
    // anything else unchanged. Replaces the "dtype <= 2 ? dtype + 2 : dtype" idiom
    // (without depending on the enum layout).
    static constexpr unsigned int to_real(unsigned int type_id) {
      check_type(type_id);
      if (type_id == ComplexDouble) return Double;
      if (type_id == ComplexFloat) return Float;
      return type_id;
    }
 
    // Complex counterpart of a dtype: Double -> ComplexDouble, Float -> ComplexFloat,
    // complex types unchanged, Void unchanged, integral/bool -> ComplexDouble.
    static constexpr unsigned int to_complex(unsigned int type_id) {
      check_type(type_id);
      if (is_complex(type_id)) return type_id;
      if (type_id == Double) return ComplexDouble;
      if (type_id == Float) return ComplexFloat;
      if (type_id == Void) return Void;
      return ComplexDouble;
    }
 
    // The dtype linalg::Norm() produces for an input of the given dtype: the
    // real counterpart for floating/complex dtypes, Double for integer/bool
    // inputs (which Norm computes in double precision). This is the single
    // home of that policy -- src/linalg/Norm.cpp and callers pre-sizing
    // norm accumulators (e.g. UniTensor::normalize_) both use it, so they
    // cannot drift apart.
    static constexpr unsigned int norm_result_dtype(unsigned int type_id) {
      check_type(type_id);
      // Void has no norm: reject it here rather than silently mapping it to
      // Double (which would let an uninitialized tensor flow onward).
      cytnx_error_msg(type_id == Void,
                      "[ERROR] norm_result_dtype: Void has no norm result dtype.%s", "\n");
      if (is_float(type_id)) return to_real(type_id);
      return Double;
    }
 
    // Find a common type for typeL and typeR
    static constexpr unsigned int type_promote(unsigned int typeL, unsigned int typeR) {
      if (typeL == Void || typeR == Void) return Void;
      // Mixed complex/real: promote the real counterparts, then re-complexify.
      // Fixes ComplexFloat + Double -> ComplexDouble (previously ComplexFloat,
      // discarding precision, because the enum interleaves complexness and
      // precision and promotion picked the lower index).
      if (is_complex(typeL) != is_complex(typeR)) {
        return to_complex(type_promote(to_real(typeL), to_real(typeR)));
      }
      if (typeL < typeR) {
        if (!is_unsigned(typeR) && is_unsigned(typeL)) {
          return typeL - 1;
        } else {
          return typeL;
        }
      } else {
        if (!is_unsigned(typeL) && is_unsigned(typeR)) {
          return typeR - 1;
        } else {
          return typeR;
        }
      }
    }
 
    // type metafunction for type promotion
    template <typename TL, typename TR>
    using type_promote_t =
      std::variant_alternative_t<Type_class::type_promote(variant_index_v<TL, Type_list>,
                                                          variant_index_v<TR, Type_list>),
                                 Type_list>;
 
    // Runtime counterpart of make_floating_point_t, for call sites that only
    // have a dtype id (not a C++ type) at the point an operation's output
    // Tensor/Storage is pre-sized, e.g. Div's out-of-place output allocation
    // before the typed visitor runs. Named distinctly from the
    // make_floating_point<T> type-trait below (a function and a class
    // template cannot share a name in the same scope).
    static constexpr unsigned int make_floating_point_dtype(unsigned int type_id) {
      check_type(type_id);
      if (is_float(type_id)) return type_id;
      return Double;
    }
 
    // The true-division output type for a promoted dtype: integral/bool dtypes
    // become cytnx_double (Python true-division semantics), existing floating
    // dtypes and complex dtypes are unchanged. Used by Div's output-type rule
    // (make_floating_point_t<type_promote_t<TL,TR>>) so int/int division
    // produces a floating result instead of truncating (#941).
    template <typename T>
    struct make_floating_point {
      using type =
        std::conditional_t<std::is_floating_point_v<T> || is_complex_v<T>, T, cytnx_double>;
    };
 
    template <typename T>
    using make_floating_point_t = typename make_floating_point<T>::type;
 
    // Helper to promote two pointer types (note does _not_ return another pointer type)
    template <typename TL, typename TR>
    struct type_promote_from_pointer {
      using type = void;
    };
 
    template <typename TL, typename TR>
    struct type_promote_from_pointer<TL*, TR*> {
      using type = type_promote_t<std::decay_t<TL>, std::decay_t<TR>>;
    };
 
    // helper typedef
    template <typename TL, typename TR>
    using type_promote_from_pointer_t = typename type_promote_from_pointer<TL, TR>::type;
 
#ifdef UNI_GPU
    // .. and we need a version where TL and TR are GPU device pointers
    template <typename TL, typename TR>
    using type_promote_gpu_t =
      std::variant_alternative_t<Type_class::type_promote(variant_index_v<TL, Type_list_gpu>,
                                                          variant_index_v<TR, Type_list_gpu>),
                                 Type_list_gpu>;
 
    template <typename TL, typename TR>
    struct type_promote_from_gpu_pointer {
      using type = void;
    };
 
    template <typename TL, typename TR>
    struct type_promote_from_gpu_pointer<TL*, TR*> {
      using type = type_promote_gpu_t<std::decay_t<TL>, std::decay_t<TR>>;
    };
 
    // helper typedef
    template <typename TL, typename TR>
    using type_promote_from_gpu_pointer_t = typename type_promote_from_gpu_pointer<TL, TR>::type;
#endif
 
  };  // Type_class
 
  constexpr Type_class Type;
 
  extern int __blasINTsize__;
 
  // User_debug is declared in cytnx_error.hpp (included above); no need to redeclare it here.
 
}  // namespace cytnx
 
#endif  // CYTNX_TYPE_H_