[WIP][Ansor][AutoTVM v2.0] Part 0: Infrastructures for Automatic Schedule Search

[jcf94]

Hi all,

In [RFC] Ansor: An Auto-scheduler for TVM (AutoTVM v2.0), we've introduced the auto-scheduler Ansor. In the RFC, we reached an agreement that we should replace AutoTVM with Ansor.
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 is the first PR according to the integration plan mentioned in the RFC.
This PR contains the infrastructure for search (the definition of state and actions) and small modifications outside the Ansor folder.

Infrastructure for search: A lightweight IR

Automatic scheduling is a search problem. For a search problem, we need to define the states and actions.
The state of schedule search is the loop structure defined by the schedule (i.e., the TVM IR created by tvm.lower). The actions are schedule primitives to manipulate the loop structures (e.g., split, reorder, fuse).

To enable flexible manipulation of the loop structures, we implemented a lightweight loop structure IR (Intermediate Representation) specifically for search. We also implemented all schedule primitives for this IR. Basically, it is a simplified TVM IR. We don't use the existing TVM IR because:

1. We want fast incremental change to the loop structures
2. We want serializable transform history for replay, backtracking, and mutation
3. We may create some macro schedule primitives

After the search is done, we will lower this IR to TVM IR with TVM schedule primitives.

Key data structures

• ComputeDAG: Compute declaration graph and its related analysis tools
  Related files: src/ansor/compute_dag.*, python/tvm/ansor/compute_dag.py
  This is the entrance data structure of Ansor. Ansor takes a compute declaration described by tvm.compute as input and converts it to this data structure for analysis.
• TransformStep: This defines the "action" for the search problem, i.e., the schedule primitives for our 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 TransformStep for every state as its transform history. After the search is done, these transform steps will be lowered with their corresponding TVM's schedule primitives.
• State: This defines the "state" for the search problem, i.e., the current loop structure and history transform steps to reach this state.
  Related files: src/ansor/loop_state.*, python/tvm/ansor/loop_state.py
  A state consists of a current loop structure and the transform history to reach its current loop structure.
  The loop structure keeps a preview of how the schedule will finally look like after lowering (how many iterators, the extent of each iter, the location of some iterators if they have been done compute_at...), which can help the search policy to make decisions during the search.
  The history is a sequence of TransformStep which will finally be mapped to schedule primitives.

Example Walkthrough

While the search policy is implemented in C++, we also provide a python API for the new IR.
This is intended to be used by users for space customization. They look very similar to existing schedule primitives, as shown in python/tvm/ansor/loop_state.py. The API design is ongoing and may get updated later.

Take tests/python/unittest/test_ansor_loop_state.py:test_split_fuse_reorder_annotation() for example, we can print out the test state s1 as:

>>> print(s1)

Placeholder: A, B
gpu.blockIdx.x i.0 (0,32)
  vthread i.1 (0,8)
    gpu.threadIdx.y i.2 (0,2)
      parallel j.0 (0,32)
        unroll j.1 (0,8)
          vectorize j.2 (0,2)
            for k (0,512)
              C = ...

The state stores all history transform steps required to reach this state. We can print the history transform steps as TVM's python API.

>>> print(dag.print_python_code_from_state(s1))

i, j, k = tuple(C.op.axis) + tuple(C.op.reduce_axis)
i_o_i, i_i = s[C].split(i, factor=2)
i_o_o, i_o_i = s[C].split(i_o_i, factor=8)
j_o, j_i_o = s[C].split(j, nparts=32)
j_i_o, j_i_i = s[C].split(j_i_o, nparts=8)
s[C].parallel(j_o)
s[C].unroll(j_i_o)
s[C].vectorize(j_i_i)
s[C].bind(i_o_o, tvm.thread_axis("blockIdx.x"))
s[C].bind(i_o_i, tvm.thread_axis("vthread"))
s[C].bind(i_i, tvm.thread_axis("threadIdx.y"))

Or replay these steps to get a schedule for tvm.lower and tvm.build.

s, args = dag.apply_steps_from_state(s1)
print(tvm.lower(s, args, simple_mode=True))

The steps of this state can be serialized into the log file as:

[["SP", 2, 0, 512, [8, 2], 1], ["SP", 2, 3, 512, [32, 8], 0], ["AN", 2, 3, 3], ["AN", 2, 4, 1], ["AN", 2, 5, 2], ["AN", 2, 0, 5], ["AN", 2, 1, 4], ["AN", 2, 2, 8]]

