tensor_types.h

Go to the documentation of this file.

 
#ifndef ATEN_CORE_TENSOR_TYPES_H_
#define ATEN_CORE_TENSOR_TYPES_H_
 
#include <regex>
#include <string>
#include <vector>
#include <numeric>
#include <complex>
#include <utility>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <algorithm>
#include <unordered_map>
#include <initializer_list>
 
#include "source_base/module_device/types.h"
 
#if defined(__CUDACC__)
#include <base/macros/cuda.h>
#elif defined(__HIPCC__)
#include <base/macros/rocm.h>
#endif // defined(__CUDACC__) || defined(__HIPCC__)
 
namespace container {
 
template <typename T, int Accuracy>
static inline bool element_compare(T& a, T& b) {
    if (Accuracy <= 4) {
        return (a == b) || (std::norm(a - b) < 1e-7);
    } 
    else if (Accuracy <= 8) {
        return (a == b) || (std::norm(a - b) < 1e-15);
    } 
    else {
        return (a == b);
    }
}
 
enum class DataType {
    DT_INVALID = 0, 
    DT_FLOAT = 1, 
    DT_DOUBLE = 2, 
    DT_INT = 3, 
    DT_INT64 = 4, 
    DT_COMPLEX = 5, 
    DT_COMPLEX_DOUBLE = 6, 
// ... other data types
};
 
struct DEVICE_CPU;
struct DEVICE_GPU;
 
struct DEVICE_CPU {};
struct DEVICE_GPU {};
enum class DeviceType {
    UnKnown = 0,  
    CpuDevice = 1,     
    GpuDevice = 2,     
};
 
template <typename T>
struct GetTypeReal {
    using type = T; 
};
 
template <>
struct GetTypeReal<std::complex<float>> {
    using type = float; 
};
 
template <>
struct GetTypeReal<std::complex<double>> {
    using type = double; 
};
 
template <typename T> 
struct PsiToContainer {
    using type = T; 
};
 
template <>
struct PsiToContainer<base_device::DEVICE_CPU>
{
    using type = container::DEVICE_CPU; 
};
 
template <>
struct PsiToContainer<base_device::DEVICE_GPU>
{
    using type = container::DEVICE_GPU; 
};
 
template <typename T> 
struct ContainerToPsi {
    using type = T; 
};
 
template <>
struct ContainerToPsi<container::DEVICE_CPU> {
    using type = base_device::DEVICE_CPU; 
};
 
template <>
struct ContainerToPsi<container::DEVICE_GPU> {
    using type = base_device::DEVICE_GPU; 
};
 
 
template <typename T>
struct DeviceTypeToEnum {
    static constexpr DeviceType value = {};
};
// Specializations of DeviceTypeToEnum for supported devices.
template <>
struct DeviceTypeToEnum<DEVICE_CPU> {
    static constexpr DeviceType value = DeviceType::CpuDevice;
};
template <>
struct DeviceTypeToEnum<DEVICE_GPU> {
    static constexpr DeviceType value = DeviceType::GpuDevice;
};
template <>
struct DeviceTypeToEnum<base_device::DEVICE_CPU>
{
    static constexpr DeviceType value = DeviceType::CpuDevice;
};
template <>
struct DeviceTypeToEnum<base_device::DEVICE_GPU>
{
    static constexpr DeviceType value = DeviceType::GpuDevice;
};
 
template <typename T>
struct DataTypeToEnum {
    static constexpr DataType value = {};
};
// Specializations of DataTypeToEnum for supported types.
template <>
struct DataTypeToEnum<int> {
    static constexpr DataType value = DataType::DT_INT;
};
template <>
struct DataTypeToEnum<float> {
    static constexpr DataType value = DataType::DT_FLOAT;
};
template <>
struct DataTypeToEnum<double> {
    static constexpr DataType value = DataType::DT_DOUBLE;
};
template <>
struct DataTypeToEnum<int64_t> {
    static constexpr DataType value = DataType::DT_INT64;
};
template <>
struct DataTypeToEnum<std::complex<float>> {
    static constexpr DataType value = DataType::DT_COMPLEX;
};
template <>
struct DataTypeToEnum<std::complex<double>> {
    static constexpr DataType value = DataType::DT_COMPLEX_DOUBLE;
};
 
#if defined(__CUDACC__) || defined(__HIPCC__)
template <>
struct DataTypeToEnum<thrust::complex<float>> {
    static constexpr DataType value = DataType::DT_COMPLEX;
};
 
template <>
struct DataTypeToEnum<thrust::complex<double>> {
    static constexpr DataType value = DataType::DT_COMPLEX_DOUBLE;
};
#endif // defined(__CUDACC__) || defined(__HIPCC__)
 
std::ostream& operator<<(std::ostream& os, const DataType& data_type);
 
std::ostream& operator<<(std::ostream& os, const DeviceType& memory_type);
 
} // namespace container
#endif // ATEN_CORE_TENSOR_TYPES_H_