Relay Operator Strategy¶
In order to lower Relay operators to the implementations defined in TOPI library, a compute and schedule function need to be registered to each Relay operator. However, compute and schedule functions are usually specialized for each target, and further, even for the same target, we may have multiple algorithms and implementations available. To deal with the complexity, we introduce operator strategy to allow developers to define a flexible lowering strategy for each operator and target.
Operator Strategy Design¶
The basic element in operator strategy is an OpImplementation
. It includes
the a pair of compute and schedule function, the name of the implementation,
and a priority level (the use of priority level is explained in
Select Implementation from Op Strategy).
The OpStrategy
includes a list of OpSpecialization
. Each OpSpecialization
contains a list of OpImplementation
associated with a SpecializedCondition
(see definition in include/tvm/te/schedule.h
). The SpecializedCondition
can be null, indicating the implementations are generally applicable;
otherwise, the implementations are only considered when the specialized
condition is satisfied. SpecializedCondition
consists of a list
of clauses defined in Tensor Expression in conjunctive normal form (CNF) and
only supports conditions on tensor shapes.
Last, a strategy function, or FTVMStrategy
, determines which pair(s) of
compute and schedule functions should be used given a workload, and needs to be
registered to each Relay operator. FTVMStrategy
is a generic function (see
include/tvm/target/generic_func.h
), that can be overwritten for each
target. The function signature is
OpStrategy(const Attrs& attrs, const Array<Tensor>& inputs, const Type& out_type, const Target& target)
that the function returns an OpStrategy
given the op attributes, input
tensors, output types, and target to compile to.
Write A Strategy Function¶
We recommend developers to write strategy function in Python as
most TOPI compute and schedule functions are written in Python.
In python, we provide OpStrategy
class in pyton/tvm/relay/op/op.py
.
It only has one API, which is to add an implementation to the strategy:
def add_implementation(self, compute, schedule, name="default", plevel=10)
We now take topk
as an example to explain how to write the
FTVMStrategy
function:
# add to python/tvm/relay/op/strategy/generic.py
@override_native_generic_func("topk_strategy")
def topk_strategy(attrs, inputs, out_type, target):
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_topk(topi.topk),
wrap_topi_schedule(topi.generic.schedule_topk),
name="topk.generic")
return strategy
# add to each target file in python/tvm/relay/op/strategy, e.g., x86.py, cuda.py, etc.
@topk_strategy.register(["cuda", "gpu"])
def topk_strategy_cuda(attrs, inputs, out_type, target):
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_my_new_op(topi.cuda.topk),
wrap_topi_schedule(topi.cuda.schedule_topk),
name="topk.cuda")
return strategy
In this example, we use topi.cuda.topk
and topi.cuda.schedule_topk
as the compute and schedule function for CUDA or GPU target, while use TOPI
generic compute and schedule for the rest of targets.
Note that we use two wrapper functions that wrap the topi
compute and schedule to conform with the required function signature (
see FTVMCompute
and FTVMSchedule
in include/tvm/relay/op_attr_types.h
).
Usually we need to write a customized compute wrapper function for each operator
to get different fields from op attributes.
The example above shows a very basic strategy function that only adds one implementation in the strategy. But for many complicated operators, we may need to add multiple implementations that use different algorithms. For example, we can use both direct and winograd algorithm to compute a conv2d op. In order to achieve this, we can write the strategy function as follows:
strategy.add_implementation(
wrap_compute_conv2d(topi.cuda.conv2d_nchw),
wrap_topi_schedule(topi.cuda.schedule_conv2d_nchw),
name="conv2d_nchw.cuda",
plevel=10)
if winograd_condition:
strategy.add_implementation(
wrap_compute_conv2d(topi.cuda.conv2d_nchw_winograd),
wrap_topi_schedule(topi.cuda.schedule_conv2d_nchw_winograd),
name="conv2d_nchw_winograd.cuda",
plevel=15)
In this example, we add two implementations to the conv2d strategy where
winograd algorithm is only added when winograd_condition
is true.
The implementation "conv2d_nchw_winograd.cuda"
will be used to compile
conv2d when winograd_condition
is true as it has higher
priority level (this could be changed if certain implementation is an AutoTVM
template. See Select Implementation from Op Strategy for more
details). Otherwise, "conv2d_nchw.cuda"
is used.
We can extend the example above to third party library implementation. For example, we can add the implementation that invokes kernel in the cblas library when cblas is included in the target.
if "cblas" in target.libs:
strategy.add_implementation(
wrap_compute_dense(topi.x86.dense_cblas),
wrap_topi_schedule(topi.x86.schedule_dense_cblas),
name="dense_cblas.x86",
plevel=15)
Further, we can add implementation specialized for a certain range of shapes.
The code below shows an example of dense strategy that adds an implementation
that is specialized for m
greater than 16. The main difference between
hardcode python condition like examples above and specialized condition is that
it allows TVM to generate multiple kernels when the input tensors have symbolic
shapes. The compile engine will generate a dispatch function that invokes the
specialized kernel when the corresponding condition is met; otherwise,
invoke the kernel that has no associated specialized condition (dense_common
in this example). This part is still work in progress. More details will be
provided after it is done.
def dense_strategy(attrs, inputs, out_type, target):
m = inputs[0].shape[0]
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_dense(dense_compute1),
wrap_topi_schedule(dense_schedule1),
name="dense_common")
with tvm.te.SpecializedCondition(m > 16):
strategy.add_implementation(
wrap_compute_dense(dense_compute2),
wrap_topi_schedule(dense_schedule2),
name="dense_for_large_m",
plevel=15)
return strategy
Register Strategy Function to An Operator¶
After we define the strategy function for an operator, we can now register the strategy function to this operator with
register_strategy("topk", strategy.topk_strategy)
However, it takes much effort to write a strategy function for an operator. Therefore, we provide two other methods for simpler operators.
First, for operators that have injective, broadcast, or reduction pattern, we
can call register_injective_schedule
, register_broadcast_schedule
, and
register_reduce_schedule
repsectively. The schedule function for these
patterns are already registered by each target and can be applied to these
operators. We assume the compute function should be the same across all targets,
and FTVMCompute
needs to be registered to the op before invoking register
schedule.
register_broadcast_schedule("add")
Second, for operators that doesn’t have these common patterns mentioned before,
but also have the same compute function for all targets, we can use
register_schedule
API. It is easier to write FTVMSchedule
function
as we only need to provide which schedule function to use. The following
code snippet shows FTVMSchedule
function for pooling.
# add to python/tvm/relay/op/strategy/generic.py
@generic_func
def schedule_pool(attrs, outs, target):
with target:
return topi.generic.schedule_pool(outs, attrs.layout)
# add to each target file in python/tvm/relay/op/strategy, e.g., x86.py, cuda.py, etc.
@schedule_pool.register("cpu")
def schedule_pool_cpu(attrs, outs, target):
...
After we created the FTVMSchedule
for an operator, we can
register the strategy using register_schedule
:
register_schedule("nn.max_pool2d", strategy.schedule_pool)
Register Strategies for A New Target¶
There are two ways to register strategies for a new target. The more
straightforward one is adding a new target file in the directory
python/tvm/relay/op/strategy
. You only need to customize the strategy for
ops that have been implemented for this new target and reuse the generic
strategies for the rest.
Alternatively, you can also register the strategy for the new target outside the
TVM python library. The following code snippet shows an example how to do
so. You can find more examples in vta/python/vta/top/op.py
.
@relay.op.strategy.conv2d_strategy.register("mytarget")
def conv2d_strategy_mytarget(attrs, inputs, out_type, target):
...
Select Implementation from Op Strategy¶
During the compilation, Relay compile engine needs to determine which implementation to use for an operator when there are multiple. The selection policy works as follows.
When the input tensors to an operator or a fused op all have constant shapes, the compile engine first finds the best implementation based on AutoTVM tuning logs. If there is no implementation that is an AutoTVM template or all AutoTVM templates have fallback configs, the implementation with highest priority level will then be chosen. Implementations with same priority level in this case leads to an undefined behavior, and any of them might be selected.
The selection policy for ops with symbolic input shapes is still work in progress. Currently, if any input tensor has a symbolic shape, only the implementation with highest priority level will be used for this operator. This will be updated after the implementation finishes.
For debug purpose, you can add the following lines before you compile the Relay model to learn which implementation is used for each operator.
logging.getLogger("te_compiler").setLevel(logging.INFO)
logging.getLogger("te_compiler").addHandler(logging.StreamHandler(sys.stdout))