cu_trace_estimator.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_trace_estimator.h"
#include "../_definitions/definitions.h"  // USE_OPENMP
#include "../_cu_definitions/cu_types.h" // __nv_fp8_e5m2, __nv_fp8_e4m3,
                                         // __half, __nv_bfloat16
 
#if defined(USE_OPENMP) && (USE_OPENMP == 1)
    #include <omp.h>  // omp_set_num_threads
#endif
 
#include <cmath>  // std::sqrt, std::pow
#include <cstddef>  // NULL
#include "./cu_lanczos_tridiagonalization.h"  // cu_lanczos_tridiagonalization
#include "./cu_golub_kahn_bidiagonalization.h"  // cu_golub_kahn_bidiagonali...
#include "../_random_generator/random_array_generator.h"  // RandomArrayGene...
#include "../_c_trace_estimator/diagonalization.h"  // Diagonalization
#include "../_c_trace_estimator/convergence_tools.h"  // check_convergence, ...
#include "../_cuda_utilities/cuda_timer.h"  // CudaTimer
#include "../_cuda_utilities/cuda_api.h"  // CudaAPI
#include "../_cu_arithmetics/cu_arithmetics.h" // cu_arithmetics
 
 
// ==================
// cu trace estimator
// ==================
 
 
template <typename DataType>
FlagType cuTraceEstimator<DataType>::cu_trace_estimator(
        cuLinearOperator<DataType>* A,
        DataType* parameters,
        const IndexType num_inquiries,
        const Function* matrix_function,
        const FlagType gram,
        const DataType exponent,
        const FlagType orthogonalize,
        const int64_t seed,
        const IndexType lanczos_degree,
        const DataType lanczos_tol,
        const IndexType min_num_samples,
        const IndexType max_num_samples,
        const DataType error_atol,
        const DataType error_rtol,
        const DataType confidence_level,
        const DataType outlier_significance_level,
        const IndexType num_threads,
        const IndexType num_gpu_devices,
        DataType* trace,
        DataType* error,
        DataType** samples,
        IndexType* processed_samples_indices,
        IndexType* num_samples_used,
        IndexType* num_outliers,
        FlagType* converged,
        double& alg_wall_time)
{
    // Matrix size
    IndexType matrix_size = A->get_num_rows();
 
    // Set the number of threads
    #if defined(USE_OPENMP) && (USE_OPENMP == 1)
        omp_set_num_threads(num_gpu_devices);
    #else
        num_threads = 1;
    #endif
 
    // Allocate 1D array of random vectors We only allocate a random vector
    // per parallel thread. Thus, the total size of the random vectors is
    // matrix_size*num_threads. On each iteration in parallel threads, the
    // allocated memory is reused. That is, in each iteration, a new random
    // vector is generated for that specific thread id.
    IndexType random_vectors_size = matrix_size * num_gpu_devices;
    DataType* random_vectors = new DataType[random_vectors_size];
 
    // Initialize random number generator to generate in parallel threads
    // independently.
    RandomNumberGenerator random_number_generator(num_gpu_devices, seed);
 
    // The counter of filled size of processed_samples_indices array
    // This scalar variable is defined as array to be shared among al threads
    IndexType num_processed_samples = 0;
 
    // Criterion for early termination of iterations if convergence reached
    // This scalar variable is defined as array to be shared among al threads
    FlagType all_converged = 0;
 
    // Using square-root of max possible chunk size for parallel schedules
    unsigned int chunk_size = static_cast<int>(
        std::sqrt(static_cast<double>(max_num_samples) / \
                  static_cast<double>(num_gpu_devices)));
    if (chunk_size < 1)
    {
        chunk_size = 1;
    }
 
    // Timing elapsed time of algorithm
    CudaTimer cuda_timer;
    cuda_timer.start();
 
    // Shared-memory parallelism over Monte-Carlo ensemble sampling
    IndexType i;
    #if defined(USE_OPENMP) && (USE_OPENMP == 1)
    #pragma omp parallel for \
        schedule(dynamic, chunk_size)
    #endif
    for (i=0; i < max_num_samples; ++i)
    {
        if (!static_cast<bool>(all_converged))
        {
            // Switch to a device with the same device id as the cpu thread id
            unsigned int thread_id = omp_get_thread_num();
            CudaAPI<DataType>::set_device(thread_id);
 
            // Perform one Monte-Carlo sampling to estimate trace
            cuTraceEstimator<DataType>::_cu_stochastic_lanczos_quadrature(
                    A, parameters, num_inquiries, matrix_function, gram,
                    exponent, orthogonalize, lanczos_degree, lanczos_tol,
                    random_number_generator,
                    &random_vectors[matrix_size*thread_id], converged,
                    samples[i]);
 
            // Critical section
            #if defined(USE_OPENMP) && (USE_OPENMP == 1)
            #pragma omp critical
            #endif
            {
                // Store the index of processed samples
                processed_samples_indices[num_processed_samples] = i;
                ++num_processed_samples;
 
                // Check whether convergence criterion has been met to stop.
                // This check can also be done after another parallel thread
                // set all_converged to "1", but we continue to update error.
                all_converged = ConvergenceTools<DataType>::check_convergence(
                        samples, min_num_samples, num_inquiries,
                        processed_samples_indices, num_processed_samples,
                        confidence_level, error_atol, error_rtol, error,
                        num_samples_used, converged);
            }
        }
    }
 
    // Elapsed wall time of the algorithm (computation only, not array i/o)
    cuda_timer.stop();
    alg_wall_time = cuda_timer.elapsed();
 
    // Remove outliers from trace estimates and average trace estimates
    ConvergenceTools<DataType>::average_estimates(
            confidence_level, outlier_significance_level, num_inquiries,
            max_num_samples, num_samples_used, processed_samples_indices,
            samples, num_outliers, trace, error);
 
    // Deallocate memory
    delete[] random_vectors;
 
    return all_converged;
}
 
 
// ================================
// cu stochastic lanczos quadrature
// ================================
 
 
template <typename DataType>
void cuTraceEstimator<DataType>::_cu_stochastic_lanczos_quadrature(
        cuLinearOperator<DataType>* A,
        DataType* parameters,
        const IndexType num_inquiries,
        const Function* matrix_function,
        const FlagType gram,
        const DataType exponent,
        const FlagType orthogonalize,
        const IndexType lanczos_degree,
        const DataType lanczos_tol,
        RandomNumberGenerator& random_number_generator,
        DataType* random_vector,
        FlagType* converged,
        DataType* trace_estimate)
{
    // Matrix size
    IndexType matrix_size = A->get_num_rows();
 
    // Fill random vectors with Rademacher distribution (+1, -1), normalized
    // but not orthogonalized. Setting num_threads to zero indicates to not
    // create any new threads in RandomNumbrGenerator since the current
    // function is inside a parallel thread.
    IndexType num_threads = 0;
    RandomArrayGenerator<DataType>::generate_random_array(
        random_number_generator, random_vector, matrix_size, num_threads);
 
    // Allocate diagonals (alpha) and supdiagonals (beta) of Lanczos matrix
    DataType* alpha = new DataType[lanczos_degree];
    DataType* beta = new DataType[lanczos_degree];
 
    // Define 2D arrays needed to decomposition. All these arrays are
    // defined as 1D array with Fortran ordering
    DataType* eigenvectors = NULL;
    DataType* left_singularvectors = NULL;
    DataType* right_singularvectors_transposed = NULL;
 
    // Actual number of inquiries
    IndexType required_num_inquiries = num_inquiries;
    if (A->is_eigenvalue_relation_known())
    {
        // When a relation between eigenvalues and the parameters of the linear
        // operator is known, to compute eigenvalues of for each inquiry, only
        // computing one inquiry is enough. This is because an eigenvalue for
        // one parameter setting is enough to compute eigenvalue of another set
        // of parameters.
        required_num_inquiries = 1;
    }
 
    // Allocate and initialize theta
    IndexType i;
    IndexType j;
    DataType** theta = new DataType*[num_inquiries];
    for (j=0; j < num_inquiries; ++j)
    {
        theta[j] = new DataType[lanczos_degree];
 
        // Initialize components to zero
        for (i=0; i < lanczos_degree; ++i)
        {
            theta[j][i] = 0.0;
        }
    }
 
    // Allocate and initialize tau
    DataType** tau = new DataType*[num_inquiries];
    for (j=0; j < num_inquiries; ++j)
    {
        tau[j] = new DataType[lanczos_degree];
 
        // Initialize components to zero
        for (i=0; i < lanczos_degree; ++i)
        {
            tau[j][i] = 0.0;
        }
    }
 
    // Allocate lanczos size for each inquiry. This variable keeps the non-zero
    // size of the tri-diagonal (or bi-diagonal) matrix. Ideally, this matrix
    // is of the size lanczos_degree. But, due to the early termination, this
    // size might be smaller.
    IndexType* lanczos_size = new IndexType[num_inquiries];
 
    // Number of parameters of linear operator A
    IndexType num_parameters = A->get_num_parameters();
 
    // Lanczos iterations, computes theta and tau for each inquiry parameter
    for (j=0; j < required_num_inquiries; ++j)
    {
        // If trace is already converged, do not compute on the new sample.
        // However, exclude the case where required_num_inquiries is not the
        // same as num_inquiries, since in this case, we compute one inquiry
        // for multiple parameters.
        if ((converged[j] == 1) && (required_num_inquiries == num_inquiries))
        {
            continue;
        }
 
        // Set parameter of linear operator A
        A->set_parameters(&parameters[j*num_parameters]);
 
        if (gram)
        {
            // Use Golub-Kahn-Lanczos Bi-diagonalization
            lanczos_size[j] = cu_golub_kahn_bidiagonalization(
                    A, random_vector, matrix_size, lanczos_degree, lanczos_tol,
                    orthogonalize, alpha, beta);
 
            // Allocate matrix of singular vectors (1D array, Fortran ordering)
            left_singularvectors = \
                new DataType[lanczos_size[j] * lanczos_size[j]];
            right_singularvectors_transposed = \
                new DataType[lanczos_size[j] * lanczos_size[j]];
 
            // Note: alpha is written in-place with singular values
            Diagonalization<DataType>::svd_bidiagonal(
                    alpha, beta, left_singularvectors,
                    right_singularvectors_transposed, lanczos_size[j]);
 
            // theta and tau from singular values and vectors
            for (i=0; i < lanczos_size[j]; ++i)
            {
                theta[j][i] = cu_arithmetics::mul(alpha[i], alpha[i]);
                tau[j][i] = right_singularvectors_transposed[i];
            }
        }
        else
        {
            // Use Lanczos Tri-diagonalization
            lanczos_size[j] = cu_lanczos_tridiagonalization(
                    A, random_vector, matrix_size, lanczos_degree, lanczos_tol,
                    orthogonalize, alpha, beta);
 
            // Allocate eigenvectors matrix (1D array with Fortran ordering)
            eigenvectors = new DataType[lanczos_size[j] * lanczos_size[j]];
 
            // Note: alpha is written in-place with eigenvalues
            Diagonalization<DataType>::eigh_tridiagonal(
                    alpha, beta, eigenvectors, lanczos_size[j]);
 
            // theta and tau from singular values and vectors
            for (i=0; i < lanczos_size[j]; ++i)
            {
                theta[j][i] = alpha[i];
                tau[j][i] = eigenvectors[i * lanczos_size[j]];
            }
        }
    }
 
    // If an eigenvalue relation is known, compute the rest of eigenvalues
    // using the eigenvalue relation given in the operator A for its
    // eigenvalues. If no eigenvalue relation is not known, the rest of
    // eigenvalues were already computed in the above loop and no other
    // computation is needed.
    if (A->is_eigenvalue_relation_known() && num_inquiries > 1)
    {
        // When the code execution reaches this function, at least one of the
        // inquiries is not converged, but some others might have been
        // converged already. Here, we force-update those that are even
        // converged already by setting converged to false. The extra update is
        // free of charge when a relation for the eigenvalues are known.
        for (j=0; j < num_inquiries; ++j)
        {
            converged[j] = 0;
        }
 
        // Compute theta and tau for the rest of inquiry parameters
        for (j=1; j < num_inquiries; ++j)
        {
            // Only j=0 was iterated before. Set the same size for other j-s
            lanczos_size[j] = lanczos_size[0];
 
            for (i=0; i < lanczos_size[j]; ++i)
            {
                // Shift eigenvalues by the old and new parameters
                theta[j][i] = A->get_eigenvalue(
                        &parameters[0],
                        theta[0][i],
                        &parameters[j*num_parameters]);
 
                // tau is the same (at least for the affine operator)
                tau[j][i] = tau[0][i];
            }
        }
    }
 
    // Estimate trace using quadrature method
    DataType quadrature_sum;
    for (j=0; j < num_inquiries; ++j)
    {
        // If the j-th inquiry is already converged, skip.
        if (converged[j] == 1)
        {
            continue;
        }
 
        // Initialize sum for the integral of quadrature
        quadrature_sum = 0.0;
 
        // Important: This loop should iterate till lanczos_size[j], but not
        // lanczos_degree. Otherwise the computation is wrong for certain
        // matrices, such as if the input matrix is identity, or rank
        // deficient. By using lanczos_size[j] instead of lanczos_degree, all
        // issues with special matrices will resolve.
        for (i=0; i < lanczos_size[j]; ++i)
        {
            quadrature_sum += cu_arithmetics::mul(
                tau[j][i],
                tau[j][i],
                cu_arithmetics::cast<double, DataType>(
                    matrix_function->function(std::pow(
                        cu_arithmetics::cast<DataType, double>(theta[j][i]),
                        cu_arithmetics::cast<DataType, double>(exponent)
                        )
                    )
                )
            );
        }
 
        trace_estimate[j] = \
            cu_arithmetics::mul(
                quadrature_sum,
                cu_arithmetics::cast<unsigned long long int, DataType>(
                    static_cast<unsigned long long int>(matrix_size)
                )
            );
    }
 
    // Release dynamic memory
    delete[] alpha;
    delete[] beta;
    delete[] lanczos_size;
 
    for (j=0; j < required_num_inquiries; ++j)
    {
        delete[] theta[j];
    }
    delete[] theta;
 
    for (j=0; j < required_num_inquiries; ++j)
    {
        delete[] tau[j];
    }
    delete[] tau;
 
    if (eigenvectors != NULL)
    {
        delete[] eigenvectors;
    }
 
    if (left_singularvectors != NULL)
    {
        delete[] left_singularvectors;
    }
 
    if (right_singularvectors_transposed != NULL)
    {
        delete[] right_singularvectors_transposed;
    }
}
 
 
// ===============================
// Explicit template instantiation
// ===============================
 
#if defined(USE_CUDA_FP8_E5M2) && (USE_CUDA_FP8_E5M2 == 1)
    template class cuTraceEstimator<__nv_fp8_e5m2>;
#endif
 
#if defined(USE_CUDA_FP8_E4M3) && (USE_CUDA_FP8_E4M3 == 1)
    template class cuTraceEstimator<__nv_fp8_e4m3>;
#endif
 
#if defined(USE_CUDA_FP16) && (USE_CUDA_FP16 == 1)
    template class cuTraceEstimator<__half>;
#endif
 
#if defined(USE_CUDA_BF16) && (USE_CUDA_BF16 == 1)
    template class cuTraceEstimator<__nv_bfloat16>;
#endif
 
#if defined(USE_CUDA_FP32) && (USE_CUDA_FP32 == 1)
    template class cuTraceEstimator<float>;
#endif
 
#if defined(USE_CUDA_FP64) && (USE_CUDA_FP64 == 1)
    template class cuTraceEstimator<double>;
#endif