MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / HeevdImpl

Function HeevdImpl

tensorflow/core/kernels/cuda_solvers.cc:606–629  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

604
605template <typename Scalar, typename BufSizeFnT, typename SolverFnT>
606static inline Status HeevdImpl(BufSizeFnT bufsize, SolverFnT solver,
607 CudaSolver* cuda_solver,
608 OpKernelContext* context,
609 cusolverDnHandle_t cusolver_dn_handle,
610 cusolverEigMode_t jobz, cublasFillMode_t uplo,
611 int n, Scalar* dev_A, int lda,
612 typename Eigen::NumTraits<Scalar>::Real* dev_W,
613 int* dev_lapack_info) {
614 mutex_lock lock(handle_map_mutex);
615 /* Get amount of workspace memory required. */
616 int lwork;
617 TF_RETURN_IF_CUSOLVER_ERROR(bufsize(cusolver_dn_handle, jobz, uplo, n,
618 CUDAComplex(dev_A), lda,
619 CUDAComplex(dev_W), &lwork));
620 /* Allocate device memory for workspace. */
621 auto dev_workspace =
622 cuda_solver->GetScratchSpace<Scalar>(lwork, "", /* on_host */ false);
623 /* Launch the solver kernel. */
624 TF_RETURN_IF_CUSOLVER_ERROR(
625 solver(cusolver_dn_handle, jobz, uplo, n, CUDAComplex(dev_A), lda,
626 CUDAComplex(dev_W), CUDAComplex(dev_workspace.mutable_data()),
627 lwork, dev_lapack_info));
628 return Status::OK();
629}
630
631#define HEEVD_INSTANCE(Scalar, function_prefix, type_prefix) \
632 template <> \

Callers

nothing calls this directly

Calls 2

CUDAComplexFunction · 0.85
mutable_dataMethod · 0.80

Tested by

no test coverage detected