tvm
Loading...
Searching...
No Matches
Public Member Functions | Protected Member Functions | Protected Attributes | List of all members
tvm::relax::ExprMutator Class Reference

A mutator works in normal form. More...

#include <expr_functor.h>

Inheritance diagram for tvm::relax::ExprMutator:
Collaboration diagram for tvm::relax::ExprMutator:

Public Member Functions

 ExprMutator (ffi::Optional< IRModule > mod=std::nullopt)
 
Expr VisitExpr (const Expr &expr) override
 
Expr VisitExpr_ (const VarNode *op) override
 
Expr VisitExpr_ (const DataflowVarNode *op) override
 
Expr VisitExpr_ (const FunctionNode *op) override
 
Expr VisitExpr_ (const SeqExprNode *op) override
 
Expr VisitExpr_ (const IfExprNode *op) override
 
virtual void VisitBinding (const Binding &binding)
 Generic dispatcher for bindings.
 
virtual void VisitBinding_ (const VarBindingNode *binding)
 
virtual void VisitBinding_ (const MatchCastNode *binding)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const GenericConstNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const TupleNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const VarNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const DataflowVarNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const ShapeExprNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const ExternFuncNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const GlobalVarNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const FunctionNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const CallNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const SeqExprNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const IfExprNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const OpNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const TupleGetItemNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const ExprNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const StringImmNode *val)
 
virtual void VisitBinding_ (const VarBindingNode *binding, const DataTypeImmNode *val)
 
virtual BindingBlock VisitBindingBlock (const BindingBlock &block) override
 Generic dispatcher for binding blocks.
 
virtual BindingBlock VisitBindingBlock_ (const BindingBlockNode *block)
 
virtual BindingBlock VisitBindingBlock_ (const DataflowBlockNode *block)
 
virtual Var VisitVarDef (const Var &var)
 Generic dispatcher for rewriting the var definition site.
 
virtual Var VisitVarDef_ (const VarNode *var)
 
virtual Var VisitVarDef_ (const DataflowVarNode *var)
 
Expr VisitExpr_ (const GenericConstNode *op) override
 
Expr VisitExpr_ (const TupleNode *op) override
 
Expr VisitExpr_ (const VarNode *op) override
 
Expr VisitExpr_ (const DataflowVarNode *op) override
 
Expr VisitExpr_ (const ShapeExprNode *op) override
 
Expr VisitExpr_ (const ExternFuncNode *op) override
 
Expr VisitExpr_ (const GlobalVarNode *op) override
 
Expr VisitExpr_ (const FunctionNode *op) override
 
Expr VisitExpr_ (const CallNode *op) override
 
Expr VisitExpr_ (const TensorLoadNode *op) override
 
Expr VisitExpr_ (const tvm::IntImmNode *op) override
 
Expr VisitExpr_ (const tvm::FloatImmNode *op) override
 
Expr VisitExpr_ (const SeqExprNode *op) override
 
Expr VisitExpr_ (const IfExprNode *op) override
 
Expr VisitExpr_ (const OpNode *op) override
 
Expr VisitExpr_ (const TupleGetItemNode *op) override
 
Expr VisitExpr_ (const StringImmNode *op) override
 
Expr VisitExpr_ (const DataTypeImmNode *op) override
 
Expr VisitExpr_ (const prim::AddNode *op) override
 
Expr VisitExpr_ (const prim::LShiftNode *op) override
 
Expr VisitExpr_ (const prim::RShiftNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseAndNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseOrNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseXorNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseNotNode *op) override
 
Expr VisitExpr_ (const prim::SubNode *op) override
 
Expr VisitExpr_ (const prim::MulNode *op) override
 
Expr VisitExpr_ (const prim::DivNode *op) override
 
Expr VisitExpr_ (const prim::ModNode *op) override
 
Expr VisitExpr_ (const prim::FloorDivNode *op) override
 
Expr VisitExpr_ (const prim::FloorModNode *op) override
 
Expr VisitExpr_ (const prim::MinNode *op) override
 
Expr VisitExpr_ (const prim::MaxNode *op) override
 
