| 259 | : SmallVector<size_t>{algo_workspace_size}}; |
| 260 | } |
| 261 | bool CUBLASLTMatmulDesc::get_algorithm_heuristic( |
| 262 | const SizeArgs& args, size_t ws_limit, cublasLtMatmulAlgo_t& algo) { |
| 263 | bool result; |
| 264 | int return_algo_count; |
| 265 | size_t algo_ws_limit; |
| 266 | cublasStatus_t status; |
| 267 | cublasLtMatmulPreference_t algo_pref; |
| 268 | cublasLtMatmulHeuristicResult_t algo_result{}; |
| 269 | auto&& handle = concrete_handle(args.handle); |
| 270 | auto&& cublasLt_handle = handle->cublasLt_handle(); |
| 271 | |
| 272 | size_t temp = workspace_b + workspace_a + workspace_c; |
| 273 | algo_ws_limit = (ws_limit > temp) ? (ws_limit - temp) : 0; |
| 274 | |
| 275 | /** |
| 276 | * \Note: algo_ws_limit must be zero if cublasLtGetVersion() <= 10100 |
| 277 | */ |
| 278 | // algo_ws_limit = 0; |
| 279 | if (dt_c == CUDA_R_32I) { |
| 280 | //[FIXME]: cublasLt(Version 10020) produce wrong result when k in |
| 281 | //[64*n+1 , 64*n+32] for small matrix |
| 282 | |
| 283 | //[TODO]: check if this bug is fixed in latter cublasLt. |
| 284 | size_t k_pos = (is_batched ? 1 : 0) + (args.transposeA ? 0 : 1); |
| 285 | size_t k = args.layout_a.shape[k_pos]; |
| 286 | bool flt = (k < 65 || ((k - 1) / 32) % 2 == 1); |
| 287 | if (!flt) |
| 288 | return false; |
| 289 | } |
| 290 | result = false; |
| 291 | cublas_check(cublasLtMatmulPreferenceCreate(&algo_pref)); |
| 292 | cublas_check(cublasLtMatmulPreferenceSetAttribute( |
| 293 | algo_pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &algo_ws_limit, |
| 294 | sizeof(algo_ws_limit))); |
| 295 | #if CUDA_VERSION < 11000 |
| 296 | bool is_f32_config = args.layout_a.dtype == dtype::Float32() && |
| 297 | args.layout_b.dtype == dtype::Float32() && |
| 298 | args.layout_c.dtype == dtype::Float32(); |
| 299 | if (is_f32_config) { |
| 300 | // disable HMMA tensor op matmul when inputs and output are all f32 |
| 301 | // tensors, to avoid the potential accuracy loss |
| 302 | uint32_t math_mode = CUBLAS_DEFAULT_MATH; |
| 303 | cublas_check(cublasLtMatmulPreferenceSetAttribute( |
| 304 | algo_pref, CUBLASLT_MATMUL_PREF_MATH_MODE_MASK, &math_mode, |
| 305 | sizeof(math_mode))); |
| 306 | } |
| 307 | #endif |
| 308 | status = cublasLtMatmulAlgoGetHeuristic( |
| 309 | cublasLt_handle, matmul_desc, |
| 310 | dt_c == CUDA_R_32I ? layout_trans_b : layout_b, |
| 311 | dt_c == CUDA_R_32I ? layout_trans_a : layout_a, |
| 312 | dt_c == CUDA_R_32I ? layout_trans_c : layout_c, |
| 313 | dt_c == CUDA_R_32I ? layout_trans_c : layout_c, algo_pref, 1, &algo_result, |
| 314 | &return_algo_count); |
| 315 | if (status == CUBLAS_STATUS_SUCCESS && return_algo_count > 0 && |
| 316 | // perform cublasLtAlgoCheck() to make sure the algo is correct |
| 317 | get_workspace_bundle(args, algo_result.algo).nr_workspace() > 0) { |
| 318 | result = true; |
nothing calls this directly
no test coverage detected