UpSample-nearest cuda kernel update

[jjsjann123]

updating upsampling kernel:

1. avoids atomicAdd for better fp16 performance.
2. better launch configures for 2D input.

[jjsjann123]

Perf number/scripts will be posted shortly.
I'll have another PR for bilinear upsampling coming soon as well.

cc @ngimel

[jjsjann123]

fp16 forward perf number has been observed to be all over the place, especially for tiny input. :/
Things to notice here is the speedup on the fp16 backward path.

Here's the script for the benchmark

import torch
import numpy as np

nrep = 300
sample_mode = "nearest"

def bench(size, fn, factor, half=False, b=2, c=32, dim=2):
   x=torch.ones([b*c*(size**dim)], device='cuda', dtype = torch.float)
   if half:
     x = x.half()
   if dim==1:
     x=x.view(b, c, size).requires_grad_()
   elif dim==2:
     x=x.view(b, c, size, size).requires_grad_()
   elif dim==3:
     x=x.view(b, c, size, size, size).requires_grad_()
   torch.cuda.synchronize()
   import time
   start = time.time()
   for i in range(nrep):
      out = fn(x, scale_factor=factor, mode=sample_mode)
   torch.cuda.synchronize()
   end = time.time()
   inp_size = x.size()
   out_size = out.size()
   del x, out
   return ((end-start)/nrep, inp_size, out_size)

def bench_back(size, fn, factor, half=False, b=2, c=32, dim=2):
   x=torch.ones([b*c*(size**dim)], device='cuda', dtype = torch.float)
   if half:
     x = x.half()
   if dim==1:
     x=x.view(b, c, size).requires_grad_()
   elif dim==2:
     x=x.view(b, c, size, size).requires_grad_()
   elif dim==3:
     x=x.view(b, c, size, size, size).requires_grad_()
   torch.cuda.synchronize()
   out = fn(x, scale_factor=factor, mode=sample_mode)
   grad = torch.randn_like(out)
   import time
   start = time.time()
   for i in range(nrep):
      out.backward(grad)
   torch.cuda.synchronize()
   end = time.time()
   inp_size = x.size()
   out_size = out.size()
   del x, out, grad
   return ((end-start)/nrep, inp_size, out_size)


spatial_size = [2**i for i in range(5,12)]
batch = [8]
channel = [32]
dim = [1, 2, 3]
scale_factor = [2]
bool_flag = [False, True]
cap = 2**31

from itertools import product
for d, b, c, s, f, half_flag in product(dim, batch, channel, spatial_size, scale_factor, bool_flag):
  if ((s*f)**d)*b*c*(2 if half_flag else 4) < cap:
    (fw_time, inp_size, out_size) = bench(s, torch.nn.functional.interpolate, f, half_flag, b, c, d)
    (bw_time, inp_size, out_size) = bench_back(s, torch.nn.functional.interpolate, f, half_flag, b, c, d)
    print(inp_size, f, half_flag, 1./fw_time, 1./bw_time)

[jjsjann123]

Removed the specialized 2d kernel, as the speedup is sparse. Caching seems to have done a great job saving the memory accessing pattern.

I don't think I can justify having a dedicated kernel there 😢

[ezyang]

@ngimel happy to merge this if you give it the OK

[ngimel]

Make sure there are checks against empty tensors, and make sure you are not excessively zeroing outputs. Those already might be somewhere in the code, and I might be blind, in which case it is good to go.

[jjsjann123]

Addressed review comments. Should be good to go when test passes

[ngimel]

Review comments are addressed, great job, Jie!

[facebook-github-bot]

@ezyang is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@ezyang merged this pull request in c471a63.