1try: 

2 import pycuda.gpuarray as gpuarray 

3except ImportError: 

4 gpuarray = None 

5import numpy as np 

6 

7import pystencils 

8 

9 

10class PyCudaArrayHandler: 

11 

12 def __init__(self): 

13 import pycuda.autoinit # NOQA 

14 

15 def zeros(self, shape, dtype=np.float64, order='C'): 

16 cpu_array = np.zeros(shape=shape, dtype=dtype, order=order) 

17 return self.to_gpu(cpu_array) 

18 

19 def ones(self, shape, dtype=np.float64, order='C'): 

20 cpu_array = np.ones(shape=shape, dtype=dtype, order=order) 

21 return self.to_gpu(cpu_array) 

22 

23 def empty(self, shape, dtype=np.float64, layout=None): 

24 if layout: 

25 cpu_array = pystencils.field.create_numpy_array_with_layout(shape=shape, dtype=dtype, layout=layout) 

26 return self.to_gpu(cpu_array) 

27 else: 

28 return gpuarray.empty(shape, dtype) 

29 

30 @staticmethod 

31 def to_gpu(array): 

32 return gpuarray.to_gpu(array) 

33 

34 @staticmethod 

35 def upload(array, numpy_array): 

36 array.set(numpy_array) 

37 

38 @staticmethod 

39 def download(array, numpy_array): 

40 array.get(numpy_array) 

41 

42 def randn(self, shape, dtype=np.float64): 

43 cpu_array = np.random.randn(*shape).astype(dtype) 

44 return self.to_gpu(cpu_array) 

45 

46 from_numpy = to_gpu 

47 

48 

49class PyCudaNotAvailableHandler: 

50 def __getattribute__(self, name): 

51 raise NotImplementedError("Unable to initiaize PyCuda! " 

52 "Try to run `import pycuda.autoinit` to check whether PyCuda is working correctly!")