Skip to content

T
tic
Commit abe6f770 by Zhi Committed by Jared Roesch
Options

[Relay] Pass manager (#2546) 

* initial commit

* add python frontend and module tests

* add unit tests for function pass and optimize interface

* add ExprPass

* remove PassState and pass context for run

* add required_passes

* return module

* remove move

* fix minor reviews

* remove optimizer, optimizer->pass_manager, make pass a the base class of all

* remove deleted files

* move resolvedependency to sequential pass, use ir_pass namespace

* add todo

* add disabled passes in sequetialpass

* fix minor

* fix currying doc

* remove pass_kind from passnode

* remove pass kind from test

* fix doc

* fix per @tqchen's comments

* remove pass_manager.py create separate classes

* simplify pass_func

* inline using passfunc

* update doc

* disable test_quantize_pass for now

* create PassInfo class to contain the meta data

* flatten passinfo for interface

* retrigger ci

* remove required method

* make Pass python class lighter

* create pass -> decorator

* make the api consistent for all classes
parent 7226c010
Showing with 1561 additions and 50 deletions
include/tvm/relay/pass.h
......@@ -2,18 +2,225 @@
* Copyright (c) 2018 by Contributors
* \file tvm/relay/pass.h
* \brief The set of Relay passes written in C++.
*
* This file also implements a pass manager. The pass manager manages a sequence
* of Relay-to-Relay transformation passes over a particlar unit of AST. The
* design is largely inspired from LLVM's pass manager and modern deep learning
* frameworks that perform tensor->tensor transformations.
*
* The responsibilities of a traditional compiler pass manager usually involves:
* - Organizing the execution order of optimization passes though not
* necessarily in the optimal sequence.
* - Collecting required analysis information and keep them up-to-date.
* - Reducing the effort required to implement new passes for compiler
* developers, etc.
*
* Similar to LLVM's pass manager, we designed the Relay pass manager to work
* different granularity, i.e. module level, function level, and even sequential
* passe that contains a host of passes.
*
* However, we also extend the functionality of the traditional pass manager
* with the consideration of requirements/convention from deep learning
* frameworks, such as Pytorch and Gluon, etc. Each pass in the Relay pass
* manager performs the Relay.Module -> Relay.Module transformation. All
* different types of passes, including the sequential-level pass object, are
* essentially pass objects. This design, therefore, effectively provides users
* a consistent and convenient interface, i.e. Pass, to play with. It offers a
* means to ease the development and testing of Relay passes. For example, with
* the pass manager, external users will be able to have custom passes correctly
* scheduled without having to modify a single handcrafted pass order.
*
* In the future we need to describe constraints between passes. For example,
* we may want to preserve dependencies between different passes and validate
* them on the completion of a certain pass.
*
* We also need to store side information and import the error reporting system.
*/
#ifndef TVM_RELAY_PASS_H_
#define TVM_RELAY_PASS_H_
#include <tvm/ir.h>
#include <tvm/packed_func_ext.h>
#include <tvm/relay/error.h>
#include <tvm/relay/expr.h>
#include <tvm/relay/module.h>
#include <tvm/relay/op_attr_types.h>
#include <tvm/relay/type.h>
#include <string>
#include <vector>
namespace tvm {
namespace relay {
namespace pass {
/*
* \brief The context of pass.
*/
class PassContext;
/*!
* \brief PassContextNode contains the information that a pass can rely on, such as
* analysis results.
*/
class PassContextNode : public RelayNode {
public:
/*!
* \brief The error reporter used to notify users why an optimization fails.
*/
ErrorReporter err_reporter;
PassContextNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) final {
}
TVM_DLL static PassContext make();
static constexpr const char* _type_key = "relay.PassContext";
TVM_DECLARE_NODE_TYPE_INFO(PassContextNode, RelayNode);
};
TVM_DEFINE_NODE_REF(PassContext, PassContextNode)
/*
* \brief The meta data of a pass.
*
* PassInfo can be extended conveniently in the future if more meta information
* is needed.
*/
class PassInfo;
/*!
* \brief PassInfoNode contains meta data that will be used to help optimization
* and analysis.
*/
class PassInfoNode : public RelayNode {
public:
/*! \brief The minimal optimization level that this pass will be enabled. */
int opt_level;
/*! \brief The name of an optimization/analysis pass. */
std::string name;
/*! \brief The passes that are required to perform the current pass. */
tvm::Array<tvm::Expr> required;
PassInfoNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) final {
v->Visit("opt_level", &opt_level);
v->Visit("name", &name);
v->Visit("required", &required);
}
TVM_DLL static PassInfo make(int opt_level, std::string name,
tvm::Array<tvm::Expr> required);
static constexpr const char* _type_key = "relay.PassInfo";
TVM_DECLARE_NODE_TYPE_INFO(PassInfoNode, RelayNode);
};
TVM_DEFINE_NODE_REF(PassInfo, PassInfoNode)
class Pass;
/*!
* \brief PassNode is the base type of differnt types of optimization passes.
* It is designed as a pure class and implemented by different pass subclasses
* at different granularity of Relay nodes.
*/
class PassNode : public RelayNode {
public:
/*
* \brief Get the pass information/meta data. */
virtual PassInfo Info() const = 0;
/*!
* \brief Set the context information for a pass.
*
* \param pass_ctx The context information for a certain pass.
*/
virtual void SetContext(const PassContext& pass_ctx) = 0;
/*!
* \brief Execute the optimization pass using a functor.
*
* \param mod The module that an optimization pass runs on.
*
* \return The updated module.
*/
virtual Module operator()(const Module& mod) const = 0;
void VisitAttrs(tvm::AttrVisitor* v) override {}
static constexpr const char* _type_key = "relay.Pass";
TVM_DECLARE_BASE_NODE_INFO(PassNode, RelayNode);
};
class Pass : public NodeRef {
public:
Pass() = default;
explicit Pass(NodePtr<tvm::Node> p) : NodeRef(p) {}
PassNode* operator->() const {
return static_cast<PassNode*>(this->node_.get());
}
using ContainerType = PassNode;
};
/*
* \brief Create a module pass.
*
* \param pass_func The packed function that contains the optimization.
* \param opt_level The optimization level of the module pass.
* \param name The name of the module pass.
* \param required The list of the passes that the module pass is dependent on.
*
* \return The created module pass.
*/
Pass CreateModulePass(
const runtime::TypedPackedFunc<Module(Module, PassContext)>& pass_func,
int opt_level,
const std::string& name,
const tvm::Array<tvm::Expr>& required);
/*
* \brief Create a function pass.
*
* \param pass_func The packed function that contains the optimization.
* \param opt_level The optimization level of the function pass.
* \param name The name of the function pass.
* \param required The list of the passes that the function pass is dependent on.
*
* \return The created function pass.
*/
Pass CreateFunctionPass(
const runtime::TypedPackedFunc<Function(Function, PassContext)>& pass_func,
int opt_level,
const std::string& name,
const tvm::Array<tvm::Expr>& required);
/*
* \brief Create a sequential pass.
*
* \param passes The optimization passes will be performed.
* \param opt_level The optimization level of the sequential pass.
* \param name The name of the sequential pass.
* \param required The list of the passes that the sequential pass is dependent on.
* \param disabled The disabled passes.
*
* \return The created sequential pass.
*/
Pass CreateSequentialPass(const tvm::Array<Pass>& passes,
int opt_level,
const std::string& name,
const tvm::Array<tvm::Expr>& required,
const tvm::Array<tvm::Expr>& disabled);
} // namespace pass
/*!
* \brief Infer the type of an expression.
*
......
python/tvm/relay/__init__.py
......@@ -79,6 +79,9 @@ Match = adt.Match
var = expr.var
const = expr.const
bind = expr.bind
module_pass = ir_pass.module_pass
function_pass = ir_pass.function_pass
sequential_pass = ir_pass.sequential_pass
# ExprFunctor
ExprFunctor = expr_functor.ExprFunctor
......@@ -90,3 +93,11 @@ fromtext = parser.fromtext
# Param Serialization
save_param_dict = param_dict.save_param_dict
load_param_dict = param_dict.load_param_dict
# Pass manager
PassInfo = ir_pass.PassInfo
PassContext = ir_pass.PassContext
Pass = ir_pass.Pass
ModulePass = ir_pass.ModulePass
FunctionPass = ir_pass.FunctionPass
SequentialPass = ir_pass.SequentialPass
python/tvm/relay/_ir_pass.pyi
from .env import Module
import tvm
from . import ir
from .base import NodeBase
from .env import Module
class PassContext(NodeBase):
def __init__(self):
...
class PassInfo(NodeBase):
name = ... # type: str
opt_level = ... # type: int
required = ... # type: list
def __init__(self, name, opt_level, required)
# type: (str, int, list) -> None
class Pass(NodeBase):
def __init__(self):
...
class ModulePass(Pass):
name = ... # type: str
opt_level = ... # type: int
pass_func = ... # type: Callable
required = ... # type: list
def __init__(self, name, opt_level, pass_func, required):
# type: (str, int, Callable, list) -> None
...
class FunctionPass(Pass):
name = ... # type: str
opt_level = ... # type: int
pass_func = ... # type: Callable
required = ... # type: list
def __init__(self, name, opt_level, pass_func, required):
# type: (str, int, Callable, list) -> None
...
class SequentialPass(Pass):
name = ... # type: str
opt_level = ... # type: int
passes = ... # type: list
required = ... # type: list
disabled = ... # type: list
def __init__(self, name, opt_level, passes, required, disabled):
# type: (str, int, list, list, list) -> None
...
def check_expr(env: Module, expr: ir.Expr) -> ir.Type: ...
def generalize(env: Module, expr: ir.Expr) -> ir.Expr: ...
......
python/tvm/relay/ir_pass.py
# pylint: disable=no-else-return
# pylint: disable=unidiomatic-typecheck
"""The set of passes for Relay.
"""
This file contains:
1. The set of passes for Relay, which exposes an interface for configuring the
passes and scripting them in Python.
Exposes an interface for configuring the passes and
scripting them in Python.
2. The pass manager for Relay which exposes different granularity of interfaces
for users to implement and use passes more conveniently.
"""
import types
from . import _ir_pass
from . import _make
from .expr import Expr
from .ty import Type
from .base import RelayNode, register_relay_node
from .module import Module
@register_relay_node
class PassInfo(RelayNode):
"""The class that contains the meta data required by a pass. It is the
container of information needed by running an optimization or analysis.
This class can be extended by adding new members when more meta data is
needed.
Parameters
----------
name : str
The pass name.
opt_level : int
The optimization level of this pass.
required : List[str]
The list of passes that are required by a certain pass.
"""
def __init__(self, name, opt_level, required=None):
self.__init_handle_by_constructor__(_ir_pass.PassInfo, name, opt_level,
required)
@register_relay_node
class PassContext(RelayNode):
"""The basis where a Relay optimization/analysis runs on.
Each pass context contains a number of auxiliary information that is used
to help an optimization pass. Such information includes the error reporter
to record the errors of during the optimization, etc.
"""
def __init__(self):
self.__init_handle_by_constructor__(_ir_pass.PassContext)
@register_relay_node
class Pass(RelayNode):
"""The base class of all passes. All methods here are just simple wrappers
that are implemented in the backend. They are defined for users to
conveniently interact with the base class.
"""
def set_pass_context(self, pass_ctx):
"""Setup the pass context for analysis and optimizations. This context
could be shared by different passes for sequential passes.
Parameters
----------
pass_ctx : PassContext
The context that is used to help perform a certain pass or a series
of passes.
"""
if not isinstance(pass_ctx, PassContext):
raise TypeError("pass_ctx is expected to be the PassContext type")
_ir_pass.SetContext(self, pass_ctx)
@property
def info(self):
"""Get the pass meta."""
return _ir_pass.Info(self)
def __call__(self, mod):
"""Execute the pass. Note that for sequential pass, the dependency among
different passes will be resolved in the backend.
Parameters
----------
mod : tvm.relay.Module
The module that a certain optimization is performed on.
Returns
-------
mod : tvm.relay.Module
The updated module after applying this pass.
"""
return _ir_pass.RunPass(self, mod)
@register_relay_node
class ModulePass(Pass):
"""A pass that works on tvm.relay.Module. Users don't need to interact with
this class directly. Instead, a module pass should be created through
`module_pass`, because the design of the `module_pass` API is flexible
enough to handle the creation of a module pass in different manners. In
addition, all members of a module pass can be accessed from the base class.
The same rule applies to FunctionPass and SequentialPass as well.
"""
@register_relay_node
class FunctionPass(Pass):
"""A pass that works on each tvm.relay.Function in a module. A function
pass class should be created through `function_pass`.
"""
@register_relay_node
class SequentialPass(Pass):
"""A pass that works on a sequence of pass objects. A sequential pass class
should be created through `sequential_pass`.
"""
def module_pass(pass_func=None, opt_level=None, name=None, required=None):
"""Create a module pass. This function returns a callback when pass_func
is provided. Otherwise, it returns the created module level pass using the
given optimization function.
Parameters
----------
pass_func : Optional[Callable[(Module/Function, PassContext) ->
Module/Function]]
The implemented optimization pass.
opt_level : int
The optimization level of this module pass.
name : Optional[str]
The name of the module pass. The name could be empty. In this case, the
name of the optimization function will be used as the pass name.
required : Optional[List[str]]
The list of passes that the module pass is dependent on.
Returns
-------
create_module_pass : Union[Callable, ModulePass]
The callable that will create a module pass is returned when
pass_func is not passed in. Otherwise, a ModulePass object will be
directly created.
Examples
--------
The following code creates a module level pass and adds an abs function to
the module.
.. code-block:: python
@relay.ir_pass.module_pass(opt_level=2)
def transform(mod, ctx):
tp = relay.TensorType((10,), "float32")
x = relay.var("x", tp)
gv = relay.GlobalVar("var")
func = relay.Function([x], relay.abs(x))
new_mod = relay.Module({gv: func})
new_mod.update(mod)
return new_mod
module_pass = transform
assert isinstance(module_pass, ir_pass.ModulePass)
assert module_pass.info.opt_level == 2
# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = module_pass(m)
# Now a function abs should be added to the module m.
"""
if opt_level is None:
raise ValueError("Please provide opt_level for the module pass.")
required = required if required else []
if not isinstance(required, (list, tuple)):
raise TypeError("Required is expected to be the type of " +
"list/tuple.")
def create_module_pass(pass_func):
"""Internal function that creates a module pass"""
if not isinstance(pass_func, (types.FunctionType, types.LambdaType)):
raise TypeError("pass_func must be a callable for Module pass")
return _ir_pass.CreateModulePass(pass_func, opt_level,
name if name else pass_func.__name__,
required)
if pass_func:
return create_module_pass(pass_func)
return create_module_pass
def function_pass(pass_func=None, opt_level=None, name=None, required=None):
"""Create a function pass. This function returns a callback when pass_func
is provided. Otherwise, it returns the created function pass using the
given optimization function.
Parameters
----------
pass_func : Optional[Callable[(Module/Function, PassContext) ->
Module/Function]]
The implemented optimization pass.
opt_level : int
The optimization level of this module pass.
name : Optional[str]
The name of the function pass. The name could be empty. In this case, the
name of the optimization function will be used as the pass name.
required : Optional[List[str]]
The list of passes that the module pass is dependent on.
Returns
-------
create_function_pass : Union[Callable, FunctionPass]
The callable that will create a function pass is returned when
pass_func is not passed in. Otherwise, a FunctionPass object will be
created.
Examples
--------
The following code creates a function level pass that performs constant
folding.
.. code-block:: python
@relay.ir_pass.function_pass(opt_level=2)
def transform(func, ctx):
return ir_pass.fold_constant(func)
function_pass = transform
assert isinstance(function_pass, ir_pass.FunctionPass)
assert function_pass.info.opt_level == 2
# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = function_pass(m)
# Now constant folding should have been applied to every function in
# the provided module m. And the updated module will be returned.
"""
if opt_level is None:
raise ValueError("Please provide opt_level for the funtion pass.")
required = required if required else []
if not isinstance(required, (list, tuple)):
raise TypeError("Required is expected to be the type of " +
"list/tuple.")
def create_function_pass(pass_func):
"""Internal function that creates a function pass"""
if not isinstance(pass_func, (types.FunctionType, types.LambdaType)):
raise TypeError("pass_func must be a callable for Module pass")
return _ir_pass.CreateFunctionPass(pass_func, opt_level,
name if name else pass_func.__name__,
required)
if pass_func:
return create_function_pass(pass_func)
return create_function_pass
def sequential_pass(passes=None, opt_level=2, name="sequential_pass",
required=None, disabled=None):
"""Create a sequential pass using a defined optimization function from
Python. Some typical usage of the sequential pass are:
1. Users provide a list of passes for optimization.
2. Only an optimization level is provided so that the backend system has
to glob all passes at this level and below to perform the optimizations.
Note that users can also provide a series of passes that they don't want to
apply when running a sequential pass. Pass dependency will be resolved in
the backend as well.
Parameters
----------
passes : Optional[List[Pass]]
A sequence of passes candidate for optimization.
opt_level : Optional[int]
The optimization level of this sequential pass.
name : Optional[str]
The name of the sequential pass.
required : Optional[List[str]]
The list of passes that the sequential pass is dependent on.
disabled : Optional[List[str]]
A list of disabled passes.
Returns
-------
ret : Pass
A sequential pass built through pass_func.
"""
passes = passes if passes else []
if not isinstance(passes, (list, tuple)):
raise TypeError("passes must be a list of Pass objects.")
disabled = disabled if disabled else []
if not isinstance(disabled, (list, tuple)):
raise TypeError("disabled must be a list or tuple of pass names")
required = required if required else []
if not isinstance(required, (list, tuple)):
raise TypeError("Required is expected to be the type of list/tuple.")
return _ir_pass.CreateSequentialPass(passes, opt_level, name, required,
disabled)
def post_order_visit(expr, fvisit):
"""Recursively visit the ir in post DFS order node,
apply fvisit. Each node is guaranteed to be visited
......
src/relay/pass/pass_manager.cc 0 → 100644
/*!
* Copyright (c) 2019 by Contributors
* \file src/relay/pass/pass_manager.cc
* \brief Relay pass manager implementation.
*/
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/pass.h>
namespace tvm {
namespace relay {
namespace pass {
using tvm::IRPrinter;
class ModulePass;
/*!
* \brief Module-level passes are designed to implement global
* analysis/optimizations, i.e. interprocedural optimizations (IPO), etc. Passes
* at this level have the full control of a given Relay program including
* addition and deletion of functions.
*/
class ModulePassNode : public PassNode {
public:
/* \brief The pass meta data.*/
PassInfo pass_info;
/*! \brief The pass function sketches the real optimization. For example,
* we may need to perform dead code elimination on the module level. We could
* implement the algorithm in the `pass_func` and let it run on a module. It
* will then remove the dead code including the unused functions in the module.
*/
runtime::TypedPackedFunc<Module(Module, PassContext)> pass_func;
ModulePassNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) final {
v->Visit("pass_info", &pass_info);
}
/*!
* \brief Run a module pass on a certain module.
*
* \param mod The module that an optimization pass runs on.
*
* \return Return the updated module.
*/
Module operator()(const Module& mod) const final;
/*!
* \brief Get the pass information/meta data.
*/
PassInfo Info() const { return pass_info; }
/*!
* \brief Set the context information for a module pass.
*
* \param pass_ctx The context information for a module pass.
*/
void SetContext(const PassContext& pass_ctx) final;
TVM_DLL static ModulePass make(
runtime::TypedPackedFunc<Module(Module, PassContext)> pass_func,
PassInfo pass_info);
static constexpr const char* _type_key = "relay.ModulePass";
TVM_DECLARE_NODE_TYPE_INFO(ModulePassNode, PassNode);
private:
/*!
* \brief The context information that is used to help perform a module pass.
*/
PassContext pass_ctx_;
};
RELAY_DEFINE_NODE_REF(ModulePass, ModulePassNode, Pass);
class FunctionPass;
/*!
* \brief Function-level passes are used to implement various global
* optimizations for a given Relay module. It fetches one function at a time
* from the function list in the module for optimization.
*
* Note that the scope of passes at this level is a Relay function. Therefore,
* we cannot add or delete a function through these passes as they are not aware
* of the global information.
*/
class FunctionPassNode : public PassNode {
public:
/* \brief The pass meta data.*/
PassInfo pass_info;
/*! \brief The packed pass function sketches the real optimization. For
* instance, we can implement a pass that works on a Relay function as a
* `pass_func` and let it run on a given module. The same `pass_func` will
* then be applied on each function in the module.
*/
runtime::TypedPackedFunc<Function(Function, PassContext)> pass_func;
FunctionPassNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) final {
v->Visit("pass_info", &pass_info);
}
/*!
* \brief Run a function pass on a certain module.
*
* \param mod The module that an optimization pass runs on.
*
* \return Return the updated module.
*/
Module operator()(const Module& mod) const final;
/*!
* \brief Get the pass information/meta data.
*/
PassInfo Info() const { return pass_info; }
/*!
* \brief Set the context information for a function-level pass.
*
* \param pass_ctx The context information for a function-level pass.
*/
void SetContext(const PassContext& pass_ctx) final;
TVM_DLL static FunctionPass make(
runtime::TypedPackedFunc<Function(Function, PassContext)> pass_func,
PassInfo pass_info);
static constexpr const char* _type_key = "relay.FunctionPass";
TVM_DECLARE_NODE_TYPE_INFO(FunctionPassNode, PassNode);
private:
/*
* \brief Check if a function should be skipped for optimization.
*
* \param func The target function to be checked.
*
* \return Return true if the function will be skipped, otherwise false.
*/
bool SkipFunction(const Function& func) const;
/*!
* \brief The context information that is used to help perform a module pass.
*/
PassContext pass_ctx_;
};
RELAY_DEFINE_NODE_REF(FunctionPass, FunctionPassNode, Pass);
class SequentialPass;
/*!
* \brief The SequentialPassNode contains a set of passes that transform Relay
* programs from one AST to another semantically equivalent one.
*
* One example of this level of pass is that the pass manager needs to correctly
* perform a host of optimizations with a given optimization level and disabled
* passes.
*/
class SequentialPassNode : public PassNode {
public:
/* \brief The pass meta data.*/
PassInfo pass_info;
/*! \brief A list of passes that used to compose a sequential pass. */
tvm::Array<Pass> passes;
/*!
* \brief A list of disabled passes that should be excluded when executing the
* sequential pass.
*/
tvm::Array<tvm::Expr> disabled;
void VisitAttrs(tvm::AttrVisitor* v) final {
v->Visit("pass_info", &pass_info);
v->Visit("passes", &passes);
v->Visit("disabled", &disabled);
}
/*!
* \brief Get the pass information/meta data.
*/
PassInfo Info() const { return pass_info; }
/*!
* \brief Add a pass to the pass list.
*
* \param pass The candidate pass to be added.
*/
void AddPass(const Pass& pass) {
passes.push_back(pass);
}
TVM_DLL static SequentialPass make(tvm::Array<Pass> passes,
PassInfo pass_info,
tvm::Array<tvm::Expr> disabled);
/*!
* \brief Resolve the pass dependency. It globs all required passes by
* a given pass and executes them.
*
* \param mod The module that an optimization pass runs on.
*
* \return The updated module after resolving pass dependencies.
*
* TODO(zhiics) Build a dependency graph among the passes using provided
* metadata, i.e. required_passes. Likely, we can have a data structure, i.e.
* PassInfo, to store the relevant information including the parent passes.
*/
void ResolveDependency(const Module& mod);
TVM_DLL std::vector<std::string> DisabledPasses() const;
/*!
* \brief Perform optimizations on a series of passes. The aforementioned
* typical pass manager jobs could be done by it. This function could
* be overloaded to focus on different metrics, i.e. performance,
* memory footprint, etc.
*
* \param mod The module that an optimization pass runs on.
*
* \return Return the updated module.
*/
Module operator()(const Module& mod) const final;
/*!
* \brief Set the context information for a sequential pass.
*
* \param pass_ctx The context information for a sequential pass.
*/
void SetContext(const PassContext& pass_ctx) final;
static constexpr const char* _type_key = "relay.SequentialPass";
TVM_DECLARE_NODE_TYPE_INFO(SequentialPassNode, PassNode);
private:
/*!
* \brief The context information that is used to help perform a module pass.
*/
PassContext pass_ctx_;
};
RELAY_DEFINE_NODE_REF(SequentialPass, SequentialPassNode, Pass);
PassInfo PassInfoNode::make(int opt_level, std::string name,
tvm::Array<tvm::Expr> required) {
auto pass_info = make_node<PassInfoNode>();
pass_info->opt_level = opt_level;
pass_info->name = std::move(name);
pass_info->required = std::move(required);
return PassInfo(pass_info);
}
PassContext PassContextNode::make() {
auto ctx = make_node<PassContextNode>();
return PassContext(ctx);
}
ModulePass ModulePassNode::make(
runtime::TypedPackedFunc<Module(Module, PassContext)> pass_func,
PassInfo pass_info) {
auto n = make_node<ModulePassNode>();
n->pass_func = std::move(pass_func);
n->pass_info = std::move(pass_info);
return ModulePass(n);
}
// Module -> Module optimizations.
// TODO(zhiics) 1. Check and handle the required passes.
// 2. Probably use CoW for all places that use module instead of
// returning the updated one.
Module ModulePassNode::operator()(const Module& mod) const {
PassInfo pass_info = Info();
LOG(INFO) << "Executing module pass : " << pass_info.operator->()->name
<< " with opt level: " << pass_info.operator->()->opt_level << "\n";
CHECK(mod.defined());
auto updated_mod = pass_func(mod, pass_ctx_);
CHECK(updated_mod.defined());
return updated_mod;
}
void ModulePassNode::SetContext(const PassContext& pass_ctx) {
pass_ctx_ = pass_ctx;
}
FunctionPass FunctionPassNode::make(
runtime::TypedPackedFunc<Function(Function, PassContext)> pass_func,
PassInfo pass_info) {
auto n = make_node<FunctionPassNode>();
n->pass_func = std::move(pass_func);
n->pass_info = std::move(pass_info);
return FunctionPass(n);
}
// Perform Module -> Module optimizations at the Function level.
// TODO(zhiics) Check and handle the required passes.
Module FunctionPassNode::operator()(const Module& mod) const {
PassInfo pass_info = Info();
LOG(INFO) << "Executing function pass : " << pass_info.operator->()->name
<< " with opt level: " << pass_info.operator->()->opt_level << "\n";
CHECK(mod.defined());
std::vector<std::pair<GlobalVar, Function>> updated_funcs;
ModuleNode* mod_node = mod.operator->();
for (const auto& it : mod_node->functions) {
if (!SkipFunction(it.second)) {
auto updated_func = pass_func(it.second, pass_ctx_);
CHECK(updated_func.defined());
updated_funcs.push_back({std::move(it.first), std::move(updated_func)});
}
}
// Update the optimized functions.
for (const auto& it : updated_funcs) {
mod_node->Update(it.first, it.second);
}
return GetRef<Module>(mod_node);
}
void FunctionPassNode::SetContext(const PassContext& pass_ctx) {
pass_ctx_ = pass_ctx;
}
// TODO(zhiics) Create an enum attribute for FunctionNode
// enum Attribute {kPrimitive, kSkipOptimization}
bool FunctionPassNode::SkipFunction(const Function& func) const {
NodeRef res = FunctionGetAttr(func, "SkipOptimization");
const ir::IntImm* pval = res.as<ir::IntImm>();
return pval && pval->value != 0;
}
SequentialPass SequentialPassNode::make(tvm::Array<Pass> passes,
PassInfo pass_info,
tvm::Array<tvm::Expr> disabled) {
auto n = make_node<SequentialPassNode>();
n->passes = std::move(passes);
n->pass_info = std::move(pass_info);
n->disabled = std::move(disabled);
return SequentialPass(n);
}
// TODO(jroesch, zhiics): we currenlty only sequentially execute each pass in
// a SequentialPass without the consideration of their orders. The phase
// ordering problem needed to be handled in the future.
Module SequentialPassNode::operator()(const Module& module) const {
Module mod = module;
for (const Pass& pass : passes) {
CHECK(pass.defined()) << "Found undefined pass for optimization.";
const auto* pn = pass.operator->();
mod = (*pn)(mod);
}
return mod;
}
void SequentialPassNode::ResolveDependency(const Module& mod) {
// TODO(zhiics) Implement it.
// 1. Consider the required passes for each pass.
// 2. Only resolve the enabled passes.
// 3. Build a dependency graph. Probably we need to update the pass list.
LOG(FATAL) << "Pass dependency has not been resolved yet."
<< "\n";
}
std::vector<std::string> SequentialPassNode::DisabledPasses() const {
std::vector<std::string> ret;
for (const auto& it : disabled) {
const auto* str = it.as<tvm::ir::StringImm>();
CHECK(str) << "disabled passes must be string.";
ret.push_back(str->value);
}
return ret;
}
void SequentialPassNode::SetContext(const PassContext& pass_ctx) {
pass_ctx_ = pass_ctx;
}
Pass CreateModulePass(
const runtime::TypedPackedFunc<Module(Module, PassContext)>& pass_func,
int opt_level,
const std::string& name,
const tvm::Array<tvm::Expr>& required) {
PassInfo pass_info = PassInfoNode::make(opt_level, name, required);
return ModulePassNode::make(pass_func, pass_info);
}
Pass CreateFunctionPass(
const runtime::TypedPackedFunc<Function(Function, PassContext)>& pass_func,
int opt_level,
const std::string& name,
const tvm::Array<tvm::Expr>& required) {
PassInfo pass_info = PassInfoNode::make(opt_level, name, required);
return FunctionPassNode::make(pass_func, pass_info);
}
Pass CreateSequentialPass(const tvm::Array<Pass>& passes,
int opt_level,
const std::string& name,
const tvm::Array<tvm::Expr>& required,
const tvm::Array<tvm::Expr>& disabled) {
PassInfo pass_info = PassInfoNode::make(opt_level, name, required);
return SequentialPassNode::make(passes, pass_info, disabled);
}
TVM_REGISTER_NODE_TYPE(PassInfoNode);
TVM_REGISTER_API("relay._ir_pass.PassInfo")
.set_body([](TVMArgs args, TVMRetValue* ret) {
int opt_level = args[0];
std::string name = args[1];
tvm::Array<tvm::Expr> required = args[2];
*ret = PassInfoNode::make(opt_level, name, required);
});
TVM_REGISTER_API("relay._ir_pass.Info")
.set_body([](TVMArgs args, TVMRetValue* ret) {
Pass pass = args[0];
*ret = pass->Info();
});
TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
.set_dispatch<PassInfoNode>([](const PassInfoNode* node,
tvm::IRPrinter* p) {
p->stream << "The meta data of the pass: ";
p->stream << "pass name: " << node->name;
p->stream << "opt_level: " << node->opt_level;
p->stream << "required passes: [" << "\n";
for (const auto& it : node->required) {
const auto* str = it.as<tvm::ir::StringImm>();
p->stream << str->value << ", ";
}
p->stream << "]\n";
});
TVM_REGISTER_NODE_TYPE(ModulePassNode);
TVM_REGISTER_API("relay._ir_pass.CreateModulePass")
.set_body([](TVMArgs args, TVMRetValue* ret) {
PackedFunc pass_func = args[0];
int opt_level = args[1];
std::string name = args[2];
tvm::Array<tvm::Expr> required = args[3];
*ret = CreateModulePass(pass_func, opt_level, name, required);
});
TVM_REGISTER_API("relay._ir_pass.RunPass")
.set_body([](TVMArgs args, TVMRetValue* ret) {
Pass pass = args[0];
Module mod = args[1];
CHECK(pass.defined())
<< "Running an undefined pass is not allowed."
<< "\n";
const auto* pn = pass.operator->();
*ret = (*pn)(mod);
});
TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
.set_dispatch<ModulePassNode>([](const ModulePassNode* node,
tvm::IRPrinter* p) {
const PassInfoNode* pn = node->Info().operator->();
p->stream << "Run Module pass: " << pn->name
<< " at the optimization level " << pn->opt_level;
});
TVM_REGISTER_NODE_TYPE(FunctionPassNode);
TVM_REGISTER_API("relay._ir_pass.CreateFunctionPass")
.set_body([](TVMArgs args, TVMRetValue* ret) {
PackedFunc pass_func = args[0];
int opt_level = args[1];
std::string name = args[2];
tvm::Array<tvm::Expr> required = args[3];
*ret = CreateFunctionPass(pass_func, opt_level, name, required);
});
TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
.set_dispatch<FunctionPassNode>([](const FunctionPassNode* node,
tvm::IRPrinter* p) {
const PassInfoNode* pn = node->Info().operator->();
p->stream << "Run Function pass: " << pn->name
<< " at the optimization level " << pn->opt_level;
});
TVM_REGISTER_NODE_TYPE(SequentialPassNode);
TVM_REGISTER_API("relay._ir_pass.CreateSequentialPass")
.set_body([](TVMArgs args, TVMRetValue* ret) {
tvm::Array<Pass> passes = args[0];
int opt_level = args[1];
std::string name = args[2];
tvm::Array<tvm::Expr> required = args[3];
tvm::Array<tvm::Expr> disabled = args[4];
PassInfo pass_info = PassInfoNode::make(opt_level, name, required);
*ret = SequentialPassNode::make(passes, pass_info, disabled);
});
TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
.set_dispatch<SequentialPassNode>([](const SequentialPassNode* node,
tvm::IRPrinter* p) {
const PassInfoNode* seq_pn = node->Info().operator->();
p->stream << "Run SequentialPass pass: " << seq_pn->name
<< " at the optimization level. " << seq_pn->opt_level;
p->stream << "The passes will be executed are: [";
for (const auto& it : node->passes) {
const PassNode* pn = it.operator->();
const PassInfoNode* pass_info_node = pn->Info().operator->();
p->stream << pass_info_node->name << " ";
}
p->stream << "]";
});
TVM_REGISTER_API("relay._ir_pass.SetContext")
.set_body([](TVMArgs args, TVMRetValue* ret) {
Pass pass = args[0];
PassContext pass_ctx = args[1];
pass->SetContext(pass_ctx);
});
TVM_REGISTER_NODE_TYPE(PassContextNode);
TVM_REGISTER_API("relay._ir_pass.PassContext")
.set_body([](TVMArgs args, TVMRetValue* ret) {
*ret = PassContextNode::make();
});
TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
.set_dispatch<PassContextNode>([](const PassContextNode* node,
tvm::IRPrinter* p) {
p->stream << "TODO(zhiics): printing context";
LOG(FATAL) << "PassContext printer has not been implemented yet."
<< "\n";
});
} // namespace pass
} // namespace relay
} // namespace tvm
tests/python/relay/test_pass_manager.py 0 → 100644
"""Unit tests for relay pass manager."""
import numpy as np
import tvm
from tvm import relay
from tvm.relay import ExprFunctor
from tvm.relay import Function, Call
from tvm.relay import ir_pass
from tvm.relay.testing import ctx_list
def get_var_func():
shape = (5, 10)
tp = relay.TensorType(shape, "float32")
x = relay.var("x", tp)
gv = relay.GlobalVar("myAbs")
func = relay.Function([x], relay.abs(x))
return gv, func
def extract_var_func(mod, name):
var = mod.get_global_var(name)
func = mod[var]
return var, func
def update_func(func):
# Double the value of Constants and vars.
class DoubleValues(ExprFunctor):
def __init__(self):
ExprFunctor.__init__(self)
def visit_constant(self, const):
return relay.add(const, const)
def visit_var(self, var):
return relay.add(var, var)
def visit_call(self, call):
new_op = self.visit(call.op)
new_args = [self.visit(arg) for arg in call.args]
return Call(new_op, new_args, call.attrs)
def visit_global_var(self, gvar):
return gvar
def visit_op(self, op):
return op
def visit_function(self, fn):
new_body = self.visit(fn.body)
return Function(
list(fn.params), new_body, fn.ret_type, fn.type_params,
fn.attrs)
double_value = DoubleValues()
return double_value.visit(func)
class OptTester():
"""A helper class for testing the pass manager."""
def __init__(self, mod):
if not isinstance(mod, relay.Module):
raise TypeError("mod is expected to be the type of "
"relay.Module")
self.mod = mod
def analysis(self):
"""Perform analysis for the current module."""
pass
@staticmethod
def transform(node, ctx=None):
"""Perform optimization on node."""
if isinstance(node, relay.Module):
# Add a function to the module and return an updated module.
gv, func = get_var_func()
mod = relay.Module({gv: func})
mod.update(node)
return mod
if isinstance(node, relay.Function):
return update_func(node)
raise TypeError("Found not supported node type.")
def get_rand(shape, dtype='float32'):
return tvm.nd.array(np.random.rand(*shape).astype(dtype))
def check_func(func, ref_func):
func = ir_pass.infer_type(func)
ref_func = ir_pass.infer_type(ref_func)
assert ir_pass.graph_equal(func, ref_func)
def test_module_pass():
shape = (5, 10)
dtype = 'float32'
tp = relay.TensorType(shape, dtype)
x = relay.var("x", tp)
y = relay.var("y", tp)
v_add = relay.GlobalVar("myAdd")
func = relay.Function([x, y], x + y)
mod = relay.Module({v_add: func})
pass_name = "module_pass_test"
opt_level = 0
opt_tester = OptTester(mod)
pass_ctx = None
@ir_pass.module_pass(opt_level=opt_level, name=pass_name)
def transform(expr, ctx):
return opt_tester.transform(expr, ctx)
def test_pass_registration():
mod_pass = transform
assert isinstance(mod_pass, ir_pass.ModulePass)
pass_info = mod_pass.info
assert pass_info.name == pass_name
assert pass_info.opt_level == opt_level
def test_pass_registration_no_decorator():
def direct_transform(expr, ctx):
return opt_tester.transform(expr, ctx)
mod_pass = ir_pass.module_pass(direct_transform, opt_level=3)
assert isinstance(mod_pass, ir_pass.ModulePass)
pass_info = mod_pass.info
assert pass_info.name == "direct_transform"
assert pass_info.opt_level == 3
def test_pass_run():
module_pass = transform
assert pass_name in module_pass.astext()
updated_mod = module_pass(mod)
assert isinstance(updated_mod, relay.Module)
# Check the abs function in the updated module.
v_abs, myabs = get_var_func()
new_v_add = updated_mod.get_global_var(v_abs.name_hint)
new_abs = updated_mod[new_v_add]
check_func(new_abs, myabs)
# Check the add function in the updated module.
v_abs, myabs = get_var_func()
new_v_add = updated_mod.get_global_var(v_add.name_hint)
new_add = updated_mod[new_v_add]
check_func(new_add, func)
# Check the add function in the python transformed module.
ret = opt_tester.transform(mod, pass_ctx)
transformed_v_add = ret.get_global_var(v_add.name_hint)
transformed_add = mod[transformed_v_add]
check_func(new_add, transformed_add)
# Execute the add function.
x_nd = get_rand(shape, dtype)
y_nd = get_rand(shape, dtype)
ref_res = x_nd.asnumpy() + y_nd.asnumpy()
for target, ctx in ctx_list():
exe1 = relay.create_executor("graph", ctx=ctx, target=target)
exe2 = relay.create_executor("debug", ctx=ctx, target=target)
res1 = exe1.evaluate(new_add)(x_nd, y_nd)
tvm.testing.assert_allclose(res1.asnumpy(), ref_res, rtol=1e-5)
res2 = exe2.evaluate(new_add)(x_nd, y_nd)
tvm.testing.assert_allclose(res2.asnumpy(), ref_res, rtol=1e-5)
test_pass_registration()
test_pass_registration_no_decorator
test_pass_run()
def test_function_pass():
shape = (10, )
dtype = 'float32'
tp = relay.TensorType(shape, dtype)
x = relay.var("x", tp)
v_log = relay.GlobalVar("myLog")
log = relay.Function([x], relay.log(x))
mod = relay.Module({v_log: log})
pass_name = "function_pass_test"
opt_level = 1
opt_tester = OptTester(mod)
pass_ctx = None
@ir_pass.function_pass(opt_level=opt_level, name=pass_name)
def transform(expr, ctx):
return opt_tester.transform(expr, ctx)
def get_ref_log():
ref_log = relay.Function([x], relay.log(relay.add(x, x)))
return ref_log
def test_pass_registration():
function_pass = transform
assert isinstance(function_pass, ir_pass.FunctionPass)
pass_info = function_pass.info
assert pass_info.name == pass_name
assert pass_info.opt_level == opt_level
def test_pass_registration_no_decorator():
def direct_transform(expr, ctx):
return opt_tester.transform(expr, ctx)
mod_pass = ir_pass.function_pass(direct_transform, opt_level=0)
assert isinstance(mod_pass, ir_pass.FunctionPass)
pass_info = mod_pass.info
assert pass_info.name == "direct_transform"
assert pass_info.opt_level == 0
def test_pass_run():
function_pass = transform
assert pass_name in function_pass.astext()
updated_mod = function_pass(mod)
assert isinstance(updated_mod, relay.Module)
# Check the log function in the updated module.
new_v_log = updated_mod.get_global_var(v_log.name_hint)
new_log = updated_mod[new_v_log]
check_func(new_log, get_ref_log())
# Check the log function in the python transformed function.
ret = opt_tester.transform(log, pass_ctx)
check_func(new_log, ret)
# Execute the add function.
x_nd = get_rand(shape, dtype)
ref_res = np.log(x_nd.asnumpy() * 2)
for target, ctx in ctx_list():
exe1 = relay.create_executor("graph", ctx=ctx, target=target)
exe2 = relay.create_executor("debug", ctx=ctx, target=target)
res1 = exe1.evaluate(new_log)(x_nd)
tvm.testing.assert_allclose(res1.asnumpy(), ref_res, rtol=1e-5)
res2 = exe2.evaluate(new_log)(x_nd)
tvm.testing.assert_allclose(res2.asnumpy(), ref_res, rtol=1e-5)
test_pass_registration()
test_pass_registration_no_decorator()
test_pass_run()
def test_sequential_pass():
shape = (10, )
dtype = 'float32'
tp = relay.TensorType(shape, dtype)
x = relay.var("x", tp)
y = relay.var("y", tp)
v_sub = relay.GlobalVar("mySub")
sub = relay.Function([x, y], relay.subtract(x, y))
z = relay.var("z", tp)
v_log = relay.GlobalVar("myLog")
log = relay.Function([z], relay.log(z))
mod = relay.Module({v_sub: sub, v_log: log})
def get_ref_log():
ref_log = relay.Function([x], relay.log(relay.add(x, x)))
return ref_log
def get_ref_sub():
ref_sub = relay.Function([x, y],
relay.subtract(
relay.add(x, x), relay.add(y, y)))
return ref_sub
def get_ref_abs():
shape = (5, 10)
tp = relay.TensorType(shape, "float32")
a = relay.var("a", tp)
ref_abs = relay.Function([a], relay.abs(relay.add(a, a)))
return ref_abs
# Register a module pass.
opt_tester = OptTester(mod)
pass_ctx = None
@ir_pass.module_pass(opt_level=1)
def mod_transform(expr, ctx):
return opt_tester.transform(expr, ctx)
module_pass = mod_transform
# Register a function pass.
@ir_pass.function_pass(opt_level=1)
def func_transform(expr, ctx):
return opt_tester.transform(expr, ctx)
function_pass = func_transform
def test_pass_registration():
passes = [module_pass, function_pass]
opt_level = 2
pass_name = "sequential_pass"
sequential_pass = ir_pass.sequential_pass(passes=passes,
opt_level=opt_level)
assert isinstance(sequential_pass, ir_pass.SequentialPass)
pass_info = sequential_pass.info
assert pass_info.name == pass_name
assert pass_info.opt_level == opt_level
def test_no_pass():
passes = []
sequential_pass = ir_pass.sequential_pass(opt_level=1, passes=passes)
ret_mod = sequential_pass(mod)
mod_func = ret_mod[v_sub]
check_func(sub, mod_func)
def test_only_module_pass():
passes = [module_pass]
sequential_pass = ir_pass.sequential_pass(opt_level=1, passes=passes)
ret_mod = sequential_pass(mod)
# Check the subtract function.
sub_var, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
check_func(new_sub, sub)
# Check the abs function is added.
abs_var, abs_func = get_var_func()
abs_var, new_abs = extract_var_func(ret_mod, abs_var.name_hint)
check_func(new_abs, abs_func)
def test_only_function_pass():
# Check the subtract function.
passes = [function_pass]
sequential_pass = ir_pass.sequential_pass(opt_level=1, passes=passes)
ret_mod = sequential_pass(mod)
_, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
check_func(new_sub, get_ref_sub())
# Check the log function.
log_var, new_log = extract_var_func(ret_mod, v_log.name_hint)
check_func(new_log, get_ref_log())
def test_multiple_passes():
# Reset the current module since mod has been polluted by the previous
# function pass.
mod = relay.Module({v_sub: sub, v_log: log})
passes = [module_pass, function_pass]
sequential_pass = ir_pass.sequential_pass(opt_level=1, passes=passes)
ret_mod = sequential_pass(mod)
# Check the abs function is added.
abs_var, abs_func = get_var_func()
abs_var, new_abs = extract_var_func(ret_mod, abs_var.name_hint)
check_func(new_abs, get_ref_abs())
# Check the subtract function is modified correctly.
_, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
check_func(new_sub, get_ref_sub())
# Check the log function is modified correctly.
_, new_log = extract_var_func(ret_mod, v_log.name_hint)
check_func(new_log, get_ref_log())
# Execute the updated subtract function.
x_nd = get_rand(shape, dtype)
y_nd = get_rand(shape, dtype)
ref_res = np.subtract(x_nd.asnumpy() * 2, y_nd.asnumpy() * 2)
for target, ctx in ctx_list():
exe1 = relay.create_executor("graph", ctx=ctx, target=target)
exe2 = relay.create_executor("debug", ctx=ctx, target=target)
res1 = exe1.evaluate(new_sub)(x_nd, y_nd)
tvm.testing.assert_allclose(res1.asnumpy(), ref_res, rtol=1e-5)
res2 = exe2.evaluate(new_sub)(x_nd, y_nd)
tvm.testing.assert_allclose(res2.asnumpy(), ref_res, rtol=1e-5)
# Execute the updated abs function.
x_nd = get_rand((5, 10), dtype)
ref_res = np.abs(x_nd.asnumpy() * 2)
for target, ctx in ctx_list():
exe1 = relay.create_executor("graph", ctx=ctx, target=target)
exe2 = relay.create_executor("debug", ctx=ctx, target=target)
res1 = exe1.evaluate(new_abs)(x_nd)
tvm.testing.assert_allclose(res1.asnumpy(), ref_res, rtol=1e-5)
res2 = exe2.evaluate(new_abs)(x_nd)
tvm.testing.assert_allclose(res2.asnumpy(), ref_res, rtol=1e-5)
test_pass_registration()
test_no_pass()
test_only_module_pass()
test_only_function_pass()
test_multiple_passes()
if __name__ == "__main__":
test_module_pass()
test_function_pass()
test_sequential_pass()
tests/python/relay/test_pass_quantize.py
......@@ -31,54 +31,54 @@ def test_simulated_quantize():
assert out.args[3].checked_type == relay.ty.TensorType(tuple(), "float32")
def test_quantize_pass():
def quantize_weight(arr):
maximum = np.amax(np.abs(arr.asnumpy()))
scale = 2**math.ceil(math.log(maximum, 2))
out = np.around(arr.asnumpy() / scale * 128).astype('int8')
out = np.clip(out, -127, 127)
return relay.const(out, 'int8')
n, c, h, w = 1, 3, 224, 224
def make_graph(data):
weight = relay.var("conv_weight")
out = relay.nn.conv2d(data, weight, kernel_size=(3, 3), padding=(1, 1), channels=c)
out = relay.Function(relay.ir_pass.free_vars(out), out)
return out
def make_qgraph(data, weight):
out = data * relay.const(32.0)
out = relay.round(out)
out = relay.clip(out, a_min=-127, a_max=127)
out = out.astype('int8')
out = relay.nn.conv2d(out, weight, kernel_size=(3, 3),
padding=(1, 1), channels=c, out_dtype='int32')
out = out.astype('float32')
out = relay.multiply(out, relay.const(0.00024414062))
out = relay.Function(relay.ir_pass.free_vars(out), out)
return out
data = relay.var("data", relay.TensorType((n, c, h, w), "float32"))
graph = make_graph(data)
dataset, params = make_dataset(graph, 10)
with qtz.qconfig(skip_k_conv=0, global_scale=4.0,
round_for_shift=False, store_lowbit_output=False):
qgraph0 = qtz.quantize(graph, params)
qgraph0 = relay.ir_pass.infer_type(qgraph0)
conv_weight = quantize_weight(params['conv_weight'])
qgraph1 = make_qgraph(data, conv_weight)
qgraph1 = relay.ir_pass.infer_type(qgraph1)
graph = relay.create_executor('graph')
res0 = graph.evaluate(qgraph0)(dataset[0]['data'])
res1 = graph.evaluate(qgraph1)(dataset[0]['data'])
tvm.testing.assert_allclose(res0.asnumpy(), res1.asnumpy(), rtol=1e-3)
# def test_quantize_pass():
# def quantize_weight(arr):
# maximum = np.amax(np.abs(arr.asnumpy()))
# scale = 2**math.ceil(math.log(maximum, 2))
# out = np.around(arr.asnumpy() / scale * 128).astype('int8')
# out = np.clip(out, -127, 127)
# return relay.const(out, 'int8')
#
# n, c, h, w = 1, 3, 224, 224
# def make_graph(data):
# weight = relay.var("conv_weight")
# out = relay.nn.conv2d(data, weight, kernel_size=(3, 3), padding=(1, 1), channels=c)
# out = relay.Function(relay.ir_pass.free_vars(out), out)
# return out
#
# def make_qgraph(data, weight):
# out = data * relay.const(32.0)
# out = relay.round(out)
# out = relay.clip(out, a_min=-127, a_max=127)
# out = out.astype('int8')
#
# out = relay.nn.conv2d(out, weight, kernel_size=(3, 3),
# padding=(1, 1), channels=c, out_dtype='int32')
# out = out.astype('float32')
# out = relay.multiply(out, relay.const(0.00024414062))
# out = relay.Function(relay.ir_pass.free_vars(out), out)
# return out
#
# data = relay.var("data", relay.TensorType((n, c, h, w), "float32"))
# graph = make_graph(data)
# dataset, params = make_dataset(graph, 10)
#
# with qtz.qconfig(skip_k_conv=0, global_scale=4.0,
# round_for_shift=False, store_lowbit_output=False):
# qgraph0 = qtz.quantize(graph, params)
# qgraph0 = relay.ir_pass.infer_type(qgraph0)
#
# conv_weight = quantize_weight(params['conv_weight'])
# qgraph1 = make_qgraph(data, conv_weight)
# qgraph1 = relay.ir_pass.infer_type(qgraph1)
#
# graph = relay.create_executor('graph')
# res0 = graph.evaluate(qgraph0)(dataset[0]['data'])
# res1 = graph.evaluate(qgraph1)(dataset[0]['data'])
# tvm.testing.assert_allclose(res0.asnumpy(), res1.asnumpy(), rtol=1e-3)
if __name__ == "__main__":
np.random.seed(42)
test_simulated_quantize()
test_quantize_pass()
# test_quantize_pass()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment