[TOPI] GPU sort IR refactor to enable sort by keys

[masahi]

This adds support for sort by key in topi cuda by refactoring the sort IR developed in #7099. By doing different initialization, the exact same sort IR can be reused for sort by key. This is because the original IR already supports argsort and argsort is just a special case of sort by key.

My motivation was to improve the default performance of scatter 1D using sorting based approach (see #7056), when thrust is not available. But it turns out that the sequential scatter kernel is much faster than doing TIR sort, so I dropped that goal for now (this also suggests that thrust's radix sort I used in #7056 is insanely fast). I am sending this PR anyway in the hope that someone would find sort by key useful. My branch that does scatter 1D using the TIR sorting kernel, along with a benchmark script, is available at https://github.com/masahi/tvm/tree/sort_ir_refactor_with_scatter

I also removed sort_nms_ir, since it is not used anymore and I think it would be much faster to sort entire rows using sort_ir anyway.

please review @mbrookhart @Laurawly @zhiics @tkonolige

[mbrookhart]

LGTM. Yes, thrust's radix sort is very good.

Can you share the benchmark numbers? I'm curious.

[mbrookhart]

I definitely didn't promise to be faster than thrust with #7099, just faster that what we had in TIR :)

[masahi]

Ok here it is. I think TIR numbers for 5000 and 10000 agree with the ones you had in #7099

size	current (sequential)	via thrust radix sort	via TIR merge sort
5000	0.169332	0.168660	5.319211
10000	0.301014	0.172381	10.534689
25000	0.730493	0.169007	22.173880
50000	1.457996	0.315582	44.466231
100000	2.943622	0.395291	101.347441
500000	24.262521	1.254097	481.400628

[mbrookhart]

Yep, this matches the scaling I would expect with that algorithm. I'm kind of surprised that sequential scatter isn't worse than it is. I'm looking for ways to improve the mergesort, but I haven't found a good idea that we can implement in TIR yet, what I've seen requries a while loop.

[zhiics]

Thanks @masahi @mbrookhart @Laurawly