julia fractal set using CUDA and complex numbers

julia.cu

// check comments for where to find these headers
#include "../common/book.h"
#include "../common/cpu_bitmap.h"

#define DIM 1000
#define SCALE 0.1f
#define STEPS 300
#define JULIA_REAL (-0.8f)
#define JULIA_IMAG (0.155f)


struct cuComplex {
    float   r;
    float   i;

    __device__ cuComplex( float a, float b ) : r(a), i(b) {}

    __device__ float magnitude2() const {
        return r * r + i * i;
    }

    __device__ cuComplex operator * (const cuComplex& rhs) const {
        return {r * rhs.r - i * rhs.i, i * rhs.r + r * rhs.i};
    }

    __device__ cuComplex operator + (const cuComplex& rhs) const {
        return {r + rhs.r, i + rhs.i};
    }
};


__device__ int julia( unsigned int x, unsigned int y ) {
    // todo make this a cmd arg
    float jx = SCALE * (float)(DIM - x) / (DIM);
    float jy = SCALE * (float)(DIM - y) / (DIM);

    cuComplex c(JULIA_REAL, JULIA_IMAG);
    cuComplex a(jx, jy);

    for (int i = 0; i < STEPS; i++) {
        a = a * a + c;
        if (a.magnitude2() > DIM) {
            return 0;
        }
    }
    return 1;
}


__global__ void kernel( unsigned char *ptr ){
    unsigned int x = blockIdx.x;
    unsigned int y = blockIdx.y;
    unsigned int offset = x + y * gridDim.x;

    int juliaValue = julia( x, y );
    ptr[offset * 4 + 0] = 0;
    ptr[offset * 4 + 1] = 0;
    ptr[offset * 4 + 2] = 255 * juliaValue;
    ptr[offset * 4 + 3] = 255;
}


int main() {
    CPUBitmap bitmap( DIM, DIM );

    unsigned char* dev_bitmap;

    HANDLE_ERROR(cudaMalloc((void**)&dev_bitmap, bitmap.image_size()));

    dim3 grid(DIM, DIM);

    kernel<<<grid, 1>>>(dev_bitmap);

    HANDLE_ERROR(cudaMemcpy(bitmap.get_ptr(),
                            dev_bitmap,
                            bitmap.image_size(),
                            cudaMemcpyDeviceToHost));

    bitmap.display_and_exit();

    cudaFree(dev_bitmap);

    return EXIT_SUCCESS;
}