Speeding up Buddhabrot calculation in PyCuda

I've just started using PyCuda with Python 3, but I have some experience in high performance computing on CPU. I've tried to port one of my old models for generating the Buddhabrot to run on my GPU instead. To briefly explain the algorithm to anyone not familiar, the idea is this:

1. Generate a random complex number $$\z_0\$$
2. Iteratively perform the calculation $$\z_{i+1} = z_i^2 + z_0\$$. Do this $$\N\$$ times
3. Check if the absolute value of the number $$\z_N\$$ is larger than 5*
4. If it is, calculate all the numbers $$\z_i, 0 \leq i \leq N\$$ again and map them to pixels.

*The absolute value only needs to be larger than 2, but due to the parameters I've chosen 5 became a better choice of limit.

These 4 steps need to be repeated at least a billion times, preferably a trillion times for a high quality image. That's why I looked into using PyCuda to speed it up. Here's my current script:

import numpy as np
import pycuda.autoinit
from pycuda.compiler import SourceModule
from pycuda.driver import Device
from pycuda import gpuarray
import time
import scipy.misc

code = """
#include <curand_kernel.h>
#include <pycuda-complex.hpp>
#include <stdio.h>

#define X_MIN -1.5f
#define X_MAX 1.5f
#define Y_MIN -3.2f
#define Y_MAX 2.0f
#define X_DIM %(XDIM)s
#define Y_DIM %(YDIM)s

typedef pycuda::complex<float> cmplx;

const int nstates = %(NGENERATORS)s;
__device__ curandState_t* states[nstates];

extern "C" { __global__ void init_kernel(int seed) {

int idx = threadIdx.x + blockIdx.x * blockDim.x;

if (idx < nstates) {
curandState_t* s = new curandState_t;
if (s != 0) {
curand_init(seed, idx, 0, s);
}

states[idx] = s;
} else {
printf("forbidden memory access %%d/%%d\\n", idx, nstates);
}
} }

__device__ void write_pixel(int idx, cmplx *nums, int *canvas) {
float px = nums[2*idx].imag();
float py = nums[2*idx].real();
px -= X_MIN;
py -= Y_MIN;
px /= X_MAX - X_MIN;
py /= Y_MAX - Y_MIN;
px *= X_DIM;
py *= Y_DIM;
int ix = (int)floorf(px);
int iy = (int)floorf(py);
if (0 <= ix & ix < X_DIM & 0 <= iy & iy < Y_DIM) {
canvas[iy*X_DIM + ix] += 1;
}
}

__device__ void generate_random_complex(float real, float imag, int idx,
cmplx *nums, float *dists, int *counts) {

real *= X_MAX-X_MIN+3;
real += X_MIN-2;
imag *= Y_MAX-Y_MIN+0;
imag += Y_MIN-0;

nums[2*idx+1] = cmplx(real, imag);
nums[2*idx] = cmplx(real, imag);
dists[idx] = 0;
counts[idx] = 0;
}

extern "C" {__global__ void buddha_kernel(int *counts, cmplx *nums,
float *dists, int *canvas) {
int idx = threadIdx.x + blockIdx.x * blockDim.x;
int i, j;
float real, imag;

if (idx < nstates) {
curandState_t s = *states[idx];
for(i = 0; i < 10000; i++) {

real = curand_uniform(&s);
imag = curand_uniform(&s);
generate_random_complex(real, imag, idx, nums, dists, counts);

while (counts[idx] < %(ITERS)s & dists[idx] < 5) {
counts[idx]++;
nums[2*idx] = nums[2*idx]*nums[2*idx] + nums[2*idx+1];
dists[idx] = abs(nums[2*idx]);
}

if (dists[idx] > 5) {
nums[2*idx] = cmplx(0,0);
for (j = 0; j < counts[idx]+1; j++) {
nums[2*idx] = nums[2*idx]*nums[2*idx] + nums[2*idx+1];
write_pixel(idx, nums, canvas);
}
}
}
*states[idx] = s;
} else {
printf("forbidden memory access %%d/%%d\\n", idx, nstates);
}
} }
"""

def print_stats(cpu_canvas, elapsed_time, x_dim, y_dim):
total_iterations = np.sum(cpu_canvas)
max_freq = np.max(cpu_canvas)
min_freq = np.min(cpu_canvas)
print("\tTotal iterations: %.5e" % total_iterations)
print("\tIterations per pixel: %.2f" % (total_iterations / (x_dim*y_dim),))
print("\tMaximum frequency: %d" % max_freq)
print("\tMinimum frequency: %d" % min_freq)
print("\tTotal time: %.2fs" % (elapsed_time,))
print("\tIterations per second: %.2e" % (total_iterations / (elapsed_time),))

def format_and_save(cpu_canvas, x_dim, y_dim, threads, iters):
cpu_canvas /= np.max(cpu_canvas)
cpu_canvas.shape = (y_dim, x_dim)
# this just makes the color gradient more visually pleasing
cpu_canvas = np.minimum(1.1*cpu_canvas, cpu_canvas*.2+.8)

file_name = "pycuda_%dx%d_%d_%d.png" % (x_dim, y_dim, iters, threads)
print("\n\tSaving %s..." % file_name)

scipy.misc.toimage(cpu_canvas, cmin=0.0, cmax=1.0).save(file_name)
print("\tImage saved!\n")

def generate_image(x_dim, y_dim, iters):

b_s = 2**10

device = Device(0)
print("\n\t" + device.name(), "\n")
context = device.make_context()

formatted_code = code % {
"XDIM" : x_dim,
"YDIM" : y_dim,
"ITERS" : iters
}

# generate kernel and setup random number generation
module = SourceModule(formatted_code, no_extern_c=True)
init_func = module.get_function("init_kernel")
fill_func = module.get_function("buddha_kernel")
seed = np.int32(np.random.randint(0, 1<<31))

# initialize all numpy arrays
samples = gpuarray.zeros(2*threads*b_s, dtype = np.complex64)
dists = gpuarray.zeros(threads*b_s, dtype = np.float32)
counts = gpuarray.zeros(threads*b_s, dtype = np.int32)
canvas = gpuarray.zeros(y_dim*x_dim, dtype = np.int32)

# start calculation
t0 = time.time()
fill_func(counts, samples, dists, canvas, block=(b_s,1,1), grid=(threads,1,1))
context.synchronize()
t1 = time.time()

# fetch buffer from gpu and save as image
cpu_canvas = canvas.get().astype(np.float64)
context.pop()
print_stats(cpu_canvas, t1-t0, x_dim, y_dim)
format_and_save(cpu_canvas, x_dim, y_dim, threads, iters)

if __name__ == "__main__":

x_dim = 1440
y_dim = 2560
iters = 20
generate_image(x_dim, y_dim, iters)

Here's a sample output when I run it on my laptop:

GeForce GTX 1050

Total iterations: 5.84026e+08
Iterations per pixel: 158.43
Maximum frequency: 917
Minimum frequency: 0
Total time: 3.85s
Iterations per second: 1.52e+08

Saving pycuda_1440x2560_20_64.png...
Image saved!

This runs pretty fast, but I'm hoping to squeeze some more performance out of it, and learn some more about GPU programming. I don't really know if I'm doing it right. The main reason that I think that this can be sped up is that I wrote a multithreaded CPU solution of this a few years ago, and it is almost as fast as this CUDA implementation (speed difference is less than a factor 4).

Any tips on making this run faster, or general things to think about when coding for CUDA would be appreciated!

EDIT: I know that there are optimizations related to this specific problem that can be implemented, mainly regarding sampling. I plan on implementering some importance sampling, but for this question I'm mostly interested in general CUDA practices.

