[RFC][Graph Tuner] Graph level auto-tuning

Motivation

Currently we can tune operator with handcrafted schedules or AutoTVM, which can give us descent kernel performance. However, usually these kernel templates involve transform operator data layout to other formats, such as conv2d for Intel and ARM CPU. In this case, a lot of layout transformation can be introduced into graph. Graph tuner considers both fast kernel schedules and extra layout transformations, generating descent end to end performance.

Design

There are two steps to achieve graph level optimal schedules. First, get a set of schedule candidates for each workload in the graph. This step can be finished by AutoTVM Second, feed schedule candidates into graph tuner and run graph level tuning.

In the second step, graph tuner will benchmark all possible layout transformations in the graph, given a set of schedule candidates, and then combine schedule and layout transformation execution time together to find the optimal schedule combination.

The current solution for this optimization problem is to model it as a Markov Decision Process and use dynamic programming to solve it. This will give us global optimal solution. However, the time/memory complexity for DP is prohibitively expensive for some networks, such as SSD. In this case, we need to use approximation methods. For now graph tuner provides a graph coloring algorithm(PBQP) in this scenario.

API

We provide a base class and built-in subclass:

class BaseGraphTuner(object):
    """Class to search schedules considering both kernel execution time and
    layout transformation time.

    Before creating a Graph Executor instance, schedule candidates for all kernels in
    graph should be provided through tensor searching. 
    """
    ......

class DPTuner(BaseGraphTuner):
    """Tuner which uses dynamic programming to solve MDP problem.

    Note: currently dynamic programming is used to solve this MDP problem. However,
    this problem is intrinsically non-polynomial. DP can't apply for more complicated
    models, such as networks with many element-wise sum operators. In this case, a simple
    greedy method greedy_schedule can be used to generate suboptimal schedules.
    """
    ......

class PBQPTuner(BaseGraphTuner):
    """Tuner using graph coloring algorithm to solve DP intractable graph.
    """

Implementation details:
One key part of graph tuner is to generate all possible layout transformations given a set of workloads and schedules. To hide any operator related information from graph tuner, we can add a new generic function bind with topi operator, which accepts workload and cfg, and returns i/o shapes/layouts. An example for conv2d:

@tvm.target.generic_func
def conv2d_infer_layout(workload, cfg):
    """Infer input/output shapes and layouts given a workload and a config.
    Returns two lists of tuple in the format of (input_shape, input_layout) and (output_shape, output_layout).
    """

Note that although graph tuner only supports target op with single input and output(conv2d, conv2d_transpose, dense, etc), we make this api generic enough to support operators with multiple io.

With this function, in graph tuner we only need a dictionary mapping topi function name to corresponding infer layout function, similar to what autotvm task extraction function does. After shape and layout info is extracted, we can use a generic graph traversal to fetch all possible layout transformations.

Performance Benchmark

Intel Xeon CPU(AWS c5.9xlarge, 18 physical cores)

Framework/Mode	resnet50_v1	vgg19_bn	inceptionv3	densenet201	SSD_resnet50_512
MXNet MKLDNN 1.2.1	12.95 ms	25.35 ms	15.08 ms	34.84 ms	67.06 ms
TVM with default schedules	8.07 ms	24.78 ms	13.88 ms	18.19 ms	44. 56 ms
TVM with history best schedules	7.54 ms	24.05 ms	14.39 ms	18.35 ms	40. 29 ms
TVM with graph tuner	5.73 ms	21.98 ms	10.68 ms	13.97 ms	29.05 ms

More benchmark data to be added.

PR: #1586

@eqy

@eqy @merrymercy please comment

Thanks for opening this RFC, graph level optimization is an important step in pushing performance in cases were we have to make inter-layer decisions (such as data layout in this case).

Looking at the code, I think that this is a good opportunity to either refine or more clearly lay out what we should have in our autotvm/general tuning API. Having a good API that is used healthily ensures that we keep things maintainable and gives us confidence that things are extensible.

To that end, it may be a good time to discuss how we should organize graph-level and layer-level tuning. We currently already have cost-model tuners that operate on search space; is it natural to extend the notion of a search space to cover possible graphs as well?

Taking a look at some of the odds and ends of the PR, here's some comments (in no good order):

