type.h

Go to the documentation of this file.

 
#pragma once
 
#include <cassert>
#include <concepts>
#include <memory>
#include <nlohmann/json.hpp>
#include <string>
#include <string_view>
#include <type_traits>
 
#include "sdfg/exceptions.h"
#include "sdfg/symbolic/symbolic.h"
 
using json = nlohmann::json;
 
namespace sdfg {
 
namespace types {

enum PrimitiveType {
    Void, 
    Bool, 
    Int8, 
    Int16, 
    Int32, 
    Int64, 
    Int128, 
    UInt8, 
    UInt16, 
    UInt32, 
    UInt64, 
    UInt128, 
    Half, 
    BFloat, 
    Float, 
    Double, 
    X86_FP80, 
    FP128, 
    PPC_FP128 
};

class StorageType {
public:
    enum AllocationType {
        Unmanaged, 
        Managed, 
    };
 
private:
    std::string value_; 
    symbolic::Expression allocation_size_; 
    AllocationType allocation_; 
    AllocationType deallocation_; 
    std::vector<symbolic::Expression> args_; 
 
public:
    StorageType(const std::string& value)
        : value_(value), allocation_size_(SymEngine::null), allocation_(Unmanaged), deallocation_(Unmanaged) {
    }

    StorageType(
        const std::string& value,
        const symbolic::Expression& allocation_size,
        AllocationType allocation,
        AllocationType deallocation
    )
        : value_(value), allocation_size_(allocation_size), allocation_(allocation), deallocation_(deallocation) {
    }

    std::string value() const {
        return value_;
    }

    void value(const std::string& value) {
        value_ = value;
    }

    symbolic::Expression allocation_size() const {
        return allocation_size_;
    }

    void allocation_size(const symbolic::Expression& allocation_size) {
        allocation_size_ = allocation_size;
    }

    [[deprecated("use args")]]
    const symbolic::Expression& arg1() const {
        return args_.at(0);
    }

    [[deprecated("use args")]]
    void arg1(const symbolic::Expression& arg) {
        if (args_.size() > 1) {
            args_.at(0) = arg;
        } else {
            args_.push_back(arg);
        }
    }
 
    std::vector<symbolic::Expression>& args() {
        return args_;
    }
 
    const std::vector<symbolic::Expression>& args() const {
        return args_;
    }
 
    void args(const std::vector<symbolic::Expression>& args) {
        args_ = args;
    }
 
    void args(std::vector<symbolic::Expression>&& args) {
        args_ = std::move(args);
    }

    AllocationType allocation() const {
        return allocation_;
    }

    void allocation(AllocationType allocation) {
        allocation_ = allocation;
    }

    AllocationType deallocation() const {
        return deallocation_;
    }

    void deallocation(AllocationType deallocation) {
        deallocation_ = deallocation;
    }

    bool operator==(const StorageType& other) const {
        if (value_ != other.value_) {
            return false;
        }
        if (allocation_ != other.allocation_) {
            return false;
        }
        if (deallocation_ != other.deallocation_) {
            return false;
        }
        if (!symbolic::null_safe_eq(allocation_size_, other.allocation_size_)) {
            return false;
        }
        if (args_.size() != other.args_.size()) {
            return false;
        } else {
            for (size_t i = 0; i < args_.size(); ++i) {
                if (!symbolic::null_safe_eq(args_[i], other.args_[i])) {
                    return false;
                }
            }
        }
        return true;
    }

    bool is_cpu_stack() const {
        return value_ == "CPU_Stack";
    }

    bool is_cpu_heap() const {
        return value_ == "CPU_Heap";
    }

    bool is_device() const {
        return !is_cpu_stack() && !is_cpu_heap();
    }

    bool is_nv_generic() const {
        return value_ == "NV_Generic";
    }

    bool is_amd_generic() const {
        return value_ == "AMD_Generic";
    }

    bool is_nv_global() const {
        return value_ == "NV_Global";
    }

    bool is_nv_shared() const {
        return value_ == "NV_Shared";
    }

    bool is_nv_constant() const {
        return value_ == "NV_Constant";
    }

    bool is_nv_symbol() const {
        return value_ == "NV_Symbol";
    }

    bool is_gpu_static() const {
        return value_ == "GPU_Static";
    }

    static StorageType CPU_Stack() {
        return StorageType("CPU_Stack", SymEngine::null, StorageType::Unmanaged, StorageType::Unmanaged);
    }

    static StorageType CPU_Heap() {
        return StorageType("CPU_Heap", SymEngine::null, StorageType::Unmanaged, StorageType::Unmanaged);
    }

