device.h

// device.h
// Device type

#ifndef TINYTENSOR_DEVICE_H_
#define TINYTENSOR_DEVICE_H_

#include <tt/exception.h>
#include <tt/export.h>

#include <format>
#include <ostream>
#include <string>

namespace tinytensor {

// Supported backend types
enum class TINYTENSOR_EXPORT Backend {
    cpu,
#ifdef TT_CUDA
    cuda,
#endif
};

TINYTENSOR_EXPORT constexpr auto to_string(Backend backend) -> std::string {
    switch (backend) {
    case Backend::cpu:
        return "cpu";
#ifdef TT_CUDA
    case Backend::cuda:
        return "cuda";
#endif
    }
    TT_EXCEPTION("Unknown device type.");
}

// Device is an enum backend + device ID (multi-device support)
struct TINYTENSOR_EXPORT Device {
#ifdef TT_CUDA
    constexpr static auto CUDA(int dev_id) -> Device {
        return {.backend = Backend::cuda, .id = dev_id};
    }
#endif

    // Operator!= provided by compiler since C++20
    constexpr auto operator==(const Device &other) const -> bool {
        return backend == other.backend && id == other.id;
    }

    Backend backend;
    int id;
};

// Shortnames
constexpr Device kCPU = Device{.backend = Backend::cpu, .id = 0};
#ifdef TT_CUDA
constexpr Device kCUDA = Device{.backend = Backend::cuda, .id = 0};
#endif

TINYTENSOR_EXPORT inline auto operator<<(std::ostream &os, const Device &device) -> std::ostream & {
    os << to_string(device.backend) << ":" << device.id;
    return os;
}
}    // namespace tinytensor

template <>
struct TINYTENSOR_EXPORT std::formatter<tinytensor::Device> : std::formatter<std::string> {
    auto format(const tinytensor::Device &device, format_context &ctx) const {
        return formatter<string>::format(std::format("{}:{}", to_string(device.backend), device.id), ctx);
    }
};

#endif    // TINYTENSOR_DEVICE_H_