Update rule for tensorize

taknevski · August 7, 2018, 3:15pm

Hi, why do we need a separate body rule and update rule for tensorize?
What is the distinction between the 2?

Thanks,
Anand

eqy · August 8, 2018, 12:25am

A common example in this case is a matrix vector tensorization primitive. Suppose you have a big matrix vector multiply that you want to build out of smaller matrix vector multiplies (tensorized); in this case the body for the tensorization is the current multiply that you are performing, and the update is the accumulation into previously computed values.

github.com

dmlc/tvm/blob/3fd55cc8b95da6d832a694da337a6c18eb1b47f0/tests/python/unittest/test_schedule_tensorize.py#L24


def intrin_gemv(m, n):
w = tvm.placeholder((m, n), name='w')
x = tvm.placeholder((n,), name='x')
k = tvm.reduce_axis((0, n), name='k')
z = tvm.compute((m,), lambda i:
                tvm.sum(w[i, k] * x[k], axis=k), name='z')
Wb = tvm.decl_buffer(w.shape, w.dtype,
                     name="W",
                     offset_factor=16,
                     strides=[tvm.var('ldw'), 1])
def intrin_func(ins, outs):
    ww, xx = ins
    zz = outs[0]
    ww_ptr = ww.access_ptr("r")
    xx_ptr = xx.access_ptr("r")
    zz_ptr = zz.access_ptr("w")
    body = tvm.call_packed(
        "gemv", ww_ptr, xx_ptr, zz_ptr, n, ww.strides[0])
    reset = tvm.call_packed(
        "fill_zero", zz_ptr, n)
    update = tvm.call_packed(

However, there can be cases where the “body” can be rolled into the “update,” such as when the “update” contains both the operation itself and the accumulation. In this case the “body” can simply be None.

taknevski · August 9, 2018, 11:44am

thanks! that makes it clear