jcf94 opened a new pull request #5962:
URL: https://github.com/apache/incubator-tvm/pull/5962

   Hi all,
   The last PR of Ansor #5883 is not clear enough for reviewers to fully 
understand our design. After some discussion, we changed our upstream plan to 
propose a minimal version of Ansor which contains a small but complete 
framework, so others can get a better understanding of the whole structure of 
   In [[RFC] Ansor: An Auto-scheduler for TVM (AutoTVM 
 we've introduced the **Ansor** auto-scheduler. And we reached an agreement 
that should replace AutoTVM with Ansor in the end.
   For most existing templates, current Ansor can directly replace them with 
better performance and less tuning time.
   For other special templates (low-precision, sparse), the plan is to 
introduce search space customization and gradually rewrite them with Ansor's 
new API.
   This PR introduces a self-contained minimum version of Ansor with most of 
the bones.
   This PR includes the interface of core data structures and an empty search 
policy that does nothing. More advanced search policy and cost models will be 
in the next few PRs.
   # Infrastructure: A Sketch IR for Schedule Searching
   Different from AutoTVM, whose tuning spaces are composed of predefined 
parameters, Ansor constructs a search space by manipulating the loop structures 
of the given compute DAG.
   To enable flexible manipulation of the loop structures, we implemented a 
lightweight loop structure IR (Intermediate Representation) based on the 
original TVM IR but specifically for schedule search. The IR is composed of the 
**state** and **action**, which are defined as follows:
   - **State**: A state of schedule search is the loop structure defined by the 
schedule (i.e., the TVM IR created by `tvm.lower`). See **LoopState** in the 
Key Data Structure for details.
   - **Action**: An action is composed of one or more schedule primitives to 
manipulate (e.g., split, reorder, fuse) a state. See **TransformStep** in the 
Key Data Structure for details.
   We don't use the existing TVM IR but to extend a new Sketch IR on it is 
   1. We want fast incremental change to the loop structures;
   2. We want serializable transform history for replay, backtracking, and 
   3. We may create some macro schedule primitives that represent the 
combination of several TVM schedule primitives
   After the search is done, we will lower this to TVM IR with TVM schedule 
   # Key Data Structures
   To help build an overview of Ansor, this shows the class relations of some 
important Ansor data structures:
   <img width="1222" alt="截屏2020-06-30 上午11 17 25" 
   - **ComputeDAG**: Compute declaration graph and its related analysis tools.
   Related source files: `src/ansor/compute_dag.*`, 
   Ansor takes a compute declaration, which could be a single operator or a 
subgraph, described by `tvm.compute` as an input and converts it to a 
   ComputeDAG implementation includes a set of analyses such as the total float 
operations, consumer/producer relations of each operation stage, whether a 
stage should be tiled/compute inlined, and so on (*some of the analysis will be 
included in the follow-up PRs*). These analyses can help the search policy to 
do specific decisions during schedule search process.
   - **LoopState**: This defines the "state" for the search problem.
   Related source files: `src/ansor/loop_state.*`, 
   Each LoopState corresponds to a specific schedule for the target ComputeDAG.
   A LoopState consists of: 1. the current loop structure; 2. the transform 
history to reach this loop structure.
   The loop structure keeps a **preview** of how the schedule will finally look 
like after lowering (number of iterators, the extent of each iterator, the 
`compute_at` locations, etc), which can help the search policy to make 
decisions during the search.
   The transform history is a sequence of `TransformStep` which will finally be 
mapped to schedule primitives.
   - **TransformStep**: This defines the "action" for the search problem, i.e., 
the schedule primitives for our sketch IR.
   Related files: `src/ansor/transform_step.*`, `python/tvm/ansor/loop_state.py`
   Each step has its corresponding `tvm.te` schedule primitives. We record all 
TransformSteps for every state as its transform history. After finishing the 
schedule search, these transform steps will be lowered with their corresponding 
TVM's schedule primitives.
   *Note*: This PR only contains a small subset of the TransformSteps. The 
complete set of transform steps will be in the next PRs.
   > ComputeDAG is also playing a role of connecting Ansor state system to TVM 
schedule system. That is, ComputeDAG is able to replay TransformSteps to the 
final TVM schedule (e.g., `ComputeDAG(state, actions)`).
   - **SearchTask**: Meta information and hardware parameters for a specific 
schedule search task.
   Related source files: `src/ansor/search_task.*`
   This structure includes the target ComputeDAG, device information as well as 
some hardware parameters obtained from the system or user inputs.
   - **SearchPolicy**: The search policy is defined for Ansor to auto-generate 
a high-performance schedule for different computations.
   Related source files: `src/ansor/search_policy/*`
   A SearchPolicy takes a `SearchTask`, system information and some tuning 
options as inputs, performs the schedule search, and returns a state with the 
best performance. The resulting state can be used to apply to TVM schedule 
     Note that in Ansor paper 
([https://arxiv.org/abs/2006.06762](https://arxiv.org/abs/2006.06762)), we 
proposed a sketch generation policy which achieves pretty good results with 
various workloads on different devices. On the other hand, in this minimum 
system PR, we only provide an `EmptyPolicy` to illustrate the search policy 
   # Ansor Minimum System
   This is a brief diagram for the Ansor system:
   <img width="1481" alt="截屏2020-06-30 上午10 44 59" 
   1. Define the target computation with TVM `te` API and create a ComputeDAG 
   2. Specify the target device, hardware parameters, tuning options and pack 
those with `ComputeDAG` to create a `SearchTask` structure.
   3. `SearchPolicy` takes the `SearchTask` as input and performs the schedule 
search. During the search process, `SearchPolicy` will generate multiple 
candidate states, each of which corresponds to a specific TVM schedule.
   4. Get the best state and use ComputeDAG API to transform it to the final 
TVM schedule.
   In the Ansor system, we use the sketch generation policy (will be brought in 
later PRs) described in the paper as the default search policy, which should be 
enough to cover most use cases. Meanwhile, we will have an RFC for a custom 
rule mechanism that enables user-defined template search to serve the same 
functionality as the current AutoTVM template. Specifically, we will provide 
Python APIs for the new IR that is intended to be used by users for sketch 
customization. They look very similar to existing schedule primitives, as shown 
in `python/tvm/ansor/loop_state.py`.
   Our goal is to make sure Ansor can cover all AutoTVM functionalities while 
achieving the same or better performance so that the community can gradually 
switch to Ansor from AutoTVM.
   # More Details on LoopState
   In this section, we illustrate how a loop state looks like and how does it 
connect to the current TVM build system.
   Take a simple State that includes Split, Fuse and Reorder steps for example:
   A, B, C = matmul_ansor_test(512, 512, 512)
   dag = ansor.ComputeDAG([A, B, C])
   state = dag.get_init_state()
   i, j, k = state[C].iters
   io, ii = state.split(C, i, [16])
   jo, ji = state.split(C, j, [8])
   state.reorder(C, [io, jo, k, ji, ii])
   fused_it = state.fuse(C, [io, jo])
   First, let's print out a state. It shows the loop structure of the 
corresponding TVM schedule, the "preview":
   >>> print(state)
   Placeholder: A, B
   for i.0@j.0@ (0,2048)
     for k (0,512)
       for j.1 (0,8)
         for i.1 (0,16)
           C = ...
   It stores all history transform steps required to reach the current state. 
We can print the history transform steps as TVM's python schedule API. This can 
be used for debugging or to apply the schedule on a former TVM version without 
Ansor support.
   >>> print(dag.print_python_code_from_state(state))
   i, j, k = tuple(C.op.axis) + tuple(C.op.reduce_axis)
   i_o, i_i = s[C].split(i, factor=16)
   j_o, j_i = s[C].split(j, factor=8)
   s[C].reorder(i_o, j_o, k, j_i, i_i)
   i_o_j_o_fused = s[C].fuse(i_o, j_o)
   We can also replay these steps to get a schedule for `tvm.lower` and 
   >>> sche, args = dag.apply_steps_from_state(state)
   >>> print(tvm.lower(sche, args, simple_mode=True))
   primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
     attr = {"global_symbol": "main", "tir.noalias": True}
     buffers = {C: Buffer(C_2: handle, float32, [512, 512], []),
                A: Buffer(A_2: handle, float32, [512, 512], []),
                B: Buffer(B_2: handle, float32, [512, 512], [])}
     buffer_map = {A_1: A, B_1: B, C_1: C} {
     for (i.outer.j.outer.fused: int32, 0, 2048) {
       for (j.inner.init: int32, 0, 8) {
         for (i.inner.init: int32, 0, 16) {
           C_2[((((floordiv(i.outer.j.outer.fused, 64)*8192) + 
(i.inner.init*512)) + (floormod(i.outer.j.outer.fused, 64)*8)) + j.inner.init)] 
= 0f32
       for (k: int32, 0, 512) {
         for (j.inner: int32, 0, 8) {
           for (i.inner: int32, 0, 16) {
             C_2[((((floordiv(i.outer.j.outer.fused, 64)*8192) + (i.inner*512)) 
+ (floormod(i.outer.j.outer.fused, 64)*8)) + j.inner)] = 
((float32*)C_2[((((floordiv(i.outer.j.outer.fused, 64)*8192) + (i.inner*512)) + 
(floormod(i.outer.j.outer.fused, 64)*8)) + j.inner)]) + 
((float32*)A_2[(((floordiv(i.outer.j.outer.fused, 64)*8192) + (i.inner*512)) + 
k)])*(float32*)B_2[(((k*512) + (floormod(i.outer.j.outer.fused, 64)*8)) + 
   The steps of this state can be serialized into the log file as:
   >>> target = tvm.target.create("llvm")
   >>> task = ansor.SearchTask(dag, "test", target)
   >>> inp = ansor.measure.MeasureInput(task, state)
   >>> res = ansor.measure.MeasureResult([0.1], 0, "", 0.2, 1)
   >>> with open("test.log", "w") as fp:
   >>>     ansor.serialization.write_measure_records_to_file(fp.name, [inp], 
   {"i": [["test", "llvm"], [[], [["SP", 2, 0, 512, [16], 1], ["SP", 2, 2, 512, 
[8], 1], ["RE", 2, [0, 2, 4, 3, 1]], ["FU", 2, [0, 1]]]]], "r": [[0.1], 0, 0.2, 
1], "v": "v0.2"}
   Ansor serializes all transform steps to the log file; while AutoTVM 
serializes parameters of a predefined template. The log format discussion would 
be based on 
   In the next few PRs, we'll introduce the complete search policy and 
tutorials for single op/ subgraph schedule search, Relay integration, and 
tutorials for end-to-end network schedule search, custom rules to support 
customized search space.
   This is a joint work by @merrymercy @jcf94 @minminsun @FrozenGene @comaniac 
@yangjunpro @yidawang .

This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:

Reply via email to