avx2 implementation for mask_store

This commit is contained in:
Saurabh
2021-06-28 23:44:48 -07:00
committed by Heng Li
parent 6da640e551
commit 558be6b729
3 changed files with 79 additions and 26 deletions

45
chain.c
View File

@@ -69,6 +69,10 @@ mm128_t *mm_chain_dp(int max_dist_x, int max_dist_y, int bw, int max_skip, int m
v = (int32_t*)kmalloc(km, n * 4);
memset(t, 0, n * 4);
#if VECTORIZED_CHAINING && (__AVX512BW__ || __AVX2__)
/* Allocation for debugging
f_avx = (uint32_t*)kmalloc(km, n * 4);
p_avx = (int32_t*)kmalloc(km, n * 4);
*/
anchor_t* anchors = (anchor_t*)malloc(n* sizeof(anchor_t));
for (i = 0; i < n; ++i) {
uint64_t ri = a[i].x;
@@ -79,8 +83,8 @@ mm128_t *mm_chain_dp(int max_dist_x, int max_dist_y, int bw, int max_skip, int m
}
num_bits_t *anchor_r, *anchor_q, *anchor_l;
create_SoA_Anchors_32_bit(anchors, n, anchor_r, anchor_q, anchor_l);
dp_chain obj(max_dist_x, max_dist_y, bw, max_skip, max_iter, gap_scale, is_cdna, n_segs);
dp_chain obj(max_dist_x, max_dist_y, bw, max_skip, max_iter, gap_scale, is_cdna, n_segs);
obj.mm_dp_vectorized(n, &anchors[0], anchor_r, anchor_q, anchor_l, f, p, v, max_dist_x, max_dist_y, NULL, NULL);
// -16 is due to extra padding at the start of arrays
@@ -138,6 +142,45 @@ mm128_t *mm_chain_dp(int max_dist_x, int max_dist_y, int bw, int max_skip, int m
f[i] = max_f, p[i] = max_j;
v[i] = max_j >= 0 && v[max_j] > max_f? v[max_j] : max_f; // v[] keeps the peak score up to i; f[] is the score ending at i, not always the peak
}
#if 0
for (i = 0; i < n; ++i) {
assert(f[i] == f_avx[i] && p[i] == p_avx[i]);
//if(! (f[i] == f_avx[i] && p[i] == p_avx[i]))
{
#if 0
fprintf(stderr, "mm2-score:\n");
for (int itt = 0; itt < n; ++itt) {
fprintf(stderr, "%ld %ld \n", f[itt], p[itt]);
}
fprintf(stderr, "mm2-simd-score:\n");
for (int itt = 0; itt < n; ++itt) {
fprintf(stderr, "%ld %ld \n", f_avx[itt], p_avx[itt]);
}
fprintf(stderr, "anchors:\n");
fprintf(stderr, "%lld\n", n);
for (int itt = 0; itt < n; ++itt) {
uint64_t ri = a[itt].x;
int32_t qi = (int32_t)a[itt].y, q_span = a[itt].y>>32&0xff; // NB: only 8 bits of span is used!!!
fprintf(stderr, "%llu %ld %ld\n", ri, qi, q_span);
}
//exit(0);
#endif
}
}
#if 0
fprintf(stderr, "%llu\n", n);
for (int itt = 0; itt < n; ++itt) {
uint64_t ri = a[itt].x;
int32_t qi = (int32_t)a[itt].y, q_span = a[itt].y>>32&0xff; // NB: only 8 bits of span is used!!!
fprintf(stderr, "%llu %ld %ld\n", ri, qi, q_span);
}
#endif
kfree(km, f_avx); kfree(km, p_avx);
#endif
#endif
// find the ending positions of chains
memset(t, 0, n * 4);

Submodule ext/TAL updated: 4c317716e0...8cdc730276

View File

@@ -664,10 +664,20 @@ void ksw_extd2_avx512(void *km, int qlen, const uint8_t *query, int tlen, const
#endif
#ifdef __AVX2__
__m256i get_mask_store(__m256i msk, void* addr){
return _mm256_or_si256(msk , _mm256_loadu_si256((__m256i*)addr));
}
__m256i get_mask_store(__m256i msk, void* addr, __m256i store_data){
return _mm256_blendv_epi8(_mm256_loadu_si256((__m256i*)addr), store_data, msk);
}
void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const uint8_t *target, int8_t m, const int8_t *mat,
int8_t q, int8_t e, int8_t q2, int8_t e2, int w, int zdrop, int end_bonus, int flag, ksw_extz_t *ez)
{
__m256i bt32_ = _mm256_setr_epi32(0,0,0,0,4,4,4,4);//8,8,8,8,12,12,12,12);
int8_t index[32] __attribute((aligned(64)));
@@ -683,7 +693,7 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
__m256i shf256a, shf256b, slli256;
__m256i ind256_slli = _mm256_load_si256((__m256i*) index);
__mmask8 mska = 0x00;//0x90
__mmask32 mskb = 0x10000;
__mmask32 mskb = 0x00010000;//0000 0000 0000 0001 0000 0000 0000 0000
__mmask32 mskc = 0x1;
__m256i mskb_v = _mm256_set_epi32(0,0,0,255,0,0,0,0);
@@ -750,7 +760,6 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
a2= _mm256_sub_epi8(a2, tmp); \
b2= _mm256_sub_epi8(b2, tmp);
// __mmask64 msk_ar[5] = {0xFFFF, 0xFFFFFFFF, 0xFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF};
__mmask32 msk_ar2[3] = {0xFFFF, 0xFFFF, 0xFFFFFFFF};
__m256i msk_ar2_v[3];
msk_ar2_v[0] = _mm256_set_epi32(0,0,0,0,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF);
@@ -890,10 +899,11 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
{
__mmask32 msk;
int ind = (en0 - t + 16)/16;
assert(ind >= 0 && ind < 3);
//assert(ind >= 0 && ind < 3);
msk = msk_ar2[ind];
__m256i msk_v = msk_ar2_v[ind];
__m256i str =_mm256_and_si256( msk_v, tmp_256);// msk_ar2_v[ind];
__m256i str = (get_mask_store(msk_v, ((int8_t*)s + t), tmp_256));// msk_ar2_v[ind];
//__m256i str =_mm256_and_si256(get_mask_store(msk_v,(int8_t*)s + t), tmp_256);// msk_ar2_v[ind];
_mm256_storeu_si256((__m256i*)((int8_t*)s + t), str);
@@ -933,8 +943,8 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
int ind = (en0 - t*32 + 16)/16;//doubt
msk = msk_ar2[ind];
// fprintf(stderr, "en0: %d, t: %d, ind: %d, msk: %d\n", en0, t, ind, msk);
_mm256_storeu_si256(&u[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(z, vt1)));
_mm256_storeu_si256(&v[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(z, ut)));
_mm256_storeu_si256(&u[t], (get_mask_store(msk_ar2_v[ind], &u[t], _mm256_sub_epi8(z, vt1))));
_mm256_storeu_si256(&v[t], (get_mask_store(msk_ar2_v[ind], &v[t], _mm256_sub_epi8(z, ut))));
tmp = _mm256_sub_epi8(z, q_);
a = _mm256_sub_epi8(a, tmp);
b = _mm256_sub_epi8(b, tmp);
@@ -960,10 +970,10 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
//assert(ind >= 0);
msk = msk_ar2[ind];
// fprintf(stderr, "en0: %d, t: %d, ind: %d, msk: %d\n", en0, t, ind, msk);
_mm256_storeu_si256(&x[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_max_epi8(a, zero_), qe_)));
_mm256_storeu_si256(&y[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_max_epi8(b, zero_), qe_)));
_mm256_storeu_si256(&x2[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_max_epi8(a2, zero_), qe2_)));
_mm256_storeu_si256(&y2[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_max_epi8(b2, zero_), qe2_)));
_mm256_storeu_si256(&x[t], (get_mask_store(msk_ar2_v[ind], &x[t], _mm256_sub_epi8(_mm256_max_epi8(a, zero_), qe_))));
_mm256_storeu_si256(&y[t], (get_mask_store(msk_ar2_v[ind],&y[t], _mm256_sub_epi8(_mm256_max_epi8(b, zero_), qe_))));
_mm256_storeu_si256(&x2[t], (get_mask_store(msk_ar2_v[ind], &x2[t] , _mm256_sub_epi8(_mm256_max_epi8(a2, zero_), qe2_))));
_mm256_storeu_si256(&y2[t], (get_mask_store(msk_ar2_v[ind], &y2[t], _mm256_sub_epi8(_mm256_max_epi8(b2, zero_), qe2_))));
}
else
{
@@ -1042,8 +1052,8 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
// //assert(ind >= 0 && ind < 5);
msk = msk_ar2[ind];
_mm256_storeu_si256(&u[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(z, vt1)));
_mm256_storeu_si256(&v[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(z, ut)));
_mm256_storeu_si256(&u[t], (get_mask_store(msk_ar2_v[ind], &u[t], _mm256_sub_epi8(z, vt1))));
_mm256_storeu_si256(&v[t], (get_mask_store(msk_ar2_v[ind],&v[t], _mm256_sub_epi8(z, ut))));
tmp = _mm256_sub_epi8(z, q_);
a = _mm256_sub_epi8(a, tmp);
b = _mm256_sub_epi8(b, tmp);
@@ -1060,19 +1070,19 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
off_end[r] -= (2-ind)*16;//doubt
tmp_mask = _mm256_cmpgt_epi8(a, zero_);
_mm256_storeu_si256(&x[t], _mm256_and_si256(msk_v, _mm256_sub_epi8(_mm256_blendv_epi8(zero_, a, tmp_mask), qe_)));
_mm256_storeu_si256(&x[t], (get_mask_store(msk_v, &x[t], _mm256_sub_epi8(_mm256_blendv_epi8(zero_, a, tmp_mask), qe_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(zero_, s1_, tmp_mask)); // d = a > 0? 1<<3 : 0
tmp_mask = _mm256_cmpgt_epi8(b, zero_);
_mm256_storeu_si256(&y[t], _mm256_and_si256(msk_v , _mm256_sub_epi8(_mm256_blendv_epi8(zero_, b, tmp_mask), qe_)));
_mm256_storeu_si256(&y[t], (get_mask_store(msk_v,&y[t] , _mm256_sub_epi8(_mm256_blendv_epi8(zero_, b, tmp_mask), qe_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(zero_, s2_, tmp_mask)); // d = b > 0? 1<<4 : 0
tmp_mask = _mm256_cmpgt_epi8(a2, zero_);
_mm256_storeu_si256(&x2[t], _mm256_and_si256(msk_v, _mm256_sub_epi8(_mm256_blendv_epi8(zero_, a2, tmp_mask), qe2_)));
_mm256_storeu_si256(&x2[t], (get_mask_store(msk_v,&x2[t], _mm256_sub_epi8(_mm256_blendv_epi8(zero_, a2, tmp_mask), qe2_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(zero_, s3_, tmp_mask)); // d = a > 0? 1<<5 : 0
tmp_mask = _mm256_cmpgt_epi8(b2, zero_);
_mm256_storeu_si256(&y2[t], _mm256_and_si256(msk_v, _mm256_sub_epi8(_mm256_blendv_epi8(zero_, b2, tmp_mask), qe2_)));
_mm256_storeu_si256(&y2[t], (get_mask_store(msk_v,&y2[t], _mm256_sub_epi8(_mm256_blendv_epi8(zero_, b2, tmp_mask), qe2_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(zero_, s4_, tmp_mask)); // d = b > 0? 1<<6 : 0
_mm256_storeu_si256(&pr[t], _mm256_and_si256(msk_v, d));
_mm256_storeu_si256(&pr[t], (get_mask_store(msk_v, &pr[t], d)));
//_mm256_mask_storeu_epi8(&pr[t], msk, d);
}
}
@@ -1169,8 +1179,8 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
msk = msk_ar2[ind];
off_end[r] -= (2-ind)*16;//doubt
_mm256_storeu_si256(&u[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(z, vt1)));
_mm256_storeu_si256(&v[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(z, ut)));
_mm256_storeu_si256(&u[t], (get_mask_store(msk_ar2_v[ind], &u[t], _mm256_sub_epi8(z, vt1))));
_mm256_storeu_si256(&v[t], (get_mask_store(msk_ar2_v[ind], &v[t], _mm256_sub_epi8(z, ut))));
tmp = _mm256_sub_epi8(z, q_);
a = _mm256_sub_epi8(a, tmp);
b = _mm256_sub_epi8(b, tmp);
@@ -1182,19 +1192,19 @@ void ksw_extd2_avx2(void *km, int qlen, const uint8_t *query, int tlen, const ui
tmp_mask = _mm256_cmpgt_epi8(zero_, a);
_mm256_storeu_si256(&x[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_blendv_epi8(a, zero_, tmp_mask), qe_)));
_mm256_storeu_si256(&x[t], (get_mask_store(msk_ar2_v[ind],&x[t], _mm256_sub_epi8(_mm256_blendv_epi8(a, zero_, tmp_mask), qe_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8( s1_, zero_, tmp_mask)); // d = a > 0? 1<<3 : 0
tmp_mask = _mm256_cmpgt_epi8(zero_, b);
_mm256_storeu_si256(&y[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_blendv_epi8(b, zero_,tmp_mask), qe_)));
_mm256_storeu_si256(&y[t], (get_mask_store(msk_ar2_v[ind], &y[t], _mm256_sub_epi8(_mm256_blendv_epi8(b, zero_,tmp_mask), qe_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(s2_, zero_,tmp_mask)); // d = b > 0? 1<<4 : 0
tmp_mask = _mm256_cmpgt_epi8(zero_, a2);
_mm256_storeu_si256(&x2[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_blendv_epi8(a2, zero_,tmp_mask), qe2_)));
_mm256_storeu_si256(&x2[t], (get_mask_store(msk_ar2_v[ind], &x2[t], _mm256_sub_epi8(_mm256_blendv_epi8(a2, zero_,tmp_mask), qe2_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(s3_, zero_,tmp_mask)); // d = a > 0? 1<<5 : 0
tmp_mask = _mm256_cmpgt_epi8(zero_, b2);
_mm256_storeu_si256(&y2[t], _mm256_and_si256(msk_ar2_v[ind], _mm256_sub_epi8(_mm256_blendv_epi8(b2, zero_,tmp_mask), qe2_)));
_mm256_storeu_si256(&y2[t], (get_mask_store(msk_ar2_v[ind], &y2[t], _mm256_sub_epi8(_mm256_blendv_epi8(b2, zero_,tmp_mask), qe2_))));
d = _mm256_or_si256(d, _mm256_blendv_epi8(s4_, zero_, tmp_mask)); // d = b > 0? 1<<6 : 0
// _mm256_storeu_si256(&pr[t], d);
_mm256_storeu_si256(&pr[t], _mm256_and_si256(msk_ar2_v[ind], d));
_mm256_storeu_si256(&pr[t], (get_mask_store(msk_ar2_v[ind], &pr[t], d)));
}