gpu.h

#ifndef GPU_H
#define GPU_H
 
#if defined(__EMSCRIPTEN__)
#include <emscripten/emscripten.h>
#include <emscripten/html5.h>
#endif
 
#include <fstream>
#include <iostream>
#include <iterator>
#include <map>
#include <sstream>
#include <string>
#include <webgpu/webgpu_cpp.h>
 
// auto generated by tools/embed_shaders.py
#include <embedded_shaders.h>
 
class GPU {
  private:
    wgpu::Instance instance;
    wgpu::Adapter adapter;
    wgpu::Device device;
    wgpu::Queue queue;
 
    bool adapter_ready = false;
    bool device_ready = false;
    bool gpu_initialized = false;
 
    GPU() = default;
 
    bool validate_device() {
        if (!device)
            return false;
 
        // Try creating a trivial buffer to ensure the device is usable
        wgpu::BufferDescriptor desc = {};
        desc.size = 4;
        desc.usage = wgpu::BufferUsage::CopyDst;
 
        wgpu::Buffer test = device.CreateBuffer(&desc);
        return test != nullptr;
    }
 
  public:
    // makes a single global instance that other files can reference
    static GPU& getClassInstance() {
        static GPU gpuInstance;
        return gpuInstance;
    }
 
    const wgpu::Device& get_device() {
        return device;
    }
 
    const wgpu::Instance& get_instance() {
        return instance;
    }
 
    const wgpu::Queue& get_queue() {
        return queue;
    }
 
    bool is_initialized() {
        return gpu_initialized;
    }
 
    // Delete copy constructor and assignment operator to prevent duplication
    GPU(const GPU&) = delete;
    GPU& operator=(const GPU&) = delete;
    GPU(GPU&&) = delete;
    GPU& operator=(GPU&&) = delete;
 
    std::string readWGSLFile(std::string_view shader_id) {
        // A simple linear search over 9 items is blazingly fast in C++
        for (const auto& entry : embedded_shaders::shaders) {
            if (entry.id == shader_id) {
                // Safely convert the string_view back to a std::string for Dawn
                return std::string(entry.source);
            }
        }
 
        // Always good to handle the "not found" case gracefully!
        return "";
    }
 
    wgpu::ComputePipeline createPipeline(const std::string& filename, const std::string& label) {
        wgpu::ShaderSourceWGSL wgsl;
        std::string shaderCode = readWGSLFile(filename);
        wgsl.code = shaderCode.c_str();
        wgpu::ShaderModuleDescriptor md = {};
        md.nextInChain = &wgsl;
        md.label = label.c_str();
        wgpu::ShaderModule sm = device.CreateShaderModule(&md);
 
        wgpu::ComputePipelineDescriptor cpd = {};
        cpd.compute.module = sm;
        cpd.compute.entryPoint = "main";
        return device.CreateComputePipeline(&cpd);
    };
 
    static uint32_t getAlignedBytesPerRow(uint32_t width, uint32_t bytesPerPixel = 4) {
        uint32_t unaligned = width * bytesPerPixel;
        uint32_t align = 256;
        return (unaligned + align - 1) & ~(align - 1);
    };
 
    static void printShaderError(wgpu::ShaderModule shaderModule) {
        shaderModule.GetCompilationInfo(
            // Callback Mode (New API Requirement)
            wgpu::CallbackMode::AllowProcessEvents,
            // Callback Lambda
            [](wgpu::CompilationInfoRequestStatus status, const wgpu::CompilationInfo* info) {
                if (status != wgpu::CompilationInfoRequestStatus::Success || !info)
                    return;
 
                for (uint32_t i = 0; i < info->messageCount; ++i) {
                    const auto& msg = info->messages[i];
                    std::cerr << "Shader Error ["
                              << (msg.type == wgpu::CompilationMessageType::Error ? "ERR" : "WARN")
                              << "]"
                              << " Line " << msg.lineNum << ":" << msg.linePos << " - "
                              << msg.message.data // .data for StringView
                              << std::endl;
                }
            }
        );
    };
 
    void init_gpu() {
        if (gpu_initialized)
            return;
 
        wgpu::InstanceDescriptor instanceDesc = {};
        instance = wgpu::CreateInstance(&instanceDesc);
 
        if (!instance) {
            std::cerr << "Fatal: WebGPU instance creation failed." << std::endl;
            return;
        }
 
        // ---------------------------------------------------------
        // 1. Get Adapter
        // ---------------------------------------------------------
        std::cout << "Requesting Adapter..." << std::endl;
        adapter_ready = false;
 
        instance.RequestAdapter(
            nullptr,
            wgpu::CallbackMode::AllowProcessEvents, // <--- ALLOW EVENTS
            [this](wgpu::RequestAdapterStatus status, wgpu::Adapter a, wgpu::StringView msg) {
                if (status == wgpu::RequestAdapterStatus::Success) {
                    adapter = std::move(a);
                    std::cout << "Adapter Acquired" << std::endl;
                } else {
                    std::cerr << "Adapter Failed: "
                              << std::string_view(msg.data ? msg.data : "", msg.length)
                              << std::endl;
                }
                adapter_ready = true; // Unblock the loop
            }
        );
 
        // WAIT LOOP: Yield to browser so it can actually find the adapter
        while (!adapter_ready) {
            instance.ProcessEvents();
#if defined(__EMSCRIPTEN__)
            emscripten_sleep(10); // Sleep 10ms
#endif
        }
 
        if (!adapter) {
            std::cerr << "Fatal: Could not get WebGPU Adapter." << std::endl;
            return;
        }
 
        // ---------------------------------------------------------
        // 2. Get Device
        // ---------------------------------------------------------
        std::cout << "Requesting Device..." << std::endl;
        device_ready = false;
 
        wgpu::DeviceDescriptor deviceDesc = {};
        deviceDesc.SetUncapturedErrorCallback([](const wgpu::Device&, wgpu::ErrorType type,
                                                 wgpu::StringView msg) {
            // 1. Safely extract the string using the provided length
            std::string err_str = (msg.data && msg.length > 0) ? std::string(msg.data, msg.length)
                                                               : "Unknown Error (Null message)";
 
            // 2. Print it safely
            std::cerr << "\n[WEBGPU FATAL ERROR] Type: " << static_cast<uint32_t>(type)
                      << " | Msg: " << err_str << "\n"
                      << std::endl;
        });
        deviceDesc.SetDeviceLostCallback(
            wgpu::CallbackMode::AllowProcessEvents,
            [](const wgpu::Device&, wgpu::DeviceLostReason reason, wgpu::StringView msg) {
                std::string err_msg = (msg.data && msg.length > 0)
                                          ? std::string(msg.data, msg.length)
                                          : "Unknown device lost reason";
                std::cerr << "[DEVICE LOST] Reason: " << static_cast<int>(reason)
                          << " Msg: " << err_msg << std::endl;
            }
        );
 
        // Set maximum possible device hardware memory
        wgpu::Limits supportedLimits;
        wgpu::Limits requiredLimits;
        if (adapter.GetLimits(&supportedLimits)) {
            // nominally the device Buffer limit is 256MB
 
            // Copy the adapter's physical limits over to your requested limits
            requiredLimits = supportedLimits;
 
            std::cout << "maxBufferSize: " << requiredLimits.maxBufferSize << std::endl;
            std::cout << "maxStorageBufferBindingSize: "
                      << requiredLimits.maxStorageBufferBindingSize << std::endl;
 
            deviceDesc.requiredLimits = &requiredLimits;
        }
 
        adapter.RequestDevice(
            &deviceDesc,
            wgpu::CallbackMode::AllowProcessEvents, // <--- ALLOW EVENTS
            [this](wgpu::RequestDeviceStatus status, wgpu::Device d, wgpu::StringView msg) {
                if (status == wgpu::RequestDeviceStatus::Success) {
                    device = std::move(d);
                    std::cout << "Device Acquired" << std::endl;
                } else {
                    std::cerr << "Device Failed: "
                              << (msg.data && msg.length > 0 ? std::string(msg.data, msg.length)
                                                             : "Unknown error")
                              << std::endl;
                }
                device_ready = true; // Unblock the loop
            }
        );
 
        // WAIT LOOP
        while (!device_ready) {
            instance.ProcessEvents();
#if defined(__EMSCRIPTEN__)
            emscripten_sleep(10);
#endif
        }
 
        if (!device) {
            std::cerr << "Fatal: Could not get WebGPU Device." << std::endl;
            return;
        }
 
        if (!validate_device()) {
            std::cerr << "Fatal: Could not get WebGPU Device." << std::endl;
            return;
        }
 
        queue = device.GetQueue();
        gpu_initialized = true;
        std::cout << "GPU Fully Initialized." << std::endl;
    };
 
    ~GPU() {
        device = nullptr;
        adapter = nullptr;
        queue = nullptr;
        instance = nullptr;
        gpu_initialized = false;
    };
};
 
#endif