Loop Partition example

main.py

import tvm
import numpy
import timeit

n = 225
A = tvm.placeholder((n, n), name='Input')
B = tvm.placeholder((7, 7), name='Filter')
di = tvm.reduce_axis((0, 7), name='di')
dj = tvm.reduce_axis((0, 7), name='dj')
C = tvm.compute(
    (n, n),
    lambda i, j: tvm.sum(A[i + di, j + dj] * B[di, dj], axis=[di, dj]),
    name='Output')
s = tvm.create_schedule(C.op)


f = s[C].fuse(di, dj)
fo, fi = s[C].split(f, nparts=14)

print(tvm.lower(s, [A, B, C], simple_mode=True))

a = tvm.lower(s, [A, B, C], simple_mode=True)
a = tvm.ir_pass.LoopPartition(a, True)
a = tvm.ir_pass.Simplify(a)

print(tvm.lower(s, [A, B, C], simple_mode=True))