CUDA optimization: using `__restrict__` whenever possible

[ThisIsIsaac]

🚀 Feature

Motivation

to increase throughput

Pitch

for every feasible CUDA kernel, take input as THCDeviceTensor<T, DIM, IndexT, RestrictPtrTraits> instead of DefaultPtrTraits in order to add __restrict__ keyword to the device tensors. This alone seems to increase throughput about 3~5%. The increase of throughput is tested with upsampling bilinear code.

Alternatives

Additional context

[ThisIsIsaac]

1. what is scalar_t? I am not sure if it is a template variable or an actual class. I was wondering if I could cast a scalar_t array scalar_t* x into const scalar_t* __restrict__ x ( assuming const and __restrict__'s conditions are met ).

2. some cuda Kernels -- like subsample -- seems to use regular device array while other kernels -- like caffe_gpu_interp2_kernel -- seems to use THCDeviceTensor. In what cases is one preferred over another?

3. CUDA_1D_KERNEL_LOOP(i, n) is equivalent to for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads; index += blockDim.x * gridDim.x) according to the answer here. Is this correct? if so, why use CUDA_1D_KERNEL_LOOP(i, n) over the for-loop?

[VitalyFedyunin]

1. scalar_t is the template variable responsible for the data type
2. We slowly migrating all code to be native and use at::Tensor https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/RangeFactories.cu
3. Because it gets more bulky in case of 2D loops. No other reasons.

Generally speaking, using restrict makes total sense, and we gladly would accept such PR.

[zou3519]

To clarify, THCDeviceTensor is a part of our legacy codebase. The "aten native" analog of it is TensorAccessor, which looks like it already uses some __restrict__ pointers

In general it is easier to read code written with TensorAccessor rather than code using device arrays and a list of strides.

[ThisIsIsaac]

@zou3519
I assume you are talking about RestrictPtrTraits struct defined here. In that case, there is only a single file using RestrictPtrTraits ( i.e. only Normalization.cuh file is using the variable name RestrictPtrTraits excluding files that use it inside comments or files that define RestrictPtrTraits). So I was going to take input not as:

THCDeviceTensor<T, DIM> // defaultPtrTraits without __restrict__

but as

THCDeviceTensor<T, DIM, int, RestrictPtrTraits > // RestrictPtrTraits with __restrict__

---

I thought about changing the default to RestrictPtrTraits and casting to DefaultPtrTraits when we need to make aliases, but maybe that is too dangerous.

[ThisIsIsaac]

@VitalyFedyunin
thanks for the information. I will look into at::Tensor and see how I can safely cast from regular pointer to __restrict__ pointer and const __restrict__.

Do you have some internal documentation regarding these different tensor types? I am spending a few days trying to figure out the different tensor libraries you use, but I am having a difficult time as there seem to be various tensor classes and some don't even use the classes. It would be a huge help and I would really appreciate it!

p.s. do you have any specific kernel in mind that is specially inefficient or takes longer than others? I would like to look into those as well.

thanks again

[ThisIsIsaac]

@zou3519
just wondering if I should still proceed with this. A few clarifications:

1. Is my understanding in the above comment correct that, although there is a type for __restrict__, most tensors are not casted to that type even though they could have been?

2. the new aten code seems to use TensorInfo as a wrapper around device arrays. From looking at struct TensorInfo from TensorInfo.cuh, it seems like there is no way to cast pointer from and to __restrict__. The first step seems to add that functionality.

3. Is there any code that makes aliases and use both pointers within the same kernel? otherwise, we could just make the default pointer have __restrict__ tag.

Let me know what you think.

Thanks as always.

[ThisIsIsaac]

@zou3519 @VitalyFedyunin
ping

[ThisIsIsaac]

@zou3519 @VitalyFedyunin
I am going to go ahead and make a pull request. Meanwhile, I would really appreciate someone could address these two points:

2. the new aten code seems to use TensorInfo as a wrapper around device arrays. From looking at struct TensorInfo from TensorInfo.cuh, it seems like there is no way to cast pointer from and to __restrict__. The first step seems to add that functionality.
3. Is there any code that makes aliases and use both pointers within the same kernel? otherwise, we could just make the default pointer have __restrict__ tag.

Thanks.