From 44670416abc5d9f43ac31da0eccd268f0d5acd27 Mon Sep 17 00:00:00 2001
From: David Koes <dkoes@pitt.edu>
Date: Tue, 14 Sep 2021 08:52:58 -0400
Subject: [PATCH] Better multi-device support

Provide helper functions for setting current device.
Throw useful error message if tensor is not on current device.
---
 python/bindings.in.cpp    |  2 ++
 python/bindings_grids.cpp |  8 ++++++++
 test/test_gridmaker.py    | 26 ++++++++++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/python/bindings.in.cpp b/python/bindings.in.cpp
index bcd6e84e..94f85096 100644
--- a/python/bindings.in.cpp
+++ b/python/bindings.in.cpp
@@ -399,6 +399,8 @@ BOOST_PYTHON_MODULE(molgrid)
       "Set if generated grids should be on GPU by default.");
   def("tofloatptr", +[](long val) { return Pointer<float>((float*)val);}, "Return integer as float *");
   def("todoubleptr", +[](long val) { return Pointer<double>((double*)val);}, "Return integer as double *");
+  def("set_gpu_device", +[](int device)->void {LMG_CUDA_CHECK(cudaSetDevice(device));}, "Set current GPU device.");
+  def("get_gpu_device", +[]()->int {int device = 0; LMG_CUDA_CHECK(cudaGetDevice(&device)); return device;}, "Get current GPU device.");
 
   //type converters
   py_pair<int, float>();
diff --git a/python/bindings_grids.cpp b/python/bindings_grids.cpp
index 63df451f..26a007fc 100644
--- a/python/bindings_grids.cpp
+++ b/python/bindings_grids.cpp
@@ -117,6 +117,14 @@ struct Grid_from_python {
         } else {
           return false; //don't recognize
         }
+        if(info.isGPU && hasattr(t,"device") && hasattr(t.attr("device"), "index")) {
+          int d = extract<int>(t.attr("device").attr("index"));
+          int currd = 0;
+          LMG_CUDA_CHECK(cudaGetDevice(&currd));
+          if(currd != d) {
+            throw std::invalid_argument("Attempt to use GPU tensor on different device ("+itoa(d)+") than current device ("+itoa(currd)+").  Change location of tensor or change current device.");
+          }
+        }
         if(hasattr(t,"is_contiguous")) {
           if(!t.attr("is_contiguous")()) {
             throw std::invalid_argument("Attempt to use non-contiguous tensor in molgrid.  Call clone first.");
diff --git a/test/test_gridmaker.py b/test/test_gridmaker.py
index 16476646..2a27b975 100644
--- a/test/test_gridmaker.py
+++ b/test/test_gridmaker.py
@@ -102,6 +102,32 @@ def test_a_grid():
     assert 2.094017 == approx(mgridout.tonumpy().max())
     assert 2.094017 == approx(mgridgpu.tonumpy().max())
 
+def test_devices():
+    fname = datadir+"/small.types"
+    e = molgrid.ExampleProvider(data_root=datadir+"/structs")
+    e.populate(fname)
+    ex = e.next()
+
+    gmaker = molgrid.GridMaker()
+    dims = gmaker.grid_dimensions(ex.num_types()) # this should be grid_dims or get_grid_dims
+
+    try:
+        torchout = torch.zeros(dims, dtype=torch.float32, device='cuda:1')
+    except RuntimeError:
+        return # can't test multiple devices because we don't have them
+    
+    try:
+        gmaker.forward(ex, torchout)
+        assert False # should not get here
+    except ValueError:
+        pass # this should return an error
+    
+    molgrid.set_gpu_device(1)
+    assert molgrid.get_gpu_device() == 1
+    ex = e.next() 
+    gmaker.forward(ex, torchout) #should work now
+    
+    
 def test_radius_multiples():
     g1 = molgrid.GridMaker(resolution=.1,dimension=6.0)
     c = np.array([[0,0,0]],np.float32)