Here are the results  with proper benchmarks! It looks like generating random permutations on the GPU is still the way to go, if you want to generate indices for random selection. However, the GPU methods do not scale quite as well as it seemed before.
I couldn’t find a good way to access the benchmark results, so I settled for timeit(N).raw_times[0], which seems to give the median time spent.

randperm760×455 47.9 KB

import numpy as np
import torch
import random
import pandas as pd
import matplotlib.pyplot as plt
import torch.utils.benchmark as benchmark

def rand(pop_size, num_samples):
    """Use random.sample to generate indices."""
    return torch.Tensor(random.sample(range(pop_size), num_samples)).to('cuda')

def mult_cpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a CPU tensor."""
    p = torch.ones(pop_size) / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False).to('cuda')

def mult_gpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a GPU tensor."""
    p = torch.ones(pop_size, device='cuda') / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False)

def perm_cpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a CPU tensor."""
    return torch.randperm(pop_size)[:num_samples].to('cuda')

def perm_gpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a GPU tensor."""
    return torch.randperm(pop_size, device='cuda')[:num_samples]

def perm_gpu_f32(pop_size, num_samples):
    """Use torch.randperm to generate indices on a 32-bit GPU tensor."""
    return torch.randperm(pop_size, dtype=torch.int32, device='cuda')[:num_samples]

def sort_rand_cpu(pop_size, num_samples):
    """Generate a random torch.Tensor (CPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size))[:num_samples]

def sort_rand_gpu(pop_size, num_samples):
    """Generate a random torch.Tensor (GPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size, device='cuda'))[:num_samples]

idx_fns = [rand, 
    mult_cpu, 
    mult_gpu, 
    perm_cpu, 
    perm_gpu,
    perm_gpu_f32,
    sort_rand_cpu,
    sort_rand_gpu]

pop_size = np.logspace(3, 7, 10, dtype=int)

d = []
for n_p in pop_size:
    print(f"Testing functions with {n_p} points.")
    for fn in idx_fns:
        num_samples = int(0.1 * n_p)

        t0 = benchmark.Timer(
            stmt=f"{fn.__name__}(n_p, num_samples)",
            setup=f"from __main__ import {fn.__name__}",
            globals={'n_p': n_p, 'num_samples': num_samples}
        )

        d.append({
            'Population': n_p,
            'Samples': num_samples,
            'Function':fn.__name__,
            'Time': t0.timeit(11).raw_times[0]
        })

df = pd.DataFrame(d)

fig, ax = plt.subplots()

for fn in idx_fns:
    idxs = df['Function'] == fn.__name__
    ax.plot(df['Population'][idxs], df['Time'][idxs], label=fn.__name__)

ax.set_title("Runtime of Index-Generating Methods for Randomly Selecting 10% of Population")
ax.set_xlabel('Population Size')
ax.set_ylabel('Time')
ax.set_yscale('log')
ax.legend(bbox_to_anchor=(1, 0.7))
plt.savefig('figs/randperm.png', bbox_inches='tight')
plt.show()