/************************************************************************ * Derived from the BSD3-licensed * LAPACK routine (version 3.7.1) -- * Univ. of Tennessee, Univ. of California Berkeley, * Univ. of Colorado Denver and NAG Ltd.. * June 2017 * Copyright (c) 2019-2022 Advanced Micro Devices, Inc. * ***********************************************************************/ #pragma once #include "../auxiliary/rocauxiliary_lacgv.hpp" #include "../auxiliary/rocauxiliary_larfg.hpp" #include "rocblas.hpp" #include "rocsolver.h" template void rocsolver_latrd_getMemorySize(const rocblas_int n, const rocblas_int k, const rocblas_int batch_count, size_t* size_scalars, size_t* size_work, size_t* size_norms, size_t* size_workArr) { // if quick return no workspace needed if(n == 0 || k == 0 || batch_count == 0) { *size_scalars = 0; *size_work = 0; *size_norms = 0; *size_workArr = 0; return; } size_t w_temp; // size of scalars (constants) for rocblas calls *size_scalars = sizeof(T) * 3; // size of array of pointers (batched cases) if(BATCHED) *size_workArr = 2 * sizeof(T*) * batch_count; else *size_workArr = 0; // extra requirements for calling larfg rocsolver_larfg_getMemorySize(n, batch_count, size_work, size_norms); // extra requirements for calling symv/hemv rocblasCall_symv_hemv_mem(n, batch_count, &w_temp); *size_work = std::max(*size_work, w_temp); } template rocblas_status rocsolver_latrd_argCheck(rocblas_handle handle, const rocblas_fill uplo, const rocblas_int n, const rocblas_int k, const rocblas_int lda, const rocblas_int ldw, T A, S E, U tau, U W, const rocblas_int batch_count = 1) { // order is important for unit tests: // 1. invalid/non-supported values if(uplo != rocblas_fill_upper && uplo != rocblas_fill_lower) return rocblas_status_invalid_value; // 2. invalid size if(n < 0 || k < 0 || k > n || lda < n || ldw < n || batch_count < 0) return rocblas_status_invalid_size; // skip pointer check if querying memory size if(rocblas_is_device_memory_size_query(handle)) return rocblas_status_continue; // 3. invalid pointers if((n && !A) || (n && !E) || (n && !tau) || (n * k && !W)) return rocblas_status_invalid_pointer; return rocblas_status_continue; } template > rocblas_status rocsolver_latrd_template(rocblas_handle handle, const rocblas_fill uplo, const rocblas_int n, const rocblas_int k, U A, const rocblas_int shiftA, const rocblas_int lda, const rocblas_stride strideA, S* E, const rocblas_stride strideE, T* tau, const rocblas_stride strideP, T* W, const rocblas_int shiftW, const rocblas_int ldw, const rocblas_stride strideW, const rocblas_int batch_count, T* scalars, T* work, T* norms, T** workArr) { ROCSOLVER_ENTER("latrd", "uplo:", uplo, "n:", n, "k:", k, "shiftA:", shiftA, "lda:", lda, "shiftW:", shiftW, "ldw:", ldw, "bc:", batch_count); // quick return if(n == 0 || k == 0 || batch_count == 0) return rocblas_status_success; hipStream_t stream; rocblas_get_stream(handle, &stream); // everything must be executed with scalars on the device rocblas_pointer_mode old_mode; rocblas_get_pointer_mode(handle, &old_mode); rocblas_set_pointer_mode(handle, rocblas_pointer_mode_device); // configure kernels rocblas_int blocks = (batch_count - 1) / BS1 + 1; dim3 grid_b(blocks, 1); dim3 threads(BS1, 1, 1); blocks = (n - 1) / BS1 + 1; dim3 grid_n(blocks, batch_count); if(uplo == rocblas_fill_lower) { // reduce the first k columns of A // main loop running forwards (for each column) for(rocblas_int j = 0; j < k; ++j) { // update column j of A with reflector computed in step j-1 if(COMPLEX) rocsolver_lacgv_template(handle, j, W, shiftW + idx2D(j, 0, ldw), ldw, strideW, batch_count); rocblasCall_gemv(handle, rocblas_operation_none, n - j, j, cast2constType(scalars), 0, A, shiftA + idx2D(j, 0, lda), lda, strideA, W, shiftW + idx2D(j, 0, ldw), ldw, strideW, cast2constType(scalars + 2), 0, A, shiftA + idx2D(j, j, lda), 1, strideA, batch_count, workArr); if(COMPLEX) { rocsolver_lacgv_template(handle, j, W, shiftW + idx2D(j, 0, ldw), ldw, strideW, batch_count); rocsolver_lacgv_template(handle, j, A, shiftA + idx2D(j, 0, lda), lda, strideA, batch_count); } rocblasCall_gemv(handle, rocblas_operation_none, n - j, j, cast2constType(scalars), 0, W, shiftW + idx2D(j, 0, ldw), ldw, strideW, A, shiftA + idx2D(j, 0, lda), lda, strideA, cast2constType(scalars + 2), 0, A, shiftA + idx2D(j, j, lda), 1, strideA, batch_count, workArr); if(COMPLEX) rocsolver_lacgv_template(handle, j, A, shiftA + idx2D(j, 0, lda), lda, strideA, batch_count); // generate Householder reflector to work on column j rocsolver_larfg_template(handle, n - j - 1, A, shiftA + idx2D(j + 1, j, lda), A, shiftA + idx2D(min(j + 2, n - 1), j, lda), 1, strideA, (tau + j), strideP, batch_count, work, norms); // copy to E(j) the corresponding off-diagonal element of A, which is set to 1 ROCSOLVER_LAUNCH_KERNEL(set_offdiag, grid_b, threads, 0, stream, batch_count, A, shiftA + idx2D(j + 1, j, lda), strideA, (E + j), strideE); // compute/update column j of W rocblasCall_symv_hemv( handle, uplo, n - 1 - j, (scalars + 2), 0, A, shiftA + idx2D(j + 1, j + 1, lda), lda, strideA, A, shiftA + idx2D(j + 1, j, lda), 1, strideA, (scalars + 1), 0, W, shiftW + idx2D(j + 1, j, ldw), 1, strideW, batch_count, work, workArr); rocblasCall_gemv(handle, rocblas_operation_conjugate_transpose, n - j - 1, j, cast2constType(scalars + 2), 0, W, shiftW + idx2D(j + 1, 0, ldw), ldw, strideW, A, shiftA + idx2D(j + 1, j, lda), 1, strideA, cast2constType(scalars + 1), 0, W, shiftW + idx2D(0, j, ldw), 1, strideW, batch_count, workArr); rocblasCall_gemv(handle, rocblas_operation_none, n - j - 1, j, cast2constType(scalars), 0, A, shiftA + idx2D(j + 1, 0, lda), lda, strideA, W, shiftW + idx2D(0, j, ldw), 1, strideW, cast2constType(scalars + 2), 0, W, shiftW + idx2D(j + 1, j, ldw), 1, strideW, batch_count, workArr); rocblasCall_gemv(handle, rocblas_operation_conjugate_transpose, n - j - 1, j, cast2constType(scalars + 2), 0, A, shiftA + idx2D(j + 1, 0, lda), lda, strideA, A, shiftA + idx2D(j + 1, j, lda), 1, strideA, cast2constType(scalars + 1), 0, W, shiftW + idx2D(0, j, ldw), 1, strideW, batch_count, workArr); rocblasCall_gemv(handle, rocblas_operation_none, n - j - 1, j, cast2constType(scalars), 0, W, shiftW + idx2D(j + 1, 0, ldw), ldw, strideW, W, shiftW + idx2D(0, j, ldw), 1, strideW, cast2constType(scalars + 2), 0, W, shiftW + idx2D(j + 1, j, ldw), 1, strideW, batch_count, workArr); rocblasCall_scal(handle, n - j - 1, (tau + j), strideP, W, shiftW + idx2D(j + 1, j, ldw), 1, strideW, batch_count); rocblasCall_dot(handle, n - 1 - j, W, shiftW + idx2D(j + 1, j, ldw), 1, strideW, A, shiftA + idx2D(j + 1, j, lda), 1, strideA, batch_count, norms, work, workArr); // (TODO: rocblas_axpy is not yet ready to be used in rocsolver. When it becomes // available, we can use it instead of the scale_axpy kernel, if it provides // better performance.) ROCSOLVER_LAUNCH_KERNEL(scale_axpy, grid_n, threads, 0, stream, n - 1 - j, norms, tau + j, strideP, A, shiftA + idx2D(j + 1, j, lda), strideA, W, shiftW + idx2D(j + 1, j, ldw), strideW); } } else { // reduce the last k columns of A // main loop running forwards (for each column) rocblas_int jw; for(rocblas_int j = n - 1; j >= n - k; --j) { jw = j - n + k; // update column j of A with reflector computed in step j-1 if(COMPLEX) rocsolver_lacgv_template(handle, n - 1 - j, W, shiftW + idx2D(j, jw + 1, ldw), ldw, strideW, batch_count); rocblasCall_gemv(handle, rocblas_operation_none, j + 1, n - 1 - j, cast2constType(scalars), 0, A, shiftA + idx2D(0, j + 1, lda), lda, strideA, W, shiftW + idx2D(j, jw + 1, ldw), ldw, strideW, cast2constType(scalars + 2), 0, A, shiftA + idx2D(0, j, lda), 1, strideA, batch_count, workArr); if(COMPLEX) { rocsolver_lacgv_template(handle, n - 1 - j, W, shiftW + idx2D(j, jw + 1, ldw), ldw, strideW, batch_count); rocsolver_lacgv_template(handle, n - 1 - j, A, shiftA + idx2D(j, j + 1, lda), lda, strideA, batch_count); } rocblasCall_gemv(handle, rocblas_operation_none, j + 1, n - 1 - j, cast2constType(scalars), 0, W, shiftW + idx2D(0, jw + 1, ldw), ldw, strideW, A, shiftA + idx2D(j, j + 1, lda), lda, strideA, cast2constType(scalars + 2), 0, A, shiftA + idx2D(0, j, lda), 1, strideA, batch_count, workArr); if(COMPLEX) rocsolver_lacgv_template(handle, n - 1 - j, A, shiftA + idx2D(j, j + 1, lda), lda, strideA, batch_count); // generate Householder reflector to work on column j rocsolver_larfg_template(handle, j, A, shiftA + idx2D(j - 1, j, lda), A, shiftA + idx2D(0, j, lda), 1, strideA, (tau + j - 1), strideP, batch_count, work, norms); // copy to E(j) the corresponding off-diagonal element of A, which is set to 1 ROCSOLVER_LAUNCH_KERNEL(set_offdiag, grid_b, threads, 0, stream, batch_count, A, shiftA + idx2D(j - 1, j, lda), strideA, (E + j - 1), strideE); // compute/update column j of W rocblasCall_symv_hemv(handle, uplo, j, (scalars + 2), 0, A, shiftA, lda, strideA, A, shiftA + idx2D(0, j, lda), 1, strideA, (scalars + 1), 0, W, shiftW + idx2D(0, jw, ldw), 1, strideW, batch_count, work, workArr); rocblasCall_gemv(handle, rocblas_operation_conjugate_transpose, j, n - 1 - j, cast2constType(scalars + 2), 0, W, shiftW + idx2D(0, jw + 1, ldw), ldw, strideW, A, shiftA + idx2D(0, j, lda), 1, strideA, cast2constType(scalars + 1), 0, W, shiftW + idx2D(j + 1, jw, ldw), 1, strideW, batch_count, workArr); rocblasCall_gemv(handle, rocblas_operation_none, j, n - 1 - j, cast2constType(scalars), 0, A, shiftA + idx2D(0, j + 1, lda), lda, strideA, W, shiftW + idx2D(j + 1, jw, ldw), 1, strideW, cast2constType(scalars + 2), 0, W, shiftW + idx2D(0, jw, ldw), 1, strideW, batch_count, workArr); rocblasCall_gemv(handle, rocblas_operation_conjugate_transpose, j, n - 1 - j, cast2constType(scalars + 2), 0, A, shiftA + idx2D(0, j + 1, lda), lda, strideA, A, shiftA + idx2D(0, j, lda), 1, strideA, cast2constType(scalars + 1), 0, W, shiftW + idx2D(j + 1, jw, ldw), 1, strideW, batch_count, workArr); rocblasCall_gemv(handle, rocblas_operation_none, j, n - 1 - j, cast2constType(scalars), 0, W, shiftW + idx2D(0, jw + 1, ldw), ldw, strideW, W, shiftW + idx2D(j + 1, jw, ldw), 1, strideW, cast2constType(scalars + 2), 0, W, shiftW + idx2D(0, jw, ldw), 1, strideW, batch_count, workArr); rocblasCall_scal(handle, j, (tau + j - 1), strideP, W, shiftW + idx2D(0, jw, ldw), 1, strideW, batch_count); rocblasCall_dot(handle, j, W, shiftW + idx2D(0, jw, ldw), 1, strideW, A, shiftA + idx2D(0, j, lda), 1, strideA, batch_count, norms, work, workArr); // (TODO: rocblas_axpy is not yet ready to be used in rocsolver. When it becomes // available, we can use it instead of the scale_axpy kernel, if it provides // better performance.) ROCSOLVER_LAUNCH_KERNEL(scale_axpy, grid_n, threads, 0, stream, j, norms, tau + j - 1, strideP, A, shiftA + idx2D(0, j, lda), strideA, W, shiftW + idx2D(0, jw, ldw), strideW); } } rocblas_set_pointer_mode(handle, old_mode); return rocblas_status_success; }