Swap detection order in randperm_out_cuda to avoid unnecessary conversion from float when the input is small.

[xuhdev]

Stack from ghstack:

• Refactor and improve randperm tests. #22121 Refactor and improve randperm tests.
• Swap detection order in randperm_out_cuda to avoid unnecessary conversion from float when the input is small. #22103 Swap detection order in randperm_out_cuda to avoid unnecessary conversion from float when the input is small.
• Support Half type in randperm. #22102 Support Half type in randperm.

Previously, when n is small and dtype is not Half, randperm on cuda
would offload to CPU with a Float type, which has been changed to Half
type in this commit.

This commit basically swaps the following two blocks:

if (result.scalar_type() == at::ScalarType::Half)

and

if (n < 30000)

Differential Revision: D16153585

[xuhdev]

@pytorchbot retest this please

[xuhdev]

@pytorchbot retest this please

[xuhdev]

@pytorchbot retest this please

[xuhdev]

@gchanan @ngimel @yf225 @li-roy Could you give a quick review? This is actually a pretty simple change, if you view the diff without whitespace changes: https://github.com/pytorch/pytorch/pull/22103/files?w=1

[ngimel]

can we have some performance benchmarks? I'm concerned about int64_t to half conversion on the CPU that's happening for small tensors as a result of this PR

pytorch/aten/src/ATen/native/TensorFactories.cpp

Line 488 in 0df39bb

r__data[i*r__stride_0] = static_cast<scalar_t>(i);

- usually those are slow, slower than doing things in float on cpu and then copying to half.
Also, initialTensorOptions() implicitly set scalar type of the tensor to float, so it works now, but the comment in the file says that it's not a stable API, and what if that scalar type ever changes? https://github.com/pytorch/pytorch/blob/4453a1ff887dec226355b375d4f1bfa1eb016728/aten/src/ATen/InitialTensorOptions.h
I'd prefer scalar type to be explicitly set to kFloat.

[xuhdev]

@ngimel The conversion from int64_t to float is checked in #22102: https://github.com/pytorch/pytorch/pull/22102/files#diff-37ce10604989fdc0a02a62d4949658b2

I have added an explicit dtype.

Is there a way to trigger the benchmark?

[ngimel]

I don't think there are ready benchmarks that you can trigger, you can use benchmarks similar to those in original issue #7606, for before and after your PR.
Don't forget to run a few warmup iterations for the sizes you are going to benchmark to make sure allocator settles, and synchronize after cuda benchmarks. Since the changes are for small tensors only, looks like you should be fine, int64_t -> half conversion should not slow you down too much compared to exposed h2d copy latencies.

[xuhdev]

Here's the perf. It looks like this patch significantly improves the performance for half on cuda. I turned off CPU turbo boost and always ran three times the benchmark (warmup) before the one that is used, so the results should be reliable.

before this patch:

randperm(10) 100000 times
cpu, half	2.8076633399905404
cpu, float	2.7852712649910245
cpu, double	2.7890169029997196
cuda, half	18.277449170011096
cuda, float	10.359675675004837
cuda, double	10.433411638994585
randperm(1000) 1000 times
cpu, half	0.09014522400684655
cpu, float	0.06891961999644991
cpu, double	0.06935312499990687
cuda, half	0.2106321010069223
cuda, float	0.14546358000370674
cuda, double	0.1470240909984568

after this patch:

randperm(10) 100000 times
cpu, half	2.798496307004825
cpu, float	2.8195796069921926
cpu, double	2.827226548004546
cuda, half	11.731637821998447
cuda, float	10.317510107008275
cuda, double	10.415388538996922
randperm(1000) 1000 times
cpu, half	0.09070123700075783
cpu, float	0.06949712400091812
cpu, double	0.06988387599994894
cuda, half	0.16773680900223553
cuda, float	0.14579828100977466
cuda, double	0.14784032000170555

[xuhdev]

I also realized that the if (result.scalar_type() == at::ScalarType::Half) block should never be reached because if n>=30000 and the type is Half, check_supported_max_int_with_precision should have reported an error. Instead of removing this block completely, I left the code here for the sake of clarity (because half in thrust is spotty, and we do not want future change unaware of this).

[facebook-github-bot]

@li-roy merged this pull request in 32709af.