cuda_experiments.py

import pycuda.autoinit
import pycuda.driver as drv
import numpy as np
import h5py
from imageio import imwrite
import matplotlib.pyplot as plt
import matplotlib

from pycuda.compiler import SourceModule

mod = SourceModule(
    '''
__device__ float kernel_func(float x) {
    float C = 2.5464790894703255;
    float kernel;
    if (x <= 0.5) {
        kernel = 1.-6.*x*x*(1.-x);
    } else {
        kernel = 2.*(1.-x)*(1.-x)*(1.-x);
    }
    return kernel * C;
}

__global__ void render_image(float *coordinates, float* smoothing_length, float* mass, float* density, long long n_particles, int* image, long long image_dim) {
    int idx = (blockIdx.x * blockDim.x) + threadIdx.x;
    if (idx > n_particles) {
        return;
    }

    int particle_x = coordinates[3*idx + 0] - 10412.19823645 + image_dim/2.0;
    int particle_y = coordinates[3*idx + 1] - 1952.631666050 + image_dim/2.0;
    //int particle_z = coordinates[3*idx + 2] - 4828.253386520 + image_dim/2.0;

    float ih_j2 = 1.0 / max(smoothing_length[idx]*smoothing_length[idx], 1.0);

    float prefactor = mass[idx] / pow(smoothing_length[idx], 2);

    int pixel_radius = (int)smoothing_length[idx];

    for (int x = particle_x - pixel_radius; x < particle_x + pixel_radius + 1; x++) {
        if (x < 0 || x >= image_dim) {
            continue;
        }
        int x_distance = (x - particle_x) * (x - particle_x);

        for (int y = particle_y - pixel_radius; y < particle_y + pixel_radius + 1; y++) {
            if (y < 0 || y >= image_dim) {
                continue;
            }

            int y_distance = (y - particle_y) * (y - particle_y);
            float distance = (x_distance + y_distance) * ih_j2;
            if (distance >= 1.0) {
                continue;
            }

            int image_index = y * image_dim + x;
            atomicAdd(&image[image_index], prefactor * kernel_func(distance) * 1e9); // evil rescaling hack instead of writing an atomic f32 add
        }
    }
}
'''
)

with h5py.File("snap_N64L16_135.hdf5", "r") as f:
    coordinates = f["PartType0/Coordinates"][:]
    smoothing_length = f["PartType0/SmoothingLength"][:]
    mass = f["PartType0/Masses"][:]
    density = f["PartType0/Density"][:]

render_image = mod.get_function("render_image")
image = np.zeros((1024, 1024), dtype=np.int32)

import time
start = time.time()
render_image(
    drv.In(coordinates),
    drv.In(smoothing_length),
    drv.In(mass),
    drv.In(density),
    np.uint64(coordinates.shape[0]),
    drv.Out(image),
    np.uint64(image.shape[0]),
    block=(64, 1, 1),
    grid=(coordinates.shape[0]//64, 1),
)
print(time.time() - start)

plt.imshow(image, norm=matplotlib.colors.LogNorm(), cmap='magma').write_png('test.png')