Prevent users from attempting to pin PyTorch non-contiguous tensors or views...

Prevent users from attempting to pin PyTorch non-contiguous tensors or views only encompassing part of tensor. (#3992)

* Disable pinning non-contiguous memory

* Prevent views from being converted for write

* Fix linting

* Add unit tests

* Improve error message for users

* Switch to pytest function

* exclude mxnet and tensorflow from inplace pinning

* Add skip

* Restrict to pytorch backend

* Use backend to retrieve device

* Fix capitalization in decorator
Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>

python/dgl/backend/pytorch/tensor.py

@@ -349,6 +349,10 @@ else:
         return nd.from_dlpack(dlpack.to_dlpack(data.contiguous()))
 
 def zerocopy_to_dgl_ndarray_for_write(input):
+    assert input.is_contiguous(), "Cannot convert non-contiguous tensors " \
+        "to dgl ndarray for write. Call .to_contiguous() first."
+    assert input.numel() == input.storage().size(), "Cannot convert view " \
+        "tensors to dgl ndarray for write."
     return zerocopy_to_dgl_ndarray(input)
 
 def zerocopy_from_dgl_ndarray(data):

python/dgl/utils/pin_memory.py

 """Utility functions related to pinned memory tensors."""
 
+from ..base import DGLError
 from .. import backend as F
 from .._ffi.function import _init_api
 
 def pin_memory_inplace(tensor):
     """Register the tensor into pinned memory in-place (i.e. without copying)."""
-    F.to_dgl_nd(tensor).pin_memory_()
+    if F.backend_name in ['mxnet', 'tensorflow']:
+        raise DGLError("The {} backend does not support pinning " \
+            "tensors in-place.".format(F.backend_name))
+
+    # needs to be writable to allow in-place modification
+    try:
+        F.zerocopy_to_dgl_ndarray_for_write(tensor).pin_memory_()
+    except Exception as e:
+        raise DGLError("Failed to pin memory in-place due to: {}".format(e))
 
 def unpin_memory_inplace(tensor):
     """Unregister the tensor from pinned memory in-place (i.e. without copying)."""
-    F.to_dgl_nd(tensor).unpin_memory_()
+    # needs to be writable to allow in-place modification
+    try:
+        F.zerocopy_to_dgl_ndarray_for_write(tensor).unpin_memory_()
+    except Exception as e:
+        raise DGLError("Failed to unpin memory in-place due to: {}".format(e))
 
 def gather_pinned_tensor_rows(tensor, rows):
     """Directly gather rows from a CPU tensor given an indices array on CUDA devices,

tests/compute/test_heterograph.py

@@ -980,6 +980,7 @@ def test_to_device2(g, idtype):
         assert g1.canonical_etypes == g.canonical_etypes
 
 @unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
+@unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="Pinning graph inplace only supported for PyTorch")
 @parametrize_dtype
 def test_pin_memory_(idtype):
     # TODO: rewrite this test case to accept different graphs so we

tests/compute/test_pin_memory.py

+import backend as F
+import dgl
+import pytest
+
+@pytest.mark.skipif(F._default_context_str == 'cpu', reason="Need gpu for this test")
+def test_pin_unpin():
+    t = F.arange(0, 100, dtype=F.int64, ctx=F.cpu())
+
+    assert not F.is_pinned(t)
+
+    if F.backend_name == 'pytorch':
+        dgl.utils.pin_memory_inplace(t)
+        assert F.is_pinned(t)
+        dgl.utils.unpin_memory_inplace(t)
+        assert not F.is_pinned(t)
+    else:
+        with pytest.raises(dgl.DGLError):
+            # tensorflow and mxnet should throw an erro
+            dgl.utils.pin_memory_inplace(t)
+
+if __name__ == "__main__":
+    test_pin_unpin()

tests/compute/test_subgraph.py

@@ -610,7 +610,6 @@ def test_khop_out_subgraph(idtype):
     assert F.array_equal(F.astype(inv['game'], idtype), F.tensor([0], idtype))
 
 @unittest.skipIf(not F.gpu_ctx(), 'only necessary with GPU')
-@unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="UVA only supported for PyTorch")
 @pytest.mark.parametrize(
     'parent_idx_device', [('cpu', F.cpu()), ('cuda', F.cuda()), ('uva', F.cpu()), ('uva', F.cuda())])
 @pytest.mark.parametrize('child_device', [F.cpu(), F.cuda()])
@@ -624,6 +623,8 @@ def test_subframes(parent_idx_device, child_device):
     if parent_device == 'cuda':
         g = g.to(F.cuda())
     elif parent_device == 'uva':
+        if F.backend_name != 'pytorch':
+            pytest.skip("UVA only supported for PyTorch")
         g = g.to(F.cpu())
         g.create_formats_()
         g.pin_memory_()
@@ -631,19 +632,20 @@ def test_subframes(parent_idx_device, child_device):
         g = g.to(F.cpu())
     idx = F.copy_to(idx, idx_device)
     sg = g.sample_neighbors(idx, 2).to(child_device)
-    assert sg.device == sg.ndata['x'].device
-    assert sg.device == sg.edata['a'].device
+    assert sg.device == F.context(sg.ndata['x'])
+    assert sg.device == F.context(sg.edata['a'])
     assert sg.device == child_device
     if parent_device != 'uva':
         sg = g.to(child_device).sample_neighbors(F.copy_to(idx, child_device), 2)
-        assert sg.device == sg.ndata['x'].device
-        assert sg.device == sg.edata['a'].device
+        assert sg.device == F.context(sg.ndata['x'])
+        assert sg.device == F.context(sg.edata['a'])
         assert sg.device == child_device
     if parent_device == 'uva':
         g.unpin_memory_()
 
 @unittest.skipIf(F._default_context_str != "gpu", reason="UVA only available on GPU")
 @pytest.mark.parametrize('device', [F.cpu(), F.cuda()])
+@unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="UVA only supported for PyTorch")
 @parametrize_dtype
 def test_uva_subgraph(idtype, device):
     g = create_test_heterograph(idtype)

tests/pytorch/test_pin_memory.py

+import backend as F
+import dgl
+import pytest
+import torch
+
+@pytest.mark.skipif(F._default_context_str == 'cpu', reason="Need gpu for this test")
+def test_pin_noncontiguous():
+    t = torch.empty([10, 100]).transpose(0, 1)
+
+    assert not t.is_contiguous()
+    assert not F.is_pinned(t)
+
+    with pytest.raises(dgl.DGLError):
+        dgl.utils.pin_memory_inplace(t)
+
+@pytest.mark.skipif(F._default_context_str == 'cpu', reason="Need gpu for this test")
+def test_pin_view():
+    t = torch.empty([100, 10])
+    v = t[10:20]
+
+    assert v.is_contiguous()
+    assert not F.is_pinned(t)
+
+    with pytest.raises(dgl.DGLError):
+        dgl.utils.pin_memory_inplace(v)
+
+
+if __name__ == "__main__":
+    test_pin_noncontiguous()
+    test_pin_view()