[CUDA] Parallel Cuda Mergesort

[mbrookhart]

@Laurawly @zhiics @icemelon9 @csullivan @tkonolige

There have been many complaints recently about stability and performance of the tir-based cuda sort kernel. I've spent a couple of days this week getting a cuda version of Parallel Mergesort. It's a stable sort, so it fixes the flakiness we've seen with argsort and argwhere, it changes the threading to support dynamic shapes, and it increases the performance significantly over the previous kernel.

This PR only addresses the core sort_ir function, extending this to other versions sort in this file is future work.

I tested performance on a variety of shapes using this script and obtained these numbers on my 1070TI. It's not as fast as Thrust, as expected, but it's much closer for all shapes tested here, and even manages to beat thrust on a few. (times are in milliseconds)

Thanks!

Shape	main	thrust	this
(2000, 2, 2)	7.77	0.58	1.67
(2, 2000, 2)	4.8	0.7	1.59
(2, 2, 2000)	3.24	0.63	1.54
(4000, 2, 2)	25.53	0.65	4.05
(2, 4000, 2)	13.78	0.62	3.3
(2, 2, 4000)	9.85	0.63	4.04
(2, 12000, 2)	369.99	0.68	13.87
(2, 2, 12000)	86.55	0.66	11.11
(12000, 2, 2)	486.65	0.66	13.69
(2000, 8, 8)	259.21	10.4	4.22
(8, 2000, 8)	111.14	8.45	3.43
(8, 8, 2000)	50.37	9.05	3.05
(4000, 8, 8)	671.53	8.24	9.58
(8, 4000, 8)	368.59	8.47	10.12
(8, 8, 4000)	171.18	8.74	6.27
(12000, 8, 8)	3571.97	15.22	42.99
(8, 12000, 8)	3517.72	15.07	45.84
(8, 8, 12000)	1417.97	15.03	27.57

[mbrookhart]

I'm hitting some very odd segfaults, just in the debug runtime with nvptx. Trying to figure out what's going on, I'll keep this as WIP until I can get that fixed.

[tkonolige]

Look great! I think the main implementation might benefit from a couple of comments describing what it is doing.

[csullivan]

Nice approach. I didn't spot an obvious reason for the segfault.. Also agree with @tkonolige on docs, some high level info on the problem flattening and index mapping / thread assignment for each slice (and as slices are merged) will help maintainability by others.

[mbrookhart]

Many thanks to @masahi for helping me find an issue with heterogeneous lowering and some overflow issues in how I was handling the threads. I think it should be ready for review now, thanks everyone!

[masahi]

For the record, segfault with nvptx was happening because the generated host code was calling intrinsics registered for nvptx, like __nv_log2 or __nv_ceil. The reason it was working on CUDA was just by coincident: there is no CUDA intrinsics registered for fp64 log2, ceil, so TVM was using the default lowering, which happens to be the right one (llvm).

This change fixes that issue.

[masahi]

@mbrookhart I think we can revive some tests that are currently disabled due to flaky sort. See

tvm/tests/python/relay/test_any.py

Lines 253 to 256 in bad149e

	# TODO(zhiics) Enable argwhere gpu test after sort is fixed. Otherwise, we have
	# to use thrust to guarantee the correct results which has been tested locally.
	# @tvm.testing.uses_gpu
	def test_any_argwhere():

tvm/tests/python/topi/python/test_topi_argwhere.py

Lines 63 to 71 in 76b4ad0

	# TODO(zhiics) Enable argwhere gpu test after sort is fixed.
	if ctx.device_type != 1:
	continue
	check_device(target, ctx)
	# TODO(zhiics) Enable argwhere gpu test after sort is fixed. Otherwise, we have
	# to use thrust to guarantee the correct results which has been tested locally.
	# @tvm.testing.uses_gpu

[masahi]

We should also remove

tvm/tests/python/topi/python/test_topi_argwhere.py

Lines 63 to 65 in 76b4ad0

	# TODO(zhiics) Enable argwhere gpu test after sort is fixed.
	if ctx.device_type != 1:
	continue

[zhiics]

thanks for the work. please fix the unit test

[mbrookhart]

Oh no! A copy-paste error! Will fix

[Laurawly]

LGTM