    static StorageType CPU_Heap(
        symbolic::Expression allocation_size,
        StorageType::AllocationType allocation,
        StorageType::AllocationType deallocation
    ) {
        return StorageType("CPU_Heap", allocation_size, allocation, deallocation);
    }

    static StorageType NV_Generic() {
        return StorageType("NV_Generic");
    }

    static StorageType NV_Global() {
        return StorageType("NV_Global");
    }

    static StorageType NV_Shared() {
        return StorageType("NV_Shared");
    }

    static StorageType NV_Constant() {
        return StorageType("NV_Constant");
    }

    static StorageType NV_Symbol() {
        return StorageType("NV_Symbol");
    }

    static StorageType GPU_Static() {
        return StorageType("GPU_Static");
    }
};

constexpr const char* primitive_type_to_string(PrimitiveType e) {
    switch (e) {
        case PrimitiveType::Void:
            return "Void";
        case PrimitiveType::Bool:
            return "Bool";
        case PrimitiveType::Int8:
            return "Int8";
        case PrimitiveType::Int16:
            return "Int16";
        case PrimitiveType::Int32:
            return "Int32";
        case PrimitiveType::Int64:
            return "Int64";
        case PrimitiveType::Int128:
            return "Int128";
        case PrimitiveType::UInt8:
            return "UInt8";
        case PrimitiveType::UInt16:
            return "UInt16";
        case PrimitiveType::UInt32:
            return "UInt32";
        case PrimitiveType::UInt64:
            return "UInt64";
        case PrimitiveType::UInt128:
            return "UInt128";
        case PrimitiveType::Half:
            return "Half";
        case PrimitiveType::BFloat:
            return "BFloat";
        case PrimitiveType::Float:
            return "Float";
        case PrimitiveType::Double:
            return "Double";
        case PrimitiveType::X86_FP80:
            return "X86_FP80";
        case PrimitiveType::FP128:
            return "FP128";
        case PrimitiveType::PPC_FP128:
            return "PPC_FP128";
    }
    throw std::invalid_argument("Invalid primitive type");
};

constexpr PrimitiveType primitive_type_from_string(std::string_view e) {
    if (e == "Void") {
        return PrimitiveType::Void;
    } else if (e == "Bool") {
        return PrimitiveType::Bool;
    } else if (e == "Int8") {
        return PrimitiveType::Int8;
    } else if (e == "Int16") {
        return PrimitiveType::Int16;
    } else if (e == "Int32") {
        return PrimitiveType::Int32;
    } else if (e == "Int64") {
        return PrimitiveType::Int64;
    } else if (e == "Int128") {
        return PrimitiveType::Int128;
    } else if (e == "UInt8") {
        return PrimitiveType::UInt8;
    } else if (e == "UInt16") {
        return PrimitiveType::UInt16;
    } else if (e == "UInt32") {
        return PrimitiveType::UInt32;
    } else if (e == "UInt64") {
        return PrimitiveType::UInt64;
    } else if (e == "UInt128") {
        return PrimitiveType::UInt128;
    } else if (e == "Half") {
        return PrimitiveType::Half;
    } else if (e == "BFloat") {
        return PrimitiveType::BFloat;
    } else if (e == "Float") {
        return PrimitiveType::Float;
    } else if (e == "Double") {
        return PrimitiveType::Double;
    } else if (e == "X86_FP80") {
        return PrimitiveType::X86_FP80;
    } else if (e == "FP128") {
        return PrimitiveType::FP128;
    } else if (e == "PPC_FP128") {
        return PrimitiveType::PPC_FP128;
    }
    throw std::invalid_argument("Invalid primitive type");
};

constexpr size_t bit_width(PrimitiveType e) {
    switch (e) {
        case PrimitiveType::Void:
            return 0;
        case PrimitiveType::Bool:
            return 1;
        case PrimitiveType::Int8:
            return 8;
        case PrimitiveType::Int16:
            return 16;
        case PrimitiveType::Int32:
            return 32;
        case PrimitiveType::Int64:
            return 64;
        case PrimitiveType::Int128:
            return 128;
        case PrimitiveType::UInt8:
            return 8;
        case PrimitiveType::UInt16:
            return 16;
        case PrimitiveType::UInt32:
            return 32;
        case PrimitiveType::UInt64:
            return 64;
        case PrimitiveType::UInt128:
            return 128;
        case PrimitiveType::Half:
            return 16;
        case PrimitiveType::BFloat:
            return 16;
        case PrimitiveType::Float:
            return 32;
        case PrimitiveType::Double:
            return 64;
        case PrimitiveType::X86_FP80:
            return 80;
        case PrimitiveType::FP128:
            return 128;
        case PrimitiveType::PPC_FP128:
            return 128;
    }
    throw std::invalid_argument("Invalid primitive type");
};

constexpr bool is_floating_point(PrimitiveType e) noexcept {
    switch (e) {
        case PrimitiveType::Half:
        case PrimitiveType::BFloat:
        case PrimitiveType::Float:
        case PrimitiveType::Double:
        case PrimitiveType::X86_FP80:
        case PrimitiveType::FP128:
        case PrimitiveType::PPC_FP128:
            return true;
        default:
            return false;
    }
};

constexpr bool is_integer(PrimitiveType e) noexcept {
    switch (e) {
        case PrimitiveType::Bool:
        case PrimitiveType::Int8:
        case PrimitiveType::Int16:
        case PrimitiveType::Int32:
        case PrimitiveType::Int64:
        case PrimitiveType::Int128:
        case PrimitiveType::UInt8:
        case PrimitiveType::UInt16:
        case PrimitiveType::UInt32:
        case PrimitiveType::UInt64:
        case PrimitiveType::UInt128:
            return true;
        default:
            return false;
    }
};

constexpr bool is_signed(PrimitiveType e) noexcept {
    switch (e) {
        case PrimitiveType::Int8:
        case PrimitiveType::Int16:
        case PrimitiveType::Int32:
        case PrimitiveType::Int64:
        case PrimitiveType::Int128:
            return true;
        default:
            return false;
    }
};

constexpr bool is_unsigned(PrimitiveType e) noexcept {
    switch (e) {
        case PrimitiveType::UInt8:
        case PrimitiveType::UInt16:
        case PrimitiveType::UInt32:
        case PrimitiveType::UInt64:
        case PrimitiveType::UInt128:
            return true;
        default:
            return false;
    }
};

constexpr PrimitiveType as_signed(PrimitiveType e) noexcept {
    switch (e) {
        case PrimitiveType::UInt8:
            return PrimitiveType::Int8;
        case PrimitiveType::UInt16:
            return PrimitiveType::Int16;
        case PrimitiveType::UInt32:
            return PrimitiveType::Int32;
        case PrimitiveType::UInt64:
            return PrimitiveType::Int64;
        case PrimitiveType::UInt128:
            return PrimitiveType::Int128;
        default:
            return e;
    }
};

constexpr PrimitiveType as_unsigned(PrimitiveType e) noexcept {
    switch (e) {
        case PrimitiveType::Int8:
            return PrimitiveType::UInt8;
        case PrimitiveType::Int16:
            return PrimitiveType::UInt16;
        case PrimitiveType::Int32:
            return PrimitiveType::UInt32;
        case PrimitiveType::Int64:
            return PrimitiveType::UInt64;
        case PrimitiveType::Int128:
            return PrimitiveType::UInt128;
        default:
            return e;
    }
};

enum class TypeID {
    Scalar, 
    Array, 
    Structure, 
    Pointer, 
    Reference, 
    Function, 
    Tensor, 
};

class IType {
protected:
    StorageType storage_type_; 
    size_t alignment_; 
    std::string initializer_; 
 
public:
    IType(StorageType storage_type = StorageType::CPU_Stack(), size_t alignment = 0, const std::string& initializer = "")
        : storage_type_(storage_type), alignment_(alignment), initializer_(initializer) {};

    virtual ~IType() = default;

    virtual TypeID type_id() const = 0;

    StorageType storage_type() const {
        return storage_type_;
    }

    virtual bool is_pointer_like() const = 0;

    StorageType& storage_type() {
        return storage_type_;
    };

    void storage_type(const StorageType& storage_type) {
        storage_type_ = storage_type;
    };

    size_t alignment() const {
        return alignment_;
    };

    void alignment(size_t alignment) {
        alignment_ = alignment;
    };

    std::string initializer() const {
        return initializer_;
    };

    void initializer(const std::string& initializer) {
        initializer_ = initializer;
    };

    virtual PrimitiveType primitive_type() const = 0;

    virtual bool is_symbol() const = 0;

    virtual bool operator==(const IType& other) const = 0;

    virtual std::unique_ptr<IType> clone() const = 0;

    virtual std::string print() const = 0;

    friend std::ostream& operator<<(std::ostream& os, const IType& type) {
        os << type.print();
        return os;
    };

    virtual void replace_symbols(const symbolic::Expression old_expression, const symbolic::Expression new_expression) = 0;
    virtual void replace_symbols(const symbolic::ExpressionMapping& replacements) = 0;
};
 
} // namespace types
} // namespace sdfg