1try: 

2 import pyopencl.array as gpuarray 

3except ImportError: 

4 gpuarray = None 

5 

6import numpy as np 

7 

8import pystencils 

9 

10 

11class PyOpenClArrayHandler: 

12 

13 def __init__(self, queue): 

14 if not queue: 

15 from pystencils.opencl.opencljit import get_global_cl_queue 

16 queue = get_global_cl_queue() 

17 assert queue, "OpenCL queue missing!\n" \ 

18 "Use `import pystencils.opencl.autoinit` if you want it to be automatically created" 

19 self.queue = queue 

20 

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

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

23 return self.to_gpu(cpu_array) 

24 

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

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

27 return self.to_gpu(cpu_array) 

28 

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

30 if layout: 

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

32 return self.to_gpu(cpu_array) 

33 else: 

34 return gpuarray.empty(self.queue, shape, dtype) 

35 

36 def to_gpu(self, array): 

37 return gpuarray.to_device(self.queue, array) 

38 

39 def upload(self, gpuarray, numpy_array): 

40 gpuarray.set(numpy_array, self.queue) 

41 

42 def download(self, gpuarray, numpy_array): 

43 gpuarray.get(self.queue, numpy_array) 

44 

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

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

47 return self.from_numpy(cpu_array) 

48 

49 from_numpy = to_gpu