cu_dense_affine_matrix_function.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_dense_affine_matrix_function.h"
#include <cstddef>  // NULL
#include <cassert>  // assert
#include "../_cu_definitions/cu_types.h" // __nv_fp8_e5m2, __nv_fp8_e4m3,
                                         // __half, __nv_bfloat16
#include "../_definitions/debugging.h"  // ASSERT
 
 
// =============
// constructor 1
// =============
 
 
template <typename DataType>
cuDenseAffineMatrixFunction<DataType>::cuDenseAffineMatrixFunction(
        const DataType* A_,
        const LongIndexType num_rows_,
        const LongIndexType num_columns_,
        const FlagType A_is_row_major_,
        const FlagType A_is_symmetric_,
        const int num_gpu_devices_):
 
    // Base class constructor
    cLinearOperatorBase(num_rows_, num_columns_),
    cuLinearOperator<DataType>(num_gpu_devices_),
 
    // Initializer list
    A(A_, num_rows_, num_columns_, A_is_row_major_, A_is_symmetric_,
      num_gpu_devices_)
{
    // This constructor is called assuming B is identity
    this->B_is_identity = true;
 
    // When B is identity, the eigenvalues of A+tB are known for any t
    this->eigenvalue_relation_known = 1;
 
    // Set gpu device
    this->initialize_cublas_handle();
}
 
 
// =============
// constructor 2
// =============
 
 
template <typename DataType>
cuDenseAffineMatrixFunction<DataType>::cuDenseAffineMatrixFunction(
        const DataType* A_,
        const LongIndexType num_rows_,
        const LongIndexType num_columns_,
        const FlagType A_is_row_major_,
        const FlagType A_is_symmetric_,
        const DataType* B_,
        const FlagType B_is_row_major_,
        const FlagType B_is_symmetric_,
        const int num_gpu_devices_):
 
    // Base class constructor
    cLinearOperatorBase(num_rows_, num_columns_),
    cuLinearOperator<DataType>(num_gpu_devices_),
 
    // Initializer list
    A(A_, num_rows_, num_columns_, A_is_row_major_, A_is_symmetric_,
      num_gpu_devices_),
    B(B_, num_rows_, num_columns_, B_is_row_major_, B_is_symmetric_,
      num_gpu_devices_)
{
    // Matrix B is assumed to be non-zero. Check if it is identity or generic
    if (this->B.is_identity_matrix())
    {
        this->B_is_identity = true;
        this->eigenvalue_relation_known = 1;
    }
 
    // Set gpu device
    this->initialize_cublas_handle();
}
 
 
// ==========
// destructor
// ==========
 
 
template <typename DataType>
cuDenseAffineMatrixFunction<DataType>::~cuDenseAffineMatrixFunction()
{
}
 
 
// ============
// set symmetry
// ============
 
 
template <typename DataType>
void cuDenseAffineMatrixFunction<DataType>::set_symmetry(
        const FlagType symmetric)
{
    if (symmetric == 1)
    {
        this->A.set_symmetry(1);
        this->B.set_symmetry(1);
    }
    else
    {
        this->A.set_symmetry(0);
        this->B.set_symmetry(0);
    }
}
 
 
// ===
// dot
// ===
 
 
template <typename DataType>
void cuDenseAffineMatrixFunction<DataType>::dot(
        const DataType* vector,
        DataType* product)
{
    // Matrix A times vector
    this->A.dot(vector, product);
    LongIndexType min_vector_size;
 
    // Matrix B times vector to be added to the product
    if (this->B_is_identity)
    {
        // Check parameter is set
        ASSERT((this->parameters != NULL), "Parameter is not set.");
 
        // Find minimum of the number of rows and columns
        min_vector_size = \
            (this->num_rows < this->num_columns) ? \
            this->num_rows : this->num_columns;
 
        // Adding input vector to product
        this->_add_scaled_vector(vector, min_vector_size,
                                 this->parameters[0], product);
    }
    else
    {
        // Check parameter is set
        ASSERT((this->parameters != NULL), "Parameter is not set.");
 
        // Adding parameter times B times input vector to the product
        this->B.dot_plus(vector, this->parameters[0], product);
    }
}
 
 
// =============
// transpose dot
// =============
 
 
template <typename DataType>
void cuDenseAffineMatrixFunction<DataType>::transpose_dot(
        const DataType* vector,
        DataType* product)
{
    // Matrix A times vector
    this->A.transpose_dot(vector, product);
    LongIndexType min_vector_size;
 
    // Matrix B times vector to be added to the product
    if (this->B_is_identity)
    {
        // Check parameter is set
        ASSERT((this->parameters != NULL), "Parameter is not set.");
 
        // Find minimum of the number of rows and columns
        min_vector_size = \
            (this->num_rows < this->num_columns) ? \
            this->num_rows : this->num_columns;
 
        // Adding input vector to product
        this->_add_scaled_vector(vector, min_vector_size,
                                 this->parameters[0], product);
    }
    else
    {
        // Check parameter is set
        ASSERT((this->parameters != NULL), "Parameter is not set.");
 
        // Adding "parameter * B * input vector" to the product
        this->B.transpose_dot_plus(vector, this->parameters[0], product);
    }
}
 
 
// ===============================
// Explicit template instantiation
// ===============================
 
#if defined(USE_CUDA_FP8_E5M2) && (USE_CUDA_FP8_E5M2 == 1)
    template class cuDenseAffineMatrixFunction<__nv_fp8_e5m2>;
#endif
 
#if defined(USE_CUDA_FP8_E4M3) && (USE_CUDA_FP8_E4M3 == 1)
    template class cuDenseAffineMatrixFunction<__nv_fp8_e4m3>;
#endif
 
#if defined(USE_CUDA_FP16) && (USE_CUDA_FP16 == 1)
    template class cuDenseAffineMatrixFunction<__half>;
#endif
 
#if defined(USE_CUDA_BF16) && (USE_CUDA_BF16 == 1)
    template class cuDenseAffineMatrixFunction<__nv_bfloat16>;
#endif
 
#if defined(USE_CUDA_FP32) && (USE_CUDA_FP32 == 1)
    template class cuDenseAffineMatrixFunction<float>;
#endif
 
#if defined(USE_CUDA_FP64) && (USE_CUDA_FP64 == 1)
    template class cuDenseAffineMatrixFunction<double>;
#endif