We should be careful when introducing a new "Executor" and what the right organization is here. In the past, autotvm has seen a lot of "Executor" designs come and go, for good reason (too clunky, too many options, duplicated code, etc.).
We should unify the definition of a "search space" so that we do not end up having a search space definition style for every hardware backend/template, but rather one that is shared across templates to simplify how tuners can be designed
Can we reuse any of the workload definition/extraction code introduced by autotvm?
Using lists of RPC hosts/ports is ad-hoc and unnecessary as we have RPC Tracker functionality
Can we reuse mic autotvm code e.g., finding all of the factors of an integer?
We should avoid introducing a custom, specialized JSON/schedule format for only one backend (autotvm has already defined a generic one).
Is the current search space as large as it should be to provide coverage to many x86 CPU types other than C5/AVX-512 e.g., (C4/AVX-2 only, AMD Ryzen ?) We should not shy away from defining a large space search space if it means we can unlock better results through tuning.

One main question:
Currently there seem to be no calls into existing autotvm code; can we not define new tasks that the current autotvm can tune to leverage the existing infrastructure?

Currently graph tuner doesn't use any autotvm code, we should definitely reuse the AutoTVM system to tune kernel. Actually graph tuner is designed to be a standalone module and doesn't couple with any specific tensor tuner.
Things might need to be changed to make this happen:

A plain benchmark mode for operator. Graph tuner need to benchmark layout transformation, given input shapes and parameters. We need an interface to run operator given a set of input shapes and parameters. I'd prefer covering this part in AutoTVM(If it's not there yet), since multi-process/RPC module can be reused.
Schedule load method. Load schedule with operator index order when calling alter_op_layout, decl and schedule function. Currently this is only done for x86 cpu.
A uniform format for schedule template. If AutoTVM already defines one, we just need to modify x86 backend.

For the search space question, currently it should be enough for AVX512/AVX2.

It is a great step. For graph level layout planning, current DP solution is nice. For operator level tuning, we can reuse measurement/tuner infrastructure in autotvm.

It seems that we developed our custom tuner systems at a same time, and there are a lot of things to do for the merge. But keeping an unified infrastructure ensures the maintainability.
I can give some guidance on porting the executor to autotvm style.

Abstractions in AutoTVM

Task: Task stores all information for tuning an operator: target, target_host, space and workload. This corresponds to one entry of your Conv2dAVXExecutor.workload_execute
ConfigEntity: It contains all parameters for a template, which is the "schedule" tuple in your code.
Workload: It is same as "workload" in your code. One difference is that autotvm uses vanilla tuple instead of namedtuple. By default, autotvm automatically translates all arguments of TOPI call to a tuple as workload. But for custom layouts, you have to construct the workload by yourself. You can find some examples in arm cpu https://github.com/dmlc/tvm/blob/54a115ef14fb6dabbf6ea8eb9e6dd85846030c72/topi/python/topi/arm_cpu/conv2d.py#L15
MeasureInput, MeasureResult: It stores all information of one measurement (workload, schedule, time cost). This corresponds to [{"schedule": AVXConvCommonFwd(1, 4, 2, True), "time": 0.04}], in your code. AutoTVM can serialize them to json and use them as the log file (see tvm/autotvm/record.py).
DispatchContext: This corresponds to _get_schedule_conv in current x86 topi. In AutoTVM, topi.nn.conv2d and topi.generic.schedule_conv2d will construct a workload and use this workload to query corresponding ConfigEntity from the dispatch context.

Previously, schedule function topi.generic.schedule_conv2d need to reconstruct the workload from dataflow by using _get_workload. Now, in compute function topi.nn.conv2d, we can attach a tuple to the compute op https://github.com/dmlc/tvm/blob/54a115ef14fb6dabbf6ea8eb9e6dd85846030c72/topi/python/topi/arm_cpu/conv2d.py#L148. Then in schedule function, we just fetch the workload by op.attrs['workload']

Steps

port operator topi/x86 into autotvm style.
- use cfg.define_split, cfg.define_knob to define your space
- For alter_op_layout, it is more complicated. you can find some examples in arm cpu https://github.com/dmlc/tvm/blob/54a115ef14fb6dabbf6ea8eb9e6dd85846030c72/topi/python/topi/arm_cpu/conv2d.py#L476.
port graph tuner
- As you pointed out, we need a plain benchmark mode. For operators, you can use GridSearchTuner for exhaustive search. For layout, you can reuse https://github.com/dmlc/tvm/blob/54a115ef14fb6dabbf6ea8eb9e6dd85846030c72/python/tvm/autotvm/measure/measure_methods.py#L179 for measurement. You can create a LayoutTuner, and feed in customized MeasureInput.
- After the above steps, all the input to DP should be in (MeasureInput, MeasureResult) pairs. Then you have all the required measurement result for DP.
- Schedule load method. Current way that uses a global index looks like a hack. You can create a custom DispatchContext for this kind of schedule(ConfigEntity) load.

