diff --git a/src/layer/arm/gru_arm.cpp b/src/layer/arm/gru_arm.cpp index a49147813..c9b70ce19 100644 --- a/src/layer/arm/gru_arm.cpp +++ b/src/layer/arm/gru_arm.cpp @@ -273,10 +273,16 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we { int ti = reverse ? T - 1 - t : t; - int q = 0; + int remain_num_output_start = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* x = bottom_blob.row(ti); // gate reset update @@ -494,7 +500,8 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we vst1q_f32(gates_data + 4, _N); } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* x = bottom_blob.row(ti); @@ -578,10 +585,15 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we float* hidden_ptr = hidden_state; - q = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q / 4); float32x4_t _U = vld1q_f32(gates_data); @@ -596,7 +608,8 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we output_data += 4; } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { #if __ARM_NEON const float* gates_data = gates.row(q / 4 + q % 4); @@ -801,9 +814,13 @@ static int gru_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M { int ti = reverse ? T - 1 - t : t; - int q = 0; - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + int remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const __fp16* x = bottom_blob.row(ti); // gate reset update @@ -984,7 +1001,8 @@ static int gru_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M vst1q_f32(gates_data, _U); vst1q_f32(gates_data + 4, _N); } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const __fp16* x = bottom_blob.row(ti); @@ -1059,9 +1077,13 @@ static int gru_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M float* hidden_ptr = hidden_state; - q = 0; - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q / 4); float32x4_t _U = vld1q_f32(gates_data); @@ -1075,7 +1097,8 @@ static int gru_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M hidden_ptr += 4; output_data += 4; } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* gates_data = gates.row(q / 4 + q % 4); @@ -1109,9 +1132,13 @@ static int gru_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const { int ti = reverse ? T - 1 - t : t; - int q = 0; - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + int remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const __fp16* x = bottom_blob.row(ti); // gate reset update @@ -1308,7 +1335,8 @@ static int gru_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const vst1q_f32(gates_data, _U32); vst1q_f32(gates_data + 4, _N32); } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const __fp16* x = bottom_blob.row(ti); @@ -1383,9 +1411,13 @@ static int gru_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const float* hidden_ptr = hidden_state; - q = 0; - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q / 4); float32x4_t _U = vld1q_f32(gates_data); @@ -1399,7 +1431,8 @@ static int gru_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const hidden_ptr += 4; output_data += 4; } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* gates_data = gates.row(q / 4 + q % 4); @@ -1895,10 +1928,16 @@ static int gru_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M { int ti = reverse ? T - 1 - t : t; - int q = 0; + int remain_num_output_start = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const unsigned short* x = bottom_blob.row(ti); // gate reset update @@ -2116,7 +2155,8 @@ static int gru_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M vst1q_f32(gates_data + 4, _N); } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const unsigned short* x = bottom_blob.row(ti); @@ -2200,10 +2240,15 @@ static int gru_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M float* hidden_ptr = hidden_state; - q = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q / 4); float32x4_t _U = vld1q_f32(gates_data); @@ -2218,7 +2263,8 @@ static int gru_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M output_data += 4; } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { #if __ARM_NEON const float* gates_data = gates.row(q / 4 + q % 4); diff --git a/src/layer/arm/lstm_arm.cpp b/src/layer/arm/lstm_arm.cpp index 185b8ff54..6dfea26eb 100644 --- a/src/layer/arm/lstm_arm.cpp +++ b/src/layer/arm/lstm_arm.cpp @@ -151,6 +151,7 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w int ti = reverse ? T - 1 - t : t; const float* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* bias_c_IFOG = (const float*)bias_c + q * 4; @@ -288,10 +289,16 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w float* cell_ptr = cell_state; float* hidden_ptr = hidden_state; - int q = 0; + int remain_num_output_start = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q); float32x4x4_t _IFOG_4x4 = vld4q_f32(gates_data); @@ -313,7 +320,8 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w output_data += 4; } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* gates_data = gates.row(q); @@ -547,6 +555,7 @@ static int lstm_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const int ti = reverse ? T - 1 - t : t; const __fp16* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const __fp16* bias_c_IFOG = (const __fp16*)bias_c + q * 4; @@ -638,9 +647,13 @@ static int lstm_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const float* cell_ptr = cell_state; float* hidden_ptr = hidden_state; - int q = 0; - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + int remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q); float32x4x4_t _IFOG_4x4 = vld4q_f32(gates_data); @@ -661,7 +674,8 @@ static int lstm_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const hidden_ptr += 4; output_data += 4; } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* gates_data = gates.row(q); @@ -711,9 +725,13 @@ static int lstm_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const int ti = reverse ? T - 1 - t : t; - int q = 0; - for (; q + 1 < num_output; q += 2) + int nn_num_output = num_output >> 1; + int remain_num_output_start = nn_num_output << 1; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 2; + const __fp16* bias_c_IFOG = (const __fp16*)bias_c + q * 4; // gate I F O G @@ -809,7 +827,8 @@ static int lstm_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const vst1q_f16(gates_data, _IFOG); } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const __fp16* bias_c_IFOG = (const __fp16*)bias_c + q * 4; @@ -919,9 +938,13 @@ static int lstm_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const float* cell_ptr = cell_state; float* hidden_ptr = hidden_state; - q = 0; - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const __fp16* gates_data = gates.row(q); float16x4x4_t _IFOG_4x4 = vld4_f16(gates_data); @@ -942,7 +965,8 @@ static int lstm_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const hidden_ptr += 4; output_data += 4; } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const __fp16* gates_data = gates.row(q); @@ -1490,6 +1514,7 @@ static int lstm_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const int ti = reverse ? T - 1 - t : t; const unsigned short* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const unsigned short* bias_c_IFOG = (const unsigned short*)bias_c + q * 4; @@ -1629,10 +1654,16 @@ static int lstm_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const float* cell_ptr = cell_state; float* hidden_ptr = hidden_state; - int q = 0; + int remain_num_output_start = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* gates_data = gates.row(q); float32x4x4_t _IFOG_4x4 = vld4q_f32(gates_data); @@ -1654,7 +1685,8 @@ static int lstm_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const output_data += 4; } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* gates_data = gates.row(q); diff --git a/src/layer/arm/rnn_arm.cpp b/src/layer/arm/rnn_arm.cpp index 5b6e61ee8..d1a5112be 100644 --- a/src/layer/arm/rnn_arm.cpp +++ b/src/layer/arm/rnn_arm.cpp @@ -160,10 +160,16 @@ static int rnn(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we const float* x = bottom_blob.row(ti); - int q = 0; + int remain_num_output_start = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const float* weight_xc_ptr = weight_xc.row(q / 4); const float* weight_hc_ptr = weight_hc.row(q / 4); @@ -243,7 +249,8 @@ static int rnn(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we vst1q_f32((float*)gates + q, _H); } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { #if __ARM_NEON const float* weight_xc_ptr = weight_xc.row(q / 4 + q % 4); @@ -274,10 +281,15 @@ static int rnn(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we float* hidden_ptr = hidden_state; - q = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + float32x4_t _H = vld1q_f32((float*)gates + q); vst1q_f32(hidden_ptr, _H); @@ -287,7 +299,8 @@ static int rnn(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we output_data += 4; } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { float H = gates[q]; @@ -485,9 +498,13 @@ static int rnn_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M const __fp16* x = bottom_blob.row(ti); - int q = 0; - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + int remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const __fp16* weight_xc_ptr = weight_xc.row(q / 4); const __fp16* weight_hc_ptr = weight_hc.row(q / 4); @@ -552,7 +569,8 @@ static int rnn_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M vst1q_f32((float*)gates + q, _H); } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const __fp16* weight_xc_ptr = weight_xc.row(q / 4 + q % 4); const __fp16* weight_hc_ptr = weight_hc.row(q / 4 + q % 4); @@ -578,9 +596,13 @@ static int rnn_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M float* hidden_ptr = hidden_state; - q = 0; - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + float32x4_t _H = vld1q_f32((float*)gates + q); vst1q_f32(hidden_ptr, _H); @@ -589,7 +611,8 @@ static int rnn_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M hidden_ptr += 4; output_data += 4; } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { float H = gates[q]; @@ -620,9 +643,13 @@ static int rnn_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const const __fp16* x = bottom_blob.row(ti); - int q = 0; - for (; q + 7 < num_output; q += 8) + int nn_num_output = num_output >> 3; + int remain_num_output_start = nn_num_output << 3; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 8; + const __fp16* weight_xc_ptr = weight_xc.row(q / 8); const __fp16* weight_hc_ptr = weight_hc.row(q / 8); @@ -689,8 +716,12 @@ static int rnn_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const vst1q_f32((float*)gates + q, _H32low); vst1q_f32((float*)gates + q + 4, _H32high); } - for (; q + 3 < num_output; q += 4) + nn_num_output = (num_output - remain_num_output_start) >> 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = remain_num_output_start + qq * 4; + const __fp16* weight_xc_ptr = weight_xc.row(q / 8 + (q % 8) / 4); const __fp16* weight_hc_ptr = weight_hc.row(q / 8 + (q % 8) / 4); @@ -755,7 +786,9 @@ static int rnn_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const vst1q_f32((float*)gates + q, _H32); } - for (; q < num_output; q++) + remain_num_output_start += nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const __fp16* weight_xc_ptr = weight_xc.row(q / 8 + (q % 8) / 4 + q % 4); const __fp16* weight_hc_ptr = weight_hc.row(q / 8 + (q % 8) / 4 + q % 4); @@ -781,9 +814,13 @@ static int rnn_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const float* hidden_ptr = hidden_state; - q = 0; - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + float32x4_t _H = vld1q_f32((float*)gates + q); vst1q_f32(hidden_ptr, _H); @@ -792,7 +829,8 @@ static int rnn_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const hidden_ptr += 4; output_data += 4; } - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { float H = gates[q]; @@ -1233,10 +1271,16 @@ static int rnn_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M const unsigned short* x = bottom_blob.row(ti); - int q = 0; + int remain_num_output_start = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + int nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + const unsigned short* weight_xc_ptr = weight_xc.row(q / 4); const unsigned short* weight_hc_ptr = weight_hc.row(q / 4); @@ -1316,7 +1360,8 @@ static int rnn_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M vst1q_f32((float*)gates + q, _H); } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { #if __ARM_NEON const unsigned short* weight_xc_ptr = weight_xc.row(q / 4 + q % 4); @@ -1347,10 +1392,15 @@ static int rnn_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M float* hidden_ptr = hidden_state; - q = 0; #if __ARM_NEON - for (; q + 3 < num_output; q += 4) + nn_num_output = num_output >> 2; + remain_num_output_start = nn_num_output << 2; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 4; + float32x4_t _H = vld1q_f32((float*)gates + q); vst1q_f32(hidden_ptr, _H); @@ -1360,7 +1410,8 @@ static int rnn_bf16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M output_data += 4; } #endif // __ARM_NEON - for (; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { float H = gates[q]; diff --git a/src/layer/gru.cpp b/src/layer/gru.cpp index 5690252ad..675aa4f3c 100644 --- a/src/layer/gru.cpp +++ b/src/layer/gru.cpp @@ -72,6 +72,7 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we int ti = reverse ? T - 1 - t : t; const float* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { float* gates_data = gates.row(q); @@ -143,6 +144,7 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we // h_t := (1 - update) .* new + update .* h_{t-1} float* output_data = top_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* gates_data = gates.row(q); diff --git a/src/layer/lstm.cpp b/src/layer/lstm.cpp index 7907e1dc0..6749f05e0 100644 --- a/src/layer/lstm.cpp +++ b/src/layer/lstm.cpp @@ -79,6 +79,7 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w int ti = reverse ? T - 1 - t : t; const float* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* bias_c_I = bias_c.row(0); @@ -138,6 +139,7 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w // c_t := f_t .* c_{t-1} + i_t .* g_t // h_t := o_t .* tanh[c_t] float* output_data = top_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* gates_data = gates.row(q); diff --git a/src/layer/riscv/gru_riscv.cpp b/src/layer/riscv/gru_riscv.cpp index b57866162..7e8e66ca5 100644 --- a/src/layer/riscv/gru_riscv.cpp +++ b/src/layer/riscv/gru_riscv.cpp @@ -44,6 +44,7 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we int ti = reverse ? T - 1 - t : t; const float* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { float* gates_data = gates.row(q); @@ -189,6 +190,7 @@ static int gru(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we // h_t := (1 - update) .* new + update .* h_{t-1} float* output_data = top_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* gates_data = gates.row(q); @@ -407,6 +409,7 @@ static int gru_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M int ti = reverse ? T - 1 - t : t; const __fp16* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { float* gates_data = gates.row(q); @@ -552,6 +555,7 @@ static int gru_fp16s(const Mat& bottom_blob, Mat& top_blob, int reverse, const M // h_t := (1 - update) .* new + update .* h_{t-1} __fp16* output_data = top_blob.row<__fp16>(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* gates_data = gates.row(q); @@ -735,6 +739,7 @@ static int gru_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const int ti = reverse ? T - 1 - t : t; const __fp16* x = bottom_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { float* gates_data = gates.row(q); @@ -869,6 +874,7 @@ static int gru_fp16sa(const Mat& bottom_blob, Mat& top_blob, int reverse, const // h_t := (1 - update) .* new + update .* h_{t-1} __fp16* output_data = top_blob.row<__fp16>(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* gates_data = gates.row(q); diff --git a/src/layer/rnn.cpp b/src/layer/rnn.cpp index 54627d5c7..57ce39885 100644 --- a/src/layer/rnn.cpp +++ b/src/layer/rnn.cpp @@ -72,11 +72,10 @@ static int rnn(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we int ti = reverse ? T - 1 - t : t; const float* x = bottom_blob.row(ti); - + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { const float* weight_xc_ptr = weight_xc.row(q); - const float* weight_hc_ptr = weight_hc.row(q); float H = bias_c[q]; @@ -97,6 +96,7 @@ static int rnn(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& we } float* output_data = top_blob.row(ti); + #pragma omp parallel for num_threads(opt.num_threads) for (int q = 0; q < num_output; q++) { float H = gates[q]; diff --git a/src/layer/x86/lstm_x86.cpp b/src/layer/x86/lstm_x86.cpp index 2877120e2..59124f790 100644 --- a/src/layer/x86/lstm_x86.cpp +++ b/src/layer/x86/lstm_x86.cpp @@ -58,9 +58,14 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w // gate_input_t := W_hc * h_conted_{t-1} + W_xc * x_t + b_c int ti = reverse ? T - 1 - t : t; - int remain_output = (num_output >> 1) << 1; - for (int q = 0; q + 1 < num_output; q += 2) + + int nn_num_output = num_output >> 1; + int remain_num_output_start = nn_num_output << 1; + #pragma omp parallel for num_threads(opt.num_threads) + for (int qq = 0; qq < nn_num_output; qq++) { + int q = qq * 2; + const float* x = bottom_blob.row(ti); const float* hidden_ptr_r = hidden_state; const float* bias_c_I = bias_c.row(0); @@ -213,8 +218,8 @@ static int lstm(const Mat& bottom_blob, Mat& top_blob, int reverse, const Mat& w gates_data_O[q + 1] = sums[6]; gates_data_G[q + 1] = sums[7]; } - - for (int q = remain_output; q < num_output; q++) + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = remain_num_output_start; q < num_output; q++) { const float* x = bottom_blob.row(ti); const float* hidden_ptr_r = hidden_state;