Repository: helmholtz-analytics/heat

Issue:
[Bug]: `heat.print` fails if communication over GPUs is required
### What happened?

Printing of a large array that's distributed over GPUs fails.

### Code snippet triggering the error

```shell
import heat

N = 2**10
a=heat.arange(N,dtype=heat.float32,device='gpu',split=0)
print(a)
```


### Error message or erroneous outcome

```shell
> srun --ntasks=2 python test_gpu_qa.py

Traceback (most recent call last):
  File "/p/project/haf/users/comito1/devel/heat/test_gpu_qa.py", line 9, in <module>
    print("BEFORE reshape: ", a)
  File "/p/project/haf/users/comito1/devel/heat/heat/core/dndarray.py", line 1806, in __str__
    return printing.__str__(self)
  File "/p/project/haf/users/comito1/devel/heat/heat/core/printing.py", line 198, in __str__
    tensor_string = _tensor_str(dndarray, __INDENT + 1)
  File "/p/project/haf/users/comito1/devel/heat/heat/core/printing.py", line 285, in _tensor_str
    torch_data = _torch_data(dndarray, summarize)
  File "/p/project/haf/users/comito1/devel/heat/heat/core/printing.py", line 252, in _torch_data
    data = torch.index_select(data, i, torch.arange(end))
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)
[1678880909.786307] [hdfmlc03:18341:0]       tag_match.c:62   UCX  WARN  unexpected tag-receive descriptor 0x6902e00 was not matched
[1678880909.786324] [hdfmlc03:18341:0]       tag_match.c:62   UCX  WARN  unexpected tag-receive descriptor 0x14780c86eb50 was not matched
srun: error: hdfmlc03: task 0: Exited with exit code 1
```


### Version

1.2.x

### Python version

3.9

### PyTorch version

1.12

### MPI version

```shell
OpenMPI 4.1.2
```



Discussion:
I could reproduce the error on 1 node with 4 GPUs as well: 

```
Traceback (most recent call last):
  File "test_print.py", line 5, in <module>
    print(a)
  File "************/heat/heat/core/dndarray.py", line 1821, in __str__
    return printing.__str__(self)
  File "************/heat/heat/core/printing.py", line 198, in __str__
    tensor_string = _tensor_str(dndarray, __INDENT + 1)
  File "************/heat/heat/core/printing.py", line 285, in _tensor_str
    torch_data = _torch_data(dndarray, summarize)
  File "************/heat/heat/core/printing.py", line 252, in _torch_data
    data = torch.index_select(data, i, torch.arange(end))
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)



srun: error: **************: task 0: Exited with exit code 1
````

Resolving patch:
diff --git a/heat/core/printing.py b/heat/core/printing.py
--- a/heat/core/printing.py
+++ b/heat/core/printing.py
@@ -238,9 +238,15 @@ def _torch_data(dndarray, summarize) -> DNDarray:
 
             # non-split dimension, can slice locally
             if i != dndarray.split:
-                start_tensor = torch.index_select(data, i, torch.arange(edgeitems + 1))
+                start_tensor = torch.index_select(
+                    data, i, torch.arange(edgeitems + 1, device=data.device)
+                )
                 end_tensor = torch.index_select(
-                    data, i, torch.arange(dndarray.lshape[i] - edgeitems, dndarray.lshape[i])
+                    data,
+                    i,
+                    torch.arange(
+                        dndarray.lshape[i] - edgeitems, dndarray.lshape[i], device=data.device
+                    ),
                 )
                 data = torch.cat([start_tensor, end_tensor], dim=i)
             # split-dimension , need to respect the global offset
@@ -249,18 +255,27 @@ def _torch_data(dndarray, summarize) -> DNDarray:
 
                 if offset < edgeitems + 1:
                     end = min(dndarray.lshape[i], edgeitems + 1 - offset)
-                    data = torch.index_select(data, i, torch.arange(end))
+                    data = torch.index_select(data, i, torch.arange(end, device=data.device))
                 elif dndarray.gshape[i] - edgeitems < offset - dndarray.lshape[i]:
                     global_start = dndarray.gshape[i] - edgeitems
                     data = torch.index_select(
-                        data, i, torch.arange(max(0, global_start - offset), dndarray.lshape[i])
+                        data,
+                        i,
+                        torch.arange(
+                            max(0, global_start - offset),
+                            dndarray.lshape[i],
+                            device=data.device,
+                        ),
                     )
         # exchange data
         received = dndarray.comm.gather(data)
         if dndarray.comm.rank == 0:
             # concatenate data along the split axis
+            # problem: CUDA-aware MPI `gather`s all `data` in a list of tensors on MPI-process no. 0, but not necessarily on the same cuda device.
+            # Indeed, `received` may be a list of tensors on cuda device 0, cuda device 1, ... therefore, we need to move all entries of the list to cuda device 0 before applying `cat`.
+            device0 = received[0].device
+            received = [tens.to(device0) for tens in received]
             data = torch.cat(received, dim=dndarray.split)
-
     return data
 
 


Test patch:
diff --git a/heat/core/tests/test_printing.py b/heat/core/tests/test_printing.py
--- a/heat/core/tests/test_printing.py
+++ b/heat/core/tests/test_printing.py
@@ -429,3 +429,11 @@ def test_split_2_above_threshold(self):
 
         if dndarray.comm.rank == 0:
             self.assertEqual(comparison, __str)
+
+
+class TestPrintingGPU(TestCase):
+    def test_print_GPU(self):
+        # this test case also includes GPU now, checking the output is not done; only test whether the routine itself works...
+        a0 = ht.arange(2**20, dtype=ht.float32).reshape((2**10, 2**10)).resplit_(0)
+        a1 = ht.arange(2**20, dtype=ht.float32).reshape((2**10, 2**10)).resplit_(1)
+        print(a0, a1)