Discussions

Currently autotvm stores parameters in another git repo https://github.com/uwsaml/tvm-distro. When we create target (e.g. tvm.target.avx()), tvm will download them. I am not sure whether you like this style.

@merrymercy

@merrymercy Thanks for suggestion! I'll make changes accordingly.

@kevinthesun

@kevinthesun do you have previously collected data on the best graph level (data layout) choices for some different C5 instance types (e.g., xlarge, 2xlarge, 4xlarge, 9xlarge) on ResNet-50?

We are planning on doing some experiments with autotvm on EC2 and those would be very valuable for us.

@eqy

@eqy https://github.com/kevinthesun/intel-benchmark This repo contains link to pre-tuned best schedules for several imagenet models. These schedules are searched on c5.9xlarge, but can be directly applied to other types of c5. Convolution schedules are stored in the ascending order of node index.

@merrymercy

@merrymercy I checked arm_cpu conv2d and have several questions for implementation detail:

autotvm.task.extract_from_graph can extract tasks from graph. What if we want to extract conv2d_NCHWc instead of conv2d from graph?
Is it possible to use one DispatchContext for tuning and another DispatchContext for compiling? For tuning we can use normal DispatchContext to generate a schedule from a workload. For compiling we need to generate a schedule from a node index.
measure_batch is an internal function. To create layout tuner, do I need to create a standalone function to mimic the behavior of measure_batch?
I'm not quite sure how to apply cfg template to conv2d_NCHWc, since the input data shape is 5-D and then the ic_bn in cfg needs to be the same as ic block length of input data. Currently autotvm task can benchmark a cfg given a fixed input shape, but here we want to benchmark a cfg given a workload and the actual input shape depends on each schedule in cfg.

