Using torch.nn.DataParallel with a custom CUDA extension
To my understanding, the built-in PyTorch operations all automatically handle batches through implicit vectorization, allowing parallelism across multiple GPUs.
However, when writing a custom operation in CUDA as per the Documentation, the LLTM example given performs operations that are batch invariant, for example computing the gradient of the Sigmoid function elementwise.
However, I have a use case that is not batch element invariant and not vectorizable. Running on a single GPU, I currently (inefficiently) loop over each element in the batch, performing a kernel launch for each, like so (written in the browser, just to demonstrate):
std::vector&lt;at::Tensor&gt; op_cuda_forward(at::Tensor input, 
                                        at::Tensor elementSpecificParam) {
    
    auto output = at::zeros(torch::CUDA(/* TYPE */), {/* DIMENSIONS */});
    
    const size_t blockDim = //
    const size_t gridDim = //
    const size_t = numBatches = //

    for (size_t i = 0; i &lt; numBatches; i++) {
        op_cuda_forward_kernel&lt;T&gt;&lt;&lt;&lt;gridDim, blockDim&gt;&gt;&gt;(input[i],
                                                         elementSpecificParam[i], 
                                                         output[i]);
    }

    return {output};
}

However, I wish to split this operation over multiple GPUs by batch element.
How would the allocation of the output Tensor work in a multi-GPU scenario?
Of course, one may create intermediate Tensors on each GPU before launching the appropriate kernel, however, the overhead of copying the input data to each GPU and back again would be problematic.
Is there a simpler way to launch the kernels without first probing the environment for GPU information (# GPU's etc)?
The end goal is to have a CUDA operation that works with torch.nn.DataParallel.

In order to use torch.nn.DataParallel with a custom CUDA extension, you can follow these steps:

Define your custom CUDA extension in a subclass of torch.autograd.Function, and implement the forward() and backward() methods for the forward and backward passes, respectively.
In the forward() method, create a new output Tensor and allocate it on the same device that the input Tensor is on using output.to(input.device).
In the backward() method, create a new gradient input Tensor and allocate it on the same device that the gradient output Tensor is on using grad_input.to(grad_output.device).
In your main PyTorch code, use torch.nn.DataParallel to wrap your model, and make sure to call the to() method on the input data and move it to the same device that your model is on.

Here's an example of how this could look like:
# Define your custom CUDA extension as a subclass of torch.autograd.Function
class MyCustomCudaExtension(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, element_specific_param):
        # Create a new output Tensor and allocate it on the same device as the input Tensor
        output = torch.zeros_like(input).to(input.device)
        # Perform the forward pass using your custom CUDA kernel
        # ...
        # Save any necessary information for the backward pass
        ctx.save_for_backward(output)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # Retrieve the saved information from the forward pass
        output = ctx.saved_tensors[0]
        # Create a new gradient input Tensor and allocate it on the same device as the gradient output Tensor
        grad_input = torch.zeros_like(output).to(grad_output.device)
        # Perform the backward pass using your custom CUDA kernel
        # ...
        return grad_input, None

# In your main PyTorch code, wrap your model in torch.nn.DataParallel
model = torch.nn.DataParallel(MyModel())

# Make sure to call the .to() method on the input data and move it to the same device as your model
input = input.to(model.device)

# Perform the forward pass
output = model(input)

This way, torch.nn.DataParallel will automatically handle the parallelism across multiple GPUs, and your custom CUDA extension will be executed on the correct device.
Note that in this example, we are assuming that the input and output Tensors are on the same device. If this is not the case, you can use the input.device and grad_output.device attributes to retrieve the devices of the input and gradient output Tensors, respectively, and allocate the output and gradient input Tensors on those devices.
