Mandelbrot on GPU (MATLAB gpuArray and Python cupy)

#mandelbrot.md

Mandelbrot implementations with utilization of GPU

1. MATLAB: mandelbrot_gpu.m (gpuArray)
2. Python mandelbrot_gpu.py (cupy)

computeMandelbrot.m

function div_time = computeMandelbrot(x, y, maxIterations)
    c = complex(x, y);
    z = c; % should I initialize with c or zero?
    div_time = 0;
    while (div_time < maxIterations) && (abs(z) <= 2)
        z = z^2 + c;
        div_time = div_time + 1;
    end
end

mandelbrot_gpu.m

maxIterations = uint16(1000);

width = 5000;
height = 5000;
xCenter = -0.5;
yCenter = 0.0;
zoom = 1.0;

xWidth = 1.5;
yHeight = 1.5 * height / width;
xFrom = xCenter - xWidth / zoom;
xTo = xCenter + xWidth / zoom;
yFrom = yCenter - yHeight / zoom;
yTo = yCenter + yHeight / zoom;

t = tic();

x = gpuArray.linspace(xFrom, xTo, width);
y = gpuArray.linspace(yFrom, yTo, height);
[xGrid,yGrid] = meshgrid( x, y );

div_time_g = arrayfun(@computeMandelbrot, xGrid, yGrid, maxIterations);
div_time = gather(div_time_g);

cpuTime = toc(t);

fig = gcf;
fig.Position = [200 200 600 600];
imagesc( x, y, div_time );
colormap( [jet(); flipud(jet()); 0 0 0]);
axis off
title(sprintf('Mandelbrot, time %1.2fsecs', cpuTime));

mandelbrot_gpu.py

import cupy as cp
import matplotlib.pyplot as plt
from timeit import default_timer as timer
from datetime import timedelta


def mandelbrot(height, width, x=-0.5, y=0.0, zoom=1, max_iterations=100):
  x_width = 1.5
  y_height = 1.5 * height / width
  x_from = x - x_width / zoom
  x_to = x + x_width / zoom
  y_from = y - y_height / zoom
  y_to = y + y_height / zoom

  x = cp.linspace(x_from, x_to, width).reshape((1, width))
  y = cp.linspace(y_from, y_to, height).reshape((height, 1))
  c = (x + 1j * y)

  calculateForElement = cp.ElementwiseKernel(
    'complex128 c, uint16 max_iterations',
    'uint16 div_time',
    '''
      complex<float> z = c;
      div_time = 0;
      while (div_time < max_iterations && abs(z) <= 2) {
        z = pow(z, 2.0) + c;
        div_time++;
      }
    ''',
    'calculateForElement'
  )

  div_time = calculateForElement(c, max_iterations)
  return cp.asnumpy(div_time)


start = timer()
result = mandelbrot(5000, 5000, max_iterations=1000)
end = timer()

print("Time: {0}".format(timedelta(seconds=end - start)))

# Default image of Mandelbrot set
plt.imshow(result, cmap='magma')
plt.show()