see point 4
Exactly, you can create many dispatch contexts. In autotvm, during tuning, we use ApplyConfig to apply the config for tuning; during compilation, we use ApplyHistoryBest. For your case, you need to change the dispatch context used during compilation.
No. You can call create_measure_batch. It will return this function for you. You can see the usage of create_measure_batch here. https://github.com/dmlc/tvm/blob/7cb85d81968cd69576d923852d812590b93cc26d/python/tvm/autotvm/tuner/tuner.py#L87
You can use normal extract_from_graph to get conv2d tasks. Then transform them into conv2d_NCHWc tasks.
The current conv2d task is defined at https://github.com/dmlc/tvm/blob/7cb85d81968cd69576d923852d812590b93cc26d/python/tvm/autotvm/task/nnvm_integration.py#L99-L105
You can define the task for conv2d_NCHWc as follows

        @register("topi_nn_conv2d_NCHWc")
        def _topi_nn_conv2d(*args, **kwargs):
            assert not kwargs, "Do not support kwargs in template function call"
            args = deserialize_args(args)
            A, W = args[:2]

            # get config here
            cfg = autotvm.get_config()   
            cfg.define_knob('tile_c', [1, 2, 4, 8, 16])

            # change shape with the value in config
            VC = cfg['tile_c'].val
            raw_shape = get_const_tuple(A.shape)
            new_shape = (raw_shape[0], raw_shape[1] // VC, raw_shape[2], raw_shape[3], VC)
            args[0] = tvm.placeholder(new_shape, A.dtype)

            C = topi.nn.conv2d_NCHWc(*args, **kwargs)
            s = topi.generic.schedule_conv2d_NCHWc([C])
            return s, [A, W, C]

@merrymercy

@merrymercy I can tune conv2d_NCHWc with autotvm now. I have one issue for logging to file. It returned JSON not serializable error: TypeError: Tensor(shape=[1, 3, 4, 4], op.name=data) is not JSON serializable.

@kevinthesun

@kevinthesun Can you check that the format of MeasureInput and MeasureResult that you use don't contain any non-serializable data structures? e.g., we only use NamedTuples and Lists

TVM Tensor will be serialized by https://github.com/dmlc/tvm/blob/b11f2a0495541cb348ae89093fd233d78eefec6e/python/tvm/autotvm/task/nnvm_integration.py#L17-L23
or
https://github.com/dmlc/tvm/blob/b11f2a0495541cb348ae89093fd233d78eefec6e/python/tvm/autotvm/task/task.py#L185-L192

Can you check whether there is something missing?

Found the issue, I need to call serialize_arg before creating conv2d_NCHWc tasks since I was not using extract_from_graph.

I tried quick solution and is able to reproduce previous collected optimal results now.

@kevinthesun

@kevinthesun @merrymercy in the interest of more recent changes of int8 schedule that might also benefit from AutoTVM x86 port, do you think if it makes sense to break the PR into two part and bring in AutoTVM x86 port in first? If so, please open a separate issue to track it

@tqchen

@tqchen I open an issue to track x86 AutoTVM related tasks.

I fond a related paper on graph layout tuning from CGO 2018 https://arxiv.org/pdf/1710.01079.pdf
I do not believe their benchmark results in the paper. (for example. simple im2col matches or outperforms mkldnn).
One thing we can learn is that they leverage a off-the-shelf PBQP solver to solve the programming problem. Because this problem seems to be a classical programming problem, I think we can also leverage some existing solvers. This can help the current "too large space" case.

@merrymercy

I understand the motivation of this RFC and the paper @merrymercy linked. But I have the following questions regarding this RFC. I appreciate if somebody could answer them.

Isn't it the case that keeping the data layout in NCHWc as much as possible, and only insert layout transform when necessary, is always the best? This is what opt-level = 3 does for x86 backend. Other operators, such as max pooling, can operate directly on NCHWc layout.
The CGO 2018 paper is based on an assumption that there is data layout transform happening between each pair of layers. I think they need to do data layout transform because they have many variant of convolution algorithms each requiring different input layouts (otherwise edge cost = 0 for all edges). But TVM x86 backend has only direct algorithm at the moment. So I don't quite understand why
Graph level auto-tuning in this RFC brings such big improvements over 'TVM with default schedules'. I think comparison between AutoTVM alone vs AutoTVM + Graph level tuning would clarify the benefit of Graph tuner.

One last point: @kevinthesun Is your mxnet built with CUDA off during cmake? Otherwise, elemwise ops are not parallelized with openmp (see here ) and MXNet results would be way worse than it should be.

@masahi

@masahi

I think the point of introducing graph level tuning (e.g., for data layout) is similar to the original motivation with AutoTVM. In the past many of our schedule configurations and currently our choices of data layout are only the result of handcrafted heuristics. If there is a chance that we are leaving some efficiency on the table, then I think it is worthwhile to pursue automated approaches that leave fewer stones unturned. Having flexibility when we eventually get models that break our assumptions about how model architectures and data layout transformations is another bonus.
My understanding is that our algorithms (i.e., the best schedule configurations for each) are very shape dependent. So unless every layer has identical input/output/weight shapes, there is a chance that a different data layout will improve performance. The situation becomes even more complicated when we need to support both AVX-2 and AVX-512 CPUs, where the best choice of layout may be different depending on the vector width of the CPU.

I have some some early experiments with AutoTVM optimizing NCHWc schedules under a few hand-tuned hardcoded layout configurations. The results are OK, but my understanding is that you can only get so far without considering changes to data layout. But to reiterate I think manually defining heuristics for data layout will be brittle and unmaintainable.

If we introduce support for graph level tuning, this introduces benefits beyond data layout tuning; we can begin to consider joint architecture-schedule tuning and other techniques that are starting to become popular.

@eqy

@eqy thanks, then I'm interested to know if there is an instance where the best NCHW schedule beats the best NCHWc schedule. My assumption is that NCHWc layout is always better for direct and winograd algorithm on x86. I care a lot about TVM performance on x86, so I'm very happy if the assumption was wrong. The difference between the best AVX2 schedule and the best AVX-512 schedule would also be interesting.

I did realize that when I think of NCHWc layout, I always have one specific layout, such as NCHW8c, in mind. So to me "data layout transform" always meant NCHW <-> NCHW8c conversion. But there are also NCHW16c, NCHW32c, etc... so there are lot more combination of possible conversion ( NCHW8c <-> NCHW16c, etc). The optimal inner channel dimension can be different for each layer, so it might be better to introduce data layout conversion of different inner channel dimension. The graph level auto-tuning in this RFC automates this decision. Is this understanding correct? @eqy

@masahi

@masahi Your are right. Actually the "default" schedule for x86 conv2d is a simple heuristic method which always choose channel factor to be 16 if possible, to minimize the number of layout transformation needed. Graph tuner is an automated way to make these decisions related to data layouts.

The data layout NCHWc is suitable for graph tuner to do such searching. However, if a certain algorithm requires different input and output data layouts, say NCHW vs NCHWc, then these data layouts can not be eliminated. In this case graph tuner wouldn't help much. Another scenario which limits graph tuner is some parameters need to be fixed to fully utilize hardware resources. This happens for intel graphics, while the output channel factor should be as close to 16 as possible to use virtual threads. I'm also trying to apply autotvm to Intel graphics, and will see how graph tuner work in this case.

@kevinthesun

thanks @kevinthesun, it makes a lot more sense now. I'm looking forward to testing this feature on my network, where all inputs are NCHW8c.

I agree with @merrymercy in that the result of CGO 2018 paper seems fishy, but I did like their problem formulation. As long as we can define a 'node cost' and 'edge cost', we can model this as a standard discrete optimization problem that can be solved by off the shelf solvers. We can incorporate different algorithms each with different layout preference straightforwardly.

In the section 7 of the paper, they stated the following. To me, this sounds a lot like what AutoTVM + Graph tuning will enable to do. Very interesting!

A viable future approach might be to use code generators and auto-tuners to generate the code
and data layouts for given layers and use our approach to combine these code  segments to 
implement an entire DNN.

@merrymercy

@merrymercy @eqy I have an issue dealing with FallbackContext for x86. Currently x86 conv2d will automatically generate default schedules if no pre-tuned schedules are provided. I think this should correspond to FallbackContext? However, even if I didn't specify any dispatch before compilation, when I called autotvm.task.DispatchContext.current inside x86.conv2d.py, it returns ApplyHistoryBest. Does autotvm automatically convert FallbackContext to ApplyHistoryBest somewhere?

Yes, in nnvm.build.compiler, we will load tophub context (which is an empty ApplyHistoryBest in your case)

https://github.com/dmlc/tvm/blob/01ec533e7cfb603f6114a763f03dad4c564f589d/nnvm/python/nnvm/compiler/build_module.py#L244-L246

But this is not a problem. Although current context is ApplyHistoryBest, when you query it, you cannot find the config for your workload. Then it will query its upper context, which is the root fallback context and returns a FallbackConfigEntity. You can use cfg.is_fallback to check whether it is a fallback. Don't use something like isinstance(autotvm.task.DispatchContext.current, autotvm.FallbackContext)

@kevinthesun

@kevinthesun The performance data using auto tuning or not? According comments, seems we don't apply, wish to update the performance data.

@FrozenGene

@FrozenGene The default schedule here for x86 eliminates most layout transformations. It should have similar performance with "apply_history_best". I'll update the data for "apply_history_best".

@FrozenGene

@FrozenGene Data of "apply_history_best" updated.
@yzhliu Updated some implementation details.

Navigation Menu

Repository navigation

[RFC][Graph Tuner] Graph level auto-tuning #1585

Motivation

Design

API

Performance Benchmark

Abstractions in AutoTVM

Steps

Discussions

19 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[RFC][Graph Tuner] Graph level auto-tuning #1585

Description

Motivation

Design

API

Performance Benchmark

Activity

tqchen commented on Aug 14, 2018

eqy commented on Aug 14, 2018

kevinthesun commented on Aug 14, 2018

merrymercy commented on Aug 14, 2018

Abstractions in AutoTVM

Steps

Discussions

kevinthesun commented on Aug 16, 2018

eqy commented on Aug 20, 2018

kevinthesun commented on Aug 20, 2018

kevinthesun commented on Aug 20, 2018

merrymercy commented on Aug 21, 2018

kevinthesun commented on Aug 22, 2018

eqy commented on Aug 22, 2018

merrymercy commented on Aug 22, 2018

kevinthesun commented on Aug 22, 2018

19 remaining items

kevinthesun commented on Sep 5, 2018

tqchen commented on Sep 6, 2018

kevinthesun commented on Sep 6, 2018

merrymercy commented on Sep 8, 2018

masahi commented on Sep 8, 2018

eqy commented on Sep 8, 2018

masahi commented on Sep 8, 2018

kevinthesun commented on Sep 9, 2018

masahi commented on Sep 9, 2018

kevinthesun commented on Sep 12, 2018

merrymercy commented on Sep 12, 2018

FrozenGene commented on Mar 29, 2019

kevinthesun commented on Apr 2, 2019

kevinthesun commented on Apr 9, 2019

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions