Improve performance of broadcast_axis on CPU

[access2rohit]

Description

Improves the performance of broadcast axis by reducing ALU operations and caching stride values. This operator is crucial in performance of SSD which gets slowed down by over 50%(total training time) when MXNet is built with Large Tensor Support enabled.

This PR leverages vectorization or SIMD architecture of the CPU. It is described as follows:

Vectorization

is the term for converting a scalar program to a vector program. Vectorized programs can run multiple operations from a single instruction(SIMD), whereas scalar can only operate on pairs of operands at once.

Consider the following very simple loop that adds the elements of two arrays and stores the results to a third array.

for (int i=0; i<16; ++i)
    C[i] = A[i] + B[i];

Vectorizing it, produces something like this:

for (int i=0; i<16; i+=4)
    C[i:i+3] = A[i:i+3] + B[i:i+3];

Note: The key here is that the elements need to be adjacent (so they can be loaded into vector registors. Note:avx 512 ISA does have "gatherd" intstruction for gathering sparse data from scattred locations into avx512 register but that operation is still slow compared to adjacent data being loaded into the register) and there should not be any fundamental dependency(or read after write dependency) in subsequent iterations of the loop. For e.g. the following loop cannot be vectorized because C[i+1] depends on the value of C[i] which needs to bew calculated first.

for (int i=0; i<16; ++i)
    C[i] = A[i] + C[i-1];

because if we unroll the loop and substitute values:

for (int i=0; i<16; i+=4) {
    C[i]   = A[i]   + C[i-1];
     |
     |-------------------------|
                               |
    C[i+1] = A[i+1] + C[i]; <--|
     |
     |---------------------------|
                                 |
    C[i+2] = A[i+2] + C[i+1]; <--|
     |
     |---------------------------|
                                 |
    C[i+3] = A[i+3] + C[i+2]; <--|
}

The arrows above show read after write dependency.the above statements cannot be executed in parallel by loading A[1:4] and C[1:4] into vector registers(here 1:4 is just for illustration purposes depending on the word size and register size it can load more data elements).

How is it different from loop unrolling? Unrolling transforms the loop into something that looks like this:

for (int i=0; i<16; i+=4) {
    C[i]   = A[i]   + B[i];
    C[i+1] = A[i+1] + B[i+1];
    C[i+2] = A[i+2] + B[i+2];
    C[i+3] = A[i+3] + B[i+3];
}

Side Effects

for cases where 1 or very few elements need to be broadcasted to very large no. of locations the performance is not optimized. There is a slowdown when compared to master(int32 w/o large tensor support)

Checklist

Essentials

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Performance

Opperf code used to benchmark

add_res = run_performance_test(nd.broadcast_axis, run_backward=True, dtype='float32', ctx=mx.cpu(),
                               inputs=[{'data': (1000, 1, 100, 1), 'axis': (1, 3), 'size': (10, 5)}],
                               warmup=100, runs=1000, profiler='python')
print(add_res)

add_res = run_performance_test(nd.broadcast_axis, run_backward=True, dtype='float32', ctx=mx.cpu(),
                               inputs=[{'data': (1000, 1, 1, 100), 'axis': (1, 2), 'size': (10, 5)}],
                               warmup=100, runs=1000, profiler='python')
print(add_res)

add_res = run_performance_test(nd.broadcast_axis, run_backward=True, dtype='float32', ctx=mx.cpu(),
                               inputs=[{'data': (1, 1, 1, 1), 'axis': (0, 1, 2, 3), 'size': (1000, 10, 100, 5)}],
                               warmup=100, runs=1000, profiler='python')
print(add_res)

add_res = run_performance_test(nd.broadcast_axis, run_backward=True, dtype='float32', ctx=mx.cpu(),
                               inputs=[{'data': (1, 1000, 1, 100, 1), 'axis': (0, 2, 4), 'size': (2, 10, 5)}],
                               warmup=100, runs=1000, profiler='python')
print(add_res)

Results

