AutoTVM how is the search space being generated

comaniac · September 24, 2020, 5:51pm

The search space is defined in the schedule template, such as:

apache/incubator-tvm/blob/master/python/tvm/topi/cuda/conv2d_direct.py#L31


from tvm import autotvm
from ..util import get_const_tuple


def schedule_direct_cuda(cfg, s, conv):
    """schedule optimized for batch size = 1"""

    ##### space definition begin #####
    n, f, y, x = s[conv].op.axis
    rc, ry, rx = s[conv].op.reduce_axis
    cfg.define_split("tile_f", f, num_outputs=4)
    cfg.define_split("tile_y", y, num_outputs=4)
    cfg.define_split("tile_x", x, num_outputs=4)
    cfg.define_split("tile_rc", rc, num_outputs=2)
    cfg.define_split("tile_ry", ry, num_outputs=2)
    cfg.define_split("tile_rx", rx, num_outputs=2)
    cfg.define_knob("auto_unroll_max_step", [0, 512, 1500])

    target = tvm.target.Target.current()
    if target.kind.name in ["nvptx", "rocm"]:
        cfg.define_knob("unroll_explicit", [1])

This line defines a search parameter, and the candidates are the factors of the length of f.

While AutoTVM needs schedule templates defined in TOPI, auto-scheduler generates schedules from scratch. As a result, the auto-scheduler generated schedules are more flexible and expected to achieve even better performance.
The execution time is the main metric to judge the schedule quality.