mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
74 lines
3.0 KiB
C++
74 lines
3.0 KiB
C++
#ifndef TRITON_TVM_FFI_TYPE_H_
|
|
#define TRITON_TVM_FFI_TYPE_H_
|
|
|
|
#include <cstdint>
|
|
#include <tvm/ffi/object.h>
|
|
#include <tvm/ffi/string.h>
|
|
#include <type_traits>
|
|
|
|
namespace triton_tvm_ffi {
|
|
|
|
// --------------- Definitions --------------- //
|
|
|
|
#define TYPE_TABLE(V) \
|
|
V(I1, "i1", int8_t) \
|
|
V(I8, "i8", int8_t) \
|
|
V(I16, "i16", int16_t) \
|
|
V(I32, "i32", int32_t) \
|
|
V(I64, "i64", int64_t) \
|
|
V(U1, "u1", uint8_t) \
|
|
V(U8, "u8", uint8_t) \
|
|
V(U16, "u16", uint16_t) \
|
|
V(U32, "u32", uint32_t) \
|
|
V(U64, "u64", uint64_t) \
|
|
V(FP16, "fp16", double) \
|
|
V(BF16, "bf16", double) \
|
|
V(FP32, "f32", double) \
|
|
V(FP64, "fp64", double) \
|
|
V(PTR, "*?", void *) \
|
|
V(CONSTEXPR, "constexpr", void)
|
|
|
|
enum class Type : int64_t {
|
|
#define DEFINE_ENUM(type, str, ctype) type,
|
|
TYPE_TABLE(DEFINE_ENUM)
|
|
#undef DEFINE_ENUM
|
|
};
|
|
|
|
const char *TypeToString(Type type);
|
|
tvm::ffi::Optional<Type> StringToType(tvm::ffi::String str);
|
|
|
|
template <Type T> struct type_to_ctype;
|
|
#define DEFINE_TYPE_TO_CTYPE(type, str, ctype) \
|
|
template <> struct type_to_ctype<Type::type> { using t = ctype; };
|
|
TYPE_TABLE(DEFINE_TYPE_TO_CTYPE)
|
|
#undef DEFINE_TYPE_TO_CTYPE
|
|
template <Type T> using type_to_ctype_t = typename type_to_ctype<T>::t;
|
|
|
|
template <typename T, typename = void> struct type_size {
|
|
static constexpr size_t value = 0;
|
|
};
|
|
template <typename T>
|
|
struct type_size<T, std::enable_if_t<!std::is_void_v<decltype(sizeof(T))>>> {
|
|
static constexpr size_t value = sizeof(T);
|
|
};
|
|
template <typename T> constexpr size_t type_size_v = type_size<T>::value;
|
|
|
|
template <size_t... Ns> struct max;
|
|
template <size_t... Ns> constexpr size_t max_v = max<Ns...>::value;
|
|
template <size_t N> struct max<N> { static constexpr size_t value = N; };
|
|
template <size_t N, size_t... Ns> struct max<N, Ns...> {
|
|
static constexpr size_t value = N > max_v<Ns...> ? N : max_v<Ns...>;
|
|
};
|
|
|
|
static constexpr size_t kMaxOpaqueSize = max_v<
|
|
#define DEFINE_TYPE_SIZE(type, str, ctype) type_size_v<ctype>,
|
|
TYPE_TABLE(DEFINE_TYPE_SIZE)
|
|
#undef DEFINE_TYPE_SIZE
|
|
0>;
|
|
|
|
// --------------- Implementations --------------- //
|
|
|
|
} // namespace triton_tvm_ffi
|
|
|
|
#endif
|