Update masked_select and nonzero

[tongxin]

PR Category

Operator

Type of Change

Refactor & optimization

Description

This PR curbs autotuning for masked_select and nonzero, to avoid confusion during testing and debugging. masked_select is optimized in the following way:

• work requires at most two kernel launches
• compute partial sums of mask and cumsum over the partial sums in a single kernel
• write back the selected input data in the second kernel.
• no extra total sum kernel involved

Performance gains are clear for larger input sizes.

Operator: masked_select  Performance Test (dtype=torch.float16, mode=cuda,level=comprehensive)
Status       Torch Latency (ms)    Gems Latency (ms)         Gems Speedup          Torch GBPS            Gems GBPS           Size Detail
-----------------------------------------------------------------------------------------------------------------------------------------
SUCCESS               0.072576            0.085280               0.851               0.127               0.108          [torch.Size([64, 64]), torch.Size([64, 64])]
SUCCESS               0.077504            0.083840               0.924               1.893               1.750          [torch.Size([256, 256]), torch.Size([256, 256])]
SUCCESS               0.074592            0.086240               0.865              31.434              27.189          [torch.Size([1024, 1024]), torch.Size([1024, 1024])]
SUCCESS               0.574176            0.229568               2.501              65.324             163.382          [torch.Size([4096, 4096]), torch.Size([4096, 4096])]
SUCCESS               2.100032            0.718752               2.922              71.444             208.744          [torch.Size([1024, 65536]), torch.Size([1024, 65536])]
SUCCESS               0.117664            0.087264               1.348              48.628              65.569          [torch.Size([10000, 256]), torch.Size([10000, 256])]
SUCCESS              19.840832            7.250912               2.736              73.850             202.076          [torch.Size([10000, 65536]), torch.Size([10000, 65536])]


Operator: masked_select  Performance Test (dtype=torch.float32, mode=cuda,level=comprehensive)
Status       Torch Latency (ms)    Gems Latency (ms)         Gems Speedup          Torch GBPS            Gems GBPS           Size Detail
-----------------------------------------------------------------------------------------------------------------------------------------
SUCCESS               0.070848            0.080000               0.886               0.129               0.115          [torch.Size([64, 64]), torch.Size([64, 64])]
SUCCESS               0.073216            0.081664               0.897               2.000               1.793          [torch.Size([256, 256]), torch.Size([256, 256])]
SUCCESS               0.075488            0.088096               0.857              31.062              26.616          [torch.Size([1024, 1024]), torch.Size([1024, 1024])]
SUCCESS               0.601920            0.247936               2.428              62.313             151.278          [torch.Size([4096, 4096]), torch.Size([4096, 4096])]
SUCCESS               2.215040            0.788672               2.809              67.739             190.250          [torch.Size([1024, 65536]), torch.Size([1024, 65536])]
SUCCESS               0.126432            0.090624               1.395              45.270              63.157          [torch.Size([10000, 256]), torch.Size([10000, 256])]
SUCCESS              20.986303            8.212896               2.555              69.821             178.411          [torch.Size([10000, 65536]), torch.Size([10000, 65536])]


Operator: masked_select  Performance Test (dtype=torch.bfloat16, mode=cuda,level=comprehensive)
Status       Torch Latency (ms)    Gems Latency (ms)         Gems Speedup          Torch GBPS            Gems GBPS           Size Detail
-----------------------------------------------------------------------------------------------------------------------------------------
SUCCESS               0.070624            0.080832               0.874               0.130               0.114          [torch.Size([64, 64]), torch.Size([64, 64])]
SUCCESS               0.074048            0.082560               0.897               1.977               1.773          [torch.Size([256, 256]), torch.Size([256, 256])]
SUCCESS               0.082176            0.084704               0.970              28.546              27.694          [torch.Size([1024, 1024]), torch.Size([1024, 1024])]
SUCCESS               0.574752            0.228224               2.518              65.262             164.353          [torch.Size([4096, 4096]), torch.Size([4096, 4096])]
SUCCESS               2.098048            0.719136               2.917              71.514             208.640          [torch.Size([1024, 65536]), torch.Size([1024, 65536])]
SUCCESS               0.117600            0.087136               1.350              48.672              65.689          [torch.Size([10000, 256]), torch.Size([10000, 256])]
SUCCESS              19.857281            7.249600               2.739              73.786             202.107          [torch.Size([10000, 65536]), torch.Size([10000, 65536])]

Issue

#718 #712

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

[iclementine]

So this kernel computes the CTA-level exclusive-prefix-scan.

[iclementine]

LGTM

[iclementine]

Removing device guard may cause inaccessible address.