Is it possible to vectorize a reduce_axis?

[merrymercy]

I want to vectorize an axis in sum reduction, here is the code

import tvm
import numpy as np

n = 64

A = tvm.placeholder((n, n), name='A', dtype='float32')
k = tvm.reduce_axis((0, n), "k")
B = tvm.compute((n,), lambda i: tvm.sum(A[i, k], axis=k), name="C")

s = tvm.create_schedule(B.op)
k, ki = s[B].split(k, factor=8)
s[B].vectorize(ki)  # vectorize here

print(tvm.lower(s, [A, B], simple_mode=True))
func = tvm.build(s, [A, B], target='llvm')

# verify correctness
ctx = tvm.cpu(0)
a = tvm.nd.array(np.random.uniform(size=(n, n)).astype(A.dtype), ctx=ctx)
b = tvm.nd.array(np.zeros(shape=(n,), dtype=B.dtype), ctx=ctx)
func(a, b)
np.testing.assert_allclose(b.asnumpy(),  np.sum(a.asnumpy(), axis=1), rtol=1e-4)

TVM did not output any error message, but the answer is wrong. Output:

produce C {
  for (i, 0, 64) {
    C[i] = 0.000000f
    for (k.outer, 0, 8) {
      C[x8(i)] = (x8(C[i]) + A[ramp((((i*8) + k.outer)*8), 1, 8)])
    }
  }
}
...
(mismatch 100.0%)
 x: array([ 5.928585,  4.939113,  4.34326 ,  3.470788,  2.35553 ,  2.628344,
        3.791345,  4.616838,  3.214629,  2.659637,  4.374955,  5.651574,
        2.995839,  4.34803 ,  5.436972,  4.196923,  3.906123,  4.364611,...
 y: array([ 33.369698,  33.435944,  32.376495,  33.983963,  33.391079,
        26.468966,  31.47113 ,  30.853607,  31.724657,  28.669315,
        31.928965,  36.953148,  32.466213,  28.177387,  31.322906,...

[tqchen]

for now it is not possible, as vectorizing on reduction requires special logic.

[tqchen]

I think it is a good idea to raise an error when user do so for now

[nkyle04]

@tqchen there will be issues, when reduce_axis can't be vectorize.
i.e. we can't use "vector mla" instruction, which is common in neon and SSE, so can not get more performance

[merrymercy]

@nkyle04 You can do it by rfactor (see doc)

import tvm

N = 128
V = 8

A = tvm.placeholder((N, N), name='A')
k = tvm.reduce_axis((0, N), name='k')
B = tvm.compute((N,), lambda i : tvm.sum(A[i, k], axis=k), name='B')

s = tvm.create_schedule(B.op)
ko, ki = s[B].split(k, V)
BF = s.rfactor(B, ki)
s[BF].compute_at(s[B], s[B].op.axis[0])
s[BF].vectorize(s[BF].op.axis[0])

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

// attr [B.rf] storage_scope = "global"
allocate B.rf[float32 * 8 * 1]
produce B {
  for (ax0, 0, 128) {
    produce B.rf {
      B.rf[ramp(0, 1, 8)] = x8(0.000000f)
      for (k.outer, 0, 16) {
        B.rf[ramp(0, 1, 8)] = (B.rf[ramp(0, 1, 8)] + A[ramp((((ax0*16) + k.outer)*8), 1, 8)])
      }
    }
    B[ax0] = 0.000000f
    for (k.inner.v, 0, 8) {
      B[ax0] = (B[ax0] + B.rf[k.inner.v])
    }
  }
}

[nkyle04]

OK, @merrymercy thanks a lot

[merrymercy]

If you want to further vectorize the line B[ax0] = (B[ax0] + B.rf[k.inner.v]) in the code above, you may have to create an intermediate buffer so you can vectorize on ax0. This introduces a tradeoff between locality and vectorization.

import tvm

N = 128
V = 8

A = tvm.placeholder((N, N), name='A')
k = tvm.reduce_axis((0, N), name='k')
B = tvm.compute((N,), lambda i : tvm.sum(A[i, k], axis=k), name='B')

s = tvm.create_schedule(B.op)
ko, ki = s[B].split(k, V)
BF = s.rfactor(B, ki)
#s[BF].compute_at(s[B], s[B].op.axis[1])
s[BF].vectorize(s[BF].op.axis[0])

ax0 = s[B].op.axis[0]
k = s[B].op.reduce_axis[0]
ax0, vax0 = s[B].split(ax0, V)
s[B].vectorize(vax0)

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

// attr [B.rf] storage_scope = "global"
allocate B.rf[float32x8 * 8 * 16]
produce B.rf {
  for (i, 0, 128) {
    B.rf[ramp(i, 128, 8)] = x8(0.000000f)
    for (k.outer, 0, 16) {
      B.rf[ramp(i, 128, 8)] = (B.rf[ramp(i, 128, 8)] + A[ramp((((i*16) + k.outer)*8), 1, 8)])
    }
  }
}
produce B {
  for (ax0.outer, 0, 16) {
    B[ramp((ax0.outer*8), 1, 8)] = x8(0.000000f)
    for (k.inner.v, 0, 8) {
      B[ramp((ax0.outer*8), 1, 8)] = (B[ramp((ax0.outer*8), 1, 8)] + B.rf[ramp(((ax0.outer + (k.inner.v*16))*8), 1, 8)])
    }
  }
}