Repository: pytorch
Issue: TorchScript GPU Fuser Doesn't Handle In-Place Operations
Labels: module: performance,oncall: jit,triaged

The performance of a TorchScript that handles the computation in the [Adam Optimizer](https://github.com/pytorch/pytorch/blob/master/torch/optim/adam.py) is worse than using the native PyTorch operations on CUDA. 

### Background
I am trying to improve the GPU performance of the [Adam Optimizer](https://github.com/pytorch/pytorch/blob/master/torch/optim/adam.py) since it is the bottleneck in one of our workflows. 

### Results
I have written a standalone script that compares PyTorch native tensor operations, TorchScript, and Numba (three runs each):
```
native :  0.0506
native :  0.0322
native :  0.0322
TScript:  0.1917
TScript:  0.0442
TScript:  0.0442
numba  :  0.1638
numba  :  0.0145
numba  :  0.0143
```
I suspect the problem is that TorchScript does not fuse all operations into a single CUDA. It generates two CUDA kernels (see the `PYTORCH_FUSION_DEBUG` output below)

### Standalone Script
```python
import math
import torch
import time
import random
import sys

def kernel_native(param, grad, exp_avg, exp_avg_sq, beta1, beta2, step_size, eps):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float, float, float) -> None
    exp_avg.mul_(beta1).add_(1 - beta1, grad)
    exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
    denom = exp_avg_sq.sqrt().add_(eps)
    param.addcdiv_(-step_size, exp_avg, denom)   

@torch.jit.script
def kernel_jit_script(param, grad, exp_avg, exp_avg_sq, beta1, beta2, step_size, eps):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float, float, float) -> None
    exp_avg.mul_(beta1)
    exp_avg.add_((1 - beta1) + grad)
    exp_avg_sq.mul_(beta2)
    exp_avg_sq.add_((1 - beta2)*(grad*grad))
    denom = exp_avg_sq.sqrt().add_(eps)
    param += -step_size * (exp_avg / denom)

from numba import cuda
@cuda.jit
def kernel_numba_cuda(param, grad, exp_avg, exp_avg_sq, beta1, beta2, step_size, eps):
    i = cuda.grid(1)
    if i > grad.size:
        return

    exp_avg[i] = exp_avg[i] * beta1 + (1 - beta1) + grad[i]
    exp_avg_sq[i] = exp_avg_sq[i] * beta2 + (1 - beta2) *grad[i]*grad[i]
    denom = math.sqrt(exp_avg_sq[i]) + eps
    param[i] += -step_size * (exp_avg[i] / denom)

def kernel_numba(param, grad, exp_avg, exp_avg_sq, beta1, beta2, step_size, eps):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float, float, float) -> None

    import ctypes
    import numpy as np
    def get_devicendarray(t):
        assert t.type() == 'torch.cuda.FloatTensor'
        numpy_dtype = np.float32
        itemsize = 4

        ctx = cuda.cudadrv.driver.driver.get_context()
        mp = cuda.cudadrv.driver.MemoryPointer(ctx, ctypes.c_ulong(t.data_ptr()), t.numel()*itemsize)
        return cuda.cudadrv.devicearray.DeviceNDArray(t.numel(), itemsize, numpy_dtype(), gpu_data=mp, stream=torch.cuda.current_stream().cuda_stream)

    numba_param = get_devicendarray(param)
    numba_grad = get_devicendarray(grad)
    numba_exp_avg = get_devicendarray(exp_avg)
    numba_exp_avg_sq = get_devicendarray(exp_avg_sq)
    threadsperblock = 32
    blockspergrid = (param.numel() + (threadsperblock - 1)) // threadsperblock
    kernel_numba_cuda[blockspergrid, threadsperblock](numba_param, numba_grad, numba_exp_avg, numba_exp_avg_sq, beta1, beta2, step_size, eps)

kernels = [("native ", kernel_native), ("TScript", kernel_jit_script), ("numba  ", kernel_numba)]
shape = (4194304*64,)
param, grad, exp_avg, exp_avg_sq = [torch.rand(shape).to("cuda") for x in range(4)]
beta1, beta2, step_size, eps = [random.random() for x in range(4)]
torch.cuda.synchronize()

for name, func in kernels:
    for i in range(3):
        t1 = time.time()
        func(param, grad, exp_avg, exp_avg_sq, beta1, beta2, step_size, eps)
        torch.cuda.synchronize()
        t2 = time.time()
        print("%s: " % name, t2-t1)
```

### Output of PYTORCH_FUSION_DEBUG
```cu
fusion code:

typedef unsigned char uint8_t;
typedef signed char int8_t;
typedef short int  int16_t;
typedef long long int int64_t;



#define NAN __int_as_float(0x7fffffff)
#define POS_INFINITY __int_as_float(0x7f800000)
#define NEG_INFINITY __int_as_float(0xff800000)

typedef unsigned int IndexType;
template<typename T, size_t N>
struct TensorInfo {
  T* data;
  IndexType sizes[N];
  IndexType strides[N];
};
template<typename T>
struct TensorInfo<T, 0> {
  T * data;
};


extern "C" __global__
void kernel_0(IndexType totalElements, const TensorInfo<float,1> t0, double s1, const TensorInfo<float,1> t2 ) {
  
  for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
        linearIndex < totalElements;
        linearIndex += gridDim.x * blockDim.x) {
      // Convert `linearIndex` into an offset of tensor:
      IndexType t0_offset = 0;
      IndexType t0_linearIndex = linearIndex;
      
      //printf("tensor t0 sizes[0] = %d, strides[0] = %d\n", t0.sizes[0],t0.strides[0]);
      size_t t0_dimIndex0 = t0_linearIndex ;
      t0_offset += t0_dimIndex0 ;
      IndexType t2_offset = 0;
      IndexType t2_linearIndex = linearIndex;
      
      //printf("tensor t2 sizes[0] = %d, strides[0] = %d\n", t2.sizes[0],t2.strides[0]);
      size_t t2_dimIndex0 = t2_linearIndex ;
      t2_offset += t2_dimIndex0 ;
      
      // calculate the results
      float n0 = __ldg(&t0.data[t0_offset]);
      double n1 = s1;
      float n2 = n0 * n0;
      float n3 = n2 * ((float) n1);
      t2.data[t2_offset] = n3;
      
    }
}

fusion code:

typedef unsigned char uint8_t;
typedef signed char int8_t;
typedef short int  int16_t;
typedef long long int int64_t;



#define NAN __int_as_float(0x7fffffff)
#define POS_INFINITY __int_as_float(0x7f800000)
#define NEG_INFINITY __int_as_float(0xff800000)

typedef unsigned int IndexType;
template<typename T, size_t N>
struct TensorInfo {
  T* data;
  IndexType sizes[N];
  IndexType strides[N];
};
template<typename T>
struct TensorInfo<T, 0> {
  T * data;
};


extern "C" __global__
void kernel_1(IndexType totalElements, const TensorInfo<float,1> t0, const TensorInfo<float,1> t1, double s2, const TensorInfo<float,1> t3 ) {
  
  for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
        linearIndex < totalElements;
        linearIndex += gridDim.x * blockDim.x) {
      // Convert `linearIndex` into an offset of tensor:
      IndexType t0_offset = 0;
      IndexType t0_linearIndex = linearIndex;
      
      //printf("tensor t0 sizes[0] = %d, strides[0] = %d\n", t0.sizes[0],t0.strides[0]);
      size_t t0_dimIndex0 = t0_linearIndex ;
      t0_offset += t0_dimIndex0 ;
      IndexType t1_offset = 0;
      IndexType t1_linearIndex = linearIndex;
      
      //printf("tensor t1 sizes[0] = %d, strides[0] = %d\n", t1.sizes[0],t1.strides[0]);
      size_t t1_dimIndex0 = t1_linearIndex ;
      t1_offset += t1_dimIndex0 ;
      IndexType t3_offset = 0;
      IndexType t3_linearIndex = linearIndex;
      
      //printf("tensor t3 sizes[0] = %d, strides[0] = %d\n", t3.sizes[0],t3.strides[0]);
      size_t t3_dimIndex0 = t3_linearIndex ;
      t3_offset += t3_dimIndex0 ;
      
      // calculate the results
      float n0 = __ldg(&t0.data[t0_offset]);
      float n1 = __ldg(&t1.data[t1_offset]);
      double n2 = s2;
      float n3 = n0 / n1;
      float n4 = n3 * ((float) n2);
      t3.data[t3_offset] = n4;
      
    }
}
```