For CUDA extension implementing element-wise function, what is the best shape for the kernel launch?

[zfzhang]

Hi,
I’m writing a Pytorch CUDA extension that implements an element-wise function.
Assuming we don’t know whether the input is 2D or 3D or 4D, but it is known to have the first dimension as batch. How do we choose the best kernel shape (grid size, block size) at launch?

Thank you!

[zfzhang]

I found a similar question with answer, may be the solution.

How to write a custom point-wise cuda kernel?

I want to write a custom point-wise cuda kernel and I can do that easily with cpp_extension, and a custom cuda kernel luncher. The problem is that I don’t know how to set a good block_size. As far as I can tell CUDA_tensor_apply2 handles those parameters and simplifies my work. But I need CUDA_tensor_apply3 which has been removed from the aten. It seems like I should use TensorIterator with gpu_kernel but I can not include it into my code. Should I include <ATen/native/CUDALoops.cuh>? If so, I…