diff --git a/src/pystencils_autodiff/backends/_pytorch.py b/src/pystencils_autodiff/backends/_pytorch.py index 5fdd1d2cadb2f8cb80cad16064f4930c8378a55d..6d2da60f799483a5dbc62f0266d6d331ee750cb6 100644 --- a/src/pystencils_autodiff/backends/_pytorch.py +++ b/src/pystencils_autodiff/backends/_pytorch.py @@ -7,8 +7,9 @@ try: import pycuda.autoinit import pycuda.gpuarray import pycuda.driver + HAS_PYCUDA = True except Exception: - pass + HAS_PYCUDA = False def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict, forward_loop, backward_loop, @@ -94,12 +95,12 @@ def numpy_dtype_to_torch(dtype): # Fails if different context/thread -# def tensor_to_gpuarray(tensor): -# if not tensor.is_cuda: -# raise ValueError( -# 'Cannot convert CPU tensor to GPUArray (call `cuda()` on it)') -# else: -# return pycuda.gpuarray.GPUArray(tensor.shape, dtype=torch_dtype_to_numpy(tensor.dtype), gpudata=tensor.data_ptr()) +def tensor_to_gpuarray(tensor): + if not tensor.is_cuda: + raise ValueError( + 'Cannot convert CPU tensor to GPUArray (call `cuda()` on it)') + else: + return pycuda.gpuarray.GPUArray(tensor.shape, dtype=torch_dtype_to_numpy(tensor.dtype), gpudata=tensor.data_ptr()) def gpuarray_to_tensor(gpuarray, context=None): @@ -125,20 +126,23 @@ def gpuarray_to_tensor(gpuarray, context=None): return out -class GpuPointerHolder(pycuda.driver.PointerHolderBase): +if HAS_PYCUDA: + class GpuPointerHolder(pycuda.driver.PointerHolderBase): - def __init__(self, tensor): - super().__init__() - self.tensor = tensor - self.gpudata = tensor.data_ptr() + def __init__(self, tensor): + super().__init__() + self.tensor = tensor + self.gpudata = tensor.data_ptr() - def get_pointer(self): - return self.tensor.data_ptr() + def get_pointer(self): + return self.tensor.data_ptr() - def __int__(self): - return self.__index__() + def __int__(self): + return self.__index__() - # without an __index__ method, arithmetic calls to the GPUArray backed by this pointer fail - # not sure why, this needs to return some integer, apparently - def __index__(self): - return self.gpudata + # without an __index__ method, arithmetic calls to the GPUArray backed by this pointer fail + # not sure why, this needs to return some integer, apparently + def __index__(self): + return self.gpudata +else: + GpuPointerHolder = None diff --git a/tests/lbm/backends/_pytorch.py b/tests/lbm/backends/_pytorch.py index 289fc0b6489d89a504441e5bb35227a1a9e2b7b3..db53c3f92a474461f904ddf0b2246ba4cf1135e1 100644 --- a/tests/lbm/backends/_pytorch.py +++ b/tests/lbm/backends/_pytorch.py @@ -11,6 +11,18 @@ try: except Exception: HAS_PYCUDA = False +# Fails if different context/thread + + +def tensor_to_gpuarray(tensor): + if not tensor.is_cuda: + raise ValueError( + 'Cannot convert CPU tensor to GPUArray (call `cuda()` on it)') + else: + return pycuda.gpuarray.GPUArray(tensor.shape, + dtype=torch_dtype_to_numpy(tensor.dtype), + gpudata=tensor.data_ptr()) + def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict, forward_loop, backward_loop, convert_tensors_to_arrays=True): @@ -92,17 +104,9 @@ def numpy_dtype_to_torch(dtype): return getattr(torch, dtype_name) -# Fails if different context/thread -# def tensor_to_gpuarray(tensor): -# if not tensor.is_cuda: -# raise ValueError( -# 'Cannot convert CPU tensor to GPUArray (call `cuda()` on it)') -# else: -# return pycuda.gpuarray.GPUArray(tensor.shape, dtype=torch_dtype_to_numpy(tensor.dtype), gpudata=tensor.data_ptr()) - - def gpuarray_to_tensor(gpuarray, context=None): - '''Convert a :class:`pycuda.gpuarray.GPUArray` to a :class:`torch.Tensor`. The underlying + """ + Convert a :class:`pycuda.gpuarray.GPUArray` to a :class:`torch.Tensor`. The underlying storage will NOT be shared, since a new copy must be allocated. Parameters ---------- @@ -110,7 +114,7 @@ def gpuarray_to_tensor(gpuarray, context=None): Returns ------- torch.Tensor - ''' + """ if not context: context = pycuda.autoinit.context shape = gpuarray.shape