| 139 | template <typename LhsScalar, typename RhsScalar, typename AccumScalar, |
| 140 | typename DstScalar, QuantizationFlavor quantization_flavor> |
| 141 | bool CustomGemv( |
| 142 | const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data, |
| 143 | const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data, |
| 144 | const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data, |
| 145 | const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params, |
| 146 | CpuBackendContext* context) { |
| 147 | gemmlowp::ScopedProfilingLabel label("cpu_backend_gemm::Gemm: CustomGemv"); |
| 148 | using Impl = CustomGemvImpl<LhsScalar, RhsScalar, AccumScalar, DstScalar, |
| 149 | quantization_flavor>; |
| 150 | if (lhs_params.rows < Impl::kKernelRows) { |
| 151 | return false; |
| 152 | } |
| 153 | if (!Impl::IsSupportedGivenSufficientlyManyRows(lhs_params, rhs_params, |
| 154 | dst_params, params)) { |
| 155 | return false; |
| 156 | } |
| 157 | TFLITE_DCHECK_GE(lhs_params.rows, Impl::kKernelRows); |
| 158 | int thread_count = LegacyHowManyThreads<Impl::kKernelRows>( |
| 159 | context->max_num_threads(), dst_params.rows, dst_params.cols, |
| 160 | lhs_params.cols); |
| 161 | if (thread_count == 1) { |
| 162 | Impl::Run(lhs_params, lhs_data, rhs_params, rhs_data, dst_params, dst_data, |
| 163 | params, 0, lhs_params.rows); |
| 164 | } else { |
| 165 | using Task = CustomGemvTask<LhsScalar, RhsScalar, AccumScalar, DstScalar, |
| 166 | quantization_flavor>; |
| 167 | std::vector<Task> tasks; |
| 168 | tasks.reserve(thread_count); |
| 169 | const int kRowsPerThread = |
| 170 | RoundUp<Impl::kKernelRows>(CeilQuotient(dst_params.rows, thread_count)); |
| 171 | int row_start = 0; |
| 172 | for (int i = 0; i < thread_count; i++) { |
| 173 | int row_end = std::min(dst_params.rows, row_start + kRowsPerThread); |
| 174 | tasks.emplace_back(lhs_params, lhs_data, rhs_params, rhs_data, dst_params, |
| 175 | dst_data, params, row_start, row_end); |
| 176 | row_start = row_end; |
| 177 | } |
| 178 | cpu_backend_threadpool::Execute(tasks.size(), tasks.data(), context); |
| 179 | } |
| 180 | return true; |
| 181 | } |
| 182 | |
| 183 | // USE_NEON still allows for x86 where we may be using the arm_neon_sse.h |
| 184 | // wrapper implementing NEON intrinsics on top of SSE4 intrinsics. |
no test coverage detected