SSE2 Smith-Waterman; unfinished

ksw.c

@@ -0,0 +1,158 @@
+#include <stdlib.h>
+#include <stdint.h>
+#include <emmintrin.h>
+#include "ksw.h"
+
+#include <stdio.h>
+
+struct _ksw_query_t {
+	int slen;
+	uint8_t shift;
+	uint8_t *mem;
+	__m128i *qp, *H0, *H1, *E;
+};
+
+ksw_query_t *ksw_qinit(int qlen, const uint8_t *query, int p, int m, const int8_t *mat)
+{
+	ksw_query_t *q;
+	uint8_t *aligned;
+	int8_t *t;
+	int qlenp, slen, a, tmp;
+
+	q = calloc(1, sizeof(ksw_query_t));
+	q->slen = slen = (qlen + p - 1) / p; // length of segment
+	qlenp = slen * p; // qlenp can be divided by p
+	q->mem = malloc(64 + qlenp * (m + 2));
+	aligned = (uint8_t*)(((size_t)q->mem + 15) & ~0xful); // align memory
+	q->qp = (__m128i*)aligned;
+	q->H0 = q->qp + qlenp * m;
+	q->H1 = q->H0 + qlenp;
+	q->E  = q->H1 + qlenp;
+	// compute shift
+	tmp = m * m;
+	for (a = 0, q->shift = 127; a < tmp; ++a) // find the minimum score (should be negative)
+		if (mat[a] < (int8_t)q->shift) q->shift = mat[a];
+	q->shift = 256 - q->shift;
+	// An example: p=8, qlen=19, slen=3 and segmentation:
+	//  {{0,3,6,9,12,15,18,-1},{1,4,7,10,13,16,-1,-1},{2,5,8,11,14,17,-1,-1}}
+	t = (int8_t*)q->qp;
+	for (a = 0; a < m; ++a) {
+		int i, k;
+		const int8_t *ma = mat + a * m;
+		for (i = 0; i < slen; ++i)
+			for (k = i; k < qlenp; k += slen) { // p iterations
+				*t++ = (k >= qlen? 0 : ma[query[k]]) + q->shift;
+				//printf("%d,%d,%d,%d\n", a, i, k, *(t-1));
+			}
+	}
+	return q;
+}
+
+void ksw_qdestroy(ksw_query_t *q)
+{
+	if (!q) return;
+	free(q->mem); free(q);
+}
+
+int ksw_sse2_16(ksw_query_t *q, int tlen, const uint8_t *target, unsigned _o, unsigned _e) // the first gap costs -(_o+_e)
+{
+	int slen, i, score;
+	__m128i zero, gapo, gape, shift, gmax, *H0, *H1, *E;
+
+#define __set_16(ret, xx) do { \
+		uint16_t t16 = ((uint16_t)(xx) << 8) | ((uint16_t)(xx) & 0x00ff); \
+		(ret) = _mm_insert_epi16((ret), t16, 0); \
+		(ret) = _mm_shufflelo_epi16((ret), 0); \
+		(ret) = _mm_shuffle_epi32((ret), 0); \
+	} while (0)
+
+#define __max_16(ret, xx) do { \
+		__m128i t; \
+		t = _mm_srli_si128((xx), 8); \
+		(xx) = _mm_max_epu8((xx), t); \
+		t = _mm_srli_si128((xx), 4); \
+		(xx) = _mm_max_epu8((xx), t); \
+		t = _mm_srli_si128((xx), 2); \
+		(xx) = _mm_max_epu8((xx), t); \
+		t = _mm_srli_si128((xx), 1); \
+		(xx) = _mm_max_epu8((xx), t); \
+    	(ret) = _mm_extract_epi16((xx), 0) & 0x00ff; \
+	} while (0)
+
+	// initialization
+	zero = _mm_xor_si128(zero, zero);
+	gmax = _mm_xor_si128(gmax, gmax);
+	__set_16(gapo, _o + _e);
+	__set_16(gape, _e);
+	H0 = q->H0; H1 = q->H1; E = q->E;
+	slen = q->slen;
+	__set_16(shift, q->shift);
+	for (i = 0; i < slen; ++i) {
+		_mm_store_si128(E + i, zero);
+		_mm_store_si128(H0 + i, zero);
+	}
+	// the core loop
+	for (i = 0; i < tlen; ++i) {
+		int j, cmp;
+		__m128i e, h, f, max, t, *S = q->qp + target[i] * slen; // s is the 1st score vector
+		max = _mm_xor_si128(max, max);
+		f = _mm_xor_si128(f, f);
+		h = _mm_load_si128(H0 + slen - 1); // h={2,5,8,11,14,17,-1,-1} in the above example
+		h = _mm_slli_si128(h, 1); // h=H(i-1,-1); << instead of >> because x86 is little-endian
+		for (score=0;score<16;++score)printf("%d ", ((int8_t*)&S[0])[score]);printf("\n");
+		for (j = 0; j < slen; ++j) {
+			// at the beginning, h=H'(i-1,j-1)
+			h = _mm_adds_epu8(h, S[j]); h = _mm_subs_epu8(h, shift); // h=H'(i-1,j-1)+S(i,j)
+			e = _mm_load_si128(E + j); // e=E'(i,j)
+			h = _mm_max_epu8(h, e);
+			h = _mm_max_epu8(h, f); // h=H'(i,j)
+			max = _mm_max_epu8(max, h); // set max
+			for (score=0;score<16;++score)printf("%d ", ((int8_t*)&h)[score]);printf("\n");
+			_mm_store_si128(H1 + j, h); // save to H'(i,j)
+			// now compute E(i+1,j)
+			h = _mm_subs_epu8(h, gapo); // h=H'(i,j)-gapo
+			e = _mm_subs_epu8(e, gape); // e=E'(i,j)-gape
+			e = _mm_max_epu8(e, h); // e=E'(i+1,j)
+			_mm_store_si128(E + j, e); // save to E'(i+1,j)
+			// now compute F'(i,j+1)
+			f = _mm_subs_epu8(f, gape);
+			f = _mm_max_epu8(f, h);
+			// get H'(i-1,j) and prepare for the next j
+			h = _mm_load_si128(H0 + j); // h=H'(i-1,j)
+		}
+		gmax = _mm_max_epu8(gmax, max);
+		j = 0;
+		h = _mm_load_si128(H1); // h=H(i,0); h->{0,3,6,9,12,15,18,-1}
+		f = _mm_slli_si128(f, 1); // f->{-1,3,6,9,12,15,18,-1}
+		t = _mm_subs_epu8(h, gapo); t = _mm_subs_epu8(f, t);
+		t = _mm_cmpeq_epi8(t, zero); cmp = _mm_movemask_epi8(t); // test whether s>0
+		while (cmp != 0xffff) { // the lazy-F loop
+			h = _mm_max_epu8(h, f);
+			_mm_store_si128(H1 + j, h);
+			e = _mm_load_si128(E + j); // e=E'(i,j)
+			h = _mm_subs_epu8(h, gapo); // h=H(i,j)-gapo
+			e = _mm_max_epu8(e, h); // e=E(i,j)
+			_mm_store_si128(E + j, e); // save to E(i,j)
+			f = _mm_subs_epu8(f, gape); // f-=gape
+			if (++j >= slen) j = 0, f = _mm_slli_si128(f, 1); // restart
+			h = _mm_load_si128(H1 + j);
+			t = _mm_subs_epu8(h, gapo); t = _mm_subs_epu8(f, t);
+			t = _mm_cmpeq_epi8(t, zero); cmp = _mm_movemask_epi8(t);
+		}
+		S = H1; H1 = H0; H0 = S; // swap H0 and H1
+	}
+	__max_16(score, gmax);
+	return score;
+}
+
+int main()
+{
+	int8_t mat[] = {1, -3, -3, 1};
+	uint8_t *t = (uint8_t*)"\1\0\1\1\0\0";
+	uint8_t *q = (uint8_t*)"\1\1";
+	ksw_query_t *qq = ksw_qinit(2, q, 16, 2, mat);
+	int s = ksw_sse2_16(qq, 6, t, 5, 2);
+	ksw_qdestroy(qq);
+	printf("%d\n", s);
+	return 0;
+}

ksw.h

@@ -0,0 +1,19 @@
+#ifndef __AC_KSW_H
+#define __AC_KSW_H
+
+struct _ksw_query_t;
+typedef struct _ksw_query_t ksw_query_t;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+ksw_query_t *ksw_qinit(int qlen, const uint8_t *query, int p, int m, const int8_t *mat);
+void ksw_qdestroy(ksw_query_t *q);
+int ksw_sse2_16(ksw_query_t *q, int tlen, const uint8_t *target, unsigned o, unsigned e); // first gap costs -(o+e)
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif