JinjieLiu/triton-tvm-ffi
1
0
Fork 0
mirror of https://github.com/sgjzfzzf/triton-tvm-ffi.git synced 2026-07-01 08:51:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

support typedvalue

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-29 16:46:00 +08:00
parent fecf48e403
commit 781a5396cc
15 changed files with 367 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/type.h
+53
View File
@@ -0,0 +1,53 @@
#ifndef TRITON_TVM_FFI_TYPE_H_
#define TRITON_TVM_FFI_TYPE_H_
#include <cstdint>
#include <tvm/ffi/object.h>
#include <tvm/ffi/string.h>
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)
enum class Type : int64_t {
#define DEFINE_ENUM(type, str, ctype) type,
TYPE_TABLE(DEFINE_ENUM)
#undef DEFINE_ENUM
PTR,
CONSTEXPR,
};
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 <> struct type_to_ctype<Type::PTR> { using t = void *; };
// TODO: check whether CUtensorMap* is correct
template <> struct type_to_ctype<Type::CONSTEXPR> { using t = void; };
template <Type T> using type_to_ctype_t = typename type_to_ctype<T>::t;
// --------------- Implementations --------------- //
} // namespace triton_tvm_ffi
#endif