tvm
binding_rewrite.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 
25 #ifndef TVM_RELAX_BINDING_REWRITE_H_
26 
27 #include <tvm/ffi/reflection/registry.h>
28 #include <tvm/ir/name_supply.h>
29 #include <tvm/relax/analysis.h>
30 #include <tvm/relax/expr.h>
31 
32 #include <map>
33 #include <set>
34 #include <type_traits>
35 #include <utility>
36 #include <vector>
37 
38 namespace tvm {
39 namespace relax {
40 
42 class DataflowBlockRewriteNode : public Object {
43  public:
45  void ReplaceAllUses(Var old_var, Var new_var);
47  void Add(Binding binding);
49  void Add(String var_name, Expr expr, bool is_dfvar = false) {
50  auto var = is_dfvar ? DataflowVar(var_name, GetStructInfo(expr)) //
51  : Var(var_name, GetStructInfo(expr));
52  Add(VarBinding(std::move(var), std::move(expr)));
53  }
55  void Add(Expr expr, bool is_dfvar = false) {
56  Add(name_supply_->FreshName("tmp"), expr, is_dfvar);
57  }
59  void RemoveUnused(Var unused, bool allow_undef = false);
62 
66  Function MutatedFunc() { return root_fn_.value(); }
69 
71  static void RegisterReflection() {
72  namespace refl = tvm::ffi::reflection;
73  refl::ObjectDef<DataflowBlockRewriteNode>()
74  .def_ro("dfb", &DataflowBlockRewriteNode::dfb_)
75  .def_ro("root_fn", &DataflowBlockRewriteNode::root_fn_);
76  }
77 
78  static constexpr const char* _type_key = "relax.DataflowBlockRewrite";
80 
81  protected:
82  friend class DataflowBlockRewrite;
83 
85  Optional<Function> root_fn_;
87  Map<Var, Array<Var>> to_users_;
88  Array<Var> fn_outputs_;
89 
90  private:
91  NameSupply name_supply_;
92 };
93 
98 class DataflowBlockRewrite : public ObjectRef {
99  public:
100  TVM_DLL explicit DataflowBlockRewrite(DataflowBlock dfb, Function root_fn);
101 
107  ICHECK(get() != nullptr);
108  return static_cast<DataflowBlockRewriteNode*>(get_mutable());
109  }
110 
112 };
113 
114 } // namespace relax
115 } // namespace tvm
116 
117 #define TVM_RELAX_BINDING_REWRITE_H_
118 #endif // TVM_RELAX_BINDING_REWRITE_H_
Managed reference class to IRModuleNode.
Definition: module.h:257
Managed reference class to NameSupplyNode.
Definition: name_supply.h:110
Managed reference to RelaxExprNode.
Definition: expr.h:446
Definition: expr.h:601
Statement rewriter for relax.DataflowBlock.
Definition: binding_rewrite.h:42
DataflowBlock MutatedDataflowBlock()
The rewritten dataflow block.
Definition: binding_rewrite.h:64
const FunctionNode * original_fn_ptr_
Pointer to the original function.
Definition: binding_rewrite.h:86
Function MutatedFunc()
The rewritten function.
Definition: binding_rewrite.h:66
void Add(String var_name, Expr expr, bool is_dfvar=false)
Insert an expression as VarBinding with variable name.
Definition: binding_rewrite.h:49
static void RegisterReflection()
Visit attributes.
Definition: binding_rewrite.h:71
Map< Var, Array< Var > > to_users_
Map from variable to its users.
Definition: binding_rewrite.h:87
TVM_DECLARE_FINAL_OBJECT_INFO(DataflowBlockRewriteNode, Object)
static constexpr const char * _type_key
Definition: binding_rewrite.h:78
void Add(Expr expr, bool is_dfvar=false)
Insert an expression as VarBinding with automatic variable name.
Definition: binding_rewrite.h:55
Array< Var > fn_outputs_
Variables required by function outputs.
Definition: binding_rewrite.h:88
void RemoveUnused(Var unused, bool allow_undef=false)
Remove the definition statement of an unused variable.
IRModule MutateIRModule(IRModule irmod)
The rewritten IRModule.
void ReplaceAllUses(Var old_var, Var new_var)
Replace all uses of old_var with new_var.
void RemoveAllUnused()
Remove the definition statements of all unused variables.
DataflowBlock dfb_
The rewritten dataflow block.
Definition: binding_rewrite.h:84
void Add(Binding binding)
Insert a Binding statement.
Optional< Function > root_fn_
The rewritten function.
Definition: binding_rewrite.h:85
A statement rewriter for relax.DataflowBlock.
Definition: binding_rewrite.h:98
TVM_DEFINE_OBJECT_REF_METHODS(DataflowBlockRewrite, ObjectRef, DataflowBlockRewriteNode)
DataflowBlockRewriteNode * operator->()
mutable accessor.
Definition: binding_rewrite.h:106
DataflowBlockRewrite(DataflowBlock dfb, Function root_fn)
Definition: expr.h:720
Definition: expr.h:414
A Relax function.
Definition: expr.h:837
Definition: expr.h:862
Definition: expr.h:674
Definition: expr.h:387
NameSupply that can be used to generate unique variable names.
Definition: repr_printer.h:91
StructInfo GetStructInfo(const Expr &expr)
Get the underlying structure info of expr.
Definition: struct_info.h:401
Var var(std::string name_hint, DataType t=DataType::Int(32))
Construct a new Var expression.
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
The set of Relax specific analysis on IR.