Hi guys,
I have been tracking this topic for a long time, but still couldn’t find a good solution.
So, I tried to implement a customized pytorch cuda extension to totally avoid memory copy on cuda device. However, my solution is very specific in use case.
In my case, I combined Concat + Conv2d into a single CatConv2d kernel, which can significantly reduce the latency in some cases (small batch, small in/out channels).
for example, the operations of concat of a list of Tensor followed by a conv2d before, now can be done in a single operator with single cuda kernel:
    x_list = [x0, x1, x2, x3,...]
    x = torch.cat(x_list, 1)
    x = conv.forward(x)

    # equivalent to 
    x = CatConv2d(x_list)

Here is the repo:
  
      

      GitHub
  

  
    

GitHub - PingoLH/CatConv2d: Concat+Conv2d fused all-in-one CUDA kernel... 10

  Concat+Conv2d fused all-in-one CUDA kernel extension for Pytorch - GitHub - PingoLH/CatConv2d: Concat+Conv2d fused all-in-one CUDA kernel extension for Pytorch


  

  
    
    
  

  


The backward path is still missing (forward only) in this repo, and there are many limitations, but I’m just trying to give a sense that concat+some_op combination is very promising in terms of performance and worth to work on.