| 565 | |
| 566 | template <typename Scalar, typename BufSizeFnT, typename SolverFnT> |
| 567 | static inline Status UngqrImpl(BufSizeFnT bufsize, SolverFnT solver, |
| 568 | CudaSolver* cuda_solver, |
| 569 | OpKernelContext* context, |
| 570 | cusolverDnHandle_t cusolver_dn_handle, int m, |
| 571 | int n, int k, Scalar* dev_a, int lda, |
| 572 | const Scalar* dev_tau, int* dev_lapack_info) { |
| 573 | mutex_lock lock(handle_map_mutex); |
| 574 | /* Get amount of workspace memory required. */ |
| 575 | int lwork; |
| 576 | TF_RETURN_IF_CUSOLVER_ERROR(bufsize(cusolver_dn_handle, m, n, k, |
| 577 | CUDAComplex(dev_a), lda, |
| 578 | CUDAComplex(dev_tau), &lwork)); |
| 579 | /* Allocate device memory for workspace. */ |
| 580 | auto dev_workspace = |
| 581 | cuda_solver->GetScratchSpace<Scalar>(lwork, "", /* on_host */ false); |
| 582 | /* Launch the solver kernel. */ |
| 583 | TF_RETURN_IF_CUSOLVER_ERROR( |
| 584 | solver(cusolver_dn_handle, m, n, k, CUDAComplex(dev_a), lda, |
| 585 | CUDAComplex(dev_tau), CUDAComplex(dev_workspace.mutable_data()), |
| 586 | lwork, dev_lapack_info)); |
| 587 | return Status::OK(); |
| 588 | } |
| 589 | |
| 590 | #define UNGQR_INSTANCE(Scalar, function_prefix, type_prefix) \ |
| 591 | template <> \ |
nothing calls this directly
no test coverage detected