cu_orthogonalization.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.
 */
 
 
// =======
// Imports
// =======
 
#include "./cu_orthogonalization.h"
#include "../_cu_definitions/cu_types.h" // __nv_fp8_e5m2, __nv_fp8_e4m3,
                                         // __half, __nv_bfloat16
#include <cstdlib>  // abort, NULL
#include <iostream>  // std::cerr, std::endl
#include <cmath>  // std::sqrt, std::fabs
#include <limits>  // std::numeric_limits
#include "../_cu_basic_algebra/cu_vector_operations.h"  // cuVectorOperations
#include "../_random_generator/random_array_generator.h"  // RandomArrayGene...
#include "../_random_generator/random_number_generator.h"  // RandomNumberGe...
#include "../_cuda_utilities/cuda_api.h"  // CudaAPI
#include "../_cu_arithmetics/cu_arithmetics.h" // cu_arithmetics
 
 
// ====================
// gram schmidt process
// ====================
 
 
template <typename DataType>
void cuOrthogonalization<DataType>::gram_schmidt_process(
        cublasHandle_t cublas_handle,
        const DataType* V,
        const LongIndexType vector_size,
        const IndexType num_vectors,
        const IndexType last_vector,
        const FlagType num_ortho,
        DataType* v)
{
    // Determine how many previous vectors to orthogonalize against
    IndexType num_steps;
    if ((num_ortho == 0) || (num_vectors < 2))
    {
        // No orthogonalization is performed
        return;
    }
    else if ((num_ortho < 0) ||
             (num_ortho > static_cast<FlagType>(num_vectors)))
    {
        // Orthogonalize against all vectors
        num_steps = num_vectors;
    }
    else
    {
        // Orthogonalize against only the last num_ortho vectors
        num_steps = num_ortho;
    }
 
    // Vectors can be orthogonalized at most to the full basis of the vector
    // space. Thus, num_steps cannot be larger than the dimension of vector
    // space, which is vector_size.
    if (num_steps > static_cast<IndexType>(vector_size))
    {
        num_steps = vector_size;
    }
 
    IndexType i;
    DataType inner_prod;
    DataType norm;
    DataType norm_v;
    DataType epsilon = cu_arithmetics::epsilon<DataType>();
    DataType distance2;
 
    // Iterate over vectors
    for (IndexType step=0; step < num_steps; ++step)
    {
        // i is the index of a column vector in V to orthogonalize v against it
        if ((last_vector % num_vectors) >= step)
        {
            i = (last_vector % num_vectors) - step;
        }
        else
        {
            // Wrap around negative indices from the end of column index
            i = (last_vector % num_vectors) - step + num_vectors;
        }
 
        // Norm of j-th vector
        norm = cuVectorOperations<DataType>::euclidean_norm(
                cublas_handle, &V[vector_size*i], vector_size);
 
        // Check norm
        if (norm < \
                cu_arithmetics::mul(
                    epsilon,
                    cu_arithmetics::cast<double, DataType>(
                        std::sqrt(vector_size))
                )
           )
        {
            std::cerr << "WARNING: norm of the given vector is too small. " \
                      << "Cannot orthogonalize against zero vector. " \
                      << "Skipping." << std::endl;
            continue;
        }
 
        // Projection
        inner_prod = cuVectorOperations<DataType>::inner_product(
                cublas_handle, &V[vector_size*i], v, vector_size);
 
        // scale for subtraction
        DataType scale = cu_arithmetics::div(
                inner_prod,
                cu_arithmetics::mul(norm, norm));
        DataType scale_abs = cu_arithmetics::abs(scale);
        DataType one = cu_arithmetics::cast<double, DataType>(1.0);
        DataType two = cu_arithmetics::cast<double, DataType>(2.0);
 
        // If scale is 1, it is possible that vector v and j-th vector are
        // identical (or close).
        if (cu_arithmetics::abs(cu_arithmetics::sub(scale_abs, one)) <= \
            cu_arithmetics::mul(two, epsilon))
        {
            // Norm of the vector v
            norm_v = cuVectorOperations<DataType>::euclidean_norm(
                    cublas_handle, v, vector_size);
 
            // Compute distance between the j-th vector and vector v
            distance2 = \
                cu_arithmetics::add(
                    cu_arithmetics::mul(norm_v, norm_v),
                    cu_arithmetics::mul(-two, inner_prod),
                    cu_arithmetics::mul(norm, norm)
                );
 
            // If distance is zero, do not re-orthogonalize i-th against
            // the j-th vector.
            if (distance2 < \
                    cu_arithmetics::mul(
                        two,
                        epsilon,
                        cu_arithmetics::cast<unsigned long long int, DataType>(
                            static_cast<unsigned long long int>(vector_size))
                    )
               )
            {
                continue;
            }
        }
 
        // Subtraction
        cuVectorOperations<DataType>::subtract_scaled_vector(
                cublas_handle, &V[vector_size*i], vector_size, scale, v);
    }
}
 
 
// =====================
// orthogonalize vectors
// =====================
 
 
template <typename DataType>
void cuOrthogonalization<DataType>::orthogonalize_vectors(
        cublasHandle_t cublas_handle,
        DataType* vectors,
        const LongIndexType vector_size,
        const IndexType num_vectors)
{
    // Do nothing if there is only one vector
    if (num_vectors < 2)
    {
        return;
    }
 
    IndexType i = 0;
    IndexType j;
    IndexType start_j;
    DataType inner_prod;
    DataType norm_j;
    DataType norm_i;
    DataType epsilon = cu_arithmetics::epsilon<DataType>();
    IndexType success = 1;
    IndexType max_num_trials = 20;
    IndexType num_trials = 0;
    IndexType num_threads = 1;
    RandomNumberGenerator random_number_generator(num_threads);
    DataType* buffer = NULL;
 
    while (i < num_vectors)
    {
        if ((success == 0) && (num_trials >= max_num_trials))
        {
            std::cerr << "ERROR: Cannot orthogonalize vectors after " \
                      << num_trials << " trials. Aborting." \
                      << std::endl;
            abort();
        }
 
        // Reset on new trial (if it was set to 0 before to start a new trial)
        success = 1;
 
        // j iterates on previous vectors in a window of at most vector_size
        if (static_cast<LongIndexType>(i) > vector_size)
        {
            // When vector_size is smaller than i, it is fine to cast to signed
            start_j = i - static_cast<IndexType>(vector_size);
        }
        else
        {
            start_j = 0;
        }
 
        // Re-orthogonalize against previous vectors
        for (j=start_j; j < i; ++j)
        {
            // Norm of the j-th vector
            norm_j = cuVectorOperations<DataType>::euclidean_norm(
                    cublas_handle, &vectors[j*vector_size], vector_size);
 
            // Check norm
            if (norm_j < \
                    cu_arithmetics::mul(
                        epsilon,
                        cu_arithmetics::cast<double, DataType>(
                            std::sqrt(vector_size))
                    )
               )
            {
                std::cerr << "WARNING: norm of the given vector is too " \
                          << " small. Cannot re-orthogonalize against zero" \
                          << "vector. Skipping."
                          << std::endl;
                continue;
            }
 
            // Projecting i-th vector to j-th vector
            inner_prod = cuVectorOperations<DataType>::inner_product(
                    cublas_handle, &vectors[i*vector_size],
                    &vectors[j*vector_size], vector_size);
 
            // Scale of subtraction
            DataType scale = cu_arithmetics::div(
                    inner_prod,
                    cu_arithmetics::mul(norm_j, norm_j));
 
            // Subtraction
            cuVectorOperations<DataType>::subtract_scaled_vector(
                    cublas_handle, &vectors[vector_size*j], vector_size, scale,
                    &vectors[vector_size*i]);
 
            // Norm of the i-th vector
            norm_i = cuVectorOperations<DataType>::euclidean_norm(
                    cublas_handle, &vectors[i*vector_size], vector_size);
 
            // If the norm is too small, regenerate the i-th vector randomly
            if (norm_i < \
                    cu_arithmetics::mul(
                        epsilon,
                        cu_arithmetics::cast<double, DataType>(
                            std::sqrt(vector_size))
                    )
                )
            {
                // Allocate buffer
                if (buffer == NULL)
                {
                    buffer = new DataType[vector_size];
                }
 
                // Regenerate new random vector on buffer
                RandomArrayGenerator<DataType>::generate_random_array(
                        random_number_generator, buffer,
                        vector_size, num_threads);
 
                // Copy buffer to the i-th vector on device
                CudaAPI<DataType>::copy_to_device(
                        buffer, vector_size, &vectors[i*vector_size]);
 
                // Repeat the re-orthogonalization for i-th vector against
                // all previous vectors again.
                success = 0;
                ++num_trials;
                break;
            }
        }
 
        if (success == 1)
        {
            ++i;
 
            // Reset if num_trials was incremented before.
            num_trials = 0;
        }
    }
 
    // Deallocate buffer
    if (buffer != NULL)
    {
        delete[] buffer;
        buffer = NULL;
    }
}
 
 
// ===============================
// Explicit template instantiation
// ===============================
 
#if defined(USE_CUDA_FP8_E5M2) && (USE_CUDA_FP8_E5M2 == 1)
    template class cuOrthogonalization<__nv_fp8_e5m2>;
#endif
 
#if defined(USE_CUDA_FP8_E4M3) && (USE_CUDA_FP8_E4M3 == 1)
    template class cuOrthogonalization<__nv_fp8_e4m3>;
#endif
    
#if defined(USE_CUDA_FP16) && (USE_CUDA_FP16 == 1)
    template class cuOrthogonalization<__half>;
#endif
 
#if defined(USE_CUDA_BF16) && (USE_CUDA_BF16 == 1)
    template class cuOrthogonalization<__nv_bfloat16>;
#endif
 
#if defined(USE_CUDA_FP32) && (USE_CUDA_FP32 == 1)
    template class cuOrthogonalization<float>;
#endif
 
#if defined(USE_CUDA_FP64) && (USE_CUDA_FP64 == 1)
    template class cuOrthogonalization<double>;
#endif