Autotvm matmul

matmul.py

import sys
import numpy
import timeit
import logging
import tvm
from tvm import autotvm

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

    s = tvm.create_schedule(C.op)
    CC = s.cache_write(C, 'global')

    cfg = autotvm.get_config()
    
    y, x = s[C].op.axis
    cfg.define_split('tile_y', y, num_outputs=2)
    cfg.define_split('tile_x', x, num_outputs=2)
    cfg.define_split('tile_k', k, num_outputs=2)

    yo, yi = cfg['tile_y'].apply(s, C, y)
    xo, xi = cfg['tile_x'].apply(s, C, x)
    s[C].reorder(yo, xo, yi, xi)
    s[CC].compute_at(s[C], xo)

    yc, xc = s[CC].op.axis
    k, = s[CC].op.reduce_axis
    ko, ki = cfg['tile_k'].apply(s, CC, k)
    s[CC].reorder(ko, yc, ki, xc)
    
    cfg.define_reorder('inner_order', [yc, ki], policy='all')
    cfg.define_reorder('outer_order', [yo, xo], policy='all')
    cfg['inner_order'].apply(s, CC, [yc, ki])
    cfg['outer_order'].apply(s, C, [yo, xo])

    return s, [A, B, C]


if __name__ == '__main__':
    M = 512
    K = 512
    N = 512

    task = autotvm.task.create(matmul, args=(M, K, N), target='llvm')

    # logging config (for printing tuning log to the screen)
    logging.getLogger('autotvm').setLevel(logging.DEBUG)
    logging.getLogger('autotvm').addHandler(logging.StreamHandler(sys.stdout))

    measure_option = autotvm.measure_option(
        builder='local',
        runner=autotvm.LocalRunner(number=5))

    # begin tuning, log records to file `matmul.log`
    n_trial = 100
    tuner = autotvm.tuner.XGBTuner(task)
    logfile = 'matmul.log'
    tuner.tune(n_trial=n_trial,
               measure_option=measure_option,
               callbacks=[autotvm.callback.log_to_file(logfile)])