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_lcs_simd.go

@@ -22,59 +22,93 @@ const (
 	lcsLenBase   = 63
 )
 
-func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) {
+func simdLanes() int {
 	var probe simd.Uint16s
-	lanes := probe.Len()
+	return probe.Len()
+}
 
-	var qbase [32]uint16
-	for i := 0; i < k; i++ {
-		qbase[i] = uint16(query>>(2*uint(k-1-i))) & 3
-	}
-
-	// 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))]
-
-		for lane, p := range chunk {
-			for j := 0; j < k; j++ {
-				pbase[j*lanes+lane] = uint16(p>>(2*uint(k-1-j))) & 3
-			}
-		}
-
-		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])
-			diag := simd.LoadUint16s(row)
-			left := simd.BroadcastUint16s(uint16(lcsLenBase - i))
-			left.Store(row)
-
-			for j := 1; j <= k; j++ {
-				up := simd.LoadUint16s(row[j*lanes:])
-				pv := simd.LoadUint16s(pbase[(j-1)*lanes:])
-				match := qv.Equal(pv)
-				best := diag.Add(bonus.Masked(match)).Max(up).Max(left).Sub(one)
-				best.Store(row[j*lanes:])
-				diag = up
-				left = best
-			}
-		}
-
-		last := row[k*lanes:]
-		for lane := range chunk {
-			v := int(last[lane])
-			out[start+lane] = KmerLCSResult{
-				LCS:      v / lcsScoreUnit,
-				AlignLen: lcsLenBase - v%lcsScoreUnit,
-			}
+// 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++ {
+			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(uint16(query>>(2*uint(k-i))) & 3)
+		diag := simd.LoadUint16s(row)
+		left := simd.BroadcastUint16s(uint16(lcsLenBase - i))
+		left.Store(row)
+
+		for j := 1; j <= k; j++ {
+			up := simd.LoadUint16s(row[j*lanes:])
+			pv := simd.LoadUint16s(pbase[(j-1)*lanes:])
+			match := qv.Equal(pv)
+			best := diag.Add(bonus.Masked(match)).Max(up).Max(left).Sub(one)
+			best.Store(row[j*lanes:])
+			diag = up
+			left = best
+		}
+	}
+
+	last := row[k*lanes:]
+	for lane := range out {
+		v := int(last[lane])
+		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])
+	}
+}