| 544 | |
| 545 | template <typename InpT, typename FuncT, typename... Args> |
| 546 | bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream, |
| 547 | bool pointer_mode_host, bool err_on_failure, |
| 548 | Args... args) { |
| 549 | absl::MutexLock lock(&mu_); |
| 550 | |
| 551 | CHECK(blas_ != nullptr); |
| 552 | if (!SetStream(stream)) { |
| 553 | return false; |
| 554 | } |
| 555 | |
| 556 | #if CUDA_VERSION >= 9000 |
| 557 | ScopedCublasMathMode math_mode{blas_}; |
| 558 | #if CUBLAS_VER_MAJOR < 11 |
| 559 | // Pre cublas 11, no TF32 is available, so we can |
| 560 | // enable tensor cores in all cases. |
| 561 | cublasMath_t math_type = CUBLAS_TENSOR_OP_MATH; |
| 562 | #else |
| 563 | cublasMath_t math_type; |
| 564 | if (tensorflow::tensor_float_32_execution_enabled() && |
| 565 | (CUDADataType<InpT>::type == CUDA_R_32F || |
| 566 | CUDADataType<InpT>::type == CUDA_C_32F)) { |
| 567 | math_type = CUBLAS_TF32_TENSOR_OP_MATH; |
| 568 | } else { |
| 569 | // In cublas 11+ DEFAULT_MATH enables tensorcores in FP16 cases. |
| 570 | math_type = CUBLAS_DEFAULT_MATH; |
| 571 | } |
| 572 | #endif |
| 573 | if (!math_mode.Init(math_type)) { |
| 574 | return false; |
| 575 | } |
| 576 | #endif |
| 577 | |
| 578 | gpu::ScopedActivateExecutorContext sac{parent_}; |
| 579 | ScopedCublasPointerMode pointer_mode{blas_}; |
| 580 | if (!pointer_mode.Init(pointer_mode_host ? CUBLAS_POINTER_MODE_HOST |
| 581 | : CUBLAS_POINTER_MODE_DEVICE)) { |
| 582 | return false; |
| 583 | } |
| 584 | cublasStatus_t ret = cublas_func(blas_, args...); |
| 585 | if ((err_on_failure || VLOG_IS_ON(3)) && ret != CUBLAS_STATUS_SUCCESS) { |
| 586 | LOG(ERROR) << "failed to run cuBLAS routine: " << ToString(ret); |
| 587 | } |
| 588 | return ret == CUBLAS_STATUS_SUCCESS; |
| 589 | } |
| 590 | |
| 591 | bool CUDABlas::DoBlasAsum(Stream *stream, uint64 elem_count, |
| 592 | const DeviceMemory<float> &x, int incx, |
nothing calls this directly
no test coverage detected