gpu.h

#ifndef GPU_H
#define GPU_H
 
#if defined(__EMSCRIPTEN__)
#include <emscripten/emscripten.h>
#include <emscripten/html5.h>
#endif
 
#include <fstream>
#include <iterator>
#include <map>
#include <sstream>
#include <string>
#include <webgpu/webgpu_cpp.h>
 
// auto generated by tools/embed_shaders.py
#include "internal/log.h"
 
#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];
                    IMG2NUM_LOG_INFO(
                        "Shader Error [{}] Line {}:{} - {}",
                        msg.type == wgpu::CompilationMessageType::Error ? "ERR" : "WARN",
                        msg.lineNum, msg.linePos,
                        std::string_view(
                            msg.message.data ? msg.message.data : "", msg.message.length
                        )
                    );
                }
            }
        );
    };
 
    void init_gpu() {
        if (gpu_initialized)
            return;
 
        wgpu::InstanceDescriptor instanceDesc = {};
        instance = wgpu::CreateInstance(&instanceDesc);
 
        if (!instance) {
            IMG2NUM_LOG_INFO("Fatal: WebGPU instance creation failed.");
            return;
        }
 
        // ---------------------------------------------------------
        // 1. Get Adapter
        // ---------------------------------------------------------
        IMG2NUM_LOG_INFO("Requesting Adapter...");
        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);
                    IMG2NUM_LOG_INFO("Adapter Acquired");
                } else {
                    IMG2NUM_LOG_INFO(
                        "Adapter Failed: {}", std::string_view(msg.data ? msg.data : "", msg.length)
                    );
                }
                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) {
            IMG2NUM_LOG_INFO("Fatal: Could not get WebGPU Adapter.");
            return;
        }
 
        // ---------------------------------------------------------
        // 2. Get Device
        // ---------------------------------------------------------
        IMG2NUM_LOG_INFO("Requesting Device...");
        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
            IMG2NUM_LOG_INFO(
                "\n[WEBGPU FATAL ERROR] Type: {} | Msg: {}\n", static_cast<uint32_t>(type), err_str
            );
        });
        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";
                // 1. Ignore teardown cancellations so we don't log during static destruction
                if (reason == wgpu::DeviceLostReason::CallbackCancelled) {
                    return;
                }
                IMG2NUM_LOG_INFO(
                    "[DEVICE LOST] Reason: {} Msg: {}", static_cast<int>(reason), err_msg
                );
            }
        );
 
        // 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;
 
            IMG2NUM_LOG_INFO("maxBufferSize: {}", requiredLimits.maxBufferSize);
            IMG2NUM_LOG_INFO(
                "maxStorageBufferBindingSize: {}", requiredLimits.maxStorageBufferBindingSize
            );
 
            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);
                    IMG2NUM_LOG_INFO("Device Acquired");
                } else {
                    IMG2NUM_LOG_INFO(
                        "Device Failed: {}", msg.data && msg.length > 0
                                                 ? std::string(msg.data, msg.length)
                                                 : "Unknown error"
                    );
                }
                device_ready = true; // Unblock the loop
            }
        );
 
        // WAIT LOOP
        while (!device_ready) {
            instance.ProcessEvents();
#if defined(__EMSCRIPTEN__)
            emscripten_sleep(10);
#endif
        }
 
        if (!device) {
            IMG2NUM_LOG_INFO("Fatal: Could not get WebGPU Device.");
            return;
        }
 
        if (!validate_device()) {
            IMG2NUM_LOG_INFO("Fatal: Could not get WebGPU Device.");
            return;
        }
 
        queue = device.GetQueue();
        gpu_initialized = true;
        IMG2NUM_LOG_INFO("GPU Fully Initialized.");
    };
 
    ~GPU() {
        device = nullptr;
        adapter = nullptr;
        queue = nullptr;
        instance = nullptr;
        gpu_initialized = false;
    };
};
 
#endif