| mandelbrot_cuda.py |
| 1 | #################################################### |
| 2 | # Calculate Mandelbrot set and save it as a bmp image |
| 3 | # |
| 4 | # Data parallel version using Pycuda |
| 5 | # Create string with cuda code and let the graphics |
| 6 | # card farm out the work to each warp. |
| 7 | # |
| 8 | #################################################### |
| 9 | |
| 10 | import bmp |
| 11 | import pycuda.driver as drv |
| 12 | import pycuda.tools |
| 13 | import pycuda.autoinit |
| 14 | from pycuda.compiler import SourceModule |
| 15 | import pycuda.gpuarray as gpuarray |
| 16 | import numpy as nm |
| 17 | |
| 18 | # maximal number of iterations to compute a pixel |
| 19 | MAX_ITER = 256 |
| 20 | |
| 21 | # image dimensions |
| 22 | from sizes import nx,ny |
| 23 | |
| 24 | from pycuda.elementwise import ElementwiseKernel |
| 25 | complex_gpu = ElementwiseKernel( |
| 26 | "int nx, int ny, int maxiter, int *iteration", |
| 27 | """ |
| 28 | float zr, zi, z2; |
| 29 | float jf = 1.0f*(i%ny); |
| 30 | float rowf = 1.0f*(i/ny); |
| 31 | float nxf = 1.0f*nx; |
| 32 | float nyf = 1.0f*ny; |
| 33 | float qif = 4.0f*rowf/nyf-2.0f; |
| 34 | float qrf = 4.0f*jf/nxf-2.0f; |
| 35 | iteration[i] = maxiter; |
| 36 | zr = 0.0f; |
| 37 | zi = 0.0f; |
| 38 | for(int n=0;n < maxiter;n++) { |
| 39 | float nzr = zr*zr - zi*zi + qrf; |
| 40 | float nzi = 2.0*zr*zi + qif; |
| 41 | zi = nzi; |
| 42 | zr = nzr; |
| 43 | z2 = zr*zr+zi*zi; |
| 44 | if(z2 > 4.0f) { |
| 45 | iteration[i] = n; |
| 46 | break; |
| 47 | } |
| 48 | } |
| 49 | """, |
| 50 | "mandlebrot_gpu",) |
| 51 | |
| 52 | # allocate a numpy array |
| 53 | iterations = nm.zeros(nx*ny).astype(nm.int32) |
| 54 | |
| 55 | # allocate a gpu array |
| 56 | iterations_gpu = gpuarray.to_gpu(iterations) |
| 57 | |
| 58 | # perform the gpu calculation |
| 59 | complex_gpu(nm.int16(nx), nm.int16(ny), nm.int16(MAX_ITER), iterations_gpu) |
| 60 | |
| 61 | # copy data from the gpu array to the numpy array |
| 62 | iterations_gpu.get(iterations) |
| 63 | |
| 64 | # reshape the array to look the way we want |
| 65 | image = iterations.reshape(nx,ny) |
| 66 | |
| 67 | bmp.write_image('image.bmp', nx, ny, image, MAX_ITER) |