• Are you sure everything is correct in your code? Just looked at the last edit: 1) In write_pixel(...) you pass temp but overwrite it immediately. You pass ix and iy and never use their passed values. You could just pass z to to_pixel(...) and create a temp there or even just use the passed values to do the calculations. 2) In your main kernel you have for(i = 0; i < 1; i++)... 3) This -> for (coord.x = 0; coord.x < 1; coord.x += 1/(float)blockDim.x) may not do what you think it does (e.g. for blockDim.x = 192 on my machine). 4) You never set ix/iy in your main kernel. – Shadow Oct 23 '18 at 9:55
• @Shadow I'm not an expert when it comes to cuda, but my thinking was that pre-allocating variables could help with memory performance. It might be a terrible micro-optimization though, but I have attempted multiple solutions to squeeze every last bit of performance out of this. – maxb Oct 23 '18 at 12:48
• @Shadow could you elaborate on your comment about for (coord.x = 0; coord.x < 1; coord.x += 1/(float)blockDim.x)? I have changed that part a bit in the final iteration, but why wouldn't that part do what I think it does? – maxb Oct 25 '18 at 7:40
• Could you post the new optimized code as a self-answer instead, please? Adding improved code to the question is not allowed. – t3chb0t Oct 25 '18 at 8:09
• @t3chb0t Thanks for the info! I'll restructure the question and write an answer instead. – maxb Oct 25 '18 at 8:24

Why using binary and instead of logical and in write_pixel:

if (0 <= ix & ix < X_DIM & 0 <= iy & iy < Y_DIM) {
canvas[iy*X_DIM + ix] += 1;
}

And why not moving (and changing a bit) the check before calculations?

__device__
void write_pixel(float2 temp, int ix, int iy,
float4 z, unsigned int *canvas) {
if (X_MIN <= z.x && z.x <= X_MAX && Y_MIN <= z.y && z.y <= Y_MAX  ) {
temp.x = z.y;
temp.y = z.x;
to_pixel(temp, ix, iy);
atomicAdd(&(canvas[iy*X_DIM + ix]), 1);
}
}

Did you tried to inlining computations in to_pixel :

__device__ void to_pixel(float2 &temp, int &ix, int &iy) {
ix = __float2int_rd((temp.x - X_MIN) / (X_MAX - X_MIN) *  X_DIM);
iy = __float2int_rd((temp.y - Y_MIN) / (Y_MAX - Y_MIN) *  Y_DIM);
}

Why dont pass directly the two floats to to_pixel instead of using a float2?

PS: I dont know too much PyCuda (and more generally, interoperate Python and C), does it disallow const ?

• Thanks for the reply! I'll try all of these things out and benchmark them, and then I'll write a proper response to your comment. I really like your idea of checking against X_MIN, X_MAX etc. before calculating. – maxb Oct 23 '18 at 12:53
• Also, pycuda uses nvcc for compilation, so I think it should be fully possible to use const and any other C/C++ features. The code gets passed as a string to nvcc, and then you define hooks to the global functions. – maxb Oct 23 '18 at 13:22
• I've done some benchmarking now, and it seems that the changes above affect the speed of the script by less than 2%. As the changes are so small, it's difficult to draw conclusions from them, but it seems as if the logical and is a tiny bit slower than the binary and. My guess is that the loop prediction works better with a single boolean statement compared to multiple ones, but I'm only speculating. And from my understanding, using #define and const gives the same machine code when using -O3. I still think that the biggest bottleneck by far is the writing to global memory. – maxb Oct 24 '18 at 11:05
• And since the canvas array is at least 1MB, I can't store a copy in the shared memory to coalesce writes further. If I run with -O0 and comment the line atomicAdd(&(canvas[iy*X_DIM + ix]), 1);, the entire program runs almost 100 times faster. – maxb Oct 24 '18 at 11:11

Packaging floats

I have implemented some selective sampling that I found on a forum, and I've changed the container for my complex numbers from a pycuda::complex<float>* to a float4*. Both of these things have contributed to a 76% performance increase. Here's the updated code:

