| """Generates the specialized gemv functions.""" |
| |
| import mul_1x8_Mx8_neon |
| import mul_Nx8_Mx8_neon |
| import qnt_Nx8_neon |
| import zip_Nx8_neon |
| |
| _QUANTIZED_8BIT = 'quantized_8bit' |
| _FULL_32BIT = 'full_32bit' |
| _FULL_FLOAT = 'full_float' |
| |
| |
| class Error(Exception): |
| """Module level error.""" |
| |
| |
| class ConfigurationError(Error): |
| """Runtime configuration error.""" |
| |
| |
| def GenerateCommonTempsCountersAndConsts(emitter): |
| """Generates common gemv boilerplate variables.""" |
| emitter.EmitDeclare('const std::int32_t', 'col_chunks', 'n / 8') |
| emitter.EmitDeclare('const std::int32_t', 'padded_k', '((k + 7) / 8) * 8') |
| emitter.EmitDeclare('const std::int32_t', 'chunk_size', 'k * 4') |
| emitter.EmitDeclare('const std::int32_t', 'zipped_chunk_size', |
| '(padded_k + 16) * 4') |
| emitter.EmitDeclare('const std::uint8_t*', 'rhs_chunk', 'rhs') |
| emitter.EmitDeclare('std::uint8_t*', 'zipped_lhs', 'scratch') |
| emitter.EmitDeclare('std::int32_t*', 'zipped_lhs_offsets', |
| 'reinterpret_cast<std::int32_t*>(zipped_lhs + padded_k)') |
| emitter.EmitDeclare('std::uint8_t*', 'zipped_rhs_1', |
| 'scratch + padded_k + 16') |
| emitter.EmitDeclare('std::uint8_t*', 'zipped_rhs_2', |
| 'zipped_rhs_1 + zipped_chunk_size') |
| emitter.EmitNewline() |
| |
| |
| def GenerateQuantized8BitTempsCountersAndConsts(emitter): |
| """Generates all the boilerplate variables for the q8 gemm function.""" |
| GenerateCommonTempsCountersAndConsts(emitter) |
| emitter.EmitDeclare('const std::int32_t', 'const_offset', |
| 'lhs_offset * rhs_offset * k + result_offset') |
| emitter.EmitDeclare('const std::int32_t', 'rounding_offset', |
| '(1 << (shift - 1))') |
| emitter.EmitDeclare('std::int32_t*', 'temp_result', |
| 'reinterpret_cast<std::int32_t*>(' |
| 'zipped_rhs_2 + zipped_chunk_size)') |
| emitter.EmitDeclare('std::int32_t*', 'mul_result_chunk', 'temp_result') |
| emitter.EmitNewline() |
| |
| |
| def GenerateFullTempsCountersAndConsts(emitter, result_type): |
| """Generates all the boilerplate variables for the int32 and float gemms.""" |
| GenerateCommonTempsCountersAndConsts(emitter) |
| emitter.EmitDeclare('const std::int32_t', 'const_offset', |
| 'lhs_offset * rhs_offset * k') |
| emitter.EmitDeclare(result_type, 'mul_result_chunk', 'result') |
| emitter.EmitNewline() |
| |
| |
| def GenerateZipVector(emitter, aligned, leftovers): |
| emitter.EmitCall( |
| zip_Nx8_neon.BuildName(1, leftovers, aligned), |
| ['lhs', 'k', 'k', 'zipped_lhs', 'rhs_offset', 0]) |
| |
| |
| def GetMul2Params(result_type): |
| params = ['zipped_lhs', 'zipped_rhs_1', 'zipped_rhs_2', 'padded_k', |
| 'mul_result_chunk'] |
| if result_type is 'float': |
| params.append('result_scale') |
| return params |
| |
| |
| def GetMulParams(result_type): |
| params = ['zipped_lhs', 'zipped_rhs_1', 'padded_k', 'mul_result_chunk', 0] |
| if result_type is 'float': |
| params.append('result_scale') |
| return params |
| |
| |
| def GenerateMulCols(emitter, result_type, lhs_add, rhs_add, aligned, cols, |
| leftovers): |
| """Emits code responsible for multiplication of one horizontal lhs strip.""" |
| emitter.EmitOpenBracket('for (int i = 0; i < col_chunks; ++i)') |
| emitter.EmitCall( |
| zip_Nx8_neon.BuildName(4, leftovers, aligned), |
| ['rhs_chunk', 'k', 'k', 'zipped_rhs_1', 'lhs_offset', 'const_offset']) |
| emitter.EmitAssignIncrement('rhs_chunk', 'chunk_size') |
| |
| emitter.EmitCall( |
| zip_Nx8_neon.BuildName(4, leftovers, aligned), |
| ['rhs_chunk', 'k', 'k', 'zipped_rhs_2', 'lhs_offset', 'const_offset']) |
| emitter.EmitAssignIncrement('rhs_chunk', 'chunk_size') |
| |
| emitter.EmitCall( |
| mul_1x8_Mx8_neon.BuildName(result_type, lhs_add, rhs_add, 8), |
| GetMul2Params(result_type)) |
| |
| emitter.EmitAssignIncrement('mul_result_chunk', 8) |
| emitter.EmitCloseBracket() |
| |
| if cols > 4: |
| emitter.EmitCall( |
| zip_Nx8_neon.BuildName(4, leftovers, aligned), |
| ['rhs_chunk', 'k', 'k', 'zipped_rhs_1', 'lhs_offset', 'const_offset']) |
| emitter.EmitAssignIncrement('rhs_chunk', 'chunk_size') |
| |
| emitter.EmitCall( |
| zip_Nx8_neon.BuildName(cols - 4, leftovers, aligned), |
| ['rhs_chunk', 'k', 'k', 'zipped_rhs_2', 'lhs_offset', 'const_offset']) |
| |
| emitter.EmitCall( |
| mul_1x8_Mx8_neon.BuildName(result_type, lhs_add, rhs_add, cols), |
| GetMul2Params(result_type)) |
| elif cols > 0: |
| emitter.EmitCall( |
| zip_Nx8_neon.BuildName(cols, leftovers, aligned), |
| ['rhs_chunk', 'k', 'k', 'zipped_rhs_1', 'lhs_offset', 'const_offset']) |
| |
| emitter.EmitCall( |
| mul_Nx8_Mx8_neon.BuildName(result_type, lhs_add, rhs_add, 1, cols), |
| GetMulParams(result_type)) |
| |
| |
| def GenerateQuantized8BitMul(emitter, aligned, cols, leftovers): |
| """Emits code for all lhs strips & leftover rows. Quantize after mul code.""" |
| GenerateMulCols(emitter, 'int32', False, True, aligned, cols, leftovers) |
| emitter.EmitCall( |
| qnt_Nx8_neon.BuildName(1, cols, aligned), |
| ['temp_result', 'n', 0, 'zipped_lhs_offsets', 'result', 0, |
| 'multiplicative_offset', 'rounding_offset', '-shift']) |
| |
| |
| def GenerateFullMul(emitter, result_type, aligned, cols, leftovers): |
| GenerateMulCols(emitter, result_type, True, True, aligned, cols, leftovers) |
| |
| |
| def BuildName(output_type, aligned, cols, leftover): |
| name = BuildMainGemvName(output_type) + '_%d_%d' % (cols, leftover) |
| if aligned: |
| name += '_aligned' |
| return name |
| |
| |
| def GetCommonGemvParameters(): |
| return [['std::uint8_t*', 'scratch'], ['const std::uint8_t*', 'lhs'], |
| ['const std::uint8_t*', 'rhs'], ['std::int32_t', 'n'], |
| ['std::int32_t', 'k'], ['std::int32_t', 'lhs_offset'], |
| ['std::int32_t', 'rhs_offset']] |
| |
| |
| def GetGemvParameters(output_type): |
| """Prepares a (type, parameter) array for the gemm functions.""" |
| params = GetCommonGemvParameters() |
| if output_type is _QUANTIZED_8BIT: |
| params += [['std::int32_t', 'result_offset'], |
| ['std::int32_t', 'multiplicative_offset'], |
| ['std::int32_t', 'shift'], ['std::uint8_t*', 'result']] |
| elif output_type is _FULL_32BIT: |
| params += [['std::int32_t*', 'result']] |
| elif output_type is _FULL_FLOAT: |
| params += [['float', 'result_scale'], ['float*', 'result']] |
| else: |
| raise ConfigurationError('Unsupported output type: %s' % output_type) |
| return params |
| |
| |
| def GenerateGemv(emitter, output_type, aligned, cols, leftovers): |
| """Build one gemm function for given col, and depth leftovers.""" |
| emitter.EmitFunctionBeginA( |
| BuildName(output_type, aligned, cols, leftovers), |
| GetGemvParameters(output_type), 'void') |
| |
| emitter.EmitAssert('n %% 8 == %d' % cols) |
| emitter.EmitAssert('k %% 8 == %d' % leftovers) |
| |
| if output_type is _QUANTIZED_8BIT: |
| GenerateQuantized8BitTempsCountersAndConsts(emitter) |
| GenerateZipVector(emitter, aligned, leftovers) |
| GenerateQuantized8BitMul(emitter, aligned, cols, leftovers) |
| elif output_type is _FULL_32BIT: |
| GenerateFullTempsCountersAndConsts(emitter, 'std::int32_t*') |
| GenerateZipVector(emitter, aligned, leftovers) |
| GenerateFullMul(emitter, 'int32', aligned, cols, leftovers) |
| elif output_type is _FULL_FLOAT: |
| GenerateFullTempsCountersAndConsts(emitter, 'float*') |
| GenerateZipVector(emitter, aligned, leftovers) |
| GenerateFullMul(emitter, 'float', aligned, cols, leftovers) |
| else: |
| raise ConfigurationError('Unknown output type: %s' % output_type) |
| |
| emitter.EmitFunctionEnd() |
| |
| |
| def GenerateGemvCall(emitter, output_type, aligned, m_mod, leftovers): |
| emitter.EmitCall( |
| emitter.Scope('internal', |
| BuildName(output_type, aligned, m_mod, leftovers)), |
| [p for (unused_t, p) in GetGemvParameters(output_type)]) |
| |
| |
| def GenerateGemvSwitch2(emitter, output_type, aligned, n_mod): |
| """Second level of main switch, choose optimized version on depth leftover.""" |
| emitter.EmitSwitch('k % 8') |
| |
| for leftovers in range(0, 8): |
| emitter.EmitCase(leftovers) |
| emitter.PushIndent() |
| GenerateGemvCall(emitter, output_type, aligned, n_mod, leftovers) |
| emitter.EmitBreak() |
| emitter.PopIndent() |
| |
| emitter.EmitSwitchEnd() |
| |
| |
| def GenerateGemvSwitch1(emitter, output_type, aligned): |
| """First level of main switch, choose optimized version on cols leftover.""" |
| emitter.EmitSwitch('n % 8') |
| |
| for n_mod in range(0, 8): |
| emitter.EmitCase(n_mod) |
| emitter.PushIndent() |
| GenerateGemvSwitch2(emitter, output_type, aligned, n_mod) |
| emitter.EmitBreak() |
| emitter.PopIndent() |
| |
| emitter.EmitSwitchEnd() |
| |
| |
| def BuildMainGemvName(output_type): |
| if output_type is _QUANTIZED_8BIT: |
| return 'gemv_q8' |
| elif output_type is _FULL_32BIT: |
| return 'gemv_i32' |
| elif output_type is _FULL_FLOAT: |
| return 'gemv_f' |
| else: |
| raise ConfigurationError('Unsupported output type: %s' % output_type) |
| |
| |
| def GenerateMainGemvFunction(emitter, output_type): |
| """Emit high level gemv function that switches between optimized versions.""" |
| emitter.EmitFunctionBeginA( |
| BuildMainGemvName(output_type), GetGemvParameters(output_type), 'void') |
| |
| emitter.EmitDeclare('const bool', 'lhs_aligned', |
| '((reinterpret_cast<std::uintptr_t>(lhs) % 8) == 0)') |
| emitter.EmitDeclare('const bool', 'rhs_aligned', |
| '((reinterpret_cast<std::uintptr_t>(rhs) % 8) == 0)') |
| emitter.EmitDeclare('const bool', 'k_aligned', '((k % 8) == 0)') |
| |
| if output_type is _QUANTIZED_8BIT: |
| emitter.EmitDeclare('const bool', 'result_aligned', |
| '((reinterpret_cast<std::uintptr_t>(result) % 8) == 0)') |
| emitter.EmitDeclare('const bool', 'aligned', |
| 'lhs_aligned && rhs_aligned && result_aligned ' |
| '&& k_aligned') |
| else: |
| emitter.EmitDeclare('const bool', 'aligned', |
| 'lhs_aligned && rhs_aligned && k_aligned') |
| |
| emitter.EmitIf('aligned') |
| GenerateGemvSwitch1(emitter, output_type, True) |
| emitter.EmitElse() |
| GenerateGemvSwitch1(emitter, output_type, False) |
| emitter.EmitEndif() |
| emitter.EmitFunctionEnd() |
| |
| |
| def GenerateInternalFunctions(emitter): |
| """Generate all the functions hidden in the internal namespace.""" |
| for output_type in [_QUANTIZED_8BIT, _FULL_32BIT, _FULL_FLOAT]: |
| for aligned in [True, False]: |
| for cols in range(0, 8): |
| for leftover in range(0, 8): |
| GenerateGemv(emitter, output_type, aligned, cols, leftover) |
| emitter.EmitNewline() |
| |
| |
| def GeneratePublicFunctions(emitter): |
| for output_type in [_QUANTIZED_8BIT, _FULL_32BIT, _FULL_FLOAT]: |
| GenerateMainGemvFunction(emitter, output_type) |
| emitter.EmitNewline() |