Ansor serializes all transform steps to the log file. This is different from AutoTVM which only serializes parameters.

---

In the next few PRs, we'll introduce 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. When Ansor is able to fully support AutoTVM's features, we can gradually deprecate AutotVM.

This is a joint work by @merrymercy @jcf94 @minminsun @FrozenGene @comaniac @yangjunpro @yidawang .

Changes of original TVM code outside Ansor folders (Will later split these to separate PRs)

• CUDA device API & VerifyGPUCode pass update #5898 Add kMaxRegistersPerBlock in GPU device query (include/tvm/runtime/device_api.h, src/runtime/cuda/cuda_device_api.cc, src/runtime/opencl/opencl_device_api.cc)
• [TE] Add LegalizeInvalidAttach to legalize the compute_at location after split or fuse #5917 Add LegalizeInvalidAttach (src/te/schedule/schedule_dataflow_rewrite.cc)
• ([random] support random fill #5913) Add nd.non_empty (include/tvm/runtime/c_runtime_api.h, include/tvm/runtime/ndarray.h, python/tvm/runtime/ndarray.py, src/runtime/ndarray.cc)
• CUDA device API & VerifyGPUCode pass update #5898 Add vectorization check in VerifyGPUCode (src/tir/analysis/verify_gpu_code.cc)
• #(25) Add explicit_unroll_max_extent (src/tir/transforms/unroll_loop.cc, tests/python/unittest/test_tir_transform_unroll_loop.py)
• [Arith][GPU]Rewrite simplify fix for Vectorized Cooperative Fetching #5924 Update Index simplification to fix the vectorized cooperative fetching (src/arith/rewrite_simplify.cc)
• ([clflush] Enable x86 cpu cache flush #5914) Add flush CPU cache during tuning process to make the result more correct (src/runtime/rpc/rpc_module.cc, src/runtime/threading_backend.cc)
• #(34) Add call_all_topi_functions to RelayBuildModule (src/relay/backend/build_module.cc, python/tvm/relay/build_module.py)

[tqchen]

vector<Stage> -> Array<Stage>

[tqchen]

Let us move the schedule promitives to the StateNode instead

[tqchen]

This PR is good to get a global context. Some further comments: To make sure we get the a thoroughly review and smooth upstream, let us break this PR down further into several PRs, it also helps us to logically organize and think about the overall design architecture:

• Changes outside ansor, consider use one PR for each group of change(e.g. rewrite simplify change can be its own PR).
• A single PR for the state data structures, loop_state.h, loop_state.py
  • Clearly document all the fields, key functions
  • Discuss the implication of the design, by briefly comment on the search process in loop_state.h
• A PR contains Measurements related code for ansor
• Specific implementation of each ansor primitives(TransformStep), each as its own PR

[tqchen]

cc @merrymercy @spectrometerHBH @junrushao1994 @Hzfengsy @jwfromm

[jwfromm]

I agree with @tqchen that it would be helpful to break this up a little into separate PRs. Maybe it can be divided into the three partitions discussed in the paper: task scheduling, program sampling, and performance tuning. That would make it much more clear what we're looking it.

[merrymercy]

@jwfromm The partition of implementation is different from the organization of the paper. We cannot upstream code according to the paper. We listed the integration steps in our RFC, and we will follow those steps.
For people who read the paper, this PR corresponds to nothing in the paper. This PR only provides the infrastructures for search (The definition of state and actions in the search). The search policy (i.e. program sampling and performance fine-tuning) and task scheduler will be in next PRs.

[tqchen]

I have outlined the proposal of a possible breakdown of this PR in the above post, please see if that makes sense

[comaniac]

Since this PR is incomplete it might be not trivial for people to review. For example, feature.cc includes all utilities we used to extract features, but the use case (XGBoost cost model) is not included yet. Another example is that loop state and transform steps are included in this PR but the search policy isn't. This might cause people hard to imagine how those data structures will be used.

If we are going to breakdown this PR, I would suggest putting everything we have to this PR and use this PR as the reference when sending other small PRs.

[merrymercy]

After some discussion, we changed our upstream plan.

We will distill a minimal version of Ansor and send it as the first PR.
This minimal version will contain a small but complete framework, so people can get a better understanding of the whole structure of Ansor.
This is different from the old upstream plan where we send separate components one by one. In the old plan, the framework is not complete during upstream, making the reviewers lose the context.