[lang] Add parallel sort utility

[AmesingFlank]

This is an implementation of parallel odd-even merge sort.

related issue = #3764

References:

https://developer.nvidia.com/gpugems/gpugems2/part-vi-simulation-and-numerical-algorithms/chapter-46-improved-gpu-sorting
https://en.wikipedia.org/wiki/Batcher_odd%E2%80%93even_mergesort

This will be very useful for #3781 as well. We likely won't be having atomics in WebGPU (sadge), but one can always use a sort to simulate atomic operations. To demonstrate this, I will open another PR to add versions of mpm99.py and pbf2d.py that are atomics-free but uses sorting.

[netlify]

✔️ Deploy Preview for jovial-fermat-aa59dc ready!

🔨 Explore the source changes: 39d1344

🔍 Inspect the deploy log: https://app.netlify.com/sites/jovial-fermat-aa59dc/deploys/61b8ebe9b393ed00078c1e2f

😎 Browse the preview: https://deploy-preview-3790--jovial-fermat-aa59dc.netlify.app

[bobcao3]

Maybe add an option where we can sort based on a sort key? Which takes two arrays (N, k), which will make this a lot more useful I'd imagine. (e.g. sorting triangles based on depth or Morton code sort key for BVH)

[qiao-bo]

Hi, thanks a lot for the quick implementation. I wonder do we have another place to put these parallel intrinsics? Since more will be added later, IMHO looking at the tools folder, they seem not fitting into the same category ;)

[AmesingFlank]

Hi, thanks a lot for the quick implementation. I wonder do we have another place to put these parallel intrinsics? Since more will be added later, IMHO looking at the tools folder, they seem not fitting into the same category ;)

Yeah I'm not sure if this is right place either. I chose the tools folder because this folder contains things that are not core to the taichi language, but rather utilities and helpers, and that's what this parallel_sort function is. I'm happy to move it to another place though.

@ailzhang @k-ye @qiao-bo suggestions on where to put this?

[qiao-bo]

Hi, thanks a lot for the quick implementation. I wonder do we have another place to put these parallel intrinsics? Since more will be added later, IMHO looking at the tools folder, they seem not fitting into the same category ;)

Yeah I'm not sure if this is right place either. I chose the tools folder because this folder contains things that are not core to the taichi language, but rather utilities and helpers, and that's what this parallel_sort function is. I'm happy to move it to another place though.

@ailzhang @k-ye @qiao-bo suggestions on where to put this?

We have some linalg implementations at python/lang/linalg_impl.py, such as ti.svd(). Maybe also in lang? But i am not sure if we will move these impl to C++ IR in the future?

[AmesingFlank]

I don't really think it belongs in lang, as it is not part of the language, but rather just a utility function. And I don't think we will be moving it to C++ IR, it's a lot less readable.

[qiao-bo]

ehmm, in this case it is a utility helper correct. I thought we could offer functions such as parallel sort, reduce, scan... as the parallel intrinsics offered by the Taichi language. Nevertheless, i am also fine with keeping it here at tools for now, unless someone has a better suggestion ;/

[AmesingFlank]

I thought we could offer functions such as parallel sort, reduce, scan... as the parallel intrinsics offered by the Taichi language.

That's a pretty good idea. We should group these into the same directory once we also have an API for reduce and scan. But for now I think let's just leave it in tools.

[AmesingFlank]

Can i please get a stamp :) @qiao-bo @ailzhang @k-ye @bobcao3

[qiao-bo]

Thanks!