pytorch

8a744c31 - Up to 48% speed up using Sklansky method for innermost prefix scan algorithm (#103314)

Up to 48% speed up using Sklansky method for innermost prefix scan algorithm (#103314) I found this algorithm (Sklansky) could provide speed-up over the previously implemented Brent-Kung (BK) algorithm. In BK algorithm, the sweeps are done twice: up-sweep and down-sweep. In up-sweep, initially all threads are working, but then half of the working threads becomes inactive in the subsequent step. Similarly for down-sweep but the other way around, where it initially starts with only 1 working thread and double the number of working threads for each sweep. This results of half of the thread is idle on average and produces `2 * log2(num_threads_x)` sweep steps. On the other hand, Sklansky algorithm only use 1 sweep and in each step of the sweep, all the threads are working. This algorithm also produces `log2(num_threads_x)` sweep steps which is half of the BK algorithm. That provides the speed up. I follow the schematics of Sklansky algorithm provided in [this paper](https://research.nvidia.com/sites/default/files/pubs/2016-03_Single-pass-Parallel-Prefix/nvr-2016-002.pdf). The same paper provides a much better algorithm (the one implemented in CUB), but I haven't got my head around it, while the Sklansky algorithm is easier to digest and implement. Here are the speed up from my experiment using `cumsum` in the innermost dimension using A100: (UPDATE: the newest commit further optimize it up to 76% on `8 x 4000` matrix) (UPDATE: added shapes with 2048 and 1M in its elements) | Shape | Torch cumsum | Custom cumsum | Speed up | |--------------|---------------------------|--------------------------|---------------------| | (2, 1000) | 4.8112869262695315e-05 | 2.849102020263672e-05 | 1.688702928870293 | | (8, 4000) | 0.00017731189727783204 | 0.0001005411148071289 | 1.7635760018970834 | | (128, 10000) | 0.0005342483520507813 | 0.00035474300384521487 | 1.5060151891928222 | | (1024, 20000)| 0.0014238595962524415 | 0.0010990619659423829 | 1.2955225823246128 | | (1024, 100000)| 0.007089591026306153 | 0.005468320846557617 | 1.296484099093993 | | (2048, 1000000)| 0.058730244636535645 | 0.0458010196685791 | 1.2822912035913994 | | (1000, 2) | 1.0919570922851562e-05 | 8.106231689453125e-06 | 1.3470588235294116 | | (4000, 8) | 9.512901306152343e-06 | 7.867813110351562e-06 | 1.209090909090909 | | (10000, 128) | 2.079010009765625e-05 | 1.6164779663085937e-05 | 1.2861356932153394 | | (20000, 1024)| 0.00024993419647216796 | 0.00017964839935302734 | 1.3912408759124086 | | (100000, 1024)| 0.0011160612106323243 | 0.0009322404861450195 | 1.1971816577581138 | | (1000000, 2048) | 0.017030668258666993 | 0.014445066452026367 | 1.178995494082889 | Pull Request resolved: https://github.com/pytorch/pytorch/pull/103314 Approved by: https://github.com/ngimel