support typedvalue

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

include/exception.h

@@ -1,6 +1,7 @@
 #ifndef TRITON_TVM_FFI_EXCEPTION_H_
 #define TRITON_TVM_FFI_EXCEPTION_H_
 
+#include "type.h"
 #include <cuda.h>
 #include <exception>
 
@@ -12,7 +13,17 @@ public:
   const char *what() const noexcept override;
 
 private:
-  const CUresult code;
+  const CUresult code_;
+};
+
+class UnknownTypeException : public std::exception {
+public:
+  UnknownTypeException(Type type);
+  UnknownTypeException(std::string_view type);
+  const char *what() const noexcept override;
+
+private:
+  const std::string message_;
 };
 
 } // namespace triton_tvm_ffi

include/macro.h

@@ -0,0 +1,8 @@
+#ifndef TRITON_TVM_FFI_MACRO_H_
+#define TRITON_TVM_FFI_MACRO_H_
+
+#if defined(__GNUC__) || defined(__clang__)
+#define TRITON_TVM_FFI_INLINE __attribute__((always_inline)) inline
+#endif
+
+#endif

include/type.h

@@ -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

include/value.h

@@ -0,0 +1,45 @@
+#ifndef TRITON_TVM_FFI_VALUE_H_
+#define TRITON_TVM_FFI_VALUE_H_
+
+#include "macro.h"
+#include "type.h"
+#include <tvm/ffi/any.h>
+#include <tvm/ffi/object.h>
+
+namespace triton_tvm_ffi {
+
+class TypedValueObj : public tvm::ffi::Object {
+public:
+  TypedValueObj(Type type, const tvm::ffi::Any &value);
+  TypedValueObj(Type type, tvm::ffi::Any &&value);
+  TypedValueObj(const TypedValueObj &other) = default;
+  TypedValueObj(TypedValueObj &&other) = default;
+  TypedValueObj &operator=(const TypedValueObj &other) = default;
+  TypedValueObj &operator=(TypedValueObj &&other) = default;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("triton_tvm_ffi.TypedValue", TypedValueObj,
+                                    tvm::ffi::Object);
+  TRITON_TVM_FFI_INLINE Type GetType() const { return type_; }
+  TRITON_TVM_FFI_INLINE const tvm::ffi::Any &GetValue() const { return value_; }
+
+private:
+  Type type_;
+  tvm::ffi::Any value_;
+};
+
+class TypedValue : public tvm::ffi::ObjectRef {
+public:
+  TypedValue(Type type, const tvm::ffi::Any &value);
+  TypedValue(Type type, tvm::ffi::Any &&value);
+  using tvm::ffi::ObjectRef::ObjectRef;
+  using tvm::ffi::ObjectRef::operator=;
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypedValue, tvm::ffi::ObjectRef,
+                                                TypedValueObj);
+  TRITON_TVM_FFI_INLINE Type GetType() const { return get()->GetType(); }
+  TRITON_TVM_FFI_INLINE const tvm::ffi::Any &GetValue() const {
+    return get()->GetValue();
+  }
+};
+
+} // namespace triton_tvm_ffi
+
+#endif