diff --git a/python/bindings.in.cpp b/python/bindings.in.cpp index bcd6e84..94f8509 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*)val);}, "Return integer as float *"); def("todoubleptr", +[](long val) { return Pointer((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(); diff --git a/python/bindings_grids.cpp b/python/bindings_grids.cpp index 63df451..26a007f 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(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 1647664..2a27b97 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)