binding_rewrite.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef TVM_RELAX_BINDING_REWRITE_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr.h>
 
#include <map>
#include <set>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace tvm {
namespace relax {
 
class DataflowBlockRewriteObj : public ffi::Object {
 public:
  explicit DataflowBlockRewriteObj(DataflowBlock value) : dfb_(std::move(value)) {}
  explicit DataflowBlockRewriteObj(ffi::UnsafeInit) : dfb_(ffi::UnsafeInit{}) {}
 
  void ReplaceAllUses(Var old_var, Var new_var);
  void Add(Binding binding);
  void Add(ffi::String var_name, Expr expr, bool is_dfvar = false) {
    auto var = is_dfvar ? DataflowVar(var_name, GetType(expr))  //
                        : Var(var_name, GetType(expr));
    Add(VarBinding(std::move(var), std::move(expr)));
  }
  void Add(Expr expr, bool is_dfvar = false) {
    Add(name_supply_->FreshName("tmp"), expr, is_dfvar);
  }
  void RemoveUnused(Var unused, bool allow_undef = false);
  void RemoveAllUnused();
 
  DataflowBlock MutatedDataflowBlock() { return dfb_; }
  Function MutatedFunc() { return root_fn_.value(); }
  IRModule MutateIRModule(IRModule irmod);
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataflowBlockRewriteObj>()
        .def_ro("dfb", &DataflowBlockRewriteObj::dfb_)
        .def_ro("root_fn", &DataflowBlockRewriteObj::root_fn_);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DataflowBlockRewrite", DataflowBlockRewriteObj,
                                    ffi::Object);
 
 protected:
  friend class DataflowBlockRewrite;
 
  DataflowBlock dfb_;                        
  ffi::Optional<Function> root_fn_;          
  const FunctionNode* original_fn_ptr_;      
  ffi::Map<Var, ffi::Array<Var>> to_users_;  
  ffi::Array<Var> fn_outputs_;               
 
 private:
  UniqueNameSupply name_supply_;  
};
 
class DataflowBlockRewrite : public ffi::ObjectRef {
 public:
  TVM_DLL explicit DataflowBlockRewrite(DataflowBlock dfb, Function root_fn);
 
  DataflowBlockRewriteObj* operator->() {
    TVM_FFI_ICHECK(get() != nullptr);
    return static_cast<DataflowBlockRewriteObj*>(get_mutable());
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowBlockRewrite, ffi::ObjectRef,
                                                DataflowBlockRewriteObj);
};
 
}  // namespace relax
}  // namespace tvm
 
#define TVM_RELAX_BINDING_REWRITE_H_
#endif  // TVM_RELAX_BINDING_REWRITE_H_