diff --git a/pkg/obikmer/kmer_distance.go b/pkg/obikmer/kmer_distance.go index 8d17644..b024bf0 100644 --- a/pkg/obikmer/kmer_distance.go +++ b/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) diff --git a/pkg/obikmer/kmer_distance_test.go b/pkg/obikmer/kmer_distance_test.go index a46835a..ea11d51 100644 --- a/pkg/obikmer/kmer_distance_test.go +++ b/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) } diff --git a/pkg/obikmer/kmer_lcs_archsimd_arm64.go b/pkg/obikmer/kmer_lcs_archsimd_arm64.go deleted file mode 100644 index 3438d24..0000000 --- a/pkg/obikmer/kmer_lcs_archsimd_arm64.go +++ /dev/null @@ -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, - } - } - } -} diff --git a/pkg/obikmer/kmer_lcs_archsimd_arm64_test.go b/pkg/obikmer/kmer_lcs_archsimd_arm64_test.go deleted file mode 100644 index 1210b94..0000000 --- a/pkg/obikmer/kmer_lcs_archsimd_arm64_test.go +++ /dev/null @@ -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) } diff --git a/pkg/obikmer/kmer_lcs_scalar.go b/pkg/obikmer/kmer_lcs_scalar.go index 831a131..0095723 100644 --- a/pkg/obikmer/kmer_lcs_scalar.go +++ b/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) +} diff --git a/pkg/obikmer/kmer_lcs_simd.go b/pkg/obikmer/kmer_lcs_simd.go index 6127c62..e2f3d91 100644 --- a/pkg/obikmer/kmer_lcs_simd.go +++ b/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]) + } +}