| 176 | |
| 177 | template <typename LaunchFunc> |
| 178 | se::blas::AlgorithmConfig AutotuneMatmul( |
| 179 | const std::vector<std::unique_ptr<se::blas::IBlasLtMatmulAlgorithm>>& |
| 180 | algorithms, |
| 181 | const BatchMatmulParameters& matmul_params, OpKernelContext* context, |
| 182 | const LaunchFunc& launch_func) { |
| 183 | // Note that algorithm_config.algorithm() here is used to refer |
| 184 | // to the index within the algorithms vector, not the algorithm |
| 185 | // itself. |
| 186 | se::blas::AlgorithmConfig algorithm_config(se::blas::kNoAlgorithm); |
| 187 | if (!BlasPlansAutotuneCacheSingleton::GetInstance()->Find( |
| 188 | matmul_params, &algorithm_config)) { |
| 189 | VLOG(4) << "Autotuning BlasLtMatmul over " << algorithms.size() |
| 190 | << " algorithms."; |
| 191 | se::blas::ProfileResult best_result; |
| 192 | se::blas::ProfileResult profile_result; |
| 193 | |
| 194 | for (size_t i = 0; i != algorithms.size(); ++i) { |
| 195 | const auto& profile_algorithm = algorithms[i]; |
| 196 | |
| 197 | // Create a new scratch allocator with every autotuning run so that |
| 198 | // scratch space is deallocated between runs. |
| 199 | BlasScratchAllocator scratch_allocator(context); |
| 200 | |
| 201 | bool cublaslt_launch_ok = launch_func( |
| 202 | &scratch_allocator, profile_algorithm.get(), &profile_result); |
| 203 | |
| 204 | VLOG(4) << " Autotune algorithm " << i |
| 205 | << " result: " << profile_result.elapsed_time_in_ms() |
| 206 | << " ms, valid=" << profile_result.is_valid() |
| 207 | << ", workspace_size=" << profile_algorithm->workspace_size(); |
| 208 | |
| 209 | if (cublaslt_launch_ok && profile_result.is_valid() && |
| 210 | profile_result.elapsed_time_in_ms() < |
| 211 | best_result.elapsed_time_in_ms()) { |
| 212 | best_result = profile_result; |
| 213 | } |
| 214 | } |
| 215 | |
| 216 | if (best_result.is_valid()) { |
| 217 | algorithm_config.set_algorithm(best_result.algorithm()); |
| 218 | } |
| 219 | // We make sure that each matmul parameter set only gets one pass of |
| 220 | // autotune. If no algorithms works, we add kNoAlgorithm to the autotune |
| 221 | // map. |
| 222 | BlasPlansAutotuneCacheSingleton::GetInstance()->Insert(matmul_params, |
| 223 | algorithm_config); |
| 224 | } |
| 225 | return algorithm_config; |
| 226 | } |
| 227 | |
| 228 | } // namespace |
| 229 |
no test coverage detected