mirror of
https://github.com/metabarcoding/obitools4.git
synced 2026-11-09 09:24:41 +00:00
Eric Coissac
a0d8187fd6
Adds support for Hamming distance, constrained k-mer generation, and optimized Longest Common Subsequence (LCS) calculations using scalar, batch, and SIMD methods. This includes new methods for k-mer distance calculation, structures for k-mer constraints, and SIMD-based batch processing for LCS to improve performance and correctness.
81 lines
2.1 KiB
Go
81 lines
2.1 KiB
Go
//go:build goexperiment.simd
|
|
|
|
package obikmer
|
|
|
|
import "simd"
|
|
|
|
// Each lane of a vector holds the state of one (query, partner) alignment as
|
|
//
|
|
// score<<6 | (63 - alignment length)
|
|
//
|
|
// so that taking the maximum of two cells ranks them by score first and, for
|
|
// equal scores, by shortest alignment. With k <= 31 the score is <= 31 and the
|
|
// alignment length <= 62, hence the value fits in 11 bits and the low field
|
|
// is never exhausted. The three DP moves then all reduce to
|
|
//
|
|
// cell(i,j) = max(diag + 64*match, up, left) - 1
|
|
//
|
|
// (every move lengthens the alignment by one; only a diagonal match adds 1
|
|
// to the score).
|
|
const (
|
|
lcsScoreUnit = 64
|
|
lcsLenBase = 63
|
|
)
|
|
|
|
func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) {
|
|
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
|
|
}
|
|
|
|
// 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,
|
|
}
|
|
}
|
|
}
|
|
}
|