Feature: Introduce k-mer distance and SIMD LCS logic

Introduces k-mer distance structures and SIMD implementation for LCS simulation.

Adds support for concurrent, read-only k-mer partners, implements data preparation for SIMD alignment, and adds comprehensive testing and benchmarking capabilities.

pkg/obikmer/kmer_distance.go

@@ -97,6 +97,38 @@ func KmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) b
 	return true
 }
 
+// KmerLCSPartners is a fixed set of k-mers prepared once to be aligned against
+// many queries with KmerLCSBatch-like speed, without redoing the per-call
+// reorganisation of the partners. It is read-only after creation and can be
+// shared between goroutines.
+type KmerLCSPartners struct {
+	prepared kmerLCSPrepared
+	n        int
+}
+
+// NewKmerLCSPartners prepares partners for repeated alignment.
+// It returns nil if k is not between 1 and 31. The slice is not copied when
+// built without GOEXPERIMENT=simd, so it must not be modified afterwards.
+func NewKmerLCSPartners(partners []uint64, k int) *KmerLCSPartners {
+	if k < 1 || k > 31 {
+		return nil
+	}
+	return &KmerLCSPartners{prepared: prepareKmerLCS(partners, k), n: len(partners)}
+}
+
+// Len returns the number of prepared partners.
+func (p *KmerLCSPartners) Len() int { return p.n }
+
+// Align computes the alignment of query against every prepared partner and
+// stores it in out. It returns false if out is shorter than Len().
+func (p *KmerLCSPartners) Align(query uint64, out []KmerLCSResult) bool {
+	if len(out) < p.n {
+		return false
+	}
+	p.prepared.align(query, out[:p.n])
+	return true
+}
+
 func kmerLCSBatchScalar(query uint64, partners []uint64, k int, out []KmerLCSResult) {
 	for i, p := range partners {
 		out[i] = KmerLCS(query, p, k)

pkg/obikmer/kmer_distance_test.go

@@ -110,3 +110,62 @@ func BenchmarkKmerLCSScalar21(b *testing.B) { benchBatch(b, 21, kmerLCSBatchScal
 func BenchmarkKmerLCSBatch21(b *testing.B)  { benchBatch(b, 21, kmerLCSBatch) }
 func BenchmarkKmerLCSScalar11(b *testing.B) { benchBatch(b, 11, kmerLCSBatchScalar) }
 func BenchmarkKmerLCSBatch11(b *testing.B)  { benchBatch(b, 11, kmerLCSBatch) }
+
+func TestKmerLCSPartnersMatchesReference(t *testing.T) {
+	rng := rand.New(rand.NewSource(6))
+	for _, k := range []int{1, 5, 21, 31} {
+		for _, n := range []int{0, 1, 7, 8, 9, 50} {
+			partners := make([]uint64, n)
+			for i := range partners {
+				partners[i] = randomKmer(rng, k)
+			}
+			prep := NewKmerLCSPartners(partners, k)
+			out := make([]KmerLCSResult, n)
+			for q := 0; q < 3; q++ { // reuse the same preparation for several queries
+				query := randomKmer(rng, k)
+				if !prep.Align(query, out) {
+					t.Fatalf("k=%d n=%d: Align refused", k, n)
+				}
+				for i, p := range partners {
+					if want := KmerLCS(query, p, k); out[i] != want {
+						t.Fatalf("k=%d n=%d partner %d: got %+v, want %+v", k, n, i, out[i], want)
+					}
+				}
+			}
+		}
+	}
+	if NewKmerLCSPartners(nil, 0) != nil || NewKmerLCSPartners(nil, 32) != nil {
+		t.Fatal("invalid k must give nil")
+	}
+}
+
+// All-pairs scenario: many queries against the same partners.
+func benchAllPairs(b *testing.B, k int, prepared bool) {
+	rng := rand.New(rand.NewSource(7))
+	partners := make([]uint64, 1024)
+	for i := range partners {
+		partners[i] = randomKmer(rng, k)
+	}
+	queries := make([]uint64, 64)
+	for i := range queries {
+		queries[i] = randomKmer(rng, k)
+	}
+	out := make([]KmerLCSResult, len(partners))
+	prep := NewKmerLCSPartners(partners, k)
+	b.ResetTimer()
+	for b.Loop() {
+		for _, q := range queries {
+			if prepared {
+				prep.Align(q, out)
+			} else {
+				KmerLCSBatch(q, partners, k, out)
+			}
+		}
+	}
+	b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N*len(partners)*len(queries)), "ns/pair")
+}
+
+func BenchmarkAllPairsBatch21(b *testing.B)    { benchAllPairs(b, 21, false) }
+func BenchmarkAllPairsPrepared21(b *testing.B) { benchAllPairs(b, 21, true) }
+func BenchmarkAllPairsBatch11(b *testing.B)    { benchAllPairs(b, 11, false) }
+func BenchmarkAllPairsPrepared11(b *testing.B) { benchAllPairs(b, 11, true) }

pkg/obikmer/kmer_lcs_archsimd_arm64.go

@@ -1,62 +0,0 @@
-//go:build goexperiment.simd && arm64
-
-package obikmer
-
-import "simd/archsimd"
-
-// neonLanes is the number of 16-bit lanes of a NEON 128-bit register.
-const neonLanes = 8
-
-// kmerLCSBatchNEON is kmerLCSBatch written directly on archsimd NEON types
-// (see kmer_lcs_simd.go for the cell encoding and recurrence).
-func kmerLCSBatchNEON(query uint64, partners []uint64, k int, out []KmerLCSResult) {
-	var qbase [32]uint16
-	for i := 0; i < k; i++ {
-		qbase[i] = uint16(query>>(2*uint(k-1-i))) & 3
-	}
-
-	pbase := make([]uint16, k*neonLanes)
-	row := make([]uint16, (k+1)*neonLanes)
-	one := archsimd.BroadcastUint16x8(1)
-	bonus := archsimd.BroadcastUint16x8(lcsScoreUnit)
-
-	for start := 0; start < len(partners); start += neonLanes {
-		chunk := partners[start:min(start+neonLanes, len(partners))]
-
-		for lane, p := range chunk {
-			for j := 0; j < k; j++ {
-				pbase[j*neonLanes+lane] = uint16(p>>(2*uint(k-1-j))) & 3
-			}
-		}
-
-		for j := 0; j <= k; j++ {
-			archsimd.BroadcastUint16x8(uint16(lcsLenBase - j)).Store(row[j*neonLanes:])
-		}
-
-		for i := 1; i <= k; i++ {
-			qv := archsimd.BroadcastUint16x8(qbase[i-1])
-			diag := archsimd.LoadUint16x8(row)
-			left := archsimd.BroadcastUint16x8(uint16(lcsLenBase - i))
-			left.Store(row)
-
-			for j := 1; j <= k; j++ {
-				up := archsimd.LoadUint16x8(row[j*neonLanes:])
-				pv := archsimd.LoadUint16x8(pbase[(j-1)*neonLanes:])
-				match := qv.Equal(pv)
-				best := diag.Add(bonus.Masked(match)).Max(up).Max(left).Sub(one)
-				best.Store(row[j*neonLanes:])
-				diag = up
-				left = best
-			}
-		}
-
-		last := row[k*neonLanes:]
-		for lane := range chunk {
-			v := int(last[lane])
-			out[start+lane] = KmerLCSResult{
-				LCS:      v / lcsScoreUnit,
-				AlignLen: lcsLenBase - v%lcsScoreUnit,
-			}
-		}
-	}
-}

pkg/obikmer/kmer_lcs_archsimd_arm64_test.go

@@ -1,31 +0,0 @@
-//go:build goexperiment.simd && arm64
-
-package obikmer
-
-import (
-	"math/rand"
-	"testing"
-)
-
-func TestKmerLCSBatchNEONMatchesReference(t *testing.T) {
-	rng := rand.New(rand.NewSource(5))
-	for _, k := range []int{1, 2, 5, 11, 21, 31} {
-		for _, n := range []int{0, 1, 7, 8, 9, 100} {
-			query := randomKmer(rng, k)
-			partners := make([]uint64, n)
-			for i := range partners {
-				partners[i] = randomKmer(rng, k)
-			}
-			out := make([]KmerLCSResult, n)
-			kmerLCSBatchNEON(query, partners, k, out)
-			for i, p := range partners {
-				if want := KmerLCS(query, p, k); out[i] != want {
-					t.Fatalf("k=%d n=%d partner %d: got %+v, want %+v", k, n, i, out[i], want)
-				}
-			}
-		}
-	}
-}
-
-func BenchmarkKmerLCSNEON21(b *testing.B) { benchBatch(b, 21, kmerLCSBatchNEON) }
-func BenchmarkKmerLCSNEON11(b *testing.B) { benchBatch(b, 11, kmerLCSBatchNEON) }

pkg/obikmer/kmer_lcs_scalar.go

@@ -5,3 +5,17 @@ package obikmer
 func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) {
 	kmerLCSBatchScalar(query, partners, k, out)
 }
+
+// kmerLCSPrepared has nothing to precompute without SIMD.
+type kmerLCSPrepared struct {
+	k        int
+	partners []uint64
+}
+
+func prepareKmerLCS(partners []uint64, k int) kmerLCSPrepared {
+	return kmerLCSPrepared{k: k, partners: partners}
+}
+
+func (p *kmerLCSPrepared) align(query uint64, out []KmerLCSResult) {
+	kmerLCSBatchScalar(query, p.partners, p.k, out)
+}

pkg/obikmer/kmer_lcs_simd.go

@@ -22,37 +22,33 @@ const (
 	lcsLenBase   = 63
 )
 
-func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) {
+func simdLanes() int {
 	var probe simd.Uint16s
-	lanes := probe.Len()
-
-	var qbase [32]uint16
-	for i := 0; i < k; i++ {
-		qbase[i] = uint16(query>>(2*uint(k-1-i))) & 3
+	return probe.Len()
 }
 
-	// partner bases, transposed: pbase[j*lanes+lane] is base j of the partner in `lane`
-	pbase := make([]uint16, k*lanes)
-	// one DP row, one vector per column
-	row := make([]uint16, (k+1)*lanes)
-	one := simd.BroadcastUint16s(1)
-	bonus := simd.BroadcastUint16s(lcsScoreUnit)
-
-	for start := 0; start < len(partners); start += lanes {
-		chunk := partners[start:min(start+lanes, len(partners))]
-
+// transposeChunk writes the bases of up to `lanes` partners in dst, laid out
+// so that dst[j*lanes+lane] is base j of the partner in `lane`.
+func transposeChunk(chunk []uint64, k, lanes int, dst []uint16) {
 	for lane, p := range chunk {
 		for j := 0; j < k; j++ {
-				pbase[j*lanes+lane] = uint16(p>>(2*uint(k-1-j))) & 3
+			dst[j*lanes+lane] = uint16(p>>(2*uint(k-1-j))) & 3
 		}
 	}
+}
+
+// alignChunk aligns the query against the partners of one transposed chunk
+// and writes one result per lane in out (len(out) <= lanes).
+func alignChunk(query uint64, pbase, row []uint16, k, lanes int, out []KmerLCSResult) {
+	one := simd.BroadcastUint16s(1)
+	bonus := simd.BroadcastUint16s(lcsScoreUnit)
 
 	for j := 0; j <= k; j++ {
 		simd.BroadcastUint16s(uint16(lcsLenBase - j)).Store(row[j*lanes:])
 	}
 
 	for i := 1; i <= k; i++ {
-			qv := simd.BroadcastUint16s(qbase[i-1])
+		qv := simd.BroadcastUint16s(uint16(query>>(2*uint(k-i))) & 3)
 		diag := simd.LoadUint16s(row)
 		left := simd.BroadcastUint16s(uint16(lcsLenBase - i))
 		left.Store(row)
@@ -69,12 +65,50 @@ func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) {
 	}
 
 	last := row[k*lanes:]
-		for lane := range chunk {
+	for lane := range out {
 		v := int(last[lane])
-			out[start+lane] = KmerLCSResult{
+		out[lane] = KmerLCSResult{
 			LCS:      v / lcsScoreUnit,
 			AlignLen: lcsLenBase - v%lcsScoreUnit,
 		}
 	}
 }
+
+func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) {
+	lanes := simdLanes()
+	pbase := make([]uint16, k*lanes)
+	row := make([]uint16, (k+1)*lanes)
+
+	for start := 0; start < len(partners); start += lanes {
+		end := min(start+lanes, len(partners))
+		transposeChunk(partners[start:end], k, lanes, pbase)
+		alignChunk(query, pbase, row, k, lanes, out[start:end])
+	}
+}
+
+// kmerLCSPrepared holds partners already transposed, chunk after chunk.
+type kmerLCSPrepared struct {
+	k, lanes, n int
+	pbase       []uint16 // chunk c at [c*k*lanes : (c+1)*k*lanes]
+}
+
+func prepareKmerLCS(partners []uint64, k int) kmerLCSPrepared {
+	lanes := simdLanes()
+	chunks := (len(partners) + lanes - 1) / lanes
+	p := kmerLCSPrepared{k: k, lanes: lanes, n: len(partners),
+		pbase: make([]uint16, chunks*k*lanes)}
+	for c := 0; c < chunks; c++ {
+		end := min((c+1)*lanes, len(partners))
+		transposeChunk(partners[c*lanes:end], k, lanes, p.pbase[c*k*lanes:(c+1)*k*lanes])
+	}
+	return p
+}
+
+func (p *kmerLCSPrepared) align(query uint64, out []KmerLCSResult) {
+	row := make([]uint16, (p.k+1)*p.lanes)
+	size := p.k * p.lanes
+	for c, start := 0, 0; start < p.n; c, start = c+1, start+p.lanes {
+		end := min(start+p.lanes, p.n)
+		alignChunk(query, p.pbase[c*size:(c+1)*size], row, p.k, p.lanes, out[start:end])
+	}
 }