Repository: teddykoker/torchsort (language: python)

Issue:
Incorrect results when running on non-default cuda device
When running `torchsort.soft_rank` or `torchsort.soft_sort` on a tensor that's not on the default cuda device (usually `cuda:0`), the results are incorrect. 

```python
import torch
import torchsort

x = torch.tensor([[9,8]], device="cuda:1")

print(torchsort.soft_rank(x))
# tensor([[9., 8.]], device='cuda:1')

print(torchsort.soft_sort(x))
# tensor([[-2., -1.]], device='cuda:1')
```

Based on the GPU memory usage, torchsort tries to do something on the default cuda device `cuda:0` instead of whichever device the input tensor is on. As a workaround, you need to either change the default cuda device with `torch.cuda.set_device` or use the context manager `torch.cuda.device`.

```python
import torch
import torchsort

x = torch.tensor([[9,8]], device="cuda:1")

with torch.cuda.device(x.device):
    print(torchsort.soft_rank(x))
    # tensor([[2., 1.]], device='cuda:1')

    print(torchsort.soft_sort(x))
    # tensor([[8., 9.]], device='cuda:1')
```


Discussion / hints:
Hi Justin, thanks for finding this.

I think it might be necessary to [add a device guard](https://discuss.pytorch.org/t/c-cuda-extension-with-multiple-gpus/91241/6) to the CUDA kernel. This should be accomplished by adding the:
```cuda
#include <c10/cuda/CUDAGuard.h>
```
import and:
```cuda
const at::cuda::OptionalCUDAGuard device_guard(device_of(y));
```
to lines 309, 335, 
and 

```cuda
const at::cuda::OptionalCUDAGuard device_guard(device_of(s));
```
to lines 360 and 382 of `isotonic_cuda.cu`

I'll try to push these changes to a new branch this evening - if you could try to run it then it would be great to see if this fixes the issue.


Fix patch:
diff --git a/setup.py b/setup.py
index 6f6069e..b3be4d2 100644
--- a/setup.py
+++ b/setup.py
@@ -51,7 +51,7 @@ def ext_modules():
 
 setup(
     name="torchsort",
-    version="0.1.8",
+    version="0.1.9",
     description="Differentiable sorting and ranking in PyTorch",
     author="Teddy Koker",
     url="https://github.com/teddykoker/torchsort",
diff --git a/torchsort/isotonic_cuda.cu b/torchsort/isotonic_cuda.cu
index 2098fad..7e68bed 100644
--- a/torchsort/isotonic_cuda.cu
+++ b/torchsort/isotonic_cuda.cu
@@ -30,9 +30,7 @@
 
 
 #include <torch/extension.h>
-// #include <cuda.h>
-// #include <cuda_runtime.h>
-// #include <iostream>
+#include <c10/cuda/CUDAGuard.h>
 
 //  Copied from fast-soft-sort (https://bit.ly/3r0gOav) with the following modifications:
 //  - replace numpy functions with torch equivalents
@@ -306,6 +304,7 @@ __global__ void isotonic_kl_backward_kernel(
 // Solves an isotonic regression problem using PAV.
 // Formally, it solves argmin_{v_1 >= ... >= v_n} 0.5 ||v - y||^2.
 torch::Tensor isotonic_l2(torch::Tensor y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(y));
     auto batch = y.size(0);
     auto n = y.size(1);
     auto sol = torch::zeros_like(y);
@@ -332,6 +331,7 @@ torch::Tensor isotonic_l2(torch::Tensor y) {
 // Solves isotonic optimization with KL divergence using PAV.
 // Formally, it solves argmin_{v_1 >= ... >= v_n} <e^{y-v}, 1> + <e^w, v>.
 torch::Tensor isotonic_kl(torch::Tensor y, torch::Tensor w) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(y));
     auto batch = y.size(0);
     auto n = y.size(1);
     auto sol = torch::zeros_like(y);
@@ -357,6 +357,7 @@ torch::Tensor isotonic_kl(torch::Tensor y, torch::Tensor w) {
 }
 
 torch::Tensor isotonic_l2_backward(torch::Tensor s, torch::Tensor sol, torch::Tensor grad_input) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(s));
     auto batch = sol.size(0);
     auto n = sol.size(1);
     auto ret = torch::zeros_like(sol);
@@ -379,6 +380,7 @@ torch::Tensor isotonic_l2_backward(torch::Tensor s, torch::Tensor sol, torch::Te
 }
 
 torch::Tensor isotonic_kl_backward(torch::Tensor s, torch::Tensor sol, torch::Tensor grad_input) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(s));
     auto batch = sol.size(0);
     auto n = sol.size(1);
     auto ret = torch::zeros_like(sol);


Test patch:
diff --git a/tests/test_ops.py b/tests/test_ops.py
index e856f2f..3e67afa 100644
--- a/tests/test_ops.py
+++ b/tests/test_ops.py
@@ -15,8 +15,9 @@
 REGULARIZATION = ["l2", "kl"]
 REGULARIZATION_STRENGTH = [1e-1, 1e0, 1e1]
 
+# use CPU, and up to two CUDA devices
 DEVICES = [torch.device("cpu")] + (
-    [torch.device("cuda")] if torch.cuda.is_available() else []
+    [torch.device(f"cuda:{d}") for d in range(min(torch.cuda.device_count(), 2))]
 )
 
 torch.manual_seed(0)
@@ -67,6 +68,7 @@ def test_vs_original(funcs, regularization, regularization_strength, device):
 @pytest.mark.parametrize("device", DEVICES)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA to test fp16")
 def test_half(function, regularization, regularization_strength, device):
+    # check half precision
     x = torch.randn(BATCH_SIZE, SEQ_LEN, requires_grad=True).cuda().half()
     f = partial(
         function,
@@ -75,3 +77,5 @@ def test_half(function, regularization, regularization_strength, device):
     )
     # don't think theres a better way of testing, tolerance must be pretty high
     assert torch.allclose(f(x), f(x.float()).half(), atol=1e-1)
+
+