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

Function UnmqrImpl

tensorflow/core/kernels/cuda_solvers.cc:521–544  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

519
520template <typename Scalar, typename BufSizeFnT, typename SolverFnT>
521static inline Status UnmqrImpl(BufSizeFnT bufsize, SolverFnT solver,
522 CudaSolver* cuda_solver,
523 OpKernelContext* context,
524 cusolverDnHandle_t cusolver_dn_handle,
525 cublasSideMode_t side, cublasOperation_t trans,
526 int m, int n, int k, const Scalar* dev_a,
527 int lda, const Scalar* dev_tau, Scalar* dev_c,
528 int ldc, int* dev_lapack_info) {
529 mutex_lock lock(handle_map_mutex);
530 /* Get amount of workspace memory required. */
531 int lwork;
532 TF_RETURN_IF_CUSOLVER_ERROR(
533 bufsize(cusolver_dn_handle, side, trans, m, n, k, CUDAComplex(dev_a), lda,
534 CUDAComplex(dev_tau), CUDAComplex(dev_c), ldc, &lwork));
535 /* Allocate device memory for workspace. */
536 auto dev_workspace =
537 cuda_solver->GetScratchSpace<Scalar>(lwork, "", /* on_host */ false);
538 /* Launch the solver kernel. */
539 TF_RETURN_IF_CUSOLVER_ERROR(solver(
540 cusolver_dn_handle, side, trans, m, n, k, CUDAComplex(dev_a), lda,
541 CUDAComplex(dev_tau), CUDAComplex(dev_c), ldc,
542 CUDAComplex(dev_workspace.mutable_data()), lwork, dev_lapack_info));
543 return Status::OK();
544}
545
546// Unfortunately the LAPACK function name differs for the real and complex case
547// (complex ones are prefixed with "UN" for "unitary"), so we instantiate each

Callers

nothing calls this directly

Calls 2

CUDAComplexFunction · 0.85
mutable_dataMethod · 0.80

Tested by

no test coverage detected