| /* Copyright 2019 Google LLC. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| #include <algorithm> |
| #include <cstdint> |
| |
| #include "ruy/check_macros.h" |
| #include "ruy/kernel_x86.h" |
| #include "ruy/opt_set.h" |
| #include "ruy/platform.h" |
| #include "ruy/profiler/instrumentation.h" |
| |
| #if RUY_PLATFORM_AVX512 && RUY_OPT(ASM) |
| #include <immintrin.h> // IWYU pragma: keep |
| #endif |
| |
| namespace ruy { |
| |
| #if !(RUY_PLATFORM_AVX512 && RUY_OPT(ASM)) |
| |
| void Kernel8bitAvx512(const KernelParams8bit<16, 16>&) { |
| // CPU-ID-based checks should disable the path that would reach this point. |
| RUY_DCHECK(false); |
| } |
| |
| void Kernel8bitAvx512SingleCol(const KernelParams8bit<16, 16>&) { |
| // CPU-ID-based checks should disable the path that would reach this point. |
| RUY_DCHECK(false); |
| } |
| |
| void KernelFloatAvx512(const KernelParamsFloat<16, 16>&) { |
| // CPU-ID-based checks should disable the path that would reach this point. |
| RUY_DCHECK(false); |
| } |
| |
| void KernelFloatAvx512SingleCol(const KernelParamsFloat<16, 16>&) { |
| // CPU-ID-based checks should disable the path that would reach this point. |
| RUY_DCHECK(false); |
| } |
| |
| #else // RUY_PLATFORM_AVX512 && RUY_OPT(ASM) |
| |
| void Kernel8bitAvx512(const KernelParams8bit<16, 16>& params) { |
| profiler::ScopeLabel label("Kernel kAvx512 8-bit"); |
| |
| std::int32_t dst_stride = 0; |
| if ((params.dst_type_id == DstTypeId<std::int8_t>::kValue) || |
| (params.dst_type_id == DstTypeId<std::uint8_t>::kValue)) { |
| dst_stride = params.dst_stride; |
| } else if (params.dst_type_id == DstTypeId<std::int16_t>::kValue) { |
| dst_stride = params.dst_stride / sizeof(std::int16_t); |
| } else if (params.dst_type_id == DstTypeId<std::int32_t>::kValue) { |
| dst_stride = params.dst_stride / sizeof(std::int32_t); |
| } else { |
| RUY_DCHECK(false); |
| } |
| |
| const std::int8_t* rhs_col_ptr = params.rhs_base_ptr; |
| void* dst_col_ptr = params.dst_base_ptr; |
| |
| for (int col = params.start_col; col <= params.last_col; col += 16) { |
| const std::int8_t* lhs_col_ptr = params.lhs_base_ptr; |
| void* dst_ptr = dst_col_ptr; |
| |
| const std::int32_t lhs_zero_point = params.lhs_zero_point; |
| const bool has_rhs_sums_offsets = |
| (params.flags & RUY_ASM_FLAG_HAS_RHS_SUMS) && lhs_zero_point; |
| std::int32_t rhs_sums_offsets[16]; |
| if (has_rhs_sums_offsets) { |
| const __m512i rhs_sums_offset_v = |
| _mm512_mullo_epi32(_mm512_set1_epi32(lhs_zero_point), |
| _mm512_loadu_si512(¶ms.rhs_sums[col])); |
| _mm512_storeu_si512(reinterpret_cast<__m512i*>(rhs_sums_offsets), |
| rhs_sums_offset_v); |
| } |
| |
| for (int row = params.start_row; row <= params.last_row; row += 16) { |
| int channel = |
| (params.flags & RUY_ASM_FLAG_CHANNEL_DIMENSION_IS_COL) ? col : row; |
| int multiplier_channel = |
| (params.flags & RUY_ASM_FLAG_HAS_PERCHANNEL) ? channel : 0; |
| |
| const int residual_rows = std::min(params.dst_rows - row, 16); |
| const int residual_cols = std::min(params.dst_cols - col, 16); |
| |
| __m512i accum_data_v0; |
| __m512i accum_data_v1; |
| __m512i accum_data_v2; |
| __m512i accum_data_v3; |
| __m512i accum_data_v4; |
| __m512i accum_data_v5; |
| __m512i accum_data_v6; |
| __m512i accum_data_v7; |
| __m512i accum_data_v8; |
| __m512i accum_data_v9; |
| __m512i accum_data_va; |
| __m512i accum_data_vb; |
| __m512i accum_data_vc; |
| __m512i accum_data_vd; |
| __m512i accum_data_ve; |
| __m512i accum_data_vf; |
| |
| const __mmask16 row_mask = |
| (static_cast<std::uint32_t>(1) << residual_rows) - 1; |
| |
| // initial_accum_data will be the initialize of each of the |
| // accum_data_* accumulator registers. We compute into it terms that are |
| // identical across columns. |
| __m512i initial_accum_data = _mm512_set1_epi32(params.prod_zp_depth); |
| |
| // In the channels-are-rows case, we can load bias here. |
| if ((params.flags & RUY_ASM_FLAG_HAS_BIAS) && |
| !(params.flags & RUY_ASM_FLAG_CHANNEL_DIMENSION_IS_COL)) { |
| initial_accum_data = _mm512_add_epi32( |
| initial_accum_data, |
| _mm512_loadu_si512( |
| reinterpret_cast<const __m512i*>(params.bias + row))); |
| } |
| |
| const std::int32_t rhs_zero_point = params.rhs_zero_point; |
| if ((params.flags & RUY_ASM_FLAG_HAS_LHS_SUMS) && rhs_zero_point) { |
| const __m512i lhs_sums_offset = |
| _mm512_mullo_epi32(_mm512_set1_epi32(rhs_zero_point), |
| _mm512_loadu_si512(¶ms.lhs_sums[row])); |
| initial_accum_data = |
| _mm512_sub_epi32(initial_accum_data, lhs_sums_offset); |
| } |
| |
| // Adjustments differing across columns. |
| if (has_rhs_sums_offsets) { |
| accum_data_v0 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[0])); |
| accum_data_v1 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[1])); |
| accum_data_v2 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[2])); |
| accum_data_v3 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[3])); |
| accum_data_v4 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[4])); |
| accum_data_v5 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[5])); |
| accum_data_v6 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[6])); |
| accum_data_v7 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[7])); |
| accum_data_v8 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[8])); |
| accum_data_v9 = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[9])); |
| accum_data_va = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[10])); |
| accum_data_vb = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[11])); |
| accum_data_vc = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[12])); |
| accum_data_vd = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[13])); |
| accum_data_ve = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[14])); |
| accum_data_vf = _mm512_sub_epi32( |
| initial_accum_data, _mm512_set1_epi32(rhs_sums_offsets[15])); |
| } else { |
| accum_data_v0 = initial_accum_data; |
| accum_data_v1 = initial_accum_data; |
| accum_data_v2 = initial_accum_data; |
| accum_data_v3 = initial_accum_data; |
| accum_data_v4 = initial_accum_data; |
| accum_data_v5 = initial_accum_data; |
| accum_data_v6 = initial_accum_data; |
| accum_data_v7 = initial_accum_data; |
| accum_data_v8 = initial_accum_data; |
| accum_data_v9 = initial_accum_data; |
| accum_data_va = initial_accum_data; |
| accum_data_vb = initial_accum_data; |
| accum_data_vc = initial_accum_data; |
| accum_data_vd = initial_accum_data; |
| accum_data_ve = initial_accum_data; |
| accum_data_vf = initial_accum_data; |
| } |
| |
| // Finally, in the channels-are-columns case, load bias data here. |
| if ((params.flags & RUY_ASM_FLAG_HAS_BIAS) && |
| (params.flags & RUY_ASM_FLAG_CHANNEL_DIMENSION_IS_COL)) { |
| const __m512i bias_data = _mm512_loadu_si512( |
| reinterpret_cast<const __m512i*>(params.bias + col)); |
| accum_data_v0 = _mm512_add_epi32( |
| accum_data_v0, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(0), bias_data)); |
| accum_data_v1 = _mm512_add_epi32( |
| accum_data_v1, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(1), bias_data)); |
| accum_data_v2 = _mm512_add_epi32( |
| accum_data_v2, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(2), bias_data)); |
| accum_data_v3 = _mm512_add_epi32( |
| accum_data_v3, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(3), bias_data)); |
| accum_data_v4 = _mm512_add_epi32( |
| accum_data_v4, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(4), bias_data)); |
| accum_data_v5 = _mm512_add_epi32( |
| accum_data_v5, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(5), bias_data)); |
| accum_data_v6 = _mm512_add_epi32( |
| accum_data_v6, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(6), bias_data)); |
| accum_data_v7 = _mm512_add_epi32( |
| accum_data_v7, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(7), bias_data)); |
| accum_data_v8 = _mm512_add_epi32( |
| accum_data_v8, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(8), bias_data)); |
| accum_data_v9 = _mm512_add_epi32( |
| accum_data_v9, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(9), bias_data)); |
| accum_data_va = _mm512_add_epi32( |
| accum_data_va, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(10), bias_data)); |
| accum_data_vb = _mm512_add_epi32( |
| accum_data_vb, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(11), bias_data)); |
| accum_data_vc = _mm512_add_epi32( |
| accum_data_vc, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(12), bias_data)); |
| accum_data_vd = _mm512_add_epi32( |
| accum_data_vd, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(13), bias_data)); |
| accum_data_ve = _mm512_add_epi32( |
| accum_data_ve, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(14), bias_data)); |
| accum_data_vf = _mm512_add_epi32( |
| accum_data_vf, |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(15), bias_data)); |
| } |
| |
| const std::int8_t* lhs_ptr = lhs_col_ptr; |
| const std::int8_t* rhs_ptr = rhs_col_ptr; |
| for (int d = 0; d < params.depth; d += 4) { |
| const __m512i lhs_data = _mm512_loadu_si512(lhs_ptr); |
| __m512i rhs_data_8bit = _mm512_loadu_si512(rhs_ptr); |
| |
| // Each "int32" is two 16-bit RHS values, sign extended from 8-bit. |
| std::int32_t rhs_data[32]; |
| const __m256i rhs_data_bottom_lane = |
| _mm512_castsi512_si256(rhs_data_8bit); |
| const __m256i rhs_data_top_lane = |
| _mm512_extracti32x8_epi32(rhs_data_8bit, 1); |
| const __m512i rhs_16_bit_dup_low = |
| _mm512_cvtepi8_epi16(rhs_data_bottom_lane); |
| const __m512i rhs_16_bit_dup_high = |
| _mm512_cvtepi8_epi16(rhs_data_top_lane); |
| // Now that we have cast the RHS data, we store it so that each value |
| // can be separately loaded in the accumulation loop. |
| _mm512_storeu_si512(reinterpret_cast<__m256i*>(rhs_data), |
| rhs_16_bit_dup_low); |
| _mm512_storeu_si512(reinterpret_cast<__m256i*>(rhs_data + 16), |
| rhs_16_bit_dup_high); |
| |
| // Take bytes 0, 1, 4, 5, 8, 9, ... and expand to 16-bit. |
| const __m512i lhs_16_bit_low = |
| _mm512_cvtepi8_epi16(_mm512_cvtepi32_epi16(lhs_data)); |
| // Take bytes 2, 3, 6, 7, 10, 11, ... and expand to 16-bit. |
| const __m512i lhs_16_bit_high = _mm512_cvtepi8_epi16( |
| _mm512_cvtepi32_epi16(_mm512_srli_epi32(lhs_data, 16))); |
| |
| auto process_column = [=](int col, __m512i& accum) { |
| const __m512i rhs_16_bit_dup_low = |
| _mm512_set1_epi32(rhs_data[2 * col]); |
| const __m512i rhs_16_bit_dup_high = |
| _mm512_set1_epi32(rhs_data[2 * col + 1]); |
| |
| accum = _mm512_add_epi32( |
| accum, _mm512_madd_epi16(lhs_16_bit_low, rhs_16_bit_dup_low)); |
| accum = _mm512_add_epi32( |
| accum, _mm512_madd_epi16(lhs_16_bit_high, rhs_16_bit_dup_high)); |
| }; |
| process_column(0, accum_data_v0); |
| process_column(1, accum_data_v1); |
| process_column(2, accum_data_v2); |
| process_column(3, accum_data_v3); |
| process_column(4, accum_data_v4); |
| process_column(5, accum_data_v5); |
| process_column(6, accum_data_v6); |
| process_column(7, accum_data_v7); |
| process_column(8, accum_data_v8); |
| process_column(9, accum_data_v9); |
| process_column(10, accum_data_va); |
| process_column(11, accum_data_vb); |
| process_column(12, accum_data_vc); |
| process_column(13, accum_data_vd); |
| process_column(14, accum_data_ve); |
| process_column(15, accum_data_vf); |
| |
| lhs_ptr += 16 * 4; |
| rhs_ptr += 16 * 4; |
| } |
| |
| if (params.dst_type_id != DstTypeId<std::int32_t>::kValue) { |
| // The non-per-channel case could equivalently be handled in the per-row |
| // or per-column code path. The per-row code path is slightly more |
| // efficient so we handle it there. |
| const bool per_column_multiplier = |
| (params.flags & RUY_ASM_FLAG_HAS_PERCHANNEL) && |
| (params.flags & RUY_ASM_FLAG_CHANNEL_DIMENSION_IS_COL); |
| |
| __m512i m_vector; |
| __m512i e_vector; |
| // Does not make use of RUY_ASM_FLAG_NEEDS_LEFT_SHIFT. |
| m_vector = _mm512_loadu_si512(reinterpret_cast<const __m512i*>( |
| params.multiplier_fixedpoint + multiplier_channel)); |
| e_vector = _mm512_loadu_si512(reinterpret_cast<const __m512i*>( |
| params.multiplier_exponent + multiplier_channel)); |
| |
| const __m512i m_64bit_low = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(m_vector, 0)); |
| const __m512i m_64bit_high = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(m_vector, 1)); |
| |
| const __m512i zero_vector = _mm512_setzero_epi32(); |
| const __m512i left_shift = _mm512_max_epi32(e_vector, zero_vector); |
| const __m512i neg_e_vector = _mm512_sub_epi32(zero_vector, e_vector); |
| const __m512i right_shift = _mm512_max_epi32(neg_e_vector, zero_vector); |
| const __m512i final_right_shift = _mm512_set1_epi32(31); |
| const __m512i right_shift_low = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(right_shift, 0)); |
| const __m512i right_shift_high = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(right_shift, 1)); |
| const __m512i final_right_shift_low = _mm512_cvtepi32_epi64( |
| _mm512_extracti32x8_epi32(final_right_shift, 0)); |
| const __m512i final_right_shift_high = _mm512_cvtepi32_epi64( |
| _mm512_extracti32x8_epi32(final_right_shift, 1)); |
| |
| // A "half" added for rounding prior to truncation of 64-bit value. |
| const __m512i offset_vector = |
| _mm512_slli_epi64(_mm512_set1_epi64(1), 30); |
| |
| auto rounding_right_shift = [=](__m512i& results, |
| const __m512i& exponent) { |
| // Construct the "nudge" value for each lane if the exponent is |
| // greater than 0. Otherwise, the nudge is 0. |
| const __m512i zeros = _mm512_setzero_si512(); |
| const auto mask_rightshift_gtz = |
| _mm512_cmpgt_epi64_mask(exponent, zeros); |
| const __m512i one_shift_exp_minus1 = _mm512_sllv_epi64( |
| _mm512_set1_epi64(1), |
| _mm512_sub_epi64(exponent, _mm512_set1_epi64(1))); |
| __m512i nudge = _mm512_mask_mov_epi64(zeros, mask_rightshift_gtz, |
| one_shift_exp_minus1); |
| // Calculate the shifted sum (results + nudge) >> exp. |
| const __m512i r_plus_nudge = _mm512_add_epi64(results, nudge); |
| const __m512i shifted_sum = _mm512_srav_epi64(r_plus_nudge, exponent); |
| |
| // Identify overflow in each lane and create mask. |
| const __m512i one_shift_31minus_exp = _mm512_sllv_epi64( |
| _mm512_set1_epi64(1), |
| _mm512_sub_epi64(_mm512_set1_epi64(31), exponent)); |
| const auto mask_num_plus_nudge_overflow = _mm512_cmpgt_epi64_mask( |
| results, _mm512_sub_epi64(_mm512_set1_epi64(0x7fffffff), nudge)); |
| // Fill results with either (results + nudge) >> exponent or |
| // 1 << (31 - exp) in the case of overflow. |
| results = _mm512_mask_mov_epi64( |
| shifted_sum, mask_num_plus_nudge_overflow, one_shift_31minus_exp); |
| }; |
| |
| if (per_column_multiplier) { |
| auto apply_multiplier = [=](__m512i& accum, int col) { |
| __m512i perm_64bit_vals = _mm512_set1_epi64(col % 8); |
| // Apply the fixed-point part of the multiplier. |
| __m512i left_shift_val = |
| _mm512_permutexvar_epi32(_mm512_set1_epi32(col), left_shift); |
| __m512i m_64bit_val = _mm512_permutexvar_epi64( |
| perm_64bit_vals, col < 8 ? m_64bit_low : m_64bit_high); |
| __m512i offset_vector_val = |
| _mm512_permutexvar_epi64(perm_64bit_vals, offset_vector); |
| __m512i final_right_shift_val = _mm512_permutexvar_epi64( |
| perm_64bit_vals, |
| col < 8 ? final_right_shift_low : final_right_shift_high); |
| __m512i right_shift_val = _mm512_permutexvar_epi64( |
| perm_64bit_vals, col < 8 ? right_shift_low : right_shift_high); |
| |
| accum = _mm512_sllv_epi32(accum, left_shift_val); |
| __m512i scaled_v_low = _mm512_mul_epi32( |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(accum, 0)), |
| m_64bit_val); |
| __m512i scaled_v_high = _mm512_mul_epi32( |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(accum, 1)), |
| m_64bit_val); |
| |
| scaled_v_low = _mm512_add_epi64(scaled_v_low, offset_vector_val); |
| scaled_v_high = _mm512_add_epi64(scaled_v_high, offset_vector_val); |
| |
| scaled_v_low = |
| _mm512_srav_epi64(scaled_v_low, final_right_shift_val); |
| scaled_v_high = |
| _mm512_srav_epi64(scaled_v_high, final_right_shift_val); |
| |
| rounding_right_shift(scaled_v_low, right_shift_val); |
| rounding_right_shift(scaled_v_high, right_shift_val); |
| |
| accum = _mm512_castsi256_si512(_mm512_cvtepi64_epi32(scaled_v_low)); |
| accum = _mm512_inserti32x8(accum, |
| _mm512_cvtepi64_epi32(scaled_v_high), 1); |
| }; |
| apply_multiplier(accum_data_v0, 0); |
| apply_multiplier(accum_data_v1, 1); |
| apply_multiplier(accum_data_v2, 2); |
| apply_multiplier(accum_data_v3, 3); |
| apply_multiplier(accum_data_v4, 4); |
| apply_multiplier(accum_data_v5, 5); |
| apply_multiplier(accum_data_v6, 6); |
| apply_multiplier(accum_data_v7, 7); |
| apply_multiplier(accum_data_v8, 8); |
| apply_multiplier(accum_data_v9, 9); |
| apply_multiplier(accum_data_va, 10); |
| apply_multiplier(accum_data_vb, 11); |
| apply_multiplier(accum_data_vc, 12); |
| apply_multiplier(accum_data_vd, 13); |
| apply_multiplier(accum_data_ve, 14); |
| apply_multiplier(accum_data_vf, 15); |
| } else { // not per-column, so per-row |
| auto apply_multiplier = [=](__m512i& accum) { |
| accum = _mm512_sllv_epi32(accum, left_shift); |
| // Apply the fixed-point part of the multiplier. |
| __m512i scaled_v_low = _mm512_mul_epi32( |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(accum, 0)), |
| m_64bit_low); |
| __m512i scaled_v_high = _mm512_mul_epi32( |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(accum, 1)), |
| m_64bit_high); |
| |
| scaled_v_low = _mm512_add_epi64(scaled_v_low, offset_vector); |
| scaled_v_high = _mm512_add_epi64(scaled_v_high, offset_vector); |
| |
| scaled_v_low = |
| _mm512_srav_epi64(scaled_v_low, final_right_shift_low); |
| scaled_v_high = |
| _mm512_srav_epi64(scaled_v_high, final_right_shift_high); |
| |
| rounding_right_shift(scaled_v_low, right_shift_low); |
| rounding_right_shift(scaled_v_high, right_shift_high); |
| accum = _mm512_castsi256_si512(_mm512_cvtepi64_epi32(scaled_v_low)); |
| accum = _mm512_inserti32x8(accum, |
| _mm512_cvtepi64_epi32(scaled_v_high), 1); |
| }; |
| apply_multiplier(accum_data_v0); |
| apply_multiplier(accum_data_v1); |
| apply_multiplier(accum_data_v2); |
| apply_multiplier(accum_data_v3); |
| apply_multiplier(accum_data_v4); |
| apply_multiplier(accum_data_v5); |
| apply_multiplier(accum_data_v6); |
| apply_multiplier(accum_data_v7); |
| apply_multiplier(accum_data_v8); |
| apply_multiplier(accum_data_v9); |
| apply_multiplier(accum_data_va); |
| apply_multiplier(accum_data_vb); |
| apply_multiplier(accum_data_vc); |
| apply_multiplier(accum_data_vd); |
| apply_multiplier(accum_data_ve); |
| apply_multiplier(accum_data_vf); |
| } |
| |
| if (params.dst_zero_point != 0) { |
| __m512i dst_zero_point = _mm512_set1_epi32(params.dst_zero_point); |
| accum_data_v0 = _mm512_add_epi32(accum_data_v0, dst_zero_point); |
| accum_data_v1 = _mm512_add_epi32(accum_data_v1, dst_zero_point); |
| accum_data_v2 = _mm512_add_epi32(accum_data_v2, dst_zero_point); |
| accum_data_v3 = _mm512_add_epi32(accum_data_v3, dst_zero_point); |
| accum_data_v4 = _mm512_add_epi32(accum_data_v4, dst_zero_point); |
| accum_data_v5 = _mm512_add_epi32(accum_data_v5, dst_zero_point); |
| accum_data_v6 = _mm512_add_epi32(accum_data_v6, dst_zero_point); |
| accum_data_v7 = _mm512_add_epi32(accum_data_v7, dst_zero_point); |
| accum_data_v8 = _mm512_add_epi32(accum_data_v8, dst_zero_point); |
| accum_data_v9 = _mm512_add_epi32(accum_data_v9, dst_zero_point); |
| accum_data_va = _mm512_add_epi32(accum_data_va, dst_zero_point); |
| accum_data_vb = _mm512_add_epi32(accum_data_vb, dst_zero_point); |
| accum_data_vc = _mm512_add_epi32(accum_data_vc, dst_zero_point); |
| accum_data_vd = _mm512_add_epi32(accum_data_vd, dst_zero_point); |
| accum_data_ve = _mm512_add_epi32(accum_data_ve, dst_zero_point); |
| accum_data_vf = _mm512_add_epi32(accum_data_vf, dst_zero_point); |
| } |
| } |
| |
| const __m512i clamp_max_v = _mm512_set1_epi32(params.clamp_max); |
| const __m512i clamp_min_v = _mm512_set1_epi32(params.clamp_min); |
| |
| const bool store_full_block = |
| (residual_rows == 16) && (residual_cols == 16); |
| |
| __m512i accum_data_v[16]; |
| |
| // In most cases we would make this conditional on (!store_full_block) and |
| // unwind the clamp-and-store loop, but the benefit appears small. |
| { |
| accum_data_v[0] = accum_data_v0; |
| accum_data_v[1] = accum_data_v1; |
| accum_data_v[2] = accum_data_v2; |
| accum_data_v[3] = accum_data_v3; |
| accum_data_v[4] = accum_data_v4; |
| accum_data_v[5] = accum_data_v5; |
| accum_data_v[6] = accum_data_v6; |
| accum_data_v[7] = accum_data_v7; |
| accum_data_v[8] = accum_data_v8; |
| accum_data_v[9] = accum_data_v9; |
| accum_data_v[10] = accum_data_va; |
| accum_data_v[11] = accum_data_vb; |
| accum_data_v[12] = accum_data_vc; |
| accum_data_v[13] = accum_data_vd; |
| accum_data_v[14] = accum_data_ve; |
| accum_data_v[15] = accum_data_vf; |
| } |
| |
| if (params.dst_type_id == DstTypeId<std::int8_t>::kValue) { |
| std::int8_t* tmp_ptr = static_cast<std::int8_t*>(dst_ptr); |
| const int block_col_offset = dst_stride; |
| if (store_full_block) { |
| for (int j = 0; j < 16; ++j) { |
| __m512i result = accum_data_v[j]; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm_storeu_si128( |
| reinterpret_cast<__m128i*>(tmp_ptr + j * block_col_offset), |
| _mm512_cvtepi32_epi8(result)); |
| } |
| } else { |
| for (int j = 0; j < residual_cols; ++j) { |
| __m512i result = accum_data_v[j]; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm_mask_storeu_epi8(tmp_ptr + j * block_col_offset, row_mask, |
| _mm512_cvtepi32_epi8(result)); |
| } |
| } |
| dst_ptr = static_cast<void*>(static_cast<std::int8_t*>(dst_ptr) + 16); |
| } else if (params.dst_type_id == DstTypeId<std::uint8_t>::kValue) { |
| std::uint8_t* tmp_ptr = static_cast<std::uint8_t*>(dst_ptr); |
| const int block_col_offset = dst_stride; |
| if (store_full_block) { |
| for (int j = 0; j < residual_cols; ++j) { |
| __m512i result = accum_data_v[j]; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm_storeu_si128( |
| reinterpret_cast<__m128i*>(tmp_ptr + j * block_col_offset), |
| _mm512_cvtepi32_epi8(result)); |
| } |
| } else { |
| for (int j = 0; j < residual_cols; ++j) { |
| __m512i result = accum_data_v[j]; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm_mask_storeu_epi8(tmp_ptr + j * block_col_offset, row_mask, |
| _mm512_cvtepi32_epi8(result)); |
| } |
| } |
| dst_ptr = static_cast<void*>(static_cast<std::uint8_t*>(dst_ptr) + 16); |
| } else if (params.dst_type_id == DstTypeId<std::int16_t>::kValue) { |
| std::int16_t* tmp_ptr = static_cast<std::int16_t*>(dst_ptr); |
| const int block_col_offset = dst_stride; |
| if (store_full_block) { |
| for (int j = 0; j < 16; ++j) { |
| __m512i result = accum_data_v[j]; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm256_storeu_si256( |
| reinterpret_cast<__m256i*>(tmp_ptr + j * block_col_offset), |
| _mm512_cvtepi32_epi16(result)); |
| } |
| } else { |
| for (int j = 0; j < residual_cols; ++j) { |
| __m512i result = accum_data_v[j]; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm256_mask_storeu_epi16(tmp_ptr + j * block_col_offset, row_mask, |
| _mm512_cvtepi32_epi16(result)); |
| } |
| } |
| dst_ptr = static_cast<void*>(static_cast<std::int16_t*>(dst_ptr) + 16); |
| } else if (params.dst_type_id == DstTypeId<std::int32_t>::kValue) { |
| if (store_full_block) { |
| std::int32_t* tmp_ptr = static_cast<std::int32_t*>(dst_ptr); |
| for (int j = 0; j < 16; ++j) { |
| _mm512_storeu_si512(tmp_ptr + j * dst_stride, accum_data_v[j]); |
| } |
| } else { |
| std::int32_t* tmp_ptr = static_cast<std::int32_t*>(dst_ptr); |
| for (int j = 0; j < residual_cols; ++j) { |
| _mm512_mask_storeu_epi32(tmp_ptr + j * dst_stride, row_mask, |
| accum_data_v[j]); |
| } |
| } |
| dst_ptr = static_cast<void*>(static_cast<std::int32_t*>(dst_ptr) + 16); |
| } else { |
| RUY_DCHECK(false); |
| } |
| |
| lhs_col_ptr += 16 * params.lhs_stride; |
| } // End row-block loop. |
| |
| dst_col_ptr = static_cast<void*>(static_cast<char*>(dst_col_ptr) + |
| 16 * params.dst_stride); |
| rhs_col_ptr += 16 * params.rhs_stride; |
| } // End col-block loop. |
| } // NOLINT(readability/fn_size) |
| |
| void Kernel8bitAvx512SingleCol(const KernelParams8bit<16, 16>& params) { |
| profiler::ScopeLabel label("Kernel kAvx512 8-bit GEMV"); |
| |
| RUY_DCHECK_EQ(params.dst_cols, 1); |
| RUY_DCHECK_EQ(params.last_col, 0); |
| RUY_DCHECK_EQ(params.start_col, 0); |
| |
| int bias_ptr_block_increment = params.flags & RUY_ASM_FLAG_HAS_BIAS ? 16 : 0; |
| |
| const std::int8_t* rhs_col_ptr = params.rhs_base_ptr; |
| void* dst_col_ptr = params.dst_base_ptr; |
| const std::int32_t* bias_col_ptr = params.bias; |
| if (params.flags & RUY_ASM_FLAG_HAS_BIAS) { |
| bias_col_ptr += params.start_row; |
| } |
| |
| const std::int8_t* lhs_col_ptr = params.lhs_base_ptr; |
| void* dst_ptr = dst_col_ptr; |
| const std::int32_t* bias_ptr = bias_col_ptr; |
| |
| const std::int32_t lhs_zero_point = params.lhs_zero_point; |
| const bool has_rhs_sums_offsets = |
| (params.flags & RUY_ASM_FLAG_HAS_RHS_SUMS) && lhs_zero_point; |
| std::int32_t rhs_sums_offsets[16]; |
| if (has_rhs_sums_offsets) { |
| const __m512i rhs_sums_offset_v = |
| _mm512_mullo_epi32(_mm512_set1_epi32(lhs_zero_point), |
| _mm512_loadu_si512(¶ms.rhs_sums[0])); |
| _mm512_storeu_si512(reinterpret_cast<__m512i*>(rhs_sums_offsets), |
| rhs_sums_offset_v); |
| } |
| |
| for (int row = params.start_row; row <= params.last_row; row += 16) { |
| const int residual_rows = std::min(params.dst_rows - row, 16); |
| |
| __m512i accum_data_v0; |
| |
| // Initialize with bias. |
| const __mmask16 row_mask = |
| (static_cast<std::uint32_t>(1) << residual_rows) - 1; |
| __m512i initial_accum_data = |
| _mm512_loadu_si512(reinterpret_cast<const __m512i*>(bias_ptr)); |
| bias_ptr += bias_ptr_block_increment; |
| |
| const std::int32_t rhs_zero_point = params.rhs_zero_point; |
| if ((params.flags & RUY_ASM_FLAG_HAS_LHS_SUMS) && rhs_zero_point) { |
| const __m512i lhs_sums_offset = |
| _mm512_mullo_epi32(_mm512_set1_epi32(rhs_zero_point), |
| _mm512_loadu_si512(¶ms.lhs_sums[row])); |
| initial_accum_data = |
| _mm512_sub_epi32(initial_accum_data, lhs_sums_offset); |
| } |
| |
| const std::int32_t prod_zp_depth = params.prod_zp_depth; |
| if (prod_zp_depth != 0) { |
| initial_accum_data = _mm512_add_epi32(initial_accum_data, |
| _mm512_set1_epi32(prod_zp_depth)); |
| } |
| |
| // Adjustments differing across columns. |
| if (has_rhs_sums_offsets) { |
| accum_data_v0 = _mm512_sub_epi32(initial_accum_data, |
| _mm512_set1_epi32(rhs_sums_offsets[0])); |
| } else { |
| accum_data_v0 = initial_accum_data; |
| } |
| |
| const std::int8_t* lhs_ptr = lhs_col_ptr; |
| const std::int8_t* rhs_ptr = rhs_col_ptr; |
| for (int d = 0; d < params.depth; d += 4) { |
| const __m512i lhs_data = _mm512_loadu_si512(lhs_ptr); |
| const __m128i rhs_data_8bit = |
| _mm_loadu_si128(reinterpret_cast<const __m128i*>(rhs_ptr)); |
| |
| // Each "int32" is two 16-bit RHS values, sign extended from 8-bit. |
| // For simplicity we load 4x the data that we need and process twice the |
| // data that we need and store only the data we need. |
| std::int32_t rhs_data[2]; |
| const __m128i rhs_16_bit_dup = _mm_cvtepi8_epi16(rhs_data_8bit); |
| // Now that we have cast the RHS data, we store it so that each value |
| // can be separately loaded in the accumulation loop. |
| _mm_storeu_si64(reinterpret_cast<__m128i*>(rhs_data), rhs_16_bit_dup); |
| |
| // Take bytes 0, 1, 4, 5, 8, 9, ... and expand to 16-bit. |
| const __m512i lhs_16_bit_low = |
| _mm512_cvtepi8_epi16(_mm512_cvtepi32_epi16(lhs_data)); |
| // Take bytes 2, 3, 6, 7, 10, 11, ... and expand to 16-bit. |
| const __m512i lhs_16_bit_high = _mm512_cvtepi8_epi16( |
| _mm512_cvtepi32_epi16(_mm512_srli_epi32(lhs_data, 16))); |
| |
| // Process column 0. |
| __m512i accum_v = accum_data_v0; |
| constexpr int index = 0; |
| |
| const __m512i rhs_16_bit_dup_low = _mm512_set1_epi32(rhs_data[index]); |
| const __m512i rhs_16_bit_dup_high = |
| _mm512_set1_epi32(rhs_data[index + 1]); |
| |
| accum_v = _mm512_add_epi32( |
| accum_v, _mm512_madd_epi16(lhs_16_bit_low, rhs_16_bit_dup_low)); |
| accum_v = _mm512_add_epi32( |
| accum_v, _mm512_madd_epi16(lhs_16_bit_high, rhs_16_bit_dup_high)); |
| accum_data_v0 = accum_v; |
| |
| lhs_ptr += 16 * 4; |
| rhs_ptr += 16 * 4; |
| } |
| |
| if (params.dst_type_id != DstTypeId<std::int32_t>::kValue) { |
| __m512i m_vector; |
| __m512i e_vector; |
| // Does not make use of RUY_ASM_FLAG_NEEDS_LEFT_SHIFT. |
| int channel = (params.flags & RUY_ASM_FLAG_HAS_PERCHANNEL) ? row : 0; |
| m_vector = _mm512_loadu_si512(reinterpret_cast<const __m512i*>( |
| params.multiplier_fixedpoint + channel)); |
| e_vector = _mm512_loadu_si512(reinterpret_cast<const __m512i*>( |
| params.multiplier_exponent + channel)); |
| |
| const __m512i m_64bit_low = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(m_vector, 0)); |
| const __m512i m_64bit_high = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(m_vector, 1)); |
| |
| const __m512i zero_vector = _mm512_setzero_epi32(); |
| const __m512i left_shift = _mm512_max_epi32(e_vector, zero_vector); |
| const __m512i neg_e_vector = _mm512_sub_epi32(zero_vector, e_vector); |
| const __m512i right_shift = _mm512_max_epi32(neg_e_vector, zero_vector); |
| const __m512i final_right_shift = _mm512_set1_epi32(31); |
| const __m512i right_shift_low = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(right_shift, 0)); |
| const __m512i right_shift_high = |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(right_shift, 1)); |
| const __m512i final_right_shift_low = _mm512_cvtepi32_epi64( |
| _mm512_extracti32x8_epi32(final_right_shift, 0)); |
| const __m512i final_right_shift_high = _mm512_cvtepi32_epi64( |
| _mm512_extracti32x8_epi32(final_right_shift, 1)); |
| |
| // A "half" added for rounding prior to truncation of 64-bit value. |
| const __m512i offset_vector = _mm512_slli_epi64(_mm512_set1_epi64(1), 30); |
| |
| auto rounding_right_shift = [=](__m512i& results, |
| const __m512i& exponent) { |
| // Construct the "nudge" value for each lane if the exponent is |
| // greater than 0. Otherwise, the nudge is 0. |
| const __m512i zeros = _mm512_setzero_si512(); |
| const auto mask_rightshift_gtz = |
| _mm512_cmpgt_epi64_mask(exponent, zeros); |
| const __m512i one_shift_exp_minus1 = |
| _mm512_sllv_epi64(_mm512_set1_epi64(1), |
| _mm512_sub_epi64(exponent, _mm512_set1_epi64(1))); |
| __m512i nudge = _mm512_mask_mov_epi64(zeros, mask_rightshift_gtz, |
| one_shift_exp_minus1); |
| // Calculate the shifted sum (results + nudge) >> exp. |
| const __m512i r_plus_nudge = _mm512_add_epi64(results, nudge); |
| const __m512i shifted_sum = _mm512_srav_epi64(r_plus_nudge, exponent); |
| |
| // Identify overflow in each lane and create mask. |
| const __m512i one_shift_31minus_exp = _mm512_sllv_epi64( |
| _mm512_set1_epi64(1), |
| _mm512_sub_epi64(_mm512_set1_epi64(31), exponent)); |
| const auto mask_num_plus_nudge_overflow = _mm512_cmpgt_epi64_mask( |
| results, _mm512_sub_epi64(_mm512_set1_epi64(0x7fffffff), nudge)); |
| // Fill results with either (results + nudge) >> exponent or |
| // 1 << (31 - exp) in the case of overflow. |
| results = _mm512_mask_mov_epi64( |
| shifted_sum, mask_num_plus_nudge_overflow, one_shift_31minus_exp); |
| }; |
| |
| // Shift and round column 0. |
| accum_data_v0 = _mm512_sllv_epi32(accum_data_v0, left_shift); |
| // Apply the fixed-point part of the multiplier. |
| __m512i scaled_v_low = _mm512_mul_epi32( |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(accum_data_v0, 0)), |
| m_64bit_low); |
| __m512i scaled_v_high = _mm512_mul_epi32( |
| _mm512_cvtepi32_epi64(_mm512_extracti32x8_epi32(accum_data_v0, 1)), |
| m_64bit_high); |
| |
| scaled_v_low = _mm512_add_epi64(scaled_v_low, offset_vector); |
| scaled_v_high = _mm512_add_epi64(scaled_v_high, offset_vector); |
| |
| scaled_v_low = _mm512_srav_epi64(scaled_v_low, final_right_shift_low); |
| scaled_v_high = _mm512_srav_epi64(scaled_v_high, final_right_shift_high); |
| |
| rounding_right_shift(scaled_v_low, right_shift_low); |
| rounding_right_shift(scaled_v_high, right_shift_high); |
| |
| accum_data_v0 = |
| _mm512_castsi256_si512(_mm512_cvtepi64_epi32(scaled_v_low)); |
| accum_data_v0 = _mm512_inserti32x8( |
| accum_data_v0, _mm512_cvtepi64_epi32(scaled_v_high), 1); |
| |
| if (params.dst_zero_point != 0) { |
| __m512i dst_zero_point = _mm512_set1_epi32(params.dst_zero_point); |
| accum_data_v0 = _mm512_add_epi32(accum_data_v0, dst_zero_point); |
| } |
| } |
| |
| const __m512i clamp_max_v = _mm512_set1_epi32(params.clamp_max); |
| const __m512i clamp_min_v = _mm512_set1_epi32(params.clamp_min); |
| |
| if (params.dst_type_id == DstTypeId<std::int8_t>::kValue) { |
| std::int8_t* tmp_ptr = static_cast<std::int8_t*>(dst_ptr); |
| __m512i result = accum_data_v0; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm_mask_storeu_epi8(tmp_ptr, row_mask, _mm512_cvtepi32_epi8(result)); |
| dst_ptr = static_cast<void*>(static_cast<std::int8_t*>(dst_ptr) + 16); |
| } else if (params.dst_type_id == DstTypeId<std::uint8_t>::kValue) { |
| std::uint8_t* tmp_ptr = static_cast<std::uint8_t*>(dst_ptr); |
| __m512i result = accum_data_v0; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm_mask_storeu_epi8(tmp_ptr, row_mask, _mm512_cvtepi32_epi8(result)); |
| dst_ptr = static_cast<void*>(static_cast<std::uint8_t*>(dst_ptr) + 16); |
| } else if (params.dst_type_id == DstTypeId<std::int16_t>::kValue) { |
| std::int16_t* tmp_ptr = static_cast<std::int16_t*>(dst_ptr); |
| __m512i result = accum_data_v0; |
| result = _mm512_min_epi32(result, clamp_max_v); |
| result = _mm512_max_epi32(result, clamp_min_v); |
| _mm256_mask_storeu_epi16(tmp_ptr, row_mask, |
| _mm512_cvtepi32_epi16(result)); |
| dst_ptr = static_cast<void*>(static_cast<std::int16_t*>(dst_ptr) + 16); |
| } else if (params.dst_type_id == DstTypeId<std::int32_t>::kValue) { |
| std::int32_t* tmp_ptr = static_cast<std::int32_t*>(dst_ptr); |
| _mm512_mask_storeu_epi32(tmp_ptr, row_mask, accum_data_v0); |
| dst_ptr = static_cast<void*>(static_cast<std::int32_t*>(dst_ptr) + 16); |
| } else { |
| RUY_DCHECK(false); |
| } |
| |
| lhs_col_ptr += 16 * params.lhs_stride; |
| } // End row-block loop. |
| } // NOLINT(readability/fn_size) |
| |
| void KernelFloatAvx512(const KernelParamsFloat<16, 16>& params) { |
| profiler::ScopeLabel label("Kernel kAvx512 float"); |
| |
| // As parameters are defined, we need to scale by sizeof(float). |
| const std::int64_t lhs_stride = params.lhs_stride >> 2; |
| const std::int64_t dst_stride = params.dst_stride >> 2; |
| const std::int64_t rhs_stride = params.rhs_stride >> 2; |
| |
| int bias_ptr_block_increment = params.flags & RUY_ASM_FLAG_HAS_BIAS ? 1 : 0; |
| const int end_row = std::min(params.dst_rows, params.last_row + 16); |
| const int end_col = std::min(params.dst_cols, params.last_col + 16); |
| |
| const float* adj_rhs_col_ptr = |
| params.rhs_base_ptr - params.start_col * rhs_stride; |
| float* adj_dst_col_ptr = |
| params.dst_base_ptr - params.start_col * dst_stride - params.start_row; |
| const float* adj_lhs_col_ptr = |
| params.lhs_base_ptr - params.start_row * lhs_stride; |
| const float* bias_ptr = params.bias; |
| |
| const __m512 clamp_max_v = _mm512_set1_ps(params.clamp_max); |
| const __m512 clamp_min_v = _mm512_set1_ps(params.clamp_min); |
| const bool channel_dimension_is_col = |
| params.flags & RUY_ASM_FLAG_CHANNEL_DIMENSION_IS_COL; |
| |
| int col = params.start_col; |
| for (; col <= end_col - 16; col += 16) { |
| const float* rhs_col_ptr = adj_rhs_col_ptr + col * rhs_stride; |
| float* dst_col_ptr = adj_dst_col_ptr + col * dst_stride; |
| |
| int row = params.start_row; |
| for (; row <= end_row - 16; row += 16) { |
| const float* lhs_col_ptr = adj_lhs_col_ptr + row * lhs_stride; |
| float* dst_ptr = dst_col_ptr + row; |
| |
| // Process block in two halves, split by columns. |
| #pragma unroll(1) |
| for (int mmm = 0; mmm < 2; ++mmm) { |
| __m512 accum_data_v0; |
| __m512 accum_data_v1; |
| __m512 accum_data_v2; |
| __m512 accum_data_v3; |
| __m512 accum_data_v4; |
| __m512 accum_data_v5; |
| __m512 accum_data_v6; |
| __m512 accum_data_v7; |
| |
| // Initialize with bias. |
| if (channel_dimension_is_col) { |
| const float* bias_elem_ptr = |
| bias_ptr + (col + 8 * mmm) * bias_ptr_block_increment; |
| accum_data_v0 = _mm512_set1_ps(bias_elem_ptr[0]); |
| accum_data_v1 = _mm512_set1_ps(bias_elem_ptr[1]); |
| accum_data_v2 = _mm512_set1_ps(bias_elem_ptr[2]); |
| accum_data_v3 = _mm512_set1_ps(bias_elem_ptr[3]); |
| accum_data_v4 = _mm512_set1_ps(bias_elem_ptr[4]); |
| accum_data_v5 = _mm512_set1_ps(bias_elem_ptr[5]); |
| accum_data_v6 = _mm512_set1_ps(bias_elem_ptr[6]); |
| accum_data_v7 = _mm512_set1_ps(bias_elem_ptr[7]); |
| } else { |
| const __m512 initial_accum_data = |
| _mm512_loadu_ps(bias_ptr + row * bias_ptr_block_increment); |
| |
| accum_data_v0 = initial_accum_data; |
| accum_data_v1 = initial_accum_data; |
| accum_data_v2 = initial_accum_data; |
| accum_data_v3 = initial_accum_data; |
| accum_data_v4 = initial_accum_data; |
| accum_data_v5 = initial_accum_data; |
| accum_data_v6 = initial_accum_data; |
| accum_data_v7 = initial_accum_data; |
| } |
| |
| const float* lhs_ptr = lhs_col_ptr; |
| const float* rhs_ptr = rhs_col_ptr + 8 * mmm; |
| for (int d = 0; d < (params.depth - 1); ++d) { |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float* rhs_data = rhs_ptr; |
| lhs_ptr += 16; |
| rhs_ptr += 16; |
| |
| // GCC and clang can fuse set1+FMA into an FMA with EVEX broadcast: |
| // https://gcc.godbolt.org/z/xbfqWYfn1. Clang is more likely to do |
| // so if given an rvalue. |
| accum_data_v0 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[0]), |
| accum_data_v0); |
| accum_data_v1 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[1]), |
| accum_data_v1); |
| accum_data_v2 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[2]), |
| accum_data_v2); |
| accum_data_v3 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[3]), |
| accum_data_v3); |
| accum_data_v4 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[4]), |
| accum_data_v4); |
| accum_data_v5 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[5]), |
| accum_data_v5); |
| accum_data_v6 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[6]), |
| accum_data_v6); |
| accum_data_v7 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[7]), |
| accum_data_v7); |
| } |
| { // nested extra blocks lead to measurable speed gains |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float* rhs_data = rhs_ptr; |
| accum_data_v0 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[0]), |
| accum_data_v0); |
| accum_data_v1 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[1]), |
| accum_data_v1); |
| accum_data_v2 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[2]), |
| accum_data_v2); |
| accum_data_v3 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[3]), |
| accum_data_v3); |
| accum_data_v4 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[4]), |
| accum_data_v4); |
| accum_data_v5 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[5]), |
| accum_data_v5); |
| accum_data_v6 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[6]), |
| accum_data_v6); |
| accum_data_v7 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[7]), |
| accum_data_v7); |
| { |
| float* block_ptr = dst_ptr + (mmm * 8 + 0) * dst_stride; |
| accum_data_v0 = _mm512_min_ps(accum_data_v0, clamp_max_v); |
| accum_data_v0 = _mm512_max_ps(accum_data_v0, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 0 * dst_stride, accum_data_v0); |
| accum_data_v1 = _mm512_min_ps(accum_data_v1, clamp_max_v); |
| accum_data_v1 = _mm512_max_ps(accum_data_v1, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 1 * dst_stride, accum_data_v1); |
| accum_data_v2 = _mm512_min_ps(accum_data_v2, clamp_max_v); |
| accum_data_v2 = _mm512_max_ps(accum_data_v2, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 2 * dst_stride, accum_data_v2); |
| accum_data_v3 = _mm512_min_ps(accum_data_v3, clamp_max_v); |
| accum_data_v3 = _mm512_max_ps(accum_data_v3, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 3 * dst_stride, accum_data_v3); |
| accum_data_v4 = _mm512_min_ps(accum_data_v4, clamp_max_v); |
| accum_data_v4 = _mm512_max_ps(accum_data_v4, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 4 * dst_stride, accum_data_v4); |
| accum_data_v5 = _mm512_min_ps(accum_data_v5, clamp_max_v); |
| accum_data_v5 = _mm512_max_ps(accum_data_v5, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 5 * dst_stride, accum_data_v5); |
| accum_data_v6 = _mm512_min_ps(accum_data_v6, clamp_max_v); |
| accum_data_v6 = _mm512_max_ps(accum_data_v6, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 6 * dst_stride, accum_data_v6); |
| accum_data_v7 = _mm512_min_ps(accum_data_v7, clamp_max_v); |
| accum_data_v7 = _mm512_max_ps(accum_data_v7, clamp_min_v); |
| _mm512_storeu_ps(block_ptr + 7 * dst_stride, accum_data_v7); |
| } |
| } |
| } |
| } // End row-block loop. |
| |
| // The unrolling within this conditional may be somewhat pointless. It |
| // depends on the kinds of models. |
| if (row < end_row) { |
| const int residual_rows = end_row - row; |
| |
| const float* lhs_col_ptr = adj_lhs_col_ptr + row * lhs_stride; |
| float* dst_ptr = dst_col_ptr + row; |
| |
| const __mmask16 row_mask = |
| (static_cast<std::uint32_t>(1) << residual_rows) - 1; |
| |
| // Process block in two halves, split by columns. |
| for (int mmm = 0; mmm < 2; ++mmm) { |
| __m512 accum_data_v0; |
| __m512 accum_data_v1; |
| __m512 accum_data_v2; |
| __m512 accum_data_v3; |
| __m512 accum_data_v4; |
| __m512 accum_data_v5; |
| __m512 accum_data_v6; |
| __m512 accum_data_v7; |
| |
| // Initialize with bias. |
| if (channel_dimension_is_col) { |
| const float* bias_elem_ptr = |
| bias_ptr + (col + 8 * mmm) * bias_ptr_block_increment; |
| accum_data_v0 = _mm512_set1_ps(bias_elem_ptr[0]); |
| accum_data_v1 = _mm512_set1_ps(bias_elem_ptr[1]); |
| accum_data_v2 = _mm512_set1_ps(bias_elem_ptr[2]); |
| accum_data_v3 = _mm512_set1_ps(bias_elem_ptr[3]); |
| accum_data_v4 = _mm512_set1_ps(bias_elem_ptr[4]); |
| accum_data_v5 = _mm512_set1_ps(bias_elem_ptr[5]); |
| accum_data_v6 = _mm512_set1_ps(bias_elem_ptr[6]); |
| accum_data_v7 = _mm512_set1_ps(bias_elem_ptr[7]); |
| } else { |
| const __m512 initial_accum_data = |
| _mm512_loadu_ps(bias_ptr + row * bias_ptr_block_increment); |
| |
| accum_data_v0 = initial_accum_data; |
| accum_data_v1 = initial_accum_data; |
| accum_data_v2 = initial_accum_data; |
| accum_data_v3 = initial_accum_data; |
| accum_data_v4 = initial_accum_data; |
| accum_data_v5 = initial_accum_data; |
| accum_data_v6 = initial_accum_data; |
| accum_data_v7 = initial_accum_data; |
| } |
| |
| const float* lhs_ptr = lhs_col_ptr; |
| const float* rhs_ptr = rhs_col_ptr + 8 * mmm; |
| for (int d = 0; d < (params.depth - 1); ++d) { |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float* rhs_data = rhs_ptr; |
| lhs_ptr += 16; |
| rhs_ptr += 16; |
| // GCC and clang can fuse set1+FMA into an FMA with EVEX broadcast: |
| // https://gcc.godbolt.org/z/xbfqWYfn1. Clang is more likely to do |
| // so if given an rvalue. |
| accum_data_v0 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[0]), |
| accum_data_v0); |
| accum_data_v1 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[1]), |
| accum_data_v1); |
| accum_data_v2 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[2]), |
| accum_data_v2); |
| accum_data_v3 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[3]), |
| accum_data_v3); |
| accum_data_v4 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[4]), |
| accum_data_v4); |
| accum_data_v5 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[5]), |
| accum_data_v5); |
| accum_data_v6 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[6]), |
| accum_data_v6); |
| accum_data_v7 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[7]), |
| accum_data_v7); |
| } |
| { |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float* rhs_data = rhs_ptr; |
| accum_data_v0 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[0]), |
| accum_data_v0); |
| accum_data_v1 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[1]), |
| accum_data_v1); |
| accum_data_v2 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[2]), |
| accum_data_v2); |
| accum_data_v3 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[3]), |
| accum_data_v3); |
| accum_data_v4 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[4]), |
| accum_data_v4); |
| accum_data_v5 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[5]), |
| accum_data_v5); |
| accum_data_v6 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[6]), |
| accum_data_v6); |
| accum_data_v7 = _mm512_fmadd_ps(lhs_data, _mm512_set1_ps(rhs_data[7]), |
| accum_data_v7); |
| { |
| float* block_ptr = dst_ptr + (mmm * 8 + 0) * dst_stride; |
| accum_data_v0 = _mm512_min_ps(accum_data_v0, clamp_max_v); |
| accum_data_v0 = _mm512_max_ps(accum_data_v0, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 0 * dst_stride, row_mask, |
| accum_data_v0); |
| accum_data_v1 = _mm512_min_ps(accum_data_v1, clamp_max_v); |
| accum_data_v1 = _mm512_max_ps(accum_data_v1, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 1 * dst_stride, row_mask, |
| accum_data_v1); |
| accum_data_v2 = _mm512_min_ps(accum_data_v2, clamp_max_v); |
| accum_data_v2 = _mm512_max_ps(accum_data_v2, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 2 * dst_stride, row_mask, |
| accum_data_v2); |
| accum_data_v3 = _mm512_min_ps(accum_data_v3, clamp_max_v); |
| accum_data_v3 = _mm512_max_ps(accum_data_v3, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 3 * dst_stride, row_mask, |
| accum_data_v3); |
| accum_data_v4 = _mm512_min_ps(accum_data_v4, clamp_max_v); |
| accum_data_v4 = _mm512_max_ps(accum_data_v4, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 4 * dst_stride, row_mask, |
| accum_data_v4); |
| accum_data_v5 = _mm512_min_ps(accum_data_v5, clamp_max_v); |
| accum_data_v5 = _mm512_max_ps(accum_data_v5, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 5 * dst_stride, row_mask, |
| accum_data_v5); |
| accum_data_v6 = _mm512_min_ps(accum_data_v6, clamp_max_v); |
| accum_data_v6 = _mm512_max_ps(accum_data_v6, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 6 * dst_stride, row_mask, |
| accum_data_v6); |
| accum_data_v7 = _mm512_min_ps(accum_data_v7, clamp_max_v); |
| accum_data_v7 = _mm512_max_ps(accum_data_v7, clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr + 7 * dst_stride, row_mask, |
| accum_data_v7); |
| } |
| } |
| } // Inner half-block loop. |
| } // Residual rows, main col-block loop. |
| } // End col-block loop. |
| |
| if (col < end_col) { |
| RUY_DCHECK_GE(end_col - col, 0); |
| RUY_DCHECK_LT(end_col - col, 16); |
| |
| __m512 accum_data_v[8]; |
| |
| const float* rhs_col_ptr = adj_rhs_col_ptr + col * rhs_stride; |
| float* dst_col_ptr = adj_dst_col_ptr + col * dst_stride; |
| |
| for (int row = params.start_row; row < end_row; row += 16) { |
| const int residual_rows = std::min(end_row - row, 16); |
| |
| const float* lhs_col_ptr = adj_lhs_col_ptr + row * lhs_stride; |
| float* dst_ptr = dst_col_ptr + row; |
| |
| const __mmask16 row_mask = |
| (static_cast<std::uint32_t>(1) << residual_rows) - 1; |
| |
| // Process block in two halves, split by columns. |
| for (int mmm = 0; mmm < 2; ++mmm) { |
| // Initialize with bias. |
| if (channel_dimension_is_col) { |
| const float* bias_elem_ptr = |
| bias_ptr + (col + 8 * mmm) * bias_ptr_block_increment; |
| for (int j = 0; j < 8; ++j) { |
| accum_data_v[j] = _mm512_set1_ps(bias_elem_ptr[j]); |
| } |
| } else { |
| const __m512 initial_accum_data = |
| _mm512_loadu_ps(bias_ptr + row * bias_ptr_block_increment); |
| for (int j = 0; j < 8; ++j) { |
| accum_data_v[j] = initial_accum_data; |
| } |
| } |
| |
| const float* lhs_ptr = lhs_col_ptr; |
| const float* rhs_ptr = rhs_col_ptr + 8 * mmm; |
| for (int d = 0; d < params.depth; ++d) { |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float* rhs_data = rhs_ptr; |
| |
| for (int j = 0; j < 8; ++j) { |
| const __m512 dup_rhs_element_j = _mm512_set1_ps(rhs_data[j]); |
| accum_data_v[j] = |
| _mm512_fmadd_ps(lhs_data, dup_rhs_element_j, accum_data_v[j]); |
| } |
| lhs_ptr += 16; |
| rhs_ptr += 16; |
| } |
| |
| const int residual_cols = std::min(end_col - col - 8 * mmm, 8); |
| |
| if (residual_rows == 16) { |
| if (residual_cols == 8) { |
| for (int j = 0; j < 8; ++j) { |
| float* block_ptr = dst_ptr + (mmm * 8 + j) * dst_stride; |
| accum_data_v[j] = _mm512_min_ps(accum_data_v[j], clamp_max_v); |
| accum_data_v[j] = _mm512_max_ps(accum_data_v[j], clamp_min_v); |
| _mm512_storeu_ps(block_ptr, accum_data_v[j]); |
| } |
| } else { |
| for (int j = 0; j < residual_cols; ++j) { |
| float* block_ptr = dst_ptr + (mmm * 8 + j) * dst_stride; |
| accum_data_v[j] = _mm512_min_ps(accum_data_v[j], clamp_max_v); |
| accum_data_v[j] = _mm512_max_ps(accum_data_v[j], clamp_min_v); |
| _mm512_storeu_ps(block_ptr, accum_data_v[j]); |
| } |
| } |
| } else { |
| for (int j = 0; j < residual_cols; ++j) { |
| float* block_ptr = dst_ptr + (mmm * 8 + j) * dst_stride; |
| accum_data_v[j] = _mm512_min_ps(accum_data_v[j], clamp_max_v); |
| accum_data_v[j] = _mm512_max_ps(accum_data_v[j], clamp_min_v); |
| _mm512_mask_storeu_ps(block_ptr, row_mask, accum_data_v[j]); |
| } |
| } |
| } // Inner half-block loop. |
| } // End row-block loop. |
| } // Residual cols. |
| } |
| |
| void KernelFloatAvx512SingleCol(const KernelParamsFloat<16, 16>& params) { |
| profiler::ScopeLabel label("Kernel kAvx512 float GEMV"); |
| |
| RUY_DCHECK_EQ(params.dst_cols, 1); |
| RUY_DCHECK_EQ(params.last_col, 0); |
| RUY_DCHECK_EQ(params.start_col, 0); |
| |
| // As parameters are defined, we need to scale by sizeof(float). |
| const std::int64_t lhs_stride = params.lhs_stride >> 2; |
| |
| int bias_ptr_block_increment = params.flags & RUY_ASM_FLAG_HAS_BIAS ? 1 : 0; |
| const int end_row = std::min(params.dst_rows, params.last_row + 16); |
| |
| float* adj_dst_col_ptr = params.dst_base_ptr - params.start_row; |
| const float* adj_lhs_col_ptr = |
| params.lhs_base_ptr - params.start_row * lhs_stride; |
| const float* bias_col_ptr = params.bias; |
| |
| const __m512 clamp_max_v = _mm512_set1_ps(params.clamp_max); |
| const __m512 clamp_min_v = _mm512_set1_ps(params.clamp_min); |
| |
| __m512 accum_data_v; |
| |
| const float* rhs_col_ptr = params.rhs_base_ptr; |
| float* dst_col_ptr = adj_dst_col_ptr; |
| |
| int row = params.start_row; |
| for (; row <= end_row - 16; row += 16) { |
| const float* lhs_col_ptr = adj_lhs_col_ptr + row * lhs_stride; |
| float* dst_ptr = dst_col_ptr + row; |
| const float* bias_ptr = bias_col_ptr + row * bias_ptr_block_increment; |
| |
| // Initialize with bias. |
| accum_data_v = _mm512_loadu_ps(bias_ptr); |
| |
| const float* lhs_ptr = lhs_col_ptr; |
| const float* rhs_ptr = rhs_col_ptr; |
| for (int d = 0; d < params.depth; ++d) { |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float rhs_data = *rhs_ptr; |
| |
| const __m512 dup_rhs_element_j = _mm512_set1_ps(rhs_data); |
| accum_data_v = _mm512_fmadd_ps(lhs_data, dup_rhs_element_j, accum_data_v); |
| lhs_ptr += 16; |
| rhs_ptr += 16; |
| } |
| |
| accum_data_v = _mm512_min_ps(accum_data_v, clamp_max_v); |
| accum_data_v = _mm512_max_ps(accum_data_v, clamp_min_v); |
| _mm512_storeu_ps(dst_ptr, accum_data_v); |
| } // End row-block loop. |
| |
| if (row < end_row) { |
| const int residual_rows = end_row - row; |
| RUY_CHECK_GE(residual_rows, 1); |
| RUY_CHECK_LT(residual_rows, 16); |
| |
| const float* lhs_col_ptr = adj_lhs_col_ptr + row * lhs_stride; |
| float* dst_ptr = dst_col_ptr + row; |
| const float* bias_ptr = bias_col_ptr + row * bias_ptr_block_increment; |
| |
| // Initialize with bias. |
| const __mmask16 row_mask = |
| (static_cast<std::uint32_t>(1) << residual_rows) - 1; |
| accum_data_v = _mm512_loadu_ps(bias_ptr); |
| |
| const float* lhs_ptr = lhs_col_ptr; |
| const float* rhs_ptr = rhs_col_ptr; |
| for (int d = 0; d < params.depth; ++d) { |
| const __m512 lhs_data = _mm512_loadu_ps(lhs_ptr); |
| const float rhs_data = *rhs_ptr; |
| |
| const __m512 dup_rhs_element_j = _mm512_set1_ps(rhs_data); |
| accum_data_v = _mm512_fmadd_ps(lhs_data, dup_rhs_element_j, accum_data_v); |
| lhs_ptr += 16; |
| rhs_ptr += 16; |
| } |
| |
| accum_data_v = _mm512_min_ps(accum_data_v, clamp_max_v); |
| accum_data_v = _mm512_max_ps(accum_data_v, clamp_min_v); |
| _mm512_mask_storeu_ps(dst_ptr, row_mask, accum_data_v); |
| } // End handling of residual rows. |
| } |
| |
| #endif // RUY_PLATFORM_AVX512 && RUY_OPT(ASM) |
| |
| } // namespace ruy |
