cu_linear_operator.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_linear_operator.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 <cstddef>  // NULL
#include <cassert>  // assert
#include <cstdlib>  // abort
#include <iostream>
#include "../_cuda_utilities/cuda_api.h"  // CudaAPI
 
 
// =============
// constructor 1
// =============
 
 
template <typename DataType>
cuLinearOperator<DataType>::cuLinearOperator():
 
    // Initializer list
    num_gpu_devices(0),
    copied_host_to_device(false),
    cublas_handle(NULL),
    cusparse_handle(NULL)
{
    // Check any gpu device exists
    this->num_gpu_devices = this->query_gpu_devices();
 
    // Regardless of using dense (cublas) or sparse (cusparse) matrices, the
    // cublas handle should be initialized, since it is needed for the methods
    // in cuVectorOperations
    this->initialize_cublas_handle();
}
 
 
// =============
// constructor 2
// =============
 
 
template <typename DataType>
cuLinearOperator<DataType>::cuLinearOperator(
        const int num_gpu_devices_):
 
    // Initializer list
    num_gpu_devices(0),
    copied_host_to_device(false),
    cublas_handle(NULL),
    cusparse_handle(NULL)
{
    // Check any gpu device exists
    int device_count = this->query_gpu_devices();
 
    // Set number of gpu devices
    if (num_gpu_devices_ == 0)
    {
        this->num_gpu_devices = device_count;
    }
    else if (num_gpu_devices_ > device_count)
    {
        std::cerr << "ERROR: Number of requested gpu devices exceeds the " \
                  << "number of available gpu devices. Nummber of detected " \
                  << "devices are " << device_count << " while the " \
                  << "requested number of devices are " << num_gpu_devices_ \
                  << "." << std::endl;
        abort();
    }
    else
    {
        this->num_gpu_devices = num_gpu_devices_;
    }
 
    // Regardless of using dense (cublas) or sparse (cusparse) matrices, the
    // cublas handle should be initialized, since it is needed for the methods
    // in cuVectorOperations
    this->initialize_cublas_handle();
}
 
 
// ==========
// destructor
// ==========
 
 
template <typename DataType>
cuLinearOperator<DataType>::~cuLinearOperator()
{
    // cublas handle
    if (this->cublas_handle != NULL)
    {
        // Set the number of threads
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
            omp_set_num_threads(this->num_gpu_devices);
        #endif
 
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
        #pragma omp parallel
        #endif
        {
            // Switch to a device with the same device id as the cpu thread id
            unsigned int thread_id;
            #if defined(USE_OPENMP) && (USE_OPENMP == 1)
                thread_id = omp_get_thread_num();
            #else
                thread_id = 0;
            #endif
 
            CudaAPI<DataType>::set_device(thread_id);
 
            cublasStatus_t status = cublasDestroy(
                    this->cublas_handle[thread_id]);
            assert(status == CUBLAS_STATUS_SUCCESS);
        }
 
        // Deallocate arrays of pointers on cpu
        delete[] this->cublas_handle;
        this->cublas_handle = NULL;
    }
 
    // cusparse handle
    if (this->cusparse_handle != NULL)
    {
        // Set the number of threads
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
            omp_set_num_threads(this->num_gpu_devices);
        #endif
 
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
        #pragma omp parallel
        #endif
        {
            // Switch to a device with the same device id as the cpu thread id
            unsigned int thread_id;
            #if defined(USE_OPENMP) && (USE_OPENMP == 1)
                thread_id = omp_get_thread_num();
            #else
                thread_id = 0;
            #endif
 
            CudaAPI<DataType>::set_device(thread_id);
 
            cusparseStatus_t status = cusparseDestroy(
                    this->cusparse_handle[thread_id]);
            assert(status == CUSPARSE_STATUS_SUCCESS);
        }
 
        // Deallocate arrays of pointers on cpu
        delete[] this->cusparse_handle;
        this->cusparse_handle = NULL;
    }
}
 
 
// =================
// get cublas handle
// =================
 
 
template <typename DataType>
cublasHandle_t cuLinearOperator<DataType>::get_cublas_handle() const
{
    // Get device id
    int device_id = CudaAPI<DataType>::get_device();
 
    return this->cublas_handle[device_id];
}
 
 
// ========================
// initialize cublas handle
// ========================
 
 
template <typename DataType>
void cuLinearOperator<DataType>::initialize_cublas_handle()
{
    if (this->cublas_handle == NULL)
    {
        // Allocate pointers for each gpu device
        this->cublas_handle = new cublasHandle_t[this->num_gpu_devices];
 
        // Set the number of threads
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
            omp_set_num_threads(this->num_gpu_devices);
        #endif
 
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
        #pragma omp parallel
        #endif
        {
            // Switch to a device with the same device id as the cpu thread id
            unsigned int thread_id;
            #if defined(USE_OPENMP) && (USE_OPENMP == 1)
                thread_id = omp_get_thread_num();
            #else
                thread_id = 0;
            #endif
 
            CudaAPI<DataType>::set_device(thread_id);
 
            cublasStatus_t status_create = cublasCreate(
                    &this->cublas_handle[thread_id]);
            assert(status_create == CUBLAS_STATUS_SUCCESS);
 
            // Set tensor core whenever possible (usually for cublasXgemm)
            cublasStatus_t status_set = cublasSetMathMode(
                    this->cublas_handle[thread_id], CUBLAS_TENSOR_OP_MATH);
            assert(status_set == CUBLAS_STATUS_SUCCESS);
        }
    }
}
 
 
// ==========================
// initialize cusparse handle
// ==========================
 
 
template <typename DataType>
void cuLinearOperator<DataType>::initialize_cusparse_handle()
{
    if (this->cusparse_handle == NULL)
    {
        // Allocate pointers for each gpu device
        this->cusparse_handle = new cusparseHandle_t[this->num_gpu_devices];
 
        // Set the number of threads
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
            omp_set_num_threads(this->num_gpu_devices);
        #endif
 
        #if defined(USE_OPENMP) && (USE_OPENMP == 1)
        #pragma omp parallel
        #endif
        {
            // Switch to a device with the same device id as the cpu thread id
            unsigned int thread_id;
            #if defined(USE_OPENMP) && (USE_OPENMP == 1)
                thread_id = omp_get_thread_num();
            #else
                thread_id = 0;
            #endif
 
            CudaAPI<DataType>::set_device(thread_id);
 
            cusparseStatus_t status = cusparseCreate(
                    &this->cusparse_handle[thread_id]);
            assert(status == CUSPARSE_STATUS_SUCCESS);
        }
    }
}
 
 
// =================
// query gpu devices
// =================
 
 
template <typename DataType>
int cuLinearOperator<DataType>::query_gpu_devices() const
{
    int device_count = 0;
    cudaError_t error = cudaGetDeviceCount(&device_count);
 
    // Error code 38 means no cuda-capable device was detected.
    if ((error != cudaSuccess) || (device_count < 1))
    {
        std::cerr << "ERROR: No cuda-capable GPU device was detected on " \
                  << "this machine. If a cuda-capable GPU device exists, " \
                  << "install its cuda driver. Alternatively, set " \
                  << "'gpu=False' to use cpu instead." \
                  << std::endl;
 
        abort();
    }
 
    return device_count;
}
 
 
// ==============
// set parameters
// ==============
 
 
template <typename DataType>
void cuLinearOperator<DataType>::set_parameters(DataType* parameters_)
{
    this->parameters = parameters_;
}
 
 
// ===============================
// Explicit template instantiation
// ===============================
 
#if defined(USE_CUDA_FP8_E5M2) && (USE_CUDA_FP8_E5M2 == 1)
    template class cuLinearOperator<__nv_fp8_e5m2>;
#endif
 
#if defined(USE_CUDA_FP8_E4M3) && (USE_CUDA_FP8_E4M3 == 1)
    template class cuLinearOperator<__nv_fp8_e4m3>;
#endif
 
#if defined(USE_CUDA_FP16) && (USE_CUDA_FP16 == 1)
    template class cuLinearOperator<__half>;
#endif
 
#if defined(USE_CUDA_BF16) && (USE_CUDA_BF16 == 1)
    template class cuLinearOperator<__nv_bfloat16>;
#endif
 
#if defined(USE_CUDA_FP32) && (USE_CUDA_FP32 == 1)
    template class cuLinearOperator<float>;
#endif
 
#if defined(USE_CUDA_FP64) && (USE_CUDA_FP64 == 1)
    template class cuLinearOperator<double>;
#endif