Accelerate the performance of topk for CPU side

[ciyongch]

Description

Optimize the performance of topk algorithm for CPU side, which mentioned in issue (#10205)

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• 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)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• topk

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

@pengzhao-intel @xinyu-intel

[asmushetzel]

Can you please elaborate a bit

• what speedups do you see. I see that there are some numbers in [Operator] Accelerate the CPU side performance of topk #10205 but they all are for a single case k=3. We should make sure that other cases (k=N/2 etc) are also good.
• what actually contributes to the speedup (it is not clear from the code what is it that makes it more efficient)
  It would be nice if you can also leave some comments in the code that indicates this in order to help other people understand why certain things are done the way they are now.

[ciyongch]

@asmushetzel thanks for your comments.

1. Actually, I did some other tests with different configuration when calling topk Op, which might help to understand how much improvement we could achieve with this enhancement. Here's more performance data I collected between optimized version and out-of-box version. The number in speedup columns are speedup factor, the bigger the better. We can see that all the cases shows good, especially the case when the axis=3(equal to -1 since the input is 4d ndarray).

		ret_type=value	ret_type=indices	ret_type=both	ret_type=mask
axis	k	speedup	speedup	speedup	speedup
0	1	2.52	2.60	2.49	1.00
0	3	2.54	2.61	2.63	1.04
0	5	2.58	2.54	2.59	1.16
0	10	2.52	2.58	2.58	1.09
1	1	3.72	3.69	3.73	1.04
1	3	3.75	3.77	3.66	1.10
2	1	2.40	2.35	2.30	1.11
2	3	2.34	2.31	2.24	1.10
2	5	2.29	2.30	2.24	1.19
2	10	2.33	2.33	2.34	1.09
2	100	2.36	2.43	2.40	1.10
2	500	2.51	2.49	2.43	1.10
2	1000	2.44	2.37	2.34	1.04
3	1	67.29	70.23	61.32	5.22
3	3	64.53	64.05	61.82	5.32
3	5	61.06	61.72	52.28	4.85
3	10	43.06	50.45	49.08	4.28
3	100	21.83	22.04	21.10	2.44
3	500	10.36	10.35	10.30	1.55
3	1000	8.54	9.10	8.08	1.35

2. Regarding the code changes, I did two minor changes and kept the main logic same as current version.

   • only do the modulus calculation to ret_indices instead of full indices after calling TopkSort() function, which could reduce many redundant calculation when k is smaller than element number. So I changed indices = F<mshadow_op::mod>(indices, element_num); to ret_indices = F<mshadow_op::mod>(ret_indices, element_num);. That's why the results from ret_type=[value, indices, both] shows better than ret_type=mask.
   • in the case of axis=-1, there's no need to do the transpose to the source(input) data, but only requires a flatten to 1D operation. FlatTo1D is much efficient than reshape since it only changes the Shape but not result in data movement. I noticed that in current CPU TopkSort() function, the work Tensor is only used to stored the temporary result and then copy its content back to sorted_data. If we pass the flattened data via work Tensor to TopkSort(), and copy the top k number of data from such flattened data into sorted_data directly, then there's no need to keep additional temporal Tensor and do the copy second time. For the case axis!=-1, both transpose and reshape are required to convert a 3D Tensor to 1D Tensor. That's why the results from axis=-1 are much better than other cases.

Hope this answers your question, I could add some comments in the code if there's no other concerns.

[pengzhao-intel]

@asmushetzel any other comments?

@marcoabreu @szha could you help take a review and merge the PR? we hope it can be involved into 1.3 and then sockeye can get lots of benefits.

[asmushetzel]

@ciyongch Nice work! Reviewed it and looks good to me.

According to the code changes, the massive speedup for axis=3 and ret_type=indices/values/both is mainly attributed to the change concerning the modulo-computation. This means that a single modulo operation on the indices somehow attributed to 90% of the entire runtime, which is pretty insane. Is this a weak spot of the CPU-design? Or because mod-operations do not benefit from AVX (just speculating)? Can you guys from Intel comment on this a bit more?

[marcoabreu]

Thanks for this great improvement! Unfortunately, I'm currently quite swamped and unable to review this pull request.

[szha]

@asmushetzel thanks for the review. @pengzhao-intel I will take a look shortly.

[pengzhao-intel]

@asmushetzel Thanks for the comments.
Actually, the modulo instruction is much slower than ADD, MUL, FMA in the popular hardware, such as CPU and GPU. In practice, we have to avoid the redundant modulo operation from the software side :)

@szha, Thanks, btw, your new head portrait is really cool and looks like the big boss :)

[ciyongch]

@asmushetzel Sorry for the late response, as @pengzhao-intel mentioned above, the modulo and division instruction are not efficient as ADD/MUL. For topk Op, both transpose and modulo are most time-consuming operations, and then the reshape operation, while the TopkSort() function depends on the value of k.
In the case of axis=3 and k=1, it's no need to do transpose, so the modulo operation could take most of the time, but its proportion will decrease as k grows.
In the case of ret_type=value/indices/both and axis=0/1/2, both transpose and modulo are executed, after removing redundant modulo operation, we can see there's some speedup but not much as axis=3 due to another time-consuming operation transpose.
Since many of the neural networks use the case of axis=3, they will get benefit a lot for this improvement.

[szha]

LGTM. We may want to add dtype template to topk when there are relevant use cases.

[ciyongch]

@szha Thanks for your review :)
Once there's dtype template for topk, then the temp_workspace Tensor on CPU could be simplified.