[dmlc/tvm] [RFC][AUTOTVM] Auto-Scheduler from Compute Decleration (#2954)

[Lianmin Zheng <no...@github.com>]

# Auto-Scheduler
TVM decouples kernel implementation into compute and schedule. The compute part is a friendly DSL that can describe algorithms intuitively. However, the schedule part still requires strong expert knowledge and time-consuming tuning to provide decent performance. The tuning process is partially automated by the existing autotvm package, but a human-engineered template is still required.

This RFC proposes a "real" autotvm, which we can call auto scheduler. It aims at removing all human efforts on schedule part.

# Proposed Design 
The auto-scheduler is built on the exsiting autotvm package. It will generate a template from compute decleration. Then this template can either be 

* Statically filled by heuristic rules and cost functions to provide reasonable performance, or
* Dynamically tuned by autotvm to provide better performance with some time budget

The auto-scheduler takes a computation graph described by tvm DSL as input, then classify the type of read/write patterns and the type of computation. It dispatches the declarations to different "meta templates". The "meta templates" generates autotvm templates from the declaration. There are four types of meta templates : simple reduction, complex reduction, direct compute, and location-tunable compute. The auto-scheduler will do parallelization, vectorization, tiling, and operator fusion.

The code is available on [my branch](https://github.com/merrymercy/tvm/tree/auto-scheduler). The current implementation is in pure python bacuse autotvm is mainly written in python. But move the whole autotvm package to c++ is within long-term plan. The code is organized as follows.
* Analysis on access pattern [python/tvm/autotvm/auto_schedule/stage_analysis.py](https://github.com/merrymercy/tvm/blob/auto-scheduler/python/tvm/autotvm/auto_schedule/stage_analysis.py)
* CPU backend [python/tvm/autotvm/auto_schedule/backend/cpu.py](https://github.com/merrymercy/tvm/blob/auto-scheduler/python/tvm/autotvm/auto_schedule/backend/cpu.py)
* GPU backend [python/tvm/autotvm/auto_schedule/backend/gpu.py](https://github.com/merrymercy/tvm/blob/auto-scheduler/python/tvm/autotvm/auto_schedule/backend/gpu.py)
* Configuration for the auto-scheduler [python/tvm/autotvm/auto_schedule/common.py](https://github.com/merrymercy/tvm/blob/auto-scheduler/python/tvm/autotvm/auto_schedule/common.py)
* Experimental auto-packing for optimizing vectorization and locality [python/tvm/autotvm/auto_schedule/auto_pack.py](https://github.com/merrymercy/tvm/blob/auto-scheduler/python/tvm/autotvm/auto_schedule/auto_packing.py)
* Test case [tests/python/unittest/test_auto_scheduler.py](https://github.com/merrymercy/tvm/blob/auto-scheduler/tests/python/unittest/test_auto_scheduler.py)

## API
There are only two user-oriented API calls

* `autotvm.AutoSchedulerOptions(**kwargs)`
This is used to configure the auto scheduler. The arguments include hardware configurations(vector lanes, number of threads, size of shared memory, etc) and tuning configurations (how many tuning knobs to generate).
* `autotvm.create_schedule(tensors)`
This is similar to `tvm.create_schedule`, but returns an already optimized schedule.

```python
A = tvm.placeholder((128,), name='A')
B = tvm.placeholder((128,), name='B')
C = tvm.compute((128,),  lambda i: A[i] + B[i] * 2)

with tvm.target.create('llvm'):
    with autotvm.AutoSchedulerOptions(vec_size=8, num_threads=16):
        s, bufs = autotvm.create_schedule([A, B, C])

# NO SCHEDULE REQUIRED

func = tvm.build(s, bufs)
```

# Examples
1. [Tutorial](https://github.com/merrymercy/tvm/blob/auto-scheduler/tutorials/autotvm/auto_scheduler.py)
   This is a tutorial on how to statically use the auto-scheduler or auto-tune it.
2. [Schedule a whole network](https://github.com/merrymercy/tvm/blob/auto-scheduler/scripts/training-with-tvm.py)
   This example is adopted from #2498. It is a LeNet like convolution neural network written purely by tvm (without graph IR). The auto-scheduler also provides basic operator fusion for it. Right now we can only run forward pass. I am working on fixing the backward pass.

# Performance
One reachable performance goal is to replace more than 90% schedule code in existing TOPI by this auto-scheduler. I haven't done the experiments, but I believe the generated templates cover the existing search space for most operators (includes conv2d, reduction, ...).

Another part of the goal is to provide reasonable static performance. In the "Schedule a whole network" example, for batched forward pass, the current performance is 1.2x slower than out-of-the-box TF + Keras, and 10x faster than naive schedule (fuse and parallel outer loop) on an Intel i7-8750H. For static usage, the input of the auto-scheduler are parameters for heuristic rules and hardware configurations. We will gather all inputs into a global config, so users can still do some quick "tuning".

# Todo list
 - [ ] Performance test and improvement to cover more than 90% schedule code in TOPI
       Improve the heuristic rules to provide better static performance, do test to make sure we covor the search space of exsting templates.
 - [ ] Improve tuning speed
       The current implementation does analysis and generates the template on the fly, which is expensive and redundant during batched tuning. We should decouple the template generation and template tuning, and explicitly cache the template.
 - [ ] (long-term) Move all autotvm related code to c++
 - [ ] Improve loop partition to better handle partial tile, vectorization.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Thierry Moreau <no...@github.com>]

Thanks @merrymercy, this is really awesome work. I second Jared's comment on work involved in adding a backend. I'd be happy to chat some more about how one would add automated compilation to different hardware accelerators including VTA.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479366759

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Yao Wang <no...@github.com>]

Thank you for opening this RFC! I have a question regarding user API. Does the hardware information needed for autotvm.AutoSchedulerOptions(**kwargs) function pre-defined for different hardware architectures? If so, how much more information does a user need to provide to differentiate between different minor types of the same device target, such as Intel Xeon Platinum vs Xeon Haswell, or Nvida K80 vs V100? Today we have a single template for minor device types. Will auto-scheduler provide different templates?

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479219911

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[eqy <no...@github.com>]

@merrymercy Do you think that this is a good time to also make schedules serializable/package them with autotvm style configs? In the past we have had issues where we did not want to merge in changes to schedules because they would break compatibility with tophub, and now it seems that the variety of schedules may also change quickly as auto-schedule is changed. Instead of forcing schedules to be schedule, we can maybe side-step this by packaging schedules together with autotvm configs.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-480441416

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Jared Roesch <no...@github.com>]

@merrymercy I'm less interested in LOC and more how much conceptual burden there is. What are the key pieces that make up a backend description is more my question.

 I looked over the code but was at SysML and have two deadlines this week so I haven't had a chance to really look it over. Look forward to landing this stuff.

One idea I've been thinking about is a combined TVM + Relay language where we can auto-extract chunks that can be lowered to the compute language, auto-schedule, then auto-tune for end-to-end perf.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479393039

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Yizhi Liu <no...@github.com>]

@merrymercy Do you think this analysis design can be easily extended to be working based on TVM Tensor AST (HalideIR) instead of ScheduleStage? Not urgent but I think eventually we will make schedule primitives work on HalideIR, so that we can unify the underlying data structure of schedule and other passes.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479736425

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Lianmin Zheng <no...@github.com>]

@jroesch Currently, it is about 500 loc per backend. I am working on improvements so it may increase.

@yzhliu 
* simple reduction: reduction ops that do not have reuse opportunity (e.g. softmax, argmin)
* complex reduction: reduction ops that have reuse opportunity (e.g. matmul, conv2d
* direct compute: broadcast, elemwise, stencil computation, (e.g. relu, add)
* location-tunable compute: the same as above. The difference is that `direct compute` computes at root, while `location-tunable compute` can computes at other nodes to increase locality.

@tmoreau89 This is doable. The problem of accelerators is that if we want the auto-scheduler to take in a hardware-independent description, then we need a special pack pass to transform the layout.



-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479379983

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Yao Wang <no...@github.com>]

@merrymercy Auto-scheduler will create another search space consists of schedule templates. For a given set of hardware parameters, it will try various schedule templates and for each template do some auto-tuning on real device. This means for each minor device type, we need to do all these steps. Do I understand it correctly?

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479679532

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Tianqi Chen <no...@github.com>]

Good discussions, I think in general we can move to summarize the common patterns and make things work for specific hardware backend. As for point bought by @yzhliu (unifying schedule with pass), eventually ScheduleStage itself(or other IR structure) can be viewed as a dialect of the IR, and we can do so after we push for such unification. 

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479744491

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[eqy <no...@github.com>]

Minor question: do we consider "injective" as a special case of "simple reduction?"

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-480437297

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Jared Roesch <no...@github.com>]

@merrymercy how much work is there per backend? looking over the code now will follow up with more questions later.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479300009

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Yizhi Liu <no...@github.com>]

@merrymercy Could you elaborate a bit about the 4 types (simple reduction, complex reduction, direct compute, and location-tunable compute) ? Also it would be helpful if you can give an example of how the DAG looks like.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479345162

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Lianmin Zheng <no...@github.com>]

@eqy "injective" is considered "direct compute". Typically they will be inlined.

Serializable Template + Serializable Config seems to be a good direction to go.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-481152553

Re: [dmlc/tvm] [RFC][AUTOTVM] Auto-Schedule from Compute Declaration (#2954)

[Lianmin Zheng <no...@github.com>]

@jroesch There is no easy description for a backend. Currently these meta-templates are mainly based on the summary of existing human schedule code in TOPI. So adding a new backend is still hard. What can be reused is the classification of compute type.

@kevinthesun There is only one template for one specific op. The auto-scheduler first creates this template. Then, for static usage, it will fill the knobs in the template according to hardware paremeters. The example shown above falls in this category. For tuning usage, the auto-scheduler won't use hardware parameters. Instead, it relies real tuning. In this case, you need to explicitly create `autotvm.Task`, `autotvm.Tuner`. An example is shown in the tutorial.

@yzhliu The `tvm.compute` dsl is much easier to analyze than general Halide IR, because of its clean dependency relations and well-defined loop structure.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/2954#issuecomment-479842520