Expr VisitExpr_ (const prim::EQNode *op) override
 
Expr VisitExpr_ (const prim::NENode *op) override
 
Expr VisitExpr_ (const prim::LTNode *op) override
 
Expr VisitExpr_ (const prim::LENode *op) override
 
Expr VisitExpr_ (const prim::GTNode *op) override
 
Expr VisitExpr_ (const prim::GENode *op) override
 
Expr VisitExpr_ (const prim::AndNode *op) override
 
Expr VisitExpr_ (const prim::OrNode *op) override
 
Expr VisitExpr_ (const prim::CastNode *op) override
 
Expr VisitExpr_ (const prim::NotNode *op) override
 
Expr VisitExpr_ (const prim::SelectNode *op) override
 
Expr VisitExpr_ (const prim::RampNode *op) override
 
Expr VisitExpr_ (const prim::BroadcastNode *op) override
 
Expr VisitExpr_ (const prim::ShuffleNode *op) override
 
- Public Member Functions inherited from tvm::relax::ExprMutatorBase
Expr VisitExpr (const Expr &expr) override
 
Expr VisitExpr_ (const GenericConstNode *op) override
 
Expr VisitExpr_ (const TupleNode *op) override
 
Expr VisitExpr_ (const VarNode *op) override
 
Expr VisitExpr_ (const DataflowVarNode *op) override
 
Expr VisitExpr_ (const ShapeExprNode *op) override
 
Expr VisitExpr_ (const ExternFuncNode *op) override
 
Expr VisitExpr_ (const GlobalVarNode *op) override
 
Expr VisitExpr_ (const FunctionNode *op) override
 
Expr VisitExpr_ (const CallNode *op) override
 
Expr VisitExpr_ (const TensorLoadNode *op) override
 
Expr VisitExpr_ (const tvm::IntImmNode *op) override
 
Expr VisitExpr_ (const tvm::FloatImmNode *op) override
 
Expr VisitExpr_ (const SeqExprNode *op) override
 
Expr VisitExpr_ (const IfExprNode *op) override
 
Expr VisitExpr_ (const OpNode *op) override
 
Expr VisitExpr_ (const TupleGetItemNode *op) override
 
Expr VisitExpr_ (const StringImmNode *op) override
 
Expr VisitExpr_ (const DataTypeImmNode *op) override
 
Expr VisitExpr_ (const prim::AddNode *op) override
 
Expr VisitExpr_ (const prim::LShiftNode *op) override
 
Expr VisitExpr_ (const prim::RShiftNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseAndNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseOrNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseXorNode *op) override
 
Expr VisitExpr_ (const prim::BitwiseNotNode *op) override
 
Expr VisitExpr_ (const prim::SubNode *op) override
 
Expr VisitExpr_ (const prim::MulNode *op) override
 
Expr VisitExpr_ (const prim::DivNode *op) override
 
Expr VisitExpr_ (const prim::ModNode *op) override
 
Expr VisitExpr_ (const prim::FloorDivNode *op) override
 
Expr VisitExpr_ (const prim::FloorModNode *op) override
 
Expr VisitExpr_ (const prim::MinNode *op) override
 
Expr VisitExpr_ (const prim::MaxNode *op) override
 
Expr VisitExpr_ (const prim::EQNode *op) override
 
Expr VisitExpr_ (const prim::NENode *op) override
 
Expr VisitExpr_ (const prim::LTNode *op) override
 
Expr VisitExpr_ (const prim::LENode *op) override
 
Expr VisitExpr_ (const prim::GTNode *op) override
 
Expr VisitExpr_ (const prim::GENode *op) override
 
Expr VisitExpr_ (const prim::AndNode *op) override
 
Expr VisitExpr_ (const prim::OrNode *op) override
 
Expr VisitExpr_ (const prim::CastNode *op) override
 
Expr VisitExpr_ (const prim::NotNode *op) override
 
Expr VisitExpr_ (const prim::SelectNode *op) override
 
Expr VisitExpr_ (const prim::RampNode *op) override
 
Expr VisitExpr_ (const prim::BroadcastNode *op) override
 
