wehub-resource-sync
108 lines
5.1 KiB
C++
108 lines
5.1 KiB
C++
/*!
|
|
* \file support/module_vtable.h
|
|
* \brief Compatibility shim providing the TVM_MODULE_VTABLE_* macros that
|
|
* previously lived in <tvm/runtime/module.h>. After the TVM runtime
|
|
* refactor the macros were moved into TVM's private
|
|
* src/runtime/vm/module_utils.h, so we vendor the surface mlc-llm uses
|
|
* here to keep the existing call sites unchanged.
|
|
*/
|
|
#ifndef MLC_LLM_SUPPORT_MODULE_VTABLE_H_
|
|
#define MLC_LLM_SUPPORT_MODULE_VTABLE_H_
|
|
|
|
#include <tvm/ffi/cast.h>
|
|
#include <tvm/ffi/extra/module.h>
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/ffi/object.h>
|
|
|
|
#include <utility>
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace module_vtable_detail {
|
|
|
|
template <typename T>
|
|
struct EntryHelper {};
|
|
|
|
template <typename T, typename R, typename... Args>
|
|
struct EntryHelper<R (T::*)(Args...) const> {
|
|
using MemFnType = R (T::*)(Args...) const;
|
|
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
|
::tvm::ffi::PackedArgs args) {
|
|
auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
|
|
::tvm::ffi::details::unpack_call<R>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
|
wrapped, args.data(), args.size(), rv);
|
|
}
|
|
};
|
|
|
|
template <typename T, typename R, typename... Args>
|
|
struct EntryHelper<R (T::*)(Args...)> {
|
|
using MemFnType = R (T::*)(Args...);
|
|
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
|
::tvm::ffi::PackedArgs args) {
|
|
auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
|
|
::tvm::ffi::details::unpack_call<R>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
|
wrapped, args.data(), args.size(), rv);
|
|
}
|
|
};
|
|
|
|
template <typename T, typename... Args>
|
|
struct EntryHelper<void (T::*)(Args...) const> {
|
|
using MemFnType = void (T::*)(Args...) const;
|
|
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
|
::tvm::ffi::PackedArgs args) {
|
|
auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
|
|
::tvm::ffi::details::unpack_call<void>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
|
wrapped, args.data(), args.size(), rv);
|
|
}
|
|
};
|
|
|
|
template <typename T, typename... Args>
|
|
struct EntryHelper<void (T::*)(Args...)> {
|
|
using MemFnType = void (T::*)(Args...);
|
|
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
|
::tvm::ffi::PackedArgs args) {
|
|
auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
|
|
::tvm::ffi::details::unpack_call<void>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
|
wrapped, args.data(), args.size(), rv);
|
|
}
|
|
};
|
|
|
|
} // namespace module_vtable_detail
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#define TVM_MODULE_VTABLE_BEGIN(TypeKey) \
|
|
const char* kind() const final { return TypeKey; } \
|
|
::tvm::ffi::Optional<::tvm::ffi::Function> GetFunction(const ::tvm::ffi::String& _name) \
|
|
override { \
|
|
using SelfPtr = std::remove_cv_t<decltype(this)>; \
|
|
::tvm::ffi::ObjectPtr<::tvm::ffi::Object> _self = \
|
|
::tvm::ffi::GetObjectPtr<::tvm::ffi::Object>(this);
|
|
#define TVM_MODULE_VTABLE_END() \
|
|
return std::nullopt; \
|
|
}
|
|
#define TVM_MODULE_VTABLE_END_WITH_DEFAULT(MemFunc) \
|
|
{ \
|
|
auto f = (MemFunc); \
|
|
return (this->*f)(_name); \
|
|
} \
|
|
}
|
|
#define TVM_MODULE_VTABLE_ENTRY(Name, MemFunc) \
|
|
if (_name == Name) { \
|
|
return ::tvm::ffi::Function::FromPacked( \
|
|
[_self](::tvm::ffi::PackedArgs args, ::tvm::ffi::Any* rv) -> void { \
|
|
using Helper = ::mlc::llm::module_vtable_detail::EntryHelper<decltype(MemFunc)>; \
|
|
SelfPtr self = static_cast<SelfPtr>(_self.get()); \
|
|
Helper::Call(rv, self, MemFunc, args); \
|
|
}); \
|
|
}
|
|
#define TVM_MODULE_VTABLE_ENTRY_PACKED(Name, MemFunc) \
|
|
if (_name == Name) { \
|
|
return ::tvm::ffi::Function( \
|
|
[_self](::tvm::ffi::PackedArgs args, ::tvm::ffi::Any* rv) -> void { \
|
|
(static_cast<SelfPtr>(_self.get())->*(MemFunc))(args, rv); \
|
|
}); \
|
|
}
|
|
|
|
#endif // MLC_LLM_SUPPORT_MODULE_VTABLE_H_
|