Add unit test for UVM APIs in FBGEMM_GPU (#516)

Summary:
Pull Request resolved: #516

Added a new class CUMemTest(unittest.TestCase) with two methods to test is_uvm_tensor and uvm_to_cpu, respectively.

Used hypothesis.given to generate lists of int as input sizes to initialize tensors for testing. The test methods follow the example usage provided.

Reviewed By: jianyuh

Differential Revision: D26490239

fbshipit-source-id: e34ee97c74d396f72df458a610f43fbcc762741d

fbgemm_gpu/test/split_table_batched_embeddings_test.py

@@ -1265,5 +1265,25 @@ def test_backward_optimizers(  # noqa C901
                 )
 
 
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class CUMemTest(unittest.TestCase):
+    @given(sizes=st.lists(st.integers(min_value=1, max_value=8), min_size=1, max_size=4))
+    @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
+    def test_is_uvm_tensor(self, sizes):
+        uvm_t = torch.ops.fb.new_managed_tensor(torch.zeros(*sizes, device="cuda:0",
+                                                        dtype=torch.float), sizes)
+        assert torch.ops.fb.is_uvm_tensor(uvm_t)
+
+    @given(sizes=st.lists(st.integers(min_value=1, max_value=8), min_size=1, max_size=4))
+    @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
+    def test_uvm_to_cpu(self, sizes):
+        uvm_t = torch.ops.fb.new_managed_tensor(torch.zeros(*sizes, device="cuda:0",
+                                                            dtype=torch.float), sizes)
+        cpu_t = torch.ops.fb.uvm_to_cpu(uvm_t)
+        assert not torch.ops.fb.is_uvm_tensor(cpu_t)
+        uvm_t.copy_(cpu_t)
+        assert torch.ops.fb.is_uvm_tensor(uvm_t)
+
+
 if __name__ == "__main__":
     unittest.main()