path: root/src/gpu/gpu.cc

#include "gpu.h"
#include "platform.h"

#include <GLFW/glfw3.h>
#include <webgpu.h>
#include <wgpu.h>

#include <algorithm>
#include <cassert>
#include <cstring>
#include <iostream>
#include <vector>

static WGPUInstance g_instance = nullptr;
static WGPUAdapter g_adapter = nullptr;
static WGPUDevice g_device = nullptr;
static WGPUQueue g_queue = nullptr;
static WGPUSurface g_surface = nullptr;
static WGPUSurfaceConfiguration g_config = {};

static WGPURenderPipeline g_render_pipeline = nullptr;
static WGPUBuffer g_uniform_buffer = nullptr;
static WGPUBindGroup g_bind_group = nullptr;

static WGPUStringView str_view(const char *str) {
  if (!str) return {nullptr, 0};
  return {str, strlen(str)};
}

static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, WGPUStringView message,
                                   void *userdata1, void *userdata2) {
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter *)userdata1) = adapter;
  } else {
    printf("Request adapter failed: %.*s\n", (int)message.length, message.data);
  }
}

static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, WGPUStringView message,
                                  void *userdata1, void *userdata2) {
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice *)userdata1) = device;
  } else {
    printf("Request device failed: %.*s\n", (int)message.length, message.data);
  }
}

static void handle_device_error(WGPUDevice const *device, WGPUErrorType type,
                                WGPUStringView message, void *userdata1,
                                void *userdata2) {
  printf("WebGPU Error: %.*s\n", (int)message.length, message.data);
}

const char *shader_wgsl_code = R"(
struct Uniforms {
    audio_peak : f32,
};

@group(0) @binding(0) var<uniform> uniforms : Uniforms;

@vertex
fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4<f32> {
    let PI = 3.14159265;
    let num_sides = 7.0;
    
    let base_scale = 0.5;
    let pulse_scale = 0.2 * uniforms.audio_peak;
    let scale = base_scale + pulse_scale;

    let tri_idx = f32(vertex_index / 3u);
    let sub_idx = vertex_index % 3u;

    if (sub_idx == 0u) {
        return vec4<f32>(0.0, 0.0, 0.0, 1.0);
    }
    
    let i = tri_idx + f32(sub_idx - 1u);
    let angle = i * 2.0 * PI / num_sides;
    let x = scale * cos(angle);
    let y = scale * sin(angle);
    
    return vec4<f32>(x, y, 0.0, 1.0);
}

@fragment
fn fs_main() -> @location(0) vec4<f32> {
    let hue = uniforms.audio_peak * 0.5;
    let r = sin(hue + 0.0) * 0.5 + 0.5;
    let g = sin(hue + 2.0) * 0.5 + 0.5;
    let b = sin(hue + 4.0) * 0.5 + 0.5;
    return vec4<f32>(r, g, b, 1.0);
}
)";

void gpu_init(GLFWwindow *window) {
  g_instance = wgpuCreateInstance(nullptr);
  assert(g_instance);

  g_surface = platform_create_wgpu_surface(g_instance);
  assert(g_surface);

  WGPURequestAdapterOptions adapter_opts = {};
  adapter_opts.compatibleSurface = g_surface;
  adapter_opts.powerPreference = WGPUPowerPreference_HighPerformance;

  wgpuInstanceRequestAdapter(
      g_instance, &adapter_opts,
      {nullptr, WGPUCallbackMode_WaitAnyOnly, handle_request_adapter, &g_adapter,
       nullptr});

  while (!g_adapter) {
    wgpuInstanceWaitAny(g_instance, 0, nullptr, 0);
  }

  WGPUDeviceDescriptor device_desc = {};
  device_desc.label = str_view("Demo Device");
  device_desc.uncapturedErrorCallbackInfo.callback = handle_device_error;

  wgpuAdapterRequestDevice(
      g_adapter, &device_desc,
      {nullptr, WGPUCallbackMode_WaitAnyOnly, handle_request_device, &g_device,
       nullptr});

  while (!g_device) {
    wgpuInstanceWaitAny(g_instance, 0, nullptr, 0);
  }

  g_queue = wgpuDeviceGetQueue(g_device);

  WGPUSurfaceCapabilities caps = {};
  wgpuSurfaceGetCapabilities(g_surface, g_adapter, &caps);
  WGPUTextureFormat swap_chain_format = caps.formats[0];

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(shader_wgsl_code);

  WGPUShaderModuleDescriptor shader_module_desc = {};
  shader_module_desc.nextInChain = &wgsl_src.chain;
  shader_module_desc.label = str_view("Demo Shader");

  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(g_device, &shader_module_desc);

  WGPUBufferDescriptor uniform_buffer_desc = {};
  uniform_buffer_desc.label = str_view("Uniform Buffer");
  uniform_buffer_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
  uniform_buffer_desc.size = sizeof(float);
  g_uniform_buffer = wgpuDeviceCreateBuffer(g_device, &uniform_buffer_desc);

  WGPUBindGroupLayoutEntry bgl_entry = {};
  bgl_entry.binding = 0;
  bgl_entry.visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
  bgl_entry.buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entry.buffer.minBindingSize = sizeof(float);

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 1;
  bgl_desc.entries = &bgl_entry;
  WGPUBindGroupLayout bind_group_layout =
      wgpuDeviceCreateBindGroupLayout(g_device, &bgl_desc);

  WGPUBindGroupEntry bg_entry = {};
  bg_entry.binding = 0;
  bg_entry.buffer = g_uniform_buffer;
  bg_entry.size = sizeof(float);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout = bind_group_layout;
  bg_desc.entryCount = 1;
  bg_desc.entries = &bg_entry;
  g_bind_group = wgpuDeviceCreateBindGroup(g_device, &bg_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bind_group_layout;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(g_device, &pl_desc);

  WGPUColorTargetState color_target = {};
  color_target.format = swap_chain_format;
  color_target.writeMask = WGPUColorWriteMask_All;

  WGPUFragmentState fragment_state = {};
  fragment_state.module = shader_module;
  fragment_state.entryPoint = str_view("fs_main");
  fragment_state.targetCount = 1;
  fragment_state.targets = &color_target;

  WGPURenderPipelineDescriptor pipeline_desc = {};
  pipeline_desc.layout = pipeline_layout;
  pipeline_desc.vertex.module = shader_module;
  pipeline_desc.vertex.entryPoint = str_view("vs_main");
  pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipeline_desc.multisample.count = 1;
  pipeline_desc.multisample.mask = 0xFFFFFFFF;
  pipeline_desc.fragment = &fragment_state;

  g_render_pipeline = wgpuDeviceCreateRenderPipeline(g_device, &pipeline_desc);

  int width, height;
  glfwGetFramebufferSize(window, &width, &height);
  g_config.device = g_device;
  g_config.format = swap_chain_format;
  g_config.width = width;
  g_config.height = height;
  g_config.presentMode = WGPUPresentMode_Fifo;
  g_config.alphaMode = WGPUCompositeAlphaMode_Opaque;

  wgpuSurfaceConfigure(g_surface, &g_config);
}

void gpu_draw(float audio_peak) {
  WGPUSurfaceTexture surface_texture;
  wgpuSurfaceGetCurrentTexture(g_surface, &surface_texture);
  if (surface_texture.status != WGPUSurfaceGetCurrentTextureStatus_SuccessOptimal &&
      surface_texture.status != WGPUSurfaceGetCurrentTextureStatus_SuccessSuboptimal)
    return;

  WGPUTextureView view = wgpuTextureCreateView(surface_texture.texture, nullptr);

  wgpuQueueWriteBuffer(g_queue, g_uniform_buffer, 0, &audio_peak, sizeof(float));

  WGPUCommandEncoderDescriptor encoder_desc = {};
  WGPUCommandEncoder encoder =
      wgpuDeviceCreateCommandEncoder(g_device, &encoder_desc);

  WGPURenderPassColorAttachment color_attachment = {};
  color_attachment.view = view;
  color_attachment.loadOp = WGPULoadOp_Clear;
  color_attachment.storeOp = WGPUStoreOp_Store;
  color_attachment.clearValue = {0.1, 0.2, 0.3, 1.0};

  WGPURenderPassDescriptor render_pass_desc = {};
  render_pass_desc.colorAttachmentCount = 1;
  render_pass_desc.colorAttachments = &color_attachment;

  WGPURenderPassEncoder pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &render_pass_desc);

  wgpuRenderPassEncoderSetPipeline(pass, g_render_pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, g_bind_group, 0, nullptr);
  // Drawing a heptagon with TriangleList: 7 triangles * 3 vertices = 21
  wgpuRenderPassEncoderDraw(pass, 21, 1, 0, 0);
  wgpuRenderPassEncoderEnd(pass);

  WGPUCommandBufferDescriptor cmd_desc = {};
  WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, &cmd_desc);
  wgpuQueueSubmit(g_queue, 1, &commands);
  wgpuSurfacePresent(g_surface);

  wgpuTextureViewRelease(view);
  wgpuTextureRelease(surface_texture.texture);
}

void gpu_shutdown() {}