pytorch

af05158c - torch.multinomial : fast-path for replacement=False (#39636)

torch.multinomial : fast-path for replacement=False (#39636) Summary: Fixes https://github.com/pytorch/pytorch/issues/39624 #11931 Based on the example by RobertoLat https://github.com/pytorch/pytorch/issues/11931#issuecomment-625882503 **Fast-path is not taken on CPU for `Half` as `log` doesn't support it.** Benchmark with same build settings on same system. gcc : version 7.5.0 (Ubuntu 7.5.0-3ubuntu1~18.04) CUDA : 10.1 GPU : 1050ti ```python import time import torch import numpy as np for n, t in [(500_000, 10), (1_000_000, 10)]: for dtype in (torch.half, torch.float, torch.double): # Input Setup p = torch.from_numpy(np.random.rand(n)).to(dtype) want = 1000 print(f'torch.multinomial(a) a.numel() == {n} for {t} times {dtype}') start = time.time() # Iterate for _ in range(t): torch.multinomial(p, want, replacement=False) print(f'Took:', time.time() - start) print('****' * 10) for n, t in [(50_000, 100), (100_000, 100)]: for dtype in (torch.half, torch.float, torch.double): # Input Setup p = torch.rand(n, device='cuda', dtype=dtype) want = 1000 print(f'torch.multinomial(a) a.numel() == {n} for {t} times {dtype}') start = time.time() # torch.cuda.synchronize() # Iterate for _ in range(t): torch.multinomial(p, want, replacement=False) # torch.cuda.synchronize() print(f'CUDA Took:', time.time() - start) ``` Before: ``` torch.multinomial(a) a.numel() == 500000 for 10 times torch.float16 Took: 80.64455389976501 torch.multinomial(a) a.numel() == 500000 for 10 times torch.float32 Took: 3.7778031826019287 torch.multinomial(a) a.numel() == 500000 for 10 times torch.float64 Took: 5.045570611953735 torch.multinomial(a) a.numel() == 1000000 for 10 times torch.float16 Took: 161.53191947937012 torch.multinomial(a) a.numel() == 1000000 for 10 times torch.float32 Took: 7.640851736068726 torch.multinomial(a) a.numel() == 1000000 for 10 times torch.float64 Took: 10.399673461914062 **************************************** torch.multinomial(a) a.numel() == 50000 for 100 times torch.float16 CUDA Took: 4.873984098434448 torch.multinomial(a) a.numel() == 50000 for 100 times torch.float32 CUDA Took: 4.713594436645508 torch.multinomial(a) a.numel() == 50000 for 100 times torch.float64 CUDA Took: 11.167185068130493 torch.multinomial(a) a.numel() == 100000 for 100 times torch.float16 CUDA Took: 7.195427417755127 torch.multinomial(a) a.numel() == 100000 for 100 times torch.float32 CUDA Took: 7.669712066650391 torch.multinomial(a) a.numel() == 100000 for 100 times torch.float64 CUDA Took: 20.20938801765442 ``` After: ``` torch.multinomial(a) a.numel() == 500000 for 10 times torch.float16 Took: 80.6487455368042 torch.multinomial(a) a.numel() == 500000 for 10 times torch.float32 Took: 0.0663309097290039 torch.multinomial(a) a.numel() == 500000 for 10 times torch.float64 Took: 0.09588909149169922 torch.multinomial(a) a.numel() == 1000000 for 10 times torch.float16 Took: 161.60748076438904 torch.multinomial(a) a.numel() == 1000000 for 10 times torch.float32 Took: 0.13187885284423828 torch.multinomial(a) a.numel() == 1000000 for 10 times torch.float64 Took: 0.17609834671020508 **************************************** torch.multinomial(a) a.numel() == 50000 for 100 times torch.float16 CUDA Took: 0.007131099700927734 torch.multinomial(a) a.numel() == 50000 for 100 times torch.float32 CUDA Took: 0.022255420684814453 torch.multinomial(a) a.numel() == 50000 for 100 times torch.float64 CUDA Took: 0.0323028564453125 torch.multinomial(a) a.numel() == 100000 for 100 times torch.float16 CUDA Took: 0.04995012283325195 torch.multinomial(a) a.numel() == 100000 for 100 times torch.float32 CUDA Took: 0.04948878288269043 torch.multinomial(a) a.numel() == 100000 for 100 times torch.float64 CUDA Took: 0.05495333671569824 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/39636 Differential Revision: D21925406 Pulled By: ngimel fbshipit-source-id: f2ee5148fa7dd88e018c461ced0e2361c3a43796