para_gemm.h

Go to the documentation of this file.

#ifndef PARA_GEMM_H
#define PARA_GEMM_H
#include "source_base/module_device/device.h"
#include "source_base/module_device/memory_op.h"
 
#include <vector>
#ifdef __MPI
#include "mpi.h"
#endif
 
namespace ModuleBase
{
template <typename T, typename Device = base_device::DEVICE_CPU>
class PGemmCN
{
  public:
    PGemmCN();
    ~PGemmCN();
 
    void set_dimension(
#ifdef __MPI
        MPI_Comm comm_col,
        MPI_Comm comm_row,
#endif
        const int ncolA,
        const int LDA,
        const int ncolB,
        const int LDB,
        const int nrow,
        const int LDC,
        const int mode = 1);
 
    void multiply(const T alpha, const T* A, const T* B, const T beta, T* C);
#ifdef __MPI
    MPI_Comm col_world = MPI_COMM_NULL; 
    MPI_Comm row_world = MPI_COMM_NULL; 
 
    int col_rank = 0;  
    int col_nproc = 1; 
    int row_rank = 0;  
    int row_nproc = 1; 
 
    std::vector<int> colA_loc; 
    int max_colA = 0;          
    std::vector<int> colB_loc; 
    int max_colB = 0;          
 
    std::vector<MPI_Request> requests; 
    std::vector<int> recv_counts;      
    std::vector<int> displs;           
    int size_C_local = 0;              
    int size_C_global = 0;             
    bool gatherC = true;               
    bool divideCrow = false;           
#endif
    int ncolA = 0; 
    int ncolB = 0; 
    int nrow = 0;  
    int LDA = 0;   
    int LDB = 0;   
    int LDC = 0;   
  private:
    void multiply_single(const T alpha, const T* A, const T* B, const T beta, T* C);
#ifdef __MPI
    void multiply_col(const T alpha, const T* A, const T* B, const T beta, T* C);
    void multiply_row(const T alpha, const T* A, const T* B, const T beta, T* C);
#endif
    using resmem_dev_op = base_device::memory::resize_memory_op<T, Device>;
    using delmem_dev_op = base_device::memory::delete_memory_op<T, Device>;
    using syncmem_dev_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
    using syncmem_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
    using syncmem_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
 
#ifdef __MPI
  private:
    std::vector<T> isend_tmp_; 
    std::vector<T> A_tmp_;     
    std::vector<T> B_tmp_;     
    std::vector<T> C_tmp_;     
    std::vector<T> C_global_tmp_; 
    T* C_local_tmp_ = nullptr;    
    T* A_tmp_device_ = nullptr;   
    T* B_tmp_device_ = nullptr;   
#endif
 
 
};
} // namespace ModuleBase
#endif