Expr VisitExpr_ (const prim::ShuffleNode *op) override
 
Expr VisitExprFallback_ (const ExprNode *op) override
 
virtual PrimExpr VisitTypePrimExprField (const PrimExpr &expr)
 Used to visit the PrimExpr inside of dependent type fields.
 
virtual Type VisitExprDepTypeField (const Type &ty)
 Visit ty that may recursively contain Expr/PrimExpr.
 

Protected Member Functions

void ReEmitBinding (const VarBindingNode *binding, Expr new_value)
 Try to remit binding and bind it to a new_value.
 
Expr VisitWithNewScope (const Expr &body_expr, ffi::Optional< ffi::Array< Var > > params=std::nullopt)
 Rewrite the expr with a new scope, used in a Function's body.
 
Expr VisitWithInnerScope (const Expr &body_expr)
 Rewrite the expr with a new scope, used in the branches of IfExpr.
 
ffi::Optional< Expr > LookupBinding (const Var &var)
 Look up the value bound to a variable.
 
template<typename T >
Expr VisitExprPostOrder_ (const T *op)
 Post-order rewrite a node and normalize.
 
Var WithType (Var var, Type ty)
 Create a new var with specified type if the original var's shape or type does not match with the specified ones.
 
- Protected Member Functions inherited from tvm::relax::ExprMutatorBase
bool VisitAndCheckTypeFieldUnchanged (const ffi::ObjectRef &ty)
 Check whether VisitExprDepTypeField change ty.
 

Protected Attributes

BlockBuilder builder_
 Internal block builder to emit bindings during rewriting.
 
std::unordered_map< Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual > var_remap_
 Remap a var to a new var in use-site.
 

Detailed Description

A mutator works in normal form.

ExprMutator expects input AST to be in the normal form, i.e., the expressions are normalized(no nesting and hence the AST is in ANF), and all ty of expressions are available.

Constructor & Destructor Documentation

◆ ExprMutator()

tvm::relax::ExprMutator::ExprMutator ( ffi::Optional< IRModule >  mod = std::nullopt)
inline

Member Function Documentation

◆ LookupBinding()

ffi::Optional< Expr > tvm::relax::ExprMutator::LookupBinding ( const Var &  var)
protected

Look up the value bound to a variable.

Parameters
varThe var to be looked up.
Returns
The value bound to the input var.
Note
For function parameters, this function returns std::nullopt.

◆ ReEmitBinding()

void tvm::relax::ExprMutator::ReEmitBinding ( const VarBindingNode *  binding,
Expr  new_value 
)
protected

Try to remit binding and bind it to a new_value.

This function is called after VisitExpr(binding->value) in VisitBinding_(const VarBinding*). It will try to reuse the current binding when the new value's shape/type matches the original binding and no changes in var is needed.

Otherwise, a new binding will be emitted to replace the var specified in the current binding.

◆ VisitBinding()

virtual void tvm::relax::ExprMutator::VisitBinding ( const Binding &  binding)
virtual

Generic dispatcher for bindings.

Parameters
bindingThe binding to be visited.

◆ VisitBinding_() [1/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const MatchCastNode *  binding)
virtual

◆ VisitBinding_() [2/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding)
virtual

◆ VisitBinding_() [3/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const CallNode *  val 
)
virtual

◆ VisitBinding_() [4/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const DataflowVarNode *  val 
)
virtual

◆ VisitBinding_() [5/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const DataTypeImmNode *  val 
)
virtual

◆ VisitBinding_() [6/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const ExprNode *  val 
)
virtual

◆ VisitBinding_() [7/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const ExternFuncNode *  val 
)
virtual

◆ VisitBinding_() [8/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const FunctionNode *  val 
)
virtual

◆ VisitBinding_() [9/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const GenericConstNode *  val 
)
virtual

◆ VisitBinding_() [10/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const GlobalVarNode *  val 
)
virtual

◆ VisitBinding_() [11/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const IfExprNode *  val 
)
virtual

◆ VisitBinding_() [12/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const OpNode *  val 
)
virtual

◆ VisitBinding_() [13/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const SeqExprNode *  val 
)
virtual

