Operator Level Optimization

Decouple compute and schedule

TVM's operator implementation adopts the decoupled compute/schedule principle from Halide.

TVM introduces a tensor expression language to support automatic code generation. The following code shows an example tensor expression to compute transposed matrix multiplication:

M = 1024
K = 1024
N = 1024

k = te.reduce_axis((0, K), 'k')
A = te.placeholder((M, K), name='A')
B = te.placeholder((K, N), name='B')
C = te.compute((M, N),
           lambda x, y: te.sum(A[x, k] * B[k, y], axis=k),
           name='C')

As shown in previous GEMM chapter, this naïve implementation is not efficient and we can optimize it by better data locality and applying hardware intrinsics like SIMD. From previous Halide chapter, we know that writing optimal codes directly is complex and Halide provides a set of API to decouple computation definition and schedule. TVM adopts the same idea and provides similar APIs for the same purpose.

Just like Halide, TVM's schedule APIs include split, tile, compute_at, etc. Check TVM's tutorial here for detailed introduction. One major difference between Halide is the tensorize API, TVM supports to customize tensorize API by python codes, so we can use different intrinsic functions for different hardware easily. Check here for details.

With these schedule APIs, we can optimize matrix multiplication like below codes. Full sample codes can be found here.

s = te.create_schedule(C.op)
# Tile and split the computation into smaller blocks
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], 32, 32)
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, ki, xi, yi)
# Parallel yi dimension with intrinsics
s[C].vectorize(yi)

AutoTVM

Like Halide, TVM also supports tuning schedule parameters automatically. There're 2 steps to do it, first is to define the search space, and then search through the space automatically.

Define search space

For previous matrix multiplication sample, we set tile factor to 32, but it may not be optimal. Instead, We can define a search space.

@autotvm.template("tutorial/matmul")
def matmul(M, K, N):
    k = te.reduce_axis((0, K), 'k')
    A = te.placeholder((M, K), name='A')
    B = te.placeholder((K, N), name='B')
    C = te.compute((M, N),
            lambda x, y: te.sum(A[x, k] * B[k, y], axis=k),
            name='C')

    cfg = autotvm.get_config()
    # define search space
    cfg.define_knob("tile_x", [4, 8, 16, 32, 64])
    cfg.define_knob("tile_y", [4, 8, 16, 32, 64])

    s = te.create_schedule(C.op)
    xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], cfg["tile_x"].val, cfg["tile_y"].val)
    ko, ki = s[C].split(k, factor=4)
    s[C].reorder(xo, yo, ko, ki, xi, yi)
    # Vectorization
    s[C].vectorize(yi)

    return s, [A, B, C]

Search through the space

TVM provides different methods to search the space, like random search and grid search. They are good enough for smaller space, but for larger space Machine Learning based methods can be more efficient.

ML model that takes the lowered loop program as input and predicts its running time on a given hardware back-end. The model, trained using runtime measurement data collected during exploration, does not require the user to input detailed hardware information. The model uses a rank objective to predict the relative order of runtime costs (A runs faster than B)[1].

In practice, TVM provides a gradient tree boosting model. On average, it does prediction in 0.67ms, thousands of times faster than running a real measurement[1].

Check here for a sample of tuning E2E model automatically.

References

[1] Tianqi Chen, et al. TVM: An automated end-to-end optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), Carlsbad, CA, 2018. USENIX Association.