OpenCL Mandelbrot fixed

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pyopencl as cl
import numpy as np
import matplotlib.pyplot as plt

ctx = cl.create_some_context(interactive=True)

def mandelbrot_opencl(q, maxiter):
    global ctx
    queue = cl.CommandQueue(ctx)

    output = np.empty(q.shape, dtype=np.uint16)
    prg = cl.Program(ctx, """
    #pragma OPENCL EXTENSION cl_khr_byte_addressable_store : enable
    __kernel void mandelbrot(__global float2 *q,
                     __global ushort *output, ushort const maxiter, int const width)
    {
        int gid = get_global_id(0)*width + get_global_id(1);
        float nreal, real = 0;
        float imag = 0;
        output[gid] = 0;
        for(int curiter = 0; curiter < maxiter; curiter++) {
            nreal = real*real - imag*imag + q[gid].x;
            imag = 2* real*imag + q[gid].y;
            real = nreal;
            if (real*real + imag*imag > 4.0f) {
                 output[gid] = curiter;
                 return;
            }
        }
        output[gid] = maxiter;
    }
    """).build()
    mf = cl.mem_flags
    q_opencl = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=q)
    output_opencl = cl.Buffer(ctx, mf.WRITE_ONLY, output.nbytes)

    prg.mandelbrot(queue, output.shape, None, q_opencl,
                   output_opencl, np.uint16(maxiter), np.int32(output.shape[1]))
    cl.enqueue_copy(queue, output, output_opencl).wait()
    return output

if __name__ == '__main__':
    width, height = 750, 500
    ys, xs = np.mgrid[1:-1:height*1j, -2:1:width*1j]
    data = (xs + 1j*ys).astype('complex64')
    out = mandelbrot_opencl(data, 100)
    plt.imshow(abs(out), cmap=plt.cm.nipy_spectral_r)
    plt.show()