wehub-resource-sync
119 lines
4.2 KiB
C++
119 lines
4.2 KiB
C++
/*
|
|
* 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.
|
|
*/
|
|
|
|
/*!
|
|
* \file tvm/relax/binding_rewrite.h
|
|
* \brief An IR rewriter to easily add/remove/replace bindings (statements).
|
|
*/
|
|
|
|
#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 {
|
|
|
|
/*! \brief Statement rewriter for relax.DataflowBlock. */
|
|
class DataflowBlockRewriteNode : public ffi::Object {
|
|
public:
|
|
/*! \brief Replace all uses of old_var with new_var. */
|
|
void ReplaceAllUses(Var old_var, Var new_var);
|
|
/*! \brief Insert a Binding statement. */
|
|
void Add(Binding binding);
|
|
/*! \brief Insert an expression as VarBinding with variable name. */
|
|
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)));
|
|
}
|
|
/*! \brief Insert an expression as VarBinding with automatic variable name. */
|
|
void Add(Expr expr, bool is_dfvar = false) {
|
|
Add(name_supply_->FreshName("tmp"), expr, is_dfvar);
|
|
}
|
|
/*! \brief Remove the definition statement of an unused variable. */
|
|
void RemoveUnused(Var unused, bool allow_undef = false);
|
|
/*! \brief Remove the definition statements of all unused variables. */
|
|
void RemoveAllUnused();
|
|
|
|
/*! \brief The rewritten dataflow block. */
|
|
DataflowBlock MutatedDataflowBlock() { return dfb_; }
|
|
/*! \brief The rewritten function. */
|
|
Function MutatedFunc() { return root_fn_.value(); }
|
|
/*! \brief The rewritten IRModule. */
|
|
IRModule MutateIRModule(IRModule irmod);
|
|
|
|
/*! \brief Visit attributes. */
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<DataflowBlockRewriteNode>()
|
|
.def_ro("dfb", &DataflowBlockRewriteNode::dfb_)
|
|
.def_ro("root_fn", &DataflowBlockRewriteNode::root_fn_);
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DataflowBlockRewrite", DataflowBlockRewriteNode,
|
|
ffi::Object);
|
|
|
|
protected:
|
|
friend class DataflowBlockRewrite;
|
|
|
|
DataflowBlock dfb_; //!< The rewritten dataflow block.
|
|
ffi::Optional<Function> root_fn_; //!< The rewritten function.
|
|
const FunctionNode* original_fn_ptr_; //!< Pointer to the original function.
|
|
ffi::Map<Var, ffi::Array<Var>> to_users_; //!< Map from variable to its users.
|
|
ffi::Array<Var> fn_outputs_; //!< Variables required by function outputs.
|
|
|
|
private:
|
|
UniqueNameSupply name_supply_; //!< Unique name supply for tracking and generating unique names.
|
|
};
|
|
|
|
/*!
|
|
* \brief A statement rewriter for relax.DataflowBlock.
|
|
* \sa DataflowBlockRewriteNode
|
|
*/
|
|
class DataflowBlockRewrite : public ffi::ObjectRef {
|
|
public:
|
|
TVM_DLL explicit DataflowBlockRewrite(DataflowBlock dfb, Function root_fn);
|
|
|
|
/*!
|
|
* \brief mutable accessor.
|
|
* \return mutable access pointer.
|
|
*/
|
|
DataflowBlockRewriteNode* operator->() {
|
|
TVM_FFI_ICHECK(get() != nullptr);
|
|
return static_cast<DataflowBlockRewriteNode*>(get_mutable());
|
|
}
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowBlockRewrite, ffi::ObjectRef,
|
|
DataflowBlockRewriteNode);
|
|
};
|
|
|
|
} // namespace relax
|
|
} // namespace tvm
|
|
|
|
#define TVM_RELAX_BINDING_REWRITE_H_
|
|
#endif // TVM_RELAX_BINDING_REWRITE_H_
|