Looks correct. @t-vi would you mind running some quick perf numbers just to characterize what the effect of the extra neq test is?

My conclusion would be measurable in isolation (in the order 10%-20% slowdown), negligible in any other context. Just like max probably wasn't the bottleneck before, doing twice as many comparisons in max doesn't really matter if you do any nontrivial stuff elsewhere.

So I run this several times (to avoid init topics, the first run takes longer than the others). I lifted this from somewhere in the pytorch/benchmark repository - if you have a better methodology, I'd be happy to apply it.

import gc, torch, time
with torch.no_grad():
    a = torch.empty(100, 1000, 1000, device='cuda')
    a.normal_()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    gc.collect()
    torch.cuda.synchronize()
    start.record()
    start_cpu_secs = time.time()
    b = a.max()
    end_cpu_secs = time.time()
    end.record()
    torch.cuda.synchronize()
    gpu_msecs = start.elapsed_time(end)
    print(torch.__version__, "msecs maxall gpu", gpu_msecs, "cpu", (end_cpu_secs - start_cpu_secs)*1000)
    if 1:
        a = torch.empty(100, 1000, 1000, device='cuda')
        a.normal_()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        gc.collect()
        torch.cuda.synchronize()
        start.record()
        start_cpu_secs = time.time()
        b = a.max(1)
        end_cpu_secs = time.time()
        end.record()
        torch.cuda.synchronize()
        gpu_msecs = start.elapsed_time(end)
        print(torch.__version__, "msecs max[1] gpu", gpu_msecs, "cpu", (end_cpu_secs - start_cpu_secs)*1000)

And I get:

0.4.0 msecs maxall gpu 1.6861120462417603 cpu 1.6646385192871094
0.4.0 msecs max[1] gpu 2.239487886428833 cpu 0.07939338684082031

vs.

0.5.0a0+497ff06 msecs maxall gpu 1.8479679822921753 cpu 1.8224716186523438
0.5.0a0+497ff06 msecs max[1] gpu 1.922752022743225 cpu 0.054836273193359375

But what's the actual fix in this PR then? If we want to have kernels return NaNs when there are NaNs, then why don't we treat CPU kernels that way

Ok, never mind my comment. I did the math again and it seems to be ok.

(Sorry, the second mail doesn't seem to have reached the issue log... :( )
Oh, it is roughly the same, CPU does the comparison first and then checks if value was NaN (then theMax is NaN because how the ">" is formulated) and if it is, it aborts and theMax will be NaN.
Previously, the Cuda kernel would put NaN into the accumulator (a in how the cuda function is called, theMax in CPU) and then overwrite it when doing the next comparison.
With the proposed fix, the cuda kernel form checks if the left hand side (a - equivalent to theMax in how it is used) is NaN and if it is, it keeps a. I think this is result-equivalent to how the CPU works.