code version	cases	avg	p50	p90
master LT	(1000, 1, 100, 1)->(1000, 10, 100, 5)	27.5	26.05	21.47
	(1000, 1, 1, 100)->(1000, 10, 5, 100)	29.68	28.59	35.93
	(1, 1, 1, 1)->(1000, 10, 100, 5)	129.80	123.30	169.04
	(1, 1000, 1, 100, 1)->(2, 1000, 10, 100, 5)	63.60	58.88	74.74
master no-LT	(1000, 1, 100, 1)->(1000, 10, 100, 5)	14.30	12.76	18.00
	(1000, 1, 1, 100)->(1000, 10, 5, 100)	13.21	12.33	15.67
	(1, 1, 1, 1)->(1000, 10, 100, 5)	66.80	60.77	83.80
	(1, 1000, 1, 100, 1)->(2, 1000, 10, 100, 5)	31.81	29.71	40.23
new LT	(1000, 1, 100, 1)->(1000, 10, 100, 5)	17.53	17.36	22.42
	(1000, 1, 1, 100)->(1000, 10, 5, 100)	15.49	14.79	17.64
	(1, 1, 1, 1)->(1000, 10, 100, 5)	127.74	126.06	131.19
	(1, 1000, 1, 100, 1)->(2, 1000, 10, 100, 5)	39.23	38.95	39.48
new no-LT	(1000, 1, 100, 1)->(1000, 10, 100, 5)	9.41	8.38	11.84
	(1000, 1, 1, 100)->(1000, 10, 5, 100)	8.29	7.65	10.61
	(1, 1, 1, 1)->(1000, 10, 100, 5)	67.54	63.63	86.17
	(1, 1000, 1, 100, 1)->(2, 1000, 10, 100, 5)	23.57	23.39	23.92

SSD training performance

Code	SSD 1 Epoch time (sec)	%age Speedup/Slowdown w.r.t Master (large tensor disabled)
Master (large tensor disabled)	226	0
Master (large tensor enabled)	335	48.23% slowdown
Master + CPU Optimized broadcast_axis (large tensor disabled)	130	42.5% speedup
Master + CPU Optimized broadcast_axis (large tensor enabled)	184	18.5% speedup

[ptrendx]

Is this going to concentrate on CPU performance? GPU performance of broadcast_axis/like/etc is really poor too (few days ago I was measuring it and got ~150 GB/s on V100 GPU out of 900 GB/s peak BW.

[ptrendx]

The biggest speedup you can get here would be from vectorization, so that you don't need to do those index calculations all the time (just once per vector).

[leezu]

@access2rohit can you rebase on latest master to see if that fixes CI?

[apeforest]

Could we also take Da's suggestion by using int32_t for GPU kernels?

[access2rohit]

Could we also take Da's suggestion by using int32_t for GPU kernels?

@apeforest
For this op only or all ? For this op only I can take a look.
If for all ops I would like to do that in a separate PR.

[access2rohit]

@access2rohit can you rebase on latest master to see if that fixes CI?

Was still failing after rebase. I have fixed the failing case and pushed new changes now.

[apeforest]

Please add more details in the PR description to explain the rationale of this change:
(1) Why make the change
(2) How this change improved performance
(3) Any side effect

[access2rohit]

Please add more details in the PR description to explain the rationale of this change:

@apeforest

1. Its there in the PR Heading. Additionally, its being done to improve BERT performance regression when using Large Tensor Build
2. Leverages vectorization. Can't explain line by line as to how. Too much text to write. I can share a link about vectorization for reviewers to read. Let me know if you meant something else. I can write 1 line each for 3 cases which is necessary
3. I didn't understand this question. Could you clarify what meant by side-effect.

[access2rohit]

@mxnet-label-bot add [pr-awaiting-review]

[access2rohit]

@apeforest can you review. I have addressed your comments

[access2rohit]

@zheng-da Can you please review?

[access2rohit]

@mxnet-label-bot update [pr-awaiting-merge]