tvm
env_func.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 
24 #ifndef TVM_IR_ENV_FUNC_H_
25 #define TVM_IR_ENV_FUNC_H_
26 
27 #include <tvm/ffi/function.h>
28 #include <tvm/ffi/reflection/registry.h>
29 #include <tvm/node/node.h>
30 
31 #include <string>
32 #include <utility>
33 
34 namespace tvm {
43 class EnvFuncNode : public Object {
44  public:
46  String name;
51 
52  static void RegisterReflection() {
53  namespace refl = tvm::ffi::reflection;
54  // func do not participate in structural equal and hash.
55  refl::ObjectDef<EnvFuncNode>()
56  .def_ro("name", &EnvFuncNode::name)
57  .def_ro("func", &EnvFuncNode::func, refl::AttachFieldFlag::SEqHashIgnore());
58  }
59 
60  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
61  static constexpr const char* _type_key = "ir.EnvFunc";
62 
64 };
65 
70 class EnvFunc : public ObjectRef {
71  public:
72  EnvFunc() {}
73  explicit EnvFunc(ObjectPtr<Object> n) : ObjectRef(n) {}
75  const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
81  template <typename... Args>
82  ffi::Any operator()(Args&&... args) const {
83  const EnvFuncNode* n = operator->();
84  ICHECK(n != nullptr);
85  return n->func(std::forward<Args>(args)...);
86  }
93  TVM_DLL static EnvFunc Get(const String& name);
96 };
97 
101 template <typename FType>
103 
113 template <typename R, typename... Args>
114 class TypedEnvFunc<R(Args...)> : public ObjectRef {
115  public:
117  using TSelf = TypedEnvFunc<R(Args...)>;
119  explicit TypedEnvFunc(ObjectPtr<Object> n) : ObjectRef(n) {}
125  TSelf& operator=(const EnvFunc& other) {
126  ObjectRef::operator=(other);
127  return *this;
128  }
130  const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
136  R operator()(Args... args) const {
137  const EnvFuncNode* n = operator->();
138  ICHECK(n != nullptr);
139  if constexpr (std::is_same_v<R, void>) {
140  n->func(std::forward<Args>(args)...);
141  } else {
142  ffi::Any res = n->func(std::forward<Args>(args)...);
143  if constexpr (std::is_same_v<R, ffi::Any>) {
144  return res;
145  } else {
146  return std::move(res).cast<R>();
147  }
148  }
149  }
152 };
153 
154 } // namespace tvm
155 #endif // TVM_IR_ENV_FUNC_H_
A serializable function backed by TVM's global environment.
Definition: env_func.h:43
String name
Unique name of the global function.
Definition: env_func.h:46
static constexpr const char * _type_key
Definition: env_func.h:61
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: env_func.h:60
ffi::Function func
The internal packed function.
Definition: env_func.h:48
TVM_DECLARE_FINAL_OBJECT_INFO(EnvFuncNode, Object)
EnvFuncNode()
constructor
Definition: env_func.h:50
static void RegisterReflection()
Definition: env_func.h:52
Managed reference to EnvFuncNode.
Definition: env_func.h:70
EnvFunc(ObjectPtr< Object > n)
Definition: env_func.h:73
const EnvFuncNode * operator->() const
Definition: env_func.h:75
EnvFunc()
Definition: env_func.h:72
static EnvFunc Get(const String &name)
Get a global function based on the name.
ffi::Any operator()(Args &&... args) const
Invoke the function.
Definition: env_func.h:82
A typed version of EnvFunc. It is backed by a GlobalFuncNode internally.
Definition: env_func.h:114
TypedEnvFunc(ObjectPtr< Object > n)
Definition: env_func.h:119
TypedEnvFunc()
Definition: env_func.h:118
TSelf & operator=(const EnvFunc &other)
Assign global function to a TypedEnvFunc.
Definition: env_func.h:125
const EnvFuncNode * operator->() const
Definition: env_func.h:130
R operator()(Args... args) const
Invoke the function.
Definition: env_func.h:136
Please refer to TypedEnvFunc<R(Args..)>.
Definition: env_func.h:102
Definition: repr_printer.h:91
tvm::relax::Function Function
Definition: transform.h:42
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Definitions and helper macros for IR/AST nodes.