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 IfNode *op) override
virtual void	VisitBinding (const Binding &binding)
	Generic dispatcher for bindings. More...
virtual void	VisitBinding_ (const VarBindingNode *binding)
virtual void	VisitBinding_ (const MatchCastNode *binding)
virtual void	VisitBinding_ (const VarBindingNode *binding, const ConstantNode *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 IfNode *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. More...
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. More...
virtual Var	VisitVarDef_ (const VarNode *var)
virtual Var	VisitVarDef_ (const DataflowVarNode *var)
Expr	VisitExpr_ (const ConstantNode *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 SeqExprNode *op) override
Expr	VisitExpr_ (const IfNode *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
Public Member Functions inherited from tvm::relax::ExprMutatorBase
Expr	VisitExpr (const Expr &expr) override
Expr	VisitExpr_ (const ConstantNode *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 SeqExprNode *op) override
Expr	VisitExpr_ (const IfNode *op) override
Expr	VisitExpr_ (const OpNode *op) override
Expr	VisitExpr_ (const TupleGetItemNode *op) override
Expr	VisitExprFallback_ (const ExprNode *op) override
Expr	VisitExpr_ (const StringImmNode *op) override
Expr	VisitExpr_ (const DataTypeImmNode *op) override
virtual PrimExpr	VisitTypePrimExprField (const PrimExpr &expr)
	Used to visit the PrimExpr inside of expressions. More...
virtual Type	VisitExprDepTypeField (const Type &ty)
	Visit ty that may recursively contain Expr/PrimExpr. More...

Protected Member Functions
void	ReEmitBinding (const VarBindingNode *binding, Expr new_value)
	Try to remit binding and bind it to a new_value. More...
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. More...
Expr	VisitWithInnerScope (const Expr &body_expr)
	Rewrite the expr with a new scope, used in the branches of If. More...
ffi::Optional< Expr >	LookupBinding (const Var &var)
	Look up the value bound to a variable. More...
template<typename T >
Expr	VisitExprPostOrder_ (const T *op)
	Post-order rewrite a node and normalize. More...
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. More...
Protected Member Functions inherited from tvm::relax::ExprMutatorBase
bool	VisitAndCheckTypeFieldUnchanged (const ffi::ObjectRef &ty)
	Check whether VisitExprDepTypeField change ty. More...

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

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

var	The 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

binding

The 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 ConstantNode *  val  )

virtual

VisitBinding_() [5/18]

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

virtual

VisitBinding_() [6/18]

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

virtual

VisitBinding_() [7/18]

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

virtual

VisitBinding_() [8/18]

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

virtual

VisitBinding_() [9/18]

virtual void tvm::relax::ExprMutator::VisitBinding_ ( const VarBindingNode *  binding, const FunctionNode *  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 IfNode *  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

block

The 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/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [2/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [3/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [4/20]

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

override

VisitExpr_() [5/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [6/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [7/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [8/20]

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

override

VisitExpr_() [9/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [10/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [11/20]

Expr tvm::relax::ExprMutator::VisitExpr_ ( const IfNode *  op )

override

VisitExpr_() [12/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [13/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [14/20]

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

override

VisitExpr_() [15/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [16/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [17/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [18/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [19/20]

Expr tvm::relax::ExprMutatorBase::VisitExpr_

override

VisitExpr_() [20/20]

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

T	The node type to be rewritten.

Parameters

op	The 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

var	The 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 If.

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

Parameters

body_expr

The 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_expr	The body to be visited.
params	Optional 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

var	The var to be updated.
ty	The 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:

• include/tvm/relax/expr_functor.h