tvm
module.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
26 #ifndef TVM_RUNTIME_MODULE_H_
27 #define TVM_RUNTIME_MODULE_H_
28 
29 #include <dmlc/io.h>
30 #include <tvm/ffi/cast.h>
31 #include <tvm/ffi/extra/module.h>
32 #include <tvm/ffi/function.h>
33 #include <tvm/ffi/memory.h>
34 #include <tvm/ffi/string.h>
35 #include <tvm/runtime/base.h>
36 #include <tvm/runtime/object.h>
37 
38 #include <utility>
39 
40 namespace tvm {
41 namespace runtime {
42 
48 TVM_DLL bool RuntimeEnabled(const String& target);
49 
51 namespace symbol {
53 constexpr const char* tvm_set_device = "__tvm_set_device";
55 constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
57 constexpr const char* tvm_prepare_global_barrier = "__tvm_prepare_global_barrier";
58 } // namespace symbol
59 
60 namespace details {
61 
62 template <typename T>
64 
65 template <typename T, typename R, typename... Args>
66 struct ModuleVTableEntryHelper<R (T::*)(Args...) const> {
67  using MemFnType = R (T::*)(Args...) const;
68  static TVM_ALWAYS_INLINE void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) {
69  auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
70  ffi::details::unpack_call<R>(std::make_index_sequence<sizeof...(Args)>{}, nullptr, wrapped,
71  args.data(), args.size(), rv);
72  }
73 };
74 
75 template <typename T, typename R, typename... Args>
76 struct ModuleVTableEntryHelper<R (T::*)(Args...)> {
77  using MemFnType = R (T::*)(Args...);
78  static TVM_ALWAYS_INLINE void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) {
79  auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
80  ffi::details::unpack_call<R>(std::make_index_sequence<sizeof...(Args)>{}, nullptr, wrapped,
81  args.data(), args.size(), rv);
82  }
83 };
84 
85 template <typename T, typename... Args>
86 struct ModuleVTableEntryHelper<void (T::*)(Args...) const> {
87  using MemFnType = void (T::*)(Args...) const;
88  static TVM_ALWAYS_INLINE void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) {
89  auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
90  ffi::details::unpack_call<void>(std::make_index_sequence<sizeof...(Args)>{}, nullptr, wrapped,
91  args.data(), args.size(), rv);
92  }
93 };
94 
95 template <typename T, typename... Args>
96 struct ModuleVTableEntryHelper<void (T::*)(Args...)> {
97  using MemFnType = void (T::*)(Args...);
98  static TVM_ALWAYS_INLINE void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) {
99  auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
100  ffi::details::unpack_call<void>(std::make_index_sequence<sizeof...(Args)>{}, nullptr, wrapped,
101  args.data(), args.size(), rv);
102  }
103 };
104 } // namespace details
105 } // namespace runtime
106 } // namespace tvm
107 
108 #define TVM_MODULE_VTABLE_BEGIN(TypeKey) \
109  const char* kind() const final { return TypeKey; } \
110  ::tvm::ffi::Optional<::tvm::ffi::Function> GetFunction(const String& _name) override { \
111  using SelfPtr = std::remove_cv_t<decltype(this)>; \
112  ::tvm::ffi::ObjectPtr<::tvm::ffi::Object> _self = \
113  ::tvm::ffi::GetObjectPtr<::tvm::ffi::Object>(this);
114 #define TVM_MODULE_VTABLE_END() \
115  return std::nullopt; \
116  }
117 #define TVM_MODULE_VTABLE_END_WITH_DEFAULT(MemFunc) \
118  { \
119  auto f = (MemFunc); \
120  return (this->*f)(_name); \
121  } \
122  } // NOLINT(*)
123 #define TVM_MODULE_VTABLE_ENTRY(Name, MemFunc) \
124  if (_name == Name) { \
125  return ffi::Function::FromPacked([_self](ffi::PackedArgs args, ffi::Any* rv) -> void { \
126  using Helper = ::tvm::runtime::details::ModuleVTableEntryHelper<decltype(MemFunc)>; \
127  SelfPtr self = static_cast<SelfPtr>(_self.get()); \
128  Helper::Call(rv, self, MemFunc, args); \
129  }); \
130  }
131 #define TVM_MODULE_VTABLE_ENTRY_PACKED(Name, MemFunc) \
132  if (_name == Name) { \
133  return ffi::Function([_self](ffi::PackedArgs args, ffi::Any* rv) -> void { \
134  (static_cast<SelfPtr>(_self.get())->*(MemFunc))(args, rv); \
135  }); \
136  }
137 
138 #endif // TVM_RUNTIME_MODULE_H_
constexpr const char * tvm_set_device
global function to set device
Definition: module.h:53
constexpr const char * tvm_global_barrier_state
Auxiliary counter to global barrier.
Definition: module.h:55
constexpr const char * tvm_prepare_global_barrier
Prepare the global barrier before kernels that uses global barrier.
Definition: module.h:57
bool RuntimeEnabled(const String &target)
Check if runtime module is enabled for target.
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.
static TVM_ALWAYS_INLINE void Call(ffi::Any *rv, T *self, MemFnType f, ffi::PackedArgs args)
Definition: module.h:68
static TVM_ALWAYS_INLINE void Call(ffi::Any *rv, T *self, MemFnType f, ffi::PackedArgs args)
Definition: module.h:78
void(T::*)(Args...) const MemFnType
Definition: module.h:87
static TVM_ALWAYS_INLINE void Call(ffi::Any *rv, T *self, MemFnType f, ffi::PackedArgs args)
Definition: module.h:88
static TVM_ALWAYS_INLINE void Call(ffi::Any *rv, T *self, MemFnType f, ffi::PackedArgs args)
Definition: module.h:98