[TOPI] Rewrite GPU argwhere using exclusive scan

[masahi]

This PR improves the implementation of GPU argwhere added in #6868, using exclusive scan (see #7303).

The current implementation of argwhere is very inefficient, because it uses atomic to update the write location. Since all threads compete for the single location, this effectively makes it a sequential kernel. Moreover, since the output indices need to be lexicographically sorted, the current implementation involves sorting along each axis.

Since argwhere is literally an instance of stream compaction, this is a perfect application of exclusive scan. Now, argwhere simply consists of

• A single call to exclusive scan on a boolean flag array to compute the write indices.
• Compaction using the write indices (just copying elements with nonzero condition).

both of which are highly parallel operation. Thus, both atomic and sort are gone, vastly simplifying the implementation. Moreover, it also brings huge speed up, as shown below.

All numbers in milli sec

Shape	current main	using thrust exclusive scan	using TIR exclusive scan
(128, 65)	0.240	0.024	0.132
(200, 500)	0.571	0.029	0.184
(1000, 50)	0.270	0.027	0.154
(100000,)	0.205	0.027	0.182
(500000,)	0.947	0.090	0.418
(1000000,)	1.63	0.170	0.848
(1000, 1000)	3.17	0.185	0.863
(1000, 5000)	15.69	0.976	3.689
(100, 100, 1000)	48.70	2.087	7.423
(256, 128, 64, 32)	397.21	15.58	51.438

please review @zhiics @Laurawly @mbrookhart @tkonolige @anijain2305 @trevor-m

[mbrookhart]

Could we add a column for the performance of the PR without thrust (i.e., TIR exclusive scan?)

[mbrookhart]

I'd like to include benchmarks without thrust in the PR for posterity, but otherwise this looks great, thanks! I'd wait to merge until @zhiics can review, since he wrote the existing kernel.

[masahi]

Ok updated the numbers to include TIR scan result.

[mbrookhart]

👍 Not as fast as thrust, as expected, but it's good to see it's still a performance improvement.

[zhiics]

Thanks for the improvement.

[masahi]

Thanks @mbrookhart @zhiics