◆ VisitBinding_() [14/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const ShapeExprNode *  val 
)
virtual

◆ VisitBinding_() [15/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const StringImmNode *  val 
)
virtual

◆ VisitBinding_() [16/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const TupleGetItemNode *  val 
)
virtual

◆ VisitBinding_() [17/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const TupleNode *  val 
)
virtual

◆ VisitBinding_() [18/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding,
const VarNode *  val 
)
virtual

◆ VisitBindingBlock()

virtual BindingBlock tvm::relax::ExprMutator::VisitBindingBlock ( const BindingBlock &  block)
overridevirtual

Generic dispatcher for binding blocks.

Parameters
blockThe binding block to be visited.
Returns
The binding block after transformation.

Reimplemented from tvm::relax::ExprMutatorBase.

◆ VisitBindingBlock_() [1/2]

virtual BindingBlock tvm::relax::ExprMutator::VisitBindingBlock_ ( const BindingBlockNode *  block)
virtual

◆ VisitBindingBlock_() [2/2]

virtual BindingBlock tvm::relax::ExprMutator::VisitBindingBlock_ ( const DataflowBlockNode *  block)
virtual

◆ VisitExpr()

Expr tvm::relax::ExprMutator::VisitExpr ( const Expr &  expr)
override

◆ VisitExpr_() [1/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const CallNode *  op)
override

◆ VisitExpr_() [2/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const DataflowVarNode *  op)
override

◆ VisitExpr_() [3/52]

Expr tvm::relax::ExprMutator::VisitExpr_ ( const DataflowVarNode *  op)
override

◆ VisitExpr_() [4/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const DataTypeImmNode *  op)
override

◆ VisitExpr_() [5/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const ExternFuncNode *  op)
override

◆ VisitExpr_() [6/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const FunctionNode *  op)
override

◆ VisitExpr_() [7/52]

Expr tvm::relax::ExprMutator::VisitExpr_ ( const FunctionNode *  op)
override

◆ VisitExpr_() [8/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const GenericConstNode *  op)
override

◆ VisitExpr_() [9/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const GlobalVarNode *  op)
override

◆ VisitExpr_() [10/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const IfExprNode *  op)
override

◆ VisitExpr_() [11/52]

Expr tvm::relax::ExprMutator::VisitExpr_ ( const IfExprNode *  op)
override

◆ VisitExpr_() [12/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const OpNode *  op)
override

◆ VisitExpr_() [13/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::AddNode *  op)
override

◆ VisitExpr_() [14/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::AndNode *  op)
override

◆ VisitExpr_() [15/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::BitwiseAndNode *  op)
override

◆ VisitExpr_() [16/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::BitwiseNotNode *  op)
override

◆ VisitExpr_() [17/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::BitwiseOrNode *  op)
override

◆ VisitExpr_() [18/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::BitwiseXorNode *  op)
override

◆ VisitExpr_() [19/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::BroadcastNode *  op)
override

◆ VisitExpr_() [20/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::CastNode *  op)
override

◆ VisitExpr_() [21/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::DivNode *  op)
override

◆ VisitExpr_() [22/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::EQNode *  op)
override

◆ VisitExpr_() [23/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::FloorDivNode *  op)
override

◆ VisitExpr_() [24/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::FloorModNode *  op)
override

◆ VisitExpr_() [25/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::GENode *  op)
override

◆ VisitExpr_() [26/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::GTNode *  op)
override

◆ VisitExpr_() [27/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::LENode *  op)
override

◆ VisitExpr_() [28/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::LShiftNode *  op)
override

◆ VisitExpr_() [29/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::LTNode *  op)
override

◆ VisitExpr_() [30/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::MaxNode *  op)
override

◆ VisitExpr_() [31/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::MinNode *  op)
override

◆ VisitExpr_() [32/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::ModNode *  op)
override

◆ VisitExpr_() [33/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::MulNode *  op)
override

◆ VisitExpr_() [34/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::NENode *  op)
override

◆ VisitExpr_() [35/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::NotNode *  op)
override

