fractal_generator.cu

#include <iostream>
#include <vector>
#include <fstream>
#include <complex> // For std::complex, though we'll implement complex arithmetic manually in kernels for performance
#include <string>
#include <limits> // For numeric_limits

// CUDA includes
#include <cuda_runtime.h>

// Macro for checking CUDA errors
#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA error at %s:%d - %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// Structure to hold RGB color data
struct Color {
    unsigned char r, g, b;
};

// --- CUDA Kernels ---

/**
 * @brief CUDA kernel to generate the Mandelbrot set.
 *
 * Each thread calculates the color for one pixel.
 * The Mandelbrot set is defined by the iteration z_{n+1} = z_n^2 + c,
 * where c is the complex number corresponding to the pixel coordinates,
 * and z_0 = 0.
 *
 * @param image_data Pointer to the device memory where color data will be stored.
 * @param width Width of the image in pixels.
 * @param height Height of the image in pixels.
 * @param max_iterations Maximum number of iterations for each point.
 * @param min_x Minimum real value for the fractal plane.
 * @param max_x Maximum real value for the fractal plane.
 * @param min_y Minimum imaginary value for the fractal plane.
 * @param max_y Maximum imaginary value for the fractal plane.
 */
__global__ void mandelbrot_kernel(Color* image_data, int width, int height, int max_iterations,
                                  double min_x, double max_x, double min_y, double max_y) {
    // Calculate global thread ID
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y;

    // Check if the current thread is within the image bounds
    if (col < width && row < height) {
        // Map pixel coordinates to complex plane coordinates
        double c_real = min_x + (col / (double)width) * (max_x - min_x);
        double c_imag = min_y + (row / (double)height) * (max_y - min_y);

        double z_real = 0.0;
        double z_imag = 0.0;

        int iteration = 0;
        // Iterate z = z*z + c
        while (z_real * z_real + z_imag * z_imag < 4.0 && iteration < max_iterations) {
            double z_real_new = z_real * z_real - z_imag * z_imag + c_real;
            double z_imag_new = 2.0 * z_real * z_imag + c_imag;
            z_real = z_real_new;
            z_imag = z_imag_new;
            iteration++;
        }

        // Determine color based on iteration count
        Color pixel_color;
        if (iteration == max_iterations) {
            // Point is likely in the set (black)
            pixel_color = {0, 0, 0};
        } else {
            // Point escaped, color based on iteration count (smooth coloring can be added for better visuals)
            // Simple coloring: gradient from blue to green to red
            double t = (double)iteration / max_iterations;
            pixel_color.r = (unsigned char)(9 * (1 - t) * t * t * t * 255);
            pixel_color.g = (unsigned char)(15 * (1 - t) * (1 - t) * t * t * 255);
            pixel_color.b = (unsigned char)(8.5 * (1 - t) * (1 - t) * (1 - t) * t * 255);
        }

        // Store the color in the image data array
        image_data[row * width + col] = pixel_color;
    }
}

/**
 * @brief CUDA kernel to generate the Julia set.
 *
 * Each thread calculates the color for one pixel.
 * The Julia set is defined by the iteration z_{n+1} = z_n^2 + c,
 * where c is a fixed complex constant, and z_0 is the complex number
 * corresponding to the pixel coordinates.
 *
 * @param image_data Pointer to the device memory where color data will be stored.
 * @param width Width of the image in pixels.
 * @param height Height of the image in pixels.
 * @param max_iterations Maximum number of iterations for each point.
 * @param min_x Minimum real value for the fractal plane.
 * @param max_x Maximum real value for the fractal plane.
 * @param min_y Minimum imaginary value for the fractal plane.
 * @param max_y Maximum imaginary value for the fractal plane.
 * @param c_real Real part of the Julia constant.
 * @param c_imag Imaginary part of the Julia constant.
 */
__global__ void julia_kernel(Color* image_data, int width, int height, int max_iterations,
                             double min_x, double max_x, double min_y, double max_y,
                             double c_real_const, double c_imag_const) {
    // Calculate global thread ID
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y;

    // Check if the current thread is within the image bounds
    if (col < width && row < height) {
        // Map pixel coordinates to initial z_0
        double z_real = min_x + (col / (double)width) * (max_x - min_x);
        double z_imag = min_y + (row / (double)height) * (max_y - min_y);

        int iteration = 0;
        // Iterate z = z*z + c_const
        while (z_real * z_real + z_imag * z_imag < 4.0 && iteration < max_iterations) {
            double z_real_new = z_real * z_real - z_imag * z_imag + c_real_const;
            double z_imag_new = 2.0 * z_real * z_imag + c_imag_const;
            z_real = z_real_new;
            z_imag = z_imag_new;
            iteration++;
        }

        // Determine color based on iteration count
        Color pixel_color;
        if (iteration == max_iterations) {
            // Point is likely in the set (black)
            pixel_color = {0, 0, 0};
        } else {
            // Point escaped, color based on iteration count (smooth coloring can be added for better visuals)
            // Simple coloring: gradient from blue to green to red
            double t = (double)iteration / max_iterations;
            pixel_color.r = (unsigned char)(9 * (1 - t) * t * t * t * 255);
            pixel_color.g = (unsigned char)(15 * (1 - t) * (1 - t) * t * t * 255);
            pixel_color.b = (unsigned char)(8.5 * (1 - t) * (1 - t) * (1 - t) * t * 255);
        }

        // Store the color in the image data array
        image_data[row * width + col] = pixel_color;
    }
}

