Skip to content

Commit b61f56b

Browse files
tfennetfenne
committed
Merge PR odelaneau#292: Bit-batch phasing_hmm helpers (10-15% perf)
2 parents 781778f + 41f8520 commit b61f56b

1 file changed

Lines changed: 106 additions & 9 deletions

File tree

‎phase/src/models/phasing_hmm.h‎

Lines changed: 106 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,20 @@ void phasing_hmm::INIT_PEAK_HET(int curr_het)
141141
{
142142
const std::array <__m256, 2 > emits = {_mm256_load_ps(&EMIT0[curr_het][0]),_mm256_load_ps(&EMIT1[curr_het][0])};
143143
__m256 _sum = _mm256_set1_ps(0.0f);
144-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER)
144+
const unsigned int n_states = C->n_states;
145+
const unsigned int n_states_full = (n_states / 8) * 8;
146+
int k = 0, i = 0;
147+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
148+
{
149+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
150+
for (int b = 0 ; b < 8 ; ++b)
151+
{
152+
const bool ah = (byte >> (7 - b)) & 1;
153+
_sum = _mm256_add_ps(_sum, emits[ah]);
154+
_mm256_store_ps(&prob[i + b * HAP_NUMBER], emits[ah]);
155+
}
156+
}
157+
for ( ; k < n_states ; ++k, i += HAP_NUMBER)
145158
{
146159
const bool ah = C->Hvar.get(curr_rel_locus, k);
147160
_sum = _mm256_add_ps(_sum, emits[ah]);
@@ -156,7 +169,20 @@ void phasing_hmm::INIT_PEAK_HOM(bool ag)
156169
{
157170
const std::array <__m256, 2 > emits = {_mm256_set1_ps(1.0f),_mm256_set1_ps(C->ed_phs/C->ee_phs)};
158171
__m256 _sum = _mm256_set1_ps(0.0f);
159-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER)
172+
const unsigned int n_states = C->n_states;
173+
const unsigned int n_states_full = (n_states / 8) * 8;
174+
int k = 0, i = 0;
175+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
176+
{
177+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
178+
for (int b = 0 ; b < 8 ; ++b)
179+
{
180+
const bool ag_ah = ((byte >> (7 - b)) & 1) != ag;
181+
_sum = _mm256_add_ps(_sum, emits[ag_ah]);
182+
_mm256_store_ps(&prob[i + b * HAP_NUMBER], emits[ag_ah]);
183+
}
184+
}
185+
for ( ; k < n_states ; ++k, i += HAP_NUMBER)
160186
{
161187
const bool ag_ah = C->Hvar.get(curr_rel_locus, k)!=ag;
162188
_sum = _mm256_add_ps(_sum, emits[ag_ah]);
@@ -180,7 +206,21 @@ void phasing_hmm::RUN_PEAK_HET(int curr_het)
180206
const __m256 _nt = _mm256_set1_ps(nt / probSumT);
181207
const std::array <__m256, 2 > emits = {_mm256_load_ps(&EMIT0[curr_het][0]),_mm256_load_ps(&EMIT1[curr_het][0])};
182208
__m256 _sum = _mm256_set1_ps(0.0f);
183-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER)
209+
const unsigned int n_states = C->n_states;
210+
const unsigned int n_states_full = (n_states / 8) * 8;
211+
int k = 0, i = 0;
212+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
213+
{
214+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
215+
for (int b = 0 ; b < 8 ; ++b)
216+
{
217+
const bool ah = (byte >> (7 - b)) & 1;
218+
const __m256 _prob_curr = _mm256_mul_ps(_mm256_fmadd_ps(_mm256_load_ps(&prob[i + b * HAP_NUMBER]), _nt, _tFreq), emits[ah]);
219+
_sum = _mm256_add_ps(_sum, _prob_curr);
220+
_mm256_store_ps(&prob[i + b * HAP_NUMBER], _prob_curr);
221+
}
222+
}
223+
for ( ; k < n_states ; ++k, i += HAP_NUMBER)
184224
{
185225
const bool ah = C->Hvar.get(curr_rel_locus, k);
186226
const __m256 _prob_curr = _mm256_mul_ps(_mm256_fmadd_ps(_mm256_load_ps(&prob[i]), _nt, _tFreq), emits[ah]);
@@ -198,7 +238,23 @@ void phasing_hmm::RUN_PEAK_HOM(bool ag)
198238
const __m256 _nt = _mm256_set1_ps(nt / probSumT);
199239
const __m256 _mism = _mm256_set1_ps(C->ed_phs/C->ee_phs);
200240
__m256 _sum = _mm256_set1_ps(0.0f);
201-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER)
241+
const unsigned int n_states = C->n_states;
242+
const unsigned int n_states_full = (n_states / 8) * 8;
243+
int k = 0, i = 0;
244+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
245+
{
246+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
247+
for (int b = 0 ; b < 8 ; ++b)
248+
{
249+
const bool ah = (byte >> (7 - b)) & 1;
250+
const __m256 _prob_prev = _mm256_load_ps(&prob[i + b * HAP_NUMBER]);
251+
__m256 _prob_curr = _mm256_fmadd_ps(_prob_prev, _nt, _tFreq);
252+
if (ag!=ah) _prob_curr = _mm256_mul_ps(_prob_curr, _mism);
253+
_sum = _mm256_add_ps(_sum, _prob_curr);
254+
_mm256_store_ps(&prob[i + b * HAP_NUMBER], _prob_curr);
255+
}
256+
}
257+
for ( ; k < n_states ; ++k, i += HAP_NUMBER)
202258
{
203259
const bool ah = C->Hvar.get(curr_rel_locus, k);
204260
const __m256 _prob_prev = _mm256_load_ps(&prob[i]);
@@ -234,7 +290,21 @@ void phasing_hmm::COLLAPSE_PEAK_HET(int curr_het)
234290
const __m256 _nt = _mm256_set1_ps(nt / probSumT);
235291
const std::array <__m256, 2 > emits = {_mm256_load_ps(&EMIT0[curr_het][0]),_mm256_load_ps(&EMIT1[curr_het][0])};
236292
__m256 _sum = _mm256_set1_ps(0.0f);
237-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER)
293+
const unsigned int n_states = C->n_states;
294+
const unsigned int n_states_full = (n_states / 8) * 8;
295+
int k = 0, i = 0;
296+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
297+
{
298+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
299+
for (int b = 0 ; b < 8 ; ++b)
300+
{
301+
const bool ah = (byte >> (7 - b)) & 1;
302+
const __m256 _prob_curr = _mm256_mul_ps(_mm256_fmadd_ps(_mm256_set1_ps(probSumK[k + b]), _nt, _tFreq), emits[ah]);
303+
_sum = _mm256_add_ps(_sum, _prob_curr);
304+
_mm256_store_ps(&prob[i + b * HAP_NUMBER], _prob_curr);
305+
}
306+
}
307+
for ( ; k < n_states ; ++k, i += HAP_NUMBER)
238308
{
239309
const bool ah = C->Hvar.get(curr_rel_locus, k);
240310
const __m256 _prob_curr = _mm256_mul_ps(_mm256_fmadd_ps(_mm256_set1_ps(probSumK[k]), _nt, _tFreq), emits[ah]);
@@ -252,8 +322,22 @@ void phasing_hmm::COLLAPSE_PEAK_HOM(bool ag)
252322
const __m256 _nt = _mm256_set1_ps(nt / probSumT);
253323
__m256 _sum = _mm256_set1_ps(0.0f);
254324
const __m256 _mism = _mm256_set1_ps(C->ed_phs/C->ee_phs);
255-
256-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER)
325+
const unsigned int n_states = C->n_states;
326+
const unsigned int n_states_full = (n_states / 8) * 8;
327+
int k = 0, i = 0;
328+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
329+
{
330+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
331+
for (int b = 0 ; b < 8 ; ++b)
332+
{
333+
const bool ah = (byte >> (7 - b)) & 1;
334+
__m256 _prob_curr = _mm256_fmadd_ps(_mm256_set1_ps(probSumK[k + b]), _nt, _tFreq);
335+
if (ag!=ah) _prob_curr = _mm256_mul_ps(_prob_curr, _mism);
336+
_sum = _mm256_add_ps(_sum, _prob_curr);
337+
_mm256_store_ps(&prob[i + b * HAP_NUMBER], _prob_curr);
338+
}
339+
}
340+
for ( ; k < n_states ; ++k, i += HAP_NUMBER)
257341
{
258342
const bool ah = C->Hvar.get(curr_rel_locus, k);
259343
__m256 _prob_curr = _mm256_fmadd_ps(_mm256_set1_ps(probSumK[k]), _nt, _tFreq);
@@ -337,8 +421,21 @@ void phasing_hmm::IMPUTE_FLAT_HET()
337421
_scaleR = _mm256_div_ps(_one, _scaleR);
338422
_scaleL = _mm256_div_ps(_one, _scaleL);
339423
std::array <__m256, 2 > sums = {_mm256_set1_ps(0.0f),_mm256_set1_ps(0.0f)};
340-
341-
for(int k = 0, i = 0 ; k != C->n_states ; ++k, i += HAP_NUMBER) {
424+
const unsigned int n_states = C->n_states;
425+
const unsigned int n_states_full = (n_states / 8) * 8;
426+
int k = 0, i = 0;
427+
for ( ; k < n_states_full ; k += 8, i += 8 * HAP_NUMBER)
428+
{
429+
const unsigned char byte = C->Hvar.getByte(curr_rel_locus, k);
430+
for (int b = 0 ; b < 8 ; ++b)
431+
{
432+
const bool ah = (byte >> (7 - b)) & 1;
433+
const __m256 _p1 = _mm256_mul_ps(_mm256_load_ps(&imputeProb[curr_missing_locus*states_haps + i + b * HAP_NUMBER]), _scaleR);
434+
const __m256 _p2 = _mm256_mul_ps(_mm256_load_ps(&prob[i + b * HAP_NUMBER]), _scaleL);
435+
sums[ah] = _mm256_add_ps(sums[ah], _mm256_mul_ps(_p1,_p2));
436+
}
437+
}
438+
for ( ; k < n_states ; ++k, i += HAP_NUMBER) {
342439
const bool ah = C->Hvar.get(curr_rel_locus, k);
343440
const __m256 _p1 = _mm256_mul_ps(_mm256_load_ps(&imputeProb[curr_missing_locus*states_haps + i]), _scaleR);
344441
const __m256 _p2 = _mm256_mul_ps(_mm256_load_ps(&prob[i]), _scaleL);

0 commit comments

Comments
 (0)