cu_golub_kahn_bidiagonalization.cu

Go to the documentation of this file.

/*
 *  SPDX-FileCopyrightText: Copyright 2021, Siavash Ameli <sameli@berkeley.edu>
 *  SPDX-License-Identifier: BSD-3-Clause
 *  SPDX-FileType: SOURCE
 *
 *  This program is free software: you can redistribute it and/or modify it
 *  under the terms of the license found in the LICENSE.txt file in the root
 *  directory of this source tree.
 */
 
 
// =======
// Headers
// =======
 
#include "./cu_golub_kahn_bidiagonalization.h"
#include "../_cu_definitions/cu_types.h" // __nv_fp8_e5m2, __nv_fp8_e4m3,
                                         // __half, __nv_bfloat16
#include <cmath>  // std::sqrt
#include "../_cu_basic_algebra/cu_vector_operations.h"  // cuVectorOperations
#include "../_cu_trace_estimator/cu_orthogonalization.h"  // cuOrthogonaliza...
#include "../_cuda_utilities/cuda_api.h"  // CudaAPI
#include "../_cu_arithmetics/cu_arithmetics.h" // cu_arithmetics
 
// Avoid CUBLAS numeration value not handled in switch [-Wswitch-enum] warning
#ifdef _MSC_VER
    #pragma warning(push, 0)  // Suppress all warnings from the followings
    #include <cublas_v2.h>  // cublasHandle_t
    #pragma warning(pop)  // Restore previous warning level
#elif defined(__INTEL_LLVM_COMPILER) || defined(__INTEL_COMPILER)
    #pragma warning(push, 0)
    #include <cublas_v2.h>  // cublasHandle_t
    #pragma warning(pop)
#elif defined(__GNUC__) || defined(__clang__)
    #pragma GCC diagnostic push
    #pragma GCC diagnostic ignored "-Wswitch-enum"
    #include <cublas_v2.h>  // cublasHandle_t
    #pragma GCC diagnostic pop
#else
    #include <cublas_v2.h>  // cublasHandle_t
#endif
 
 
// ============================
// golub-kahn bidiagonalization
// ============================
 
 
template <typename DataType>
IndexType cu_golub_kahn_bidiagonalization(
        cuLinearOperator<DataType>* A,
        const DataType* v,
        const LongIndexType n,
        const IndexType m,
        const DataType lanczos_tol,
        const FlagType orthogonalize,
        DataType* alpha,
        DataType* beta)
{
    // Get cublas handle
    cublasHandle_t cublas_handle = A->get_cublas_handle();
 
    // buffer_size is number of last orthogonal vectors to keep in buffers U, V
    IndexType buffer_size;
    if (orthogonalize == 0)
    {
        // At least two vectors must be stored in buffer for Lanczos recursion
        buffer_size = 2;
    }
    else if ((orthogonalize < 0) ||
             (orthogonalize > static_cast<FlagType>(m) - 1))
    {
        // Using full re-orthogonalization, keep all of the m vectors in buffer
        buffer_size = m;
    }
    else
    {
        // Orthogonalize with less than m vectors (0 < orthogonalize < m-1)
        // plus one vector for the latest (the j-th) vector
        buffer_size = orthogonalize + 1;
    }
 
    // Allocate 2D array (as 1D array, and coalesced row-wise) to store
    // the last buffer_size of orthogonalized vectors of length n. New vectors
    // are stored by cycling through the buffer to replace with old ones.
    DataType* device_U = CudaAPI<DataType>::alloc(n * buffer_size);
    DataType* device_V = CudaAPI<DataType>::alloc(n * buffer_size);
 
    // Normalize vector v and copy to v_old
    CudaAPI<DataType>::copy_to_device(v, n, &device_V[0]);
    cuVectorOperations<DataType>::normalize_vector_in_place(
            cublas_handle, &device_V[0], n);
 
    // Declare iterators
    IndexType j;
    IndexType lanczos_size = 0;
    IndexType num_ortho;
 
    // Golub-Kahn iteration
    for (j=0; j < m; ++j)
    {
        // Counter for the non-zero size of alpha and beta
        ++lanczos_size;
 
        // u_new = A.dot(v_old)
        A->dot(&device_V[(j % buffer_size)*n], &device_U[(j % buffer_size)*n]);
 
        // Performing: u_new[i] = u_new[i] - beta[j] * u_old[i]
        if (j > 0)
        {
            cuVectorOperations<DataType>::subtract_scaled_vector(
                    cublas_handle,
                    &device_U[((j-1) % buffer_size)*n], n, beta[j-1],
                    &device_U[(j % buffer_size)*n]);
        }
 
        // orthogonalize u_new against previous vectors
        if (orthogonalize != 0)
        {
            // Find how many column vectors are filled so far in the buffer V
            if (j < buffer_size)
            {
                num_ortho = j;
            }
            else
            {
                num_ortho = buffer_size - 1;
            }
 
            // Gram-Schmidt process
            if (j > 0)
            {
                cuOrthogonalization<DataType>::gram_schmidt_process(
                        cublas_handle, &device_U[0], n, buffer_size,
                        (j-1)%buffer_size, num_ortho,
                        &device_U[(j % buffer_size)*n]);
            }
        }
 
        // Normalize u_new and set its norm to alpha[j]
        alpha[j] = cuVectorOperations<DataType>::normalize_vector_in_place(
                cublas_handle, &device_U[(j % buffer_size)*n], n);
 
        // Performing: v_new = A.T.dot(u_new) - alpha[j] * v_old
        A->transpose_dot(&device_U[(j % buffer_size)*n],
               &device_V[((j+1) % buffer_size)*n]);
 
        // Performing: v_new[i] = v_new[i] - alpha[j] * v_old[i]
        cuVectorOperations<DataType>::subtract_scaled_vector(
                cublas_handle, &device_V[(j % buffer_size)*n], n, alpha[j],
                &device_V[((j+1) % buffer_size)*n]);
 
        // orthogonalize v_new against previous vectors
        if (orthogonalize != 0)
        {
            cuOrthogonalization<DataType>::gram_schmidt_process(
                    cublas_handle, &device_V[0], n, buffer_size, j%buffer_size,
                    num_ortho, &device_V[((j+1) % buffer_size)*n]);
        }
 
        // Update beta as the norm of v_new
        beta[j] = cuVectorOperations<DataType>::normalize_vector_in_place(
                cublas_handle, &device_V[((j+1) % buffer_size)*n], n);
 
        // Exit criterion when the vector r is zero. If each component of a
        // zero vector has the tolerance epsilon, (which is called lanczos_tol
        // here), the tolerance of norm of r is epsilon times sqrt of n.
        if (beta[j] < cu_arithmetics::mul(
                    lanczos_tol,
                    cu_arithmetics::cast<double, DataType>(
                        static_cast<double>(std::sqrt(n)))
                    )
           )
        {
            break;
        }
    }
 
    // Free dynamic memory
    CudaAPI<DataType>::del(device_U);
    CudaAPI<DataType>::del(device_V);
 
    return lanczos_size;
}
 
 
// ===============================
// Explicit template instantiation
// ===============================
 
// golub kahn bidiagonalization (__nv_fp8_e5m2)
#if defined(USE_CUDA_FP8_E5M2) && (USE_CUDA_FP8_E5M2 == 1)
    template IndexType cu_golub_kahn_bidiagonalization<__nv_fp8_e5m2>(
            cuLinearOperator<__nv_fp8_e5m2>* A,
            const __nv_fp8_e5m2* v,
            const LongIndexType n,
            const IndexType m,
            const __nv_fp8_e5m2 lanczos_tol,
            const FlagType orthogonalize,
            __nv_fp8_e5m2* alpha,
            __nv_fp8_e5m2* beta);
#endif
 
// golub kahn bidiagonalization (__nv_fp8_e5m2)
#if defined(USE_CUDA_FP8_E4M3) && (USE_CUDA_FP8_E4M3 == 1)
    template IndexType cu_golub_kahn_bidiagonalization<__nv_fp8_e4m3>(
            cuLinearOperator<__nv_fp8_e4m3>* A,
            const __nv_fp8_e4m3* v,
            const LongIndexType n,
            const IndexType m,
            const __nv_fp8_e4m3 lanczos_tol,
            const FlagType orthogonalize,
            __nv_fp8_e4m3* alpha,
            __nv_fp8_e4m3* beta);
#endif
 
// golub kahn bidiagonalization (__half)
#if defined(USE_CUDA_FP16) && (USE_CUDA_FP16 == 1)
    template IndexType cu_golub_kahn_bidiagonalization<__half>(
            cuLinearOperator<__half>* A,
            const __half* v,
            const LongIndexType n,
            const IndexType m,
            const __half lanczos_tol,
            const FlagType orthogonalize,
            __half* alpha,
            __half* beta);
#endif
 
// golub kahn bidiagonalization (__nv_bfloat16)
#if defined(USE_CUDA_BF16) && (USE_CUDA_BF16 == 1)
    template IndexType cu_golub_kahn_bidiagonalization<__nv_bfloat16>(
            cuLinearOperator<__nv_bfloat16>* A,
            const __nv_bfloat16* v,
            const LongIndexType n,
            const IndexType m,
            const __nv_bfloat16 lanczos_tol,
            const FlagType orthogonalize,
            __nv_bfloat16* alpha,
            __nv_bfloat16* beta);
#endif
 
// golub kahn bidiagonalization (float)
#if defined(USE_CUDA_FP32) && (USE_CUDA_FP32 == 1)
    template IndexType cu_golub_kahn_bidiagonalization<float>(
            cuLinearOperator<float>* A,
            const float* v,
            const LongIndexType n,
            const IndexType m,
            const float lanczos_tol,
            const FlagType orthogonalize,
            float* alpha,
            float* beta);
#endif
 
// golub kahn bidiagonalization (double)
#if defined(USE_CUDA_FP64) && (USE_CUDA_FP64 == 1)
    template IndexType cu_golub_kahn_bidiagonalization<double>(
            cuLinearOperator<double>* A,
            const double* v,
            const LongIndexType n,
            const IndexType m,
            const double lanczos_tol,
            const FlagType orthogonalize,
            double* alpha,
            double* beta);
#endif