V210 issue fix + correct packing

internal/generator/wgpu.go

@@ -22,23 +22,26 @@ const wgpuWorkgroupSize = 64
 // 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
+	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
+	activeLineSize   int
+	lineSize         int
+	compactFrameSize uint64
+	frameSize        uint64
 }
 
 var _ FrameGenerator = (*WGPUGenerator)(nil)
@@ -48,15 +51,18 @@ 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),
+		width:          int(width),
+		height:         int(height),
+		blocks:         ((int(width) + 5) / 6) * int(height),
+		activeLineSize: V210ActiveLineSize(int(width)),
+		lineSize:       V210LineSize(int(width)),
+		params:         make([]byte, 16),
 	}
 	for _, opt := range opts {
 		opt(g)
 	}
-	g.frameSize = uint64(g.blocks) * 16
+	g.compactFrameSize = uint64(g.activeLineSize * g.height)
+	g.frameSize = uint64(g.lineSize * g.height)
 	binary.LittleEndian.PutUint32(g.params[0:], uint32(width))
 	binary.LittleEndian.PutUint32(g.params[4:], uint32(height))
 
@@ -81,14 +87,14 @@ func NewWGPUGenerator(width, height uint, wgsl string, opts ...WGPUOption) (*WGP
 		return nil, fmt.Errorf("wgpu: shader: %w", err)
 	}
 	if g.out, err = g.device.CreateBuffer(&wgpu.BufferDescriptor{
-		Label: "v210-out", Size: g.frameSize,
+		Label: "v210-out", Size: g.compactFrameSize,
 		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,
+		Label: "v210-host", Size: g.compactFrameSize,
 		Usage: wgpu.BufferUsageCopyDst | wgpu.BufferUsageMapRead,
 	}); err != nil {
 		g.Close()
@@ -118,7 +124,7 @@ func NewWGPUGenerator(width, height uint, wgsl string, opts ...WGPUOption) (*WGP
 	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: 0, Buffer: g.out, Size: g.compactFrameSize},
 			{Binding: 1, Buffer: g.uniform, Size: uint64(len(g.params))},
 		},
 	}); err != nil {
@@ -164,7 +170,7 @@ func (g *WGPUGenerator) GenerateFrame(dest []byte, frameIndex int) error {
 	if err := pass.End(); err != nil {
 		return fmt.Errorf("wgpu: end pass: %w", err)
 	}
-	encoder.CopyBufferToBuffer(g.out, 0, g.host, 0, g.frameSize)
+	encoder.CopyBufferToBuffer(g.out, 0, g.host, 0, g.compactFrameSize)
 	cmd, err := encoder.Finish()
 	if err != nil {
 		return fmt.Errorf("wgpu: finish: %w", err)
@@ -175,15 +181,25 @@ func (g *WGPUGenerator) GenerateFrame(dest []byte, frameIndex int) error {
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	if err := g.host.Map(ctx, wgpu.MapModeRead, 0, g.frameSize); err != nil {
+	if err := g.host.Map(ctx, wgpu.MapModeRead, 0, g.compactFrameSize); err != nil {
 		return fmt.Errorf("wgpu: map: %w", err)
 	}
-	rng, err := g.host.MappedRange(0, g.frameSize)
+	rng, err := g.host.MappedRange(0, g.compactFrameSize)
 	if err != nil {
 		_ = g.host.Unmap()
 		return fmt.Errorf("wgpu: mapped range: %w", err)
 	}
-	copy(dest, rng.Bytes())
+	mapped := rng.Bytes()
+	if g.activeLineSize == g.lineSize {
+		copy(dest[:g.frameSize], mapped)
+	} else {
+		for y := 0; y < g.height; y++ {
+			src := mapped[y*g.activeLineSize : (y+1)*g.activeLineSize]
+			dst := dest[y*g.lineSize : (y+1)*g.lineSize]
+			copy(dst, src)
+			clear(dst[g.activeLineSize:])
+		}
+	}
 	return g.host.Unmap()
 }