◆ VisitExpr_() [36/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::OrNode *  op)
override

◆ VisitExpr_() [37/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::RampNode *  op)
override

◆ VisitExpr_() [38/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::RShiftNode *  op)
override

◆ VisitExpr_() [39/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::SelectNode *  op)
override

◆ VisitExpr_() [40/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::ShuffleNode *  op)
override

◆ VisitExpr_() [41/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const prim::SubNode *  op)
override

◆ VisitExpr_() [42/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const SeqExprNode *  op)
override

◆ VisitExpr_() [43/52]

Expr tvm::relax::ExprMutator::VisitExpr_ ( const SeqExprNode *  op)
override

◆ VisitExpr_() [44/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const ShapeExprNode *  op)
override

◆ VisitExpr_() [45/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const StringImmNode *  op)
override

◆ VisitExpr_() [46/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const TensorLoadNode *  op)
override

◆ VisitExpr_() [47/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const TupleGetItemNode *  op)
override

◆ VisitExpr_() [48/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const TupleNode *  op)
override

◆ VisitExpr_() [49/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const tvm::FloatImmNode *  op)
override

◆ VisitExpr_() [50/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const tvm::IntImmNode *  op)
override

◆ VisitExpr_() [51/52]

Expr tvm::relax::ExprMutatorBase::VisitExpr_ ( const VarNode *  op)
override

◆ VisitExpr_() [52/52]

Expr tvm::relax::ExprMutator::VisitExpr_ ( const VarNode *  op)
override

◆ VisitExprPostOrder_()

template<typename T >
Expr tvm::relax::ExprMutator::VisitExprPostOrder_ ( const T *  op)
inlineprotected

Post-order rewrite a node and normalize.

Template Parameters
TThe node type to be rewritten.
Parameters
opThe node to be rewritten.
Returns
The node after post rewritten.

◆ VisitVarDef()

virtual Var tvm::relax::ExprMutator::VisitVarDef ( const Var &  var)
virtual

Generic dispatcher for rewriting the var definition site.

Parameters
varThe var to be visited.
Returns
The var after post-order rewritten.
Note
VisitExpr_(const VarNode*) will only visit the usage site of an Var

◆ VisitVarDef_() [1/2]

virtual Var tvm::relax::ExprMutator::VisitVarDef_ ( const DataflowVarNode *  var)
virtual

◆ VisitVarDef_() [2/2]

virtual Var tvm::relax::ExprMutator::VisitVarDef_ ( const VarNode *  var)
virtual

◆ VisitWithInnerScope()

Expr tvm::relax::ExprMutator::VisitWithInnerScope ( const Expr &  body_expr)
protected

Rewrite the expr with a new scope, used in the branches of IfExpr.

Visit an expression that may access variables from the current scope, but may not export definitions into the current scope.

Parameters
body_exprThe body to be visited.
Returns
The expr after visiting.
See also
VisitWithNewScope
Note
The body_expr must be an SeqExpr in the normal form.

◆ VisitWithNewScope()

Expr tvm::relax::ExprMutator::VisitWithNewScope ( const Expr &  body_expr,
ffi::Optional< ffi::Array< Var > >  params = std::nullopt 
)
protected

Rewrite the expr with a new scope, used in a Function's body.

Visit an expression that may neither access variables from the current scope, nor may export definitions into the current scope.

Parameters
body_exprThe body to be visited.
paramsOptional parameters that are visible within the scope.
Returns
The expr after visiting.
Note
The body_expr must be an SeqExpr in the normal form.

◆ WithType()

Var tvm::relax::ExprMutator::WithType ( Var  var,
Type  ty 
)
protected

Create a new var with specified type if the original var's shape or type does not match with the specified ones.

Parameters
varThe var to be updated.
tyThe type to be updated.
Returns
The var filled with type information.

Member Data Documentation

◆ builder_

BlockBuilder tvm::relax::ExprMutator::builder_
protected

Internal block builder to emit bindings during rewriting.

◆ var_remap_

std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> tvm::relax::ExprMutator::var_remap_
protected

Remap a var to a new var in use-site.


The documentation for this class was generated from the following file: