go-mxl-pattern-generator/internal/generator/wgpu.go

// WGPU pattern gen: renders v210 frames with a wgpu (WebGPU/Vulkan)
// compute shader. Pure Go — no cgo in the GPU layer.
package generator

import (
	"context"
	"encoding/binary"
	"fmt"
	"time"

	"github.com/gogpu/gputypes"
	"github.com/gogpu/wgpu"

	// Register all available GPU backends (Vulkan, DX12, GLES, Metal, etc.)
	_ "github.com/gogpu/wgpu/hal/allbackends"
)

const wgpuWorkgroupSize = 64

// WGPUGenerator renders v210 frames on the GPU via wgpu and implements
// FrameGenerator. Per frame the kernel packs the pattern into a storage
// buffer, the GPU DMAs it into a persistent host-visible buffer, and the
// mapped contents are copied straight into the destination grain.
type WGPUGenerator struct {
	instance  *wgpu.Instance
	adapter   *wgpu.Adapter
	device    *wgpu.Device
	queue     *wgpu.Queue
	shader    *wgpu.ShaderModule
	bgl       *wgpu.BindGroupLayout
	bg        *wgpu.BindGroup
	pl        *wgpu.PipelineLayout
	pipeline  *wgpu.ComputePipeline
	out       *wgpu.Buffer
	host      *wgpu.Buffer
	uniform   *wgpu.Buffer
	params    []byte
	width     int
	height    int
	blocks    int
	frameSize uint64
}

var _ FrameGenerator = (*WGPUGenerator)(nil)

// WGPUOption customizes NewWGPUGenerator.
type WGPUOption func(*WGPUGenerator)

func NewWGPUGenerator(width, height uint, wgsl string, opts ...WGPUOption) (*WGPUGenerator, error) {
	g := &WGPUGenerator{
		width:  int(width),
		height: int(height),
		blocks: int(width*height) / 6,
		params: make([]byte, 16),
	}
	for _, opt := range opts {
		opt(g)
	}
	g.frameSize = uint64(g.blocks) * 16
	binary.LittleEndian.PutUint32(g.params[0:], uint32(width))
	binary.LittleEndian.PutUint32(g.params[4:], uint32(height))

	var err error
	if g.instance, err = wgpu.CreateInstance(nil); err != nil {
		return nil, fmt.Errorf("wgpu: create instance: %w", err)
	}
	if g.adapter, err = g.instance.RequestAdapter(nil); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: request adapter: %w", err)
	}
	if g.device, err = g.adapter.RequestDevice(nil); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: request device: %w", err)
	}
	g.queue = g.device.Queue()

	if g.shader, err = g.device.CreateShaderModule(&wgpu.ShaderModuleDescriptor{
		Label: "v210-shader", WGSL: wgsl,
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: shader: %w", err)
	}
	if g.out, err = g.device.CreateBuffer(&wgpu.BufferDescriptor{
		Label: "v210-out", Size: g.frameSize,
		Usage: wgpu.BufferUsageStorage | wgpu.BufferUsageCopySrc,
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: out buffer: %w", err)
	}
	if g.host, err = g.device.CreateBuffer(&wgpu.BufferDescriptor{
		Label: "v210-host", Size: g.frameSize,
		Usage: wgpu.BufferUsageCopyDst | wgpu.BufferUsageMapRead,
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: host buffer: %w", err)
	}
	if g.uniform, err = g.device.CreateBuffer(&wgpu.BufferDescriptor{
		Label: "v210-params", Size: uint64(len(g.params)),
		Usage: wgpu.BufferUsageUniform | wgpu.BufferUsageCopyDst,
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: params buffer: %w", err)
	}
	if err := g.queue.WriteBuffer(g.uniform, 0, g.params); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: write params: %w", err)
	}
	if g.bgl, err = g.device.CreateBindGroupLayout(&wgpu.BindGroupLayoutDescriptor{
		Label: "v210-bgl",
		Entries: []gputypes.BindGroupLayoutEntry{
			{Binding: 0, Visibility: wgpu.ShaderStageCompute, Buffer: &gputypes.BufferBindingLayout{Type: gputypes.BufferBindingTypeStorage}},
			{Binding: 1, Visibility: wgpu.ShaderStageCompute, Buffer: &gputypes.BufferBindingLayout{Type: gputypes.BufferBindingTypeUniform}},
		},
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: bind group layout: %w", err)
	}
	if g.bg, err = g.device.CreateBindGroup(&wgpu.BindGroupDescriptor{
		Label: "v210-bg", Layout: g.bgl,
		Entries: []wgpu.BindGroupEntry{
			{Binding: 0, Buffer: g.out, Size: g.frameSize},
			{Binding: 1, Buffer: g.uniform, Size: uint64(len(g.params))},
		},
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: bind group: %w", err)
	}
	if g.pl, err = g.device.CreatePipelineLayout(&wgpu.PipelineLayoutDescriptor{
		Label: "v210-pl", BindGroupLayouts: []*wgpu.BindGroupLayout{g.bgl},
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: pipeline layout: %w", err)
	}
	if g.pipeline, err = g.device.CreateComputePipeline(&wgpu.ComputePipelineDescriptor{
		Label: "v210-pipeline", Layout: g.pl, Module: g.shader, EntryPoint: "main",
	}); err != nil {
		g.Close()
		return nil, fmt.Errorf("wgpu: pipeline: %w", err)
	}
	return g, nil
}

func (g *WGPUGenerator) GenerateFrame(dest []byte, frameIndex int) error {
	if uint64(len(dest)) < g.frameSize {
		return fmt.Errorf("wgpu: dest too small: %d bytes, need %d", len(dest), g.frameSize)
	}

	binary.LittleEndian.PutUint32(g.params[8:], uint32(frameIndex))
	if err := g.queue.WriteBuffer(g.uniform, 0, g.params); err != nil {
		return fmt.Errorf("wgpu: write params: %w", err)
	}

	encoder, err := g.device.CreateCommandEncoder(nil)
	if err != nil {
		return fmt.Errorf("wgpu: encoder: %w", err)
	}
	pass, err := encoder.BeginComputePass(nil)
	if err != nil {
		return fmt.Errorf("wgpu: compute pass: %w", err)
	}
	pass.SetPipeline(g.pipeline)
	pass.SetBindGroup(0, g.bg, nil)
	pass.Dispatch(uint32((g.blocks+wgpuWorkgroupSize-1)/wgpuWorkgroupSize), 1, 1)
	if err := pass.End(); err != nil {
		return fmt.Errorf("wgpu: end pass: %w", err)
	}
	encoder.CopyBufferToBuffer(g.out, 0, g.host, 0, g.frameSize)
	cmd, err := encoder.Finish()
	if err != nil {
		return fmt.Errorf("wgpu: finish: %w", err)
	}
	if _, err := g.queue.Submit(cmd); err != nil {
		return fmt.Errorf("wgpu: submit: %w", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := g.host.Map(ctx, wgpu.MapModeRead, 0, g.frameSize); err != nil {
		return fmt.Errorf("wgpu: map: %w", err)
	}
	rng, err := g.host.MappedRange(0, g.frameSize)
	if err != nil {
		_ = g.host.Unmap()
		return fmt.Errorf("wgpu: mapped range: %w", err)
	}
	copy(dest, rng.Bytes())
	return g.host.Unmap()
}

func (g *WGPUGenerator) Close() error {
	if g.pipeline != nil {
		g.pipeline.Release()
		g.pipeline = nil
	}
	if g.pl != nil {
		g.pl.Release()
		g.pl = nil
	}
	if g.bg != nil {
		g.bg.Release()
		g.bg = nil
	}
	if g.bgl != nil {
		g.bgl.Release()
		g.bgl = nil
	}
	if g.shader != nil {
		g.shader.Release()
		g.shader = nil
	}
	if g.uniform != nil {
		g.uniform.Release()
		g.uniform = nil
	}
	if g.host != nil {
		g.host.Release()
		g.host = nil
	}
	if g.out != nil {
		g.out.Release()
		g.out = nil
	}
	if g.device != nil {
		g.device.Release()
		g.device = nil
	}
	if g.adapter != nil {
		g.adapter.Release()
		g.adapter = nil
	}
	if g.instance != nil {
		g.instance.Release()
		g.instance = nil
	}
	return nil
}