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  ffi::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;
62 };
63 
68 class EnvFunc : public ObjectRef {
69  public:
70  EnvFunc() {}
71  explicit EnvFunc(ObjectPtr<Object> n) : ObjectRef(n) {}
75  explicit EnvFunc(ffi::UnsafeInit tag) : ObjectRef(tag) {}
77  const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
83  template <typename... Args>
84  ffi::Any operator()(Args&&... args) const {
85  const EnvFuncNode* n = operator->();
86  ICHECK(n != nullptr);
87  return n->func(std::forward<Args>(args)...);
88  }
95  TVM_DLL static EnvFunc Get(const ffi::String& name);
98 };
99 
103 template <typename FType>
105 
115 template <typename R, typename... Args>
116 class TypedEnvFunc<R(Args...)> : public ObjectRef {
117  public:
119  using TSelf = TypedEnvFunc<R(Args...)>;
121  explicit TypedEnvFunc(ObjectPtr<Object> n) : ObjectRef(n) {}
125  explicit TypedEnvFunc(ffi::UnsafeInit tag) : ObjectRef(tag) {}
131  TSelf& operator=(const EnvFunc& other) {
132  ObjectRef::operator=(other);
133  return *this;
134  }
136  const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
142  R operator()(Args... args) const {
143  const EnvFuncNode* n = operator->();
144  ICHECK(n != nullptr);
145  if constexpr (std::is_same_v<R, void>) {
146  n->func(std::forward<Args>(args)...);
147  } else {
148  ffi::Any res = n->func(std::forward<Args>(args)...);
149  if constexpr (std::is_same_v<R, ffi::Any>) {
150  return res;
151  } else {
152  return std::move(res).cast<R>();
153  }
154  }
155  }
158 };
159 
160 } // namespace tvm
161 #endif // TVM_IR_ENV_FUNC_H_
A serializable function backed by TVM's global environment.
Definition: env_func.h:43
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
EnvFuncNode()
constructor
Definition: env_func.h:50
static void RegisterReflection()
Definition: env_func.h:52
ffi::String name
Unique name of the global function.
Definition: env_func.h:46
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.EnvFunc", EnvFuncNode, Object)
Managed reference to EnvFuncNode.
Definition: env_func.h:68
EnvFunc(ObjectPtr< Object > n)
Definition: env_func.h:71
EnvFunc(ffi::UnsafeInit tag)
constructor with UnsafeInit
Definition: env_func.h:75
static EnvFunc Get(const ffi::String &name)
Get a global function based on the name.
const EnvFuncNode * operator->() const
Definition: env_func.h:77
EnvFunc()
Definition: env_func.h:70
ffi::Any operator()(Args &&... args) const
Invoke the function.
Definition: env_func.h:84
A typed version of EnvFunc. It is backed by a GlobalFuncNode internally.
Definition: env_func.h:116
TypedEnvFunc(ObjectPtr< Object > n)
Definition: env_func.h:121
TypedEnvFunc()
Definition: env_func.h:120
TSelf & operator=(const EnvFunc &other)
Assign global function to a TypedEnvFunc.
Definition: env_func.h:131
const EnvFuncNode * operator->() const
Definition: env_func.h:136
R operator()(Args... args) const
Invoke the function.
Definition: env_func.h:142
TypedEnvFunc(ffi::UnsafeInit tag)
constructor with UnsafeInit
Definition: env_func.h:125
Please refer to TypedEnvFunc<R(Args..)>.
Definition: env_func.h:104
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.