import numpy as np
import pycuda.autoinit
from pycuda.compiler import SourceModule
from pycuda.driver import Device
from pycuda import gpuarray
import time
import scipy.misc

code = """
#include <curand_kernel.h>
#include <stdio.h>

#define X_MIN -1.5f
#define X_MAX 1.5f
#define Y_MIN -3.2f
#define Y_MAX 2.0f
#define X_DIM %(XDIM)s
#define Y_DIM %(YDIM)s
#define ITERS %(ITERS)s

const int nstates = %(NGENERATORS)s;
__device__ curandState_t* states[nstates];

extern "C" { __global__ void init_kernel(int seed) {

int idx = threadIdx.x + blockIdx.x * blockDim.x;

if (idx < nstates) {
curandState_t* s = new curandState_t;
if (s != 0) {
curand_init(seed, idx, 0, s);
}

states[idx] = s;
} else {
printf("forbidden memory access %%d/%%d\\n", idx, nstates);
}
} }

__device__ void to_pixel(float &px, float &py, int &ix, int &iy) {
px -= X_MIN;
py -= Y_MIN;
px /= X_MAX - X_MIN;
py /= Y_MAX - Y_MIN;
px *= X_DIM;
py *= Y_DIM;
ix = __float2int_rd(px);
iy = __float2int_rd(py);
}

__device__
void write_pixel(int idx, float px, float py, int ix, int iy,
float4 *z, unsigned int *canvas) {
px = z[idx].y;
py = z[idx].x;
to_pixel(px, py, ix, iy);
if (0 <= ix & ix < X_DIM & 0 <= iy & iy < Y_DIM) {
canvas[iy*X_DIM + ix] += 1;
}
}

__device__
void generate_random_complex(float real, float imag, int idx,
float4 *z, float *dists, unsigned int *counts) {

real *= X_MAX-X_MIN+3;
real += X_MIN-2;
imag *= Y_MAX-Y_MIN+0;
imag += Y_MIN-0;

z[idx].x = real;
z[idx].y = imag;
z[idx].z = real;
z[idx].w = imag;
dists[idx] = 0;
counts[idx] = 0;
}

__device__
bool check_bulbs(int idx, float4 *z) {
float zw2 = z[idx].w*z[idx].w;
bool main_card = !(((z[idx].z-0.25)*(z[idx].z-0.25)
+ (zw2))*(((z[idx].z-0.25)*(z[idx].z-0.25)
+ (zw2))+(z[idx].z-0.25)) < 0.25* zw2);
bool period_2 = !((z[idx].z+1.0) * (z[idx].z+1.0) + (zw2) < 0.0625);
bool smaller_bulb = !((z[idx].z+1.309)*(z[idx].z+1.309) + zw2 < 0.00345);
bool smaller_bottom = !((z[idx].z+0.125)*(z[idx].z+0.125)
+ (z[idx].w-0.744)*(z[idx].w-0.744) < 0.0088);
bool smaller_top = !((z[idx].z+0.125)*(z[idx].z+0.125)
+ (z[idx].w+0.744)*(z[idx].w+0.744) < 0.0088);
return main_card & period_2 & smaller_bulb & smaller_bottom & smaller_top;
}

extern "C" {__global__ void buddha_kernel(unsigned int *counts, float4 *z,
float *dists, unsigned int *canvas) {
int idx = threadIdx.x + blockIdx.x * blockDim.x;
int i, j, ix, iy;
float real, imag;//, temp0, temp1;

if (idx < nstates) {

curandState_t s = *states[idx];
for(i = 0; i < 10000; i++) {

real = curand_uniform(&s);
imag = curand_uniform(&s);
generate_random_complex(real, imag, idx, z, dists, counts);
if (check_bulbs(idx, z)) {
while (counts[idx] < ITERS & dists[idx] < 25) {
counts[idx]++;
real = z[idx].x*z[idx].x - z[idx].y*z[idx].y + z[idx].z;
imag = 2*z[idx].x*z[idx].y + z[idx].w;
z[idx].x = real;
z[idx].y = imag;
dists[idx] = z[idx].x*z[idx].x + z[idx].y*z[idx].y;
}

if (dists[idx] > 25) {
z[idx].x = 0;
z[idx].y = 0;
for (j = 0; j < counts[idx]+1; j++) {
real = z[idx].x*z[idx].x - z[idx].y*z[idx].y + z[idx].z;
imag = 2*z[idx].x*z[idx].y + z[idx].w;
z[idx].x = real;
z[idx].y = imag;
write_pixel(idx, real, imag, ix, iy, z, canvas);
}
}
}
}
*states[idx] = s;
} else {
printf("forbidden memory access %%d/%%d\\n", idx, nstates);
}
} }
"""