// --- Host Functions ---

/**
 * @brief Saves the image data to a PPM (Portable Pixmap) file.
 *
 * @param filename The name of the output PPM file.
 * @param image_data Pointer to the host memory containing color data.
 * @param width Width of the image.
 * @param height Height of the image.
 */
void save_ppm(const std::string& filename, const std::vector<Color>& image_data, int width, int height) {
    std::ofstream ofs(filename, std::ios::out | std::ios::binary);
    if (!ofs.is_open()) {
        std::cerr << "Error: Could not open file " << filename << " for writing." << std::endl;
        return;
    }

    ofs << "P6\n" << width << " " << height << "\n255\n"; // PPM header (P6 for binary RGB)
    ofs.write(reinterpret_cast<const char*>(image_data.data()), width * height * sizeof(Color));
    ofs.close();
    std::cout << "Image saved to " << filename << std::endl;
}

/**
 * @brief Main function to handle user input, GPU memory management, and kernel launches.
 */
int main() {
    int width, height;
    int max_iterations;
    int choice;
    std::string output_filename;

    std::cout << "--- Fractal Generator ---" << std::endl;
    std::cout << "1. Mandelbrot Set" << std::endl;
    std::cout << "2. Julia Set" << std::endl;
    std::cout << "Enter your choice (1 or 2): ";
    while (!(std::cin >> choice) || (choice != 1 && choice != 2)) {
        std::cout << "Invalid choice. Please enter 1 or 2: ";
        std::cin.clear();
        std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
    }

    std::cout << "Enter image width (e.g., 1024): ";
    while (!(std::cin >> width) || width <= 0) {
        std::cout << "Invalid width. Please enter a positive integer: ";
        std::cin.clear();
        std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
    }

    std::cout << "Enter image height (e.g., 768): ";
    while (!(std::cin >> height) || height <= 0) {
        std::cout << "Invalid height. Please enter a positive integer: ";
        std::cin.clear();
        std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
    }

    std::cout << "Enter maximum iterations (e.g., 256): ";
    while (!(std::cin >> max_iterations) || max_iterations <= 0) {
        std::cout << "Invalid max iterations. Please enter a positive integer: ";
        std::cin.clear();
        std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
    }

    double min_x, max_x, min_y, max_y;
    double c_real = 0.0, c_imag = 0.0; // For Julia set constant

    if (choice == 1) { // Mandelbrot
        std::cout << "Enter output filename (e.g., mandelbrot.ppm): ";
        std::cin >> output_filename;
        min_x = -2.0; max_x = 1.0;
        min_y = -1.5; max_y = 1.5;
    } else { // Julia
        std::cout << "Enter output filename (e.g., julia.ppm): ";
        std::cin >> output_filename;
        std::cout << "Enter real part of Julia constant (e.g., -0.7): ";
        while (!(std::cin >> c_real)) {
            std::cout << "Invalid input. Please enter a double: ";
            std::cin.clear();
            std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
        }
        std::cout << "Enter imaginary part of Julia constant (e.g., 0.27015): ";
        while (!(std::cin >> c_imag)) {
            std::cout << "Invalid input. Please enter a double: ";
            std::cin.clear();
            std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
        }
        min_x = -1.5; max_x = 1.5;
        min_y = -1.5; max_y = 1.5;
    }

    // Host memory for image data
    std::vector<Color> h_image_data(width * height);

    // Device memory for image data
    Color* d_image_data;
    size_t image_size = width * height * sizeof(Color);
    CUDA_CHECK(cudaMalloc(&d_image_data, image_size));

    // Define block and grid dimensions
    // A common practice is to use 16x16 or 32x32 threads per block
    int threads_per_block_x = 16;
    int threads_per_block_y = 16;
    dim3 block_dim(threads_per_block_x, threads_per_block_y);

    dim3 grid_dim( (width + block_dim.x - 1) / block_dim.x,
                   (height + block_dim.y - 1) / block_dim.y );

    std::cout << "Launching kernel with grid_dim(" << grid_dim.x << ", " << grid_dim.y << ")"
              << " and block_dim(" << block_dim.x << ", " << block_dim.y << ")" << std::endl;

    // Launch the appropriate kernel
    if (choice == 1) {
        mandelbrot_kernel<<<grid_dim, block_dim>>>(d_image_data, width, height, max_iterations,
                                                    min_x, max_x, min_y, max_y);
    } else {
        julia_kernel<<<grid_dim, block_dim>>>(d_image_data, width, height, max_iterations,
                                               min_x, max_x, min_y, max_y,
                                               c_real, c_imag);
    }

    // Synchronize to ensure kernel completes
    CUDA_CHECK(cudaDeviceSynchronize());

    // Copy results from device to host
    CUDA_CHECK(cudaMemcpy(h_image_data.data(), d_image_data, image_size, cudaMemcpyDeviceToHost));

    // Save the image to a PPM file
    save_ppm(output_filename, h_image_data, width, height);

    // Free device memory
    CUDA_CHECK(cudaFree(d_image_data));

    std::cout << "Fractal generation complete." << std::endl;

    return 0;
}