[TOPI][ARM] Improve injective schedule

[hlu1]

The generic injective schedule does not have vectorization and is therefore slow on ARM CPU. With vectorization, it can run 2-3x faster. For example, for a upsample_relu layer with 48 x 48 x 48 (C, H, W), the vectorized code runs at 0.003 ms/iter compared to 0.008 ms/iter on raspberry pi.

[hlu1]

@ajtulloch please review :)

[FrozenGene]

Yeah. this is very useful likelayout_transform op.

@hlu1 How about (io, ii) = s[x].split(list(s[x].op.axis)[-1], 8) if s[x].op.axis[-1] <8 and how about poor arm cpu like A9 (maybe 4 is better)? I have done it like this:

  if len(s[x].op.axis) >= 5: # it is very useful when we have NCHWxC.
        fused = s[x].fuse(s[x].op.axis[0], s[x].op.axis[1], s[x].op.axis[2])
        s[x].parallel(fused)
    elif len(s[x].op.axis) >= 3:
        fused = s[x].fuse(s[x].op.axis[0], s[x].op.axis[1])
        s[x].parallel(fused)
    else:
        s[x].parallel(s[x].op.axis[0])
    s[x].vectorize(list(s[x].op.axis)[-1])

[hlu1]

We also need to consider the case of int8/unit8. For example, when you add two int8 numbers together to produce 1 int16 number, the simd width is 128/16 = 8. I think in general 8 should be a good compromise.

[FrozenGene]

Ok. Could you add >=5 like we do it in x86? https://github.com/dmlc/tvm/blob/master/topi/python/topi/x86/injective.py#L26. This could help us in NCHWxC layout transform.

[hlu1]

That should have been covered by:

if len(s[x].op.axis) >= 3:
    fused = s[x].fuse(s[x].op.axis[0], s[x].op.axis[1], s[x].op.axis[2])
    s[x].parallel(fused)

I used len(s[x].op.axis) >= 3 because it's needed for special cases like [1, 1, 224, 224].

[FrozenGene]

Oops. I haven't noticed it.

[FrozenGene]

LGTM.

[FrozenGene]

@hlu1 Could you help to see this discussion on discuss forum? https://discuss.tvm.ai/t/relay-build-target-rasp3b-something-wrong/2195 This issue should be related with this changeset.

[hlu1]

@FrozenGene, thanks for letting me know. Fixed in #3061