def print_stats(cpu_canvas, elapsed_time, x_dim, y_dim):
total_iterations = np.sum(cpu_canvas)
max_freq = np.max(cpu_canvas)
min_freq = np.min(cpu_canvas)
print("\tTotal iterations: %.5e" % total_iterations)
print("\tIterations per pixel: %.2f" % (total_iterations / (x_dim*y_dim),))
print("\tMaximum frequency: %d" % max_freq)
print("\tMinimum frequency: %d" % min_freq)
print("\tTotal time: %.2fs" % (elapsed_time,))
print("\tIterations per second: %.2e" % (total_iterations / (elapsed_time),))

def format_and_save(cpu_canvas, x_dim, y_dim, threads, iters):
cpu_canvas /= np.max(cpu_canvas)
cpu_canvas.shape = (y_dim, x_dim)
# this just makes the color gradient more visually pleasing
cpu_canvas = np.minimum(1.1*cpu_canvas, cpu_canvas*.2+.8)

file_name = "pycuda_%dx%d_%d_%d.png" % (x_dim, y_dim, iters, threads)
print("\n\tSaving %s..." % file_name)
scipy.misc.toimage(cpu_canvas, cmin=0.0, cmax=1.0).save(file_name)
print("\tImage saved!\n")

def generate_image(x_dim, y_dim, iters):

b_s = 2**8

device = Device(0)
print("\n\t" + device.name(), "\n")
context = device.make_context()

formatted_code = code % {
"XDIM" : x_dim,
"YDIM" : y_dim,
"ITERS" : iters
}

# generate kernel and setup random number generation
module = SourceModule(
formatted_code,
no_extern_c=True,
options=['--use_fast_math', '-O3', '--ptxas-options=-O3']
)
init_func = module.get_function("init_kernel")
fill_func = module.get_function("buddha_kernel")
seed = np.int32(np.random.randint(0, 1<<31))

# initialize all numpy arrays
samples = gpuarray.zeros(threads*b_s, dtype = gpuarray.vec.float4)
dists = gpuarray.zeros(threads*b_s, dtype = np.float32)
counts = gpuarray.zeros(threads*b_s, dtype = np.uint32)
canvas = gpuarray.zeros(y_dim*x_dim, dtype = np.uint32)
t0 = time.time()
fill_func(counts, samples, dists, canvas, block=(b_s,1,1), grid=(threads,1,1))
context.synchronize()
t1 = time.time()

# fetch buffer from gpu and save as image
cpu_canvas = canvas.get().astype(np.float64)
context.pop()
print_stats(cpu_canvas, t1-t0, x_dim, y_dim)
format_and_save(cpu_canvas, x_dim, y_dim, threads, iters)

if __name__ == "__main__":

x_dim = 1440
y_dim = 2560
iters = 20
generate_image(x_dim, y_dim, iters)

And here's some sample output:

GeForce GTX 1050

Total iterations: 5.83970e+08
Iterations per pixel: 158.41
Maximum frequency: 886
Minimum frequency: 0
Total time: 2.20s
Iterations per second: 2.66e+08

Saving pycuda_1440x2560_20_256.png...
Image saved!

Don't use global memory when you don't have to

I managed to get another 30% speedup by switching to using thread-local variables instead of global arrays. Now the major bottleneck seems to be the writing to the canvas array, since that has to reside in global memory. I'm still not sure how I should handle those writes as efficiently as possible.

ALWAYS try to coalesce global memory writes

I managed to get another 320% speedup, making my final version 7.37 times faster than the original! To achieve this, I removed all arrays except canvas from the global memory, and handled all variables in (what I'm guessing is) shared memory. That provided me with a 50% speedup I think.

What made everything a whole lot faster was my method of memory access coalescing. Since the algorithm is based on choosing random complex numbers, the behavior of each thread is not easily predictable. This leads to terrible branch prediction and global memory access.

By using some logic from this question on Stack Overflow, thought of a method to "sort" the random numbers.

Grid sampling

We can visualize the area that the complex numbers are sampled from as a rectangle on the complex plane. For our algorithm to work, we must ensure that every point within the rectangle is equally likely to be chosen.

To ensure this, we split the sampling area into a grid of rectangles, like the image below. We run our algorithm for each cell in the grid, only sampling from the current cell. When we have iterated over all the cells in the grid, then all points across the entire rectangle has had an equal chance of being chosen. The advantage of sampling in grid cells sequentially is that the points within a single grid cell tend to behave similarly. Since they are close to each other, they almost follow the same orbits, which helps both with branch prediction and memory coalescing.

So with this code:

import numpy as np
import pycuda.autoinit
from pycuda.compiler import SourceModule
from pycuda.driver import Device
from pycuda import gpuarray
import time
import scipy.misc

code = """
#include <curand_kernel.h>
#include <stdio.h>

#define X_MIN -1.5f
#define X_MAX 1.5f
#define Y_MIN -3.2f
#define Y_MAX 2.0f

#define X_MIN_SAMPLE -2.1f
#define X_MAX_SAMPLE 1.1f
#define Y_MIN_SAMPLE -1.8f
#define Y_MAX_SAMPLE 1.8f

#define X_DIM %(XDIM)s
#define Y_DIM %(YDIM)s
#define ITERS %(ITERS)s

const int nstates = %(NGENERATORS)s;
__device__ curandState_t* states[nstates];

extern "C" {
__global__
void init_kernel(int seed) {

int idx = threadIdx.x + blockIdx.x * blockDim.x;

if (idx < nstates) {
curandState_t* s = new curandState_t;
if (s != 0) {
curand_init(seed, idx, 0, s);
}

states[idx] = s;
} else {
printf("forbidden memory access %%d/%%d\\n", idx, nstates);
}
}
}

__device__ void to_pixel(float2 &temp, int &ix, int &iy) {
temp.x -= X_MIN;
temp.y -= Y_MIN;
temp.x /= X_MAX - X_MIN;
temp.y /= Y_MAX - Y_MIN;
temp.x *= X_DIM;
temp.y *= Y_DIM;
ix = __float2int_rd(temp.x);
iy = __float2int_rd(temp.y);
}

__device__
void write_pixel(float2 temp, int ix, int iy,
float4 z, unsigned int *canvas) {
temp.x = z.y;
temp.y = z.x;
to_pixel(temp, ix, iy);
if (0 <= ix & ix < X_DIM & 0 <= iy & iy < Y_DIM) {
atomicAdd(&(canvas[iy*X_DIM + ix]), 1);
}
}

__device__
void generate_random_complex(float2 temp,
float4 &z, float &dist, unsigned int &count) {

temp.x *= X_MAX_SAMPLE-X_MIN_SAMPLE;
temp.x += X_MIN_SAMPLE;
temp.y *= Y_MAX_SAMPLE-Y_MIN_SAMPLE;
temp.y += Y_MIN_SAMPLE;

z.x = temp.x;
z.y = temp.y;
z.z = temp.x;
z.w = temp.y;
dist = 0;
count = 0;
}

__device__
bool check_bulbs(float4 z) {
float zw2 = z.w*z.w;
bool main_card = !(((z.z-0.25)*(z.z-0.25)
+ (zw2))*(((z.z-0.25)*(z.z-0.25)
+ (zw2))+(z.z-0.25)) < 0.25* zw2);
bool period_2 = !((z.z+1.0) * (z.z+1.0) + (zw2) < 0.0625);
bool smaller_bulb = !((z.z+1.309)*(z.z+1.309) + zw2 < 0.00345);
bool smaller_bottom = !((z.z+0.125)*(z.z+0.125)
+ (z.w-0.744)*(z.w-0.744) < 0.0088);
bool smaller_top = !((z.z+0.125)*(z.z+0.125)
+ (z.w+0.744)*(z.w+0.744) < 0.0088);
return main_card & period_2 & smaller_bulb & smaller_bottom & smaller_top;
}

extern "C" {
__global__
void buddha_kernel(unsigned int *canvas, int seed) {
int idx = blockIdx.x
+ threadIdx.x * gridDim.x
+ threadIdx.y * gridDim.x * blockDim.x;
int i, j, ix, iy;
float2 temp, coord;
unsigned int count;
float4 z;
float dist;
curandState_t s;
curand_init(seed, idx, 0, &s);

for (coord.x = 0; coord.x < 1; coord.x += 1/(float)blockDim.x) {
for (coord.y = 0; coord.y < 1; coord.y += 1/(float)blockDim.x) {

for(i = 0; i < 1; i++) {

temp.x = curand_uniform(&s);
temp.y = curand_uniform(&s);
temp.x /= (float)blockDim.x;
temp.y /= (float)blockDim.x;
temp.x += coord.x;
temp.y += coord.y;

generate_random_complex(temp, z, dist, count);
if (check_bulbs(z)) {
while (count < ITERS & dist < 25) {
count++;
temp.x = z.x*z.x - z.y*z.y + z.z;
temp.y = 2*z.x*z.y + z.w;
z.x = temp.x;
z.y = temp.y;
dist = z.x*z.x + z.y*z.y;
}

if (dist > 25) {
z.x = z.z;
z.y = z.w;
for (j = 0; j < count; j++) {
temp.x = z.x*z.x - z.y*z.y + z.z;
temp.y = 2*z.x*z.y + z.w;
z.x = temp.x;
z.y = temp.y;
write_pixel(temp, ix, iy, z, canvas);
}
}
}
}

}
}

}
}
"""

def print_stats(cpu_canvas, elapsed_time, x_dim, y_dim):
total_iterations = np.sum(cpu_canvas)
max_freq = np.max(cpu_canvas)
min_freq = np.min(cpu_canvas)
print("\tTotal iterations: %.5e" % total_iterations)
print("\tIterations per pixel: %.2f" % (total_iterations / (x_dim*y_dim),))
print("\tMaximum frequency: %d" % max_freq)
print("\tMinimum frequency: %d" % min_freq)
print("\tTotal time: %.2fs" % (elapsed_time,))
print("\tIterations per second: %.2e" % (total_iterations / (elapsed_time),))

def format_and_save(cpu_canvas, x_dim, y_dim, threads, iters):
cpu_canvas /= max(1, np.max(cpu_canvas))
# cpu_canvas.shape = (y_dim, x_dim)
# this just makes the color gradient more visually pleasing
cpu_canvas = np.minimum(2.5*cpu_canvas, cpu_canvas*.2+.8)

file_name = "pycuda_%dx%d_%d_%d.png" % (x_dim, y_dim, iters, threads)
print("\n\tSaving %s..." % file_name)
scipy.misc.toimage(cpu_canvas, cmin=0.0, cmax=1.0).save(file_name)
print("\tImage saved!\n")

def generate_image(x_dim, y_dim, iters):

b_s = 2**9

device = Device(0)
print("\n\t" + device.name(), "\n")
context = device.make_context()

formatted_code = code % {
"XDIM" : x_dim,
"YDIM" : y_dim,
"ITERS" : iters
}

# generate kernel and setup random number generation
module = SourceModule(
formatted_code,
no_extern_c=True,
options=['--use_fast_math', '-O3', '--ptxas-options=-O3']
)
fill_func = module.get_function("buddha_kernel")
seed = np.int32(np.random.randint(0, 1<<31))
canvas = gpuarray.zeros((y_dim, x_dim), dtype = np.uint32)

t0 = time.time()
fill_func(canvas, seed, block=(b_s,1,1), grid=(threads,1,1))
context.synchronize()
t1 = time.time()

# fetch buffer from gpu and save as image
cpu_canvas = canvas.get().astype(np.float64)
context.pop()
print_stats(cpu_canvas, t1-t0, x_dim, y_dim)
format_and_save(cpu_canvas, x_dim, y_dim, threads, iters)

if __name__ == "__main__":

x_dim = 1440
y_dim = 2560
iters = 20
generate_image(x_dim, y_dim, iters)

The code above shows a lot of improvement:

GeForce GTX 1050

Total iterations: 2.59801e+09
Iterations per pixel: 704.76
Maximum frequency: 7873
Minimum frequency: 0
Total time: 2.33s
Iterations per second: 1.12e+09

Saving pycuda_1440x2560_20_16.png...
Image saved!

If anyone has any suggestions for further improvements, I'm happy to test them out and see what works.

Minor optimizations

I changed the lines:

for (coord.x = 0; coord.x < 1; coord.x += 1/(float)blockDim.x) {
for (coord.y = 0; coord.y < 1; coord.y += 1/(float)blockDim.x) {

float gridSize = 1/1024.0f;
for (coord.x = 0; coord.x < 1; coord.x += gridSize) {
for (coord.y = 0; coord.y < 1; coord.y += gridSize) {

That way, I could tweak the block sizes and threads independently of the grid structure for sampling. Doing this, and increasing the number of threads, I managed to get another 43% speed increase, but the downside is that it's not possible to run quick benchmarks using this method, since each thread has to perform $$\2^{20}\$$ iterations at a minimum.

GeForce GTX 1050

Total iterations: 2.54575e+11
Iterations per pixel: 69057.85
Maximum frequency: 792006
Minimum frequency: 0
Total time: 157.91s
Iterations per second: 1.61e+09

Saving pycuda_1440x2560_20_128.png...
Image saved!

Realization that the sampling can be made more efficient

From reading this blog post, I realized that my sampling could be made more effective. Since the mandelbrot set is symmetric on the imaginary axis, I can mirror any orbit to its complex conjugate. This also means that I don't have to sample any points with a negative imaginary part, and I've halved my sampling area.

This was very easy to implement, just by breaking the second part of the orbit calculation into its own function:

__device__ __forceinline__
void write_to_image(float4 z, float2 temp, int2 ixy,
int count, unsigned int *canvas) {
z.x = z.z;
z.y = z.w;
for (int j = 0; j < count; j++) {
temp.x = z.x*z.x - z.y*z.y + z.z;
temp.y = 2*z.x*z.y + z.w;
z.x = temp.x;
z.y = temp.y;
write_pixel(temp, ixy, z, canvas);
}
}

and then modifying the last part of the main loop to become:

if (dist > 4) {
write_to_image(z, temp, ixy, count, canvas);
z.w *= -1;
write_to_image(z, temp, ixy, count, canvas);
}

With some tweaking, this results in a 65% speedup!

GeForce GTX 1050

Total iterations: 1.69717e+12
Iterations per pixel: 460385.79
Maximum frequency: 5277235
Minimum frequency: 0
Total time: 639.79s
Iterations per second: 2.65e+09

Saving pycuda_1440x2560_20_128.png...
Image saved!

With this result, I have gotten my script to run over 17x faster! I'll go ahead and consider that a success, and call it a day.

I guess that the period for review of this question is over, but If you (like me) find yourself tackling this problem, look at my current solution to get some pointers on how to make things faster.