Integrate Vulkan compute backend

internal/pkg/stitch/auto.go

@@ -95,16 +95,20 @@ type AutoStitcher struct {
 
 	seq          sequence
 	dxAbsLowPass float64
+
+	pm pmatch.Instance
 }
 
 // NewAutoStitcher creates a new AutoStitcher.
 func NewAutoStitcher(c Config) *AutoStitcher {
 	return &AutoStitcher{
 		c: c,
+
+		pm: pmatch.NewInstance(),
 	}
 }
 
-func findOffset(prev, curr *image.RGBA, maxDx int) (dx int, cos float64) {
+func (r *AutoStitcher) findOffset(prev, curr *image.RGBA, maxDx int) (dx int, cos float64) {
 	t0 := time.Now()
 	defer func() {
 		log.Trace().Dur("dur", time.Since(t0)).Msg("findOffset() duration")
@@ -153,7 +157,7 @@ func findOffset(prev, curr *image.RGBA, maxDx int) (dx int, cos float64) {
 	// We expect this x value to be found by the search if the frame has not moved.
 	xZero := sliceRect.Min.Sub(subRect.Min).X
 
-	x, _, cos := pmatch.SearchRGBAC(sub.(*image.RGBA), slice.(*image.RGBA))
+	x, _, cos := r.pm.SearchRGBA(sub.(*image.RGBA), slice.(*image.RGBA))
 	return x - xZero, cos
 }
 
@@ -256,7 +260,7 @@ func (r *AutoStitcher) Frame(frameColor image.Image, ts time.Time) *Train {
 		return nil
 	}
 
-	dx, cos := findOffset(r.prevFrameRGBA, frameRGBA, maxDx)
+	dx, cos := r.findOffset(r.prevFrameRGBA, frameRGBA, maxDx)
 	log.Debug().Uint64("prevFrameIx", r.prevFrameIx).Int("dx", dx).Float64("cos", cos).Msg("received frame")
 
 	isActive := len(r.seq.dx) > 0

pkg/avg/c.go

@@ -5,7 +5,7 @@ package avg
 // #cgo CFLAGS: -Wall -Werror -Wextra -pedantic -std=c99
 // #cgo CFLAGS: -O2
 //
-// #cgo amd64 CFLAGS: -march=native
+// #cgo amd64 CFLAGS: -march=x86-64 -mtune=generic
 //
 // #cgo arm64 CFLAGS: -mcpu=cortex-a72 -mtune=cortex-a72
 //

pkg/pmatch/c.go

@@ -16,7 +16,7 @@ package pmatch
 // #cgo CFLAGS: -Wall -Werror -Wextra -pedantic -std=c99
 // #cgo CFLAGS: -O2
 //
-// #cgo amd64 CFLAGS: -march=native
+// #cgo amd64 CFLAGS: -march=x86-64 -mtune=generic
 // #cgo amd64 CFLAGS: -fopenmp
 // #cgo amd64 LDFLAGS: -fopenmp
 //

pkg/pmatch/instance.go

@@ -0,0 +1,12 @@
+package pmatch
+
+import (
+	"image"
+)
+
+// Instance is the common interface for SearchRGBA implementations.
+type Instance interface {
+	SearchRGBA(img, pat *image.RGBA) (int, int, float64)
+	Kind() string
+	Destroy()
+}

pkg/pmatch/instance_novk.go

@@ -0,0 +1,30 @@
+//go:build !vk
+// +build !vk
+
+package pmatch
+
+import "image"
+
+// C implements Instance.
+type C struct{}
+
+// Destroy implements Instance.
+func (p *C) Destroy() {}
+
+// Kind implements Instance.
+func (p *C) Kind() string {
+	return "C"
+}
+
+// SearchRGBA implements Instance.
+func (p *C) SearchRGBA(img *image.RGBA, pat *image.RGBA) (int, int, float64) {
+	return SearchRGBAC(img, pat)
+}
+
+// Compile time interface check.
+var _ Instance = (*C)(nil)
+
+// NewInstance instantiates C.
+func NewInstance() Instance {
+	return &C{}
+}

pkg/pmatch/instance_test.go

@@ -0,0 +1,36 @@
+package pmatch
+
+import (
+	"image"
+	_ "image/jpeg"
+	_ "image/png"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"jo-m.ch/go/trainbot/pkg/imutil"
+)
+
+func Test_NewInstance(t *testing.T) {
+	img := imutil.ToRGBA(LoadTestImg())
+	pat, err := imutil.Sub(img, image.Rect(x0, y0, x0+w, y0+h))
+	require.NoError(t, err)
+
+	inst := NewInstance()
+	defer inst.Destroy()
+
+	x, y, score := inst.SearchRGBA(img, pat.(*image.RGBA))
+	assert.InDelta(t, 1., score, delta)
+	assert.Equal(t, x0, x)
+	assert.Equal(t, y0, y)
+
+	// Also resets pat bounds origin to (0,0).
+	patCopy := imutil.ToRGBA(pat.(*image.RGBA))
+
+	x, y, score = inst.SearchRGBA(img, patCopy)
+	assert.InDelta(t, 1., score, delta)
+	assert.Equal(t, x0, x)
+	assert.Equal(t, y0, y)
+
+	t.Log("Instance Kind:", inst.Kind())
+}

pkg/pmatch/instance_vk.go

@@ -0,0 +1,66 @@
+//go:build vk
+// +build vk
+
+package pmatch
+
+import (
+	"image"
+
+	"github.com/rs/zerolog/log"
+)
+
+type params struct {
+	ImgBounds, PatBounds image.Rectangle
+	ImgStride, PatStride int
+}
+
+type PMatchVk struct {
+	pool map[params]*SearchVk
+}
+
+// Destroy implements Instance.
+func (p *PMatchVk) Destroy() {
+	for _, inst := range p.pool {
+		inst.Destroy()
+	}
+}
+
+// Kind implements Instance.
+func (p *PMatchVk) Kind() string {
+	return "Vk"
+}
+
+// SearchRGBA implements Instance.
+func (p *PMatchVk) SearchRGBA(img *image.RGBA, pat *image.RGBA) (int, int, float64) {
+	params := params{img.Bounds(), pat.Bounds(), img.Stride, pat.Stride}
+
+	inst, ok := p.pool[params]
+
+	if !ok {
+		log.Warn().Interface("params", params).Msg("allocating instance") // TODO: Remove.
+
+		var err error
+		inst, err = NewSearchVk(params.ImgBounds, params.PatBounds, params.ImgStride, params.PatStride, true) // TODO: false.
+		if err != nil {
+			panic(err)
+		}
+		p.pool[params] = inst
+
+		if len(p.pool) > 50 {
+			panic("too many instances are being allocated, check your image pipeline")
+		}
+	}
+
+	maxX, maxY, maxCos, err := inst.Run(img, pat)
+	if err != nil {
+		panic(err)
+	}
+	return maxX, maxY, maxCos
+}
+
+// Compile time interface check.
+var _ Instance = (*PMatchVk)(nil)
+
+func NewInstance() Instance {
+	return &PMatchVk{pool: map[params]*SearchVk{}}
+}

pkg/pmatch/pkg.go

@@ -4,4 +4,5 @@
 // 1. Naive Go implementation - rather slow, but hopefully correct
 // 2. Slightly optimized Go version
 // 3. Cgo version - fastest
+// 4. Vulkan - even faster
 package pmatch

pkg/pmatch/vk.comp

@@ -40,7 +40,7 @@ void main() {
   uint x = gl_GlobalInvocationID.x;
   uint y = gl_GlobalInvocationID.y;
 
-  if (x > M || y > N) {
+  if (x >= M || y >= N) {
     return;
   }
 

pkg/pmatch/vk.go

@@ -165,6 +165,12 @@ func (s *SearchVk) Run(img, pat *image.RGBA) (maxX, maxY int, maxCos float64, er
 		return
 	}
 
+	// Reset results buffer to zero.
+	err = s.resultsBuf.Zero(s.h, s.resSize)
+	if err != nil {
+		return
+	}
+
 	// Write to buffers.
 	err = s.imgBuf.Write(s.h, unsafe.Pointer(&img.Pix[0]), bufsz(img)) // #nosec G103
 	if err != nil {