matrix_utils.hpp
Go to the documentation of this file.
1 /*
2  * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
3  * SPDX-License-Identifier: Apache-2.0
4  */
5 #pragma once
6 #include "data.hpp"
7 #include "part_descriptor.hpp"
8 
9 #include <raft/core/handle.hpp>
10 #include <raft/random/rng.cuh>
11 
12 namespace MLCommon {
13 namespace Matrix {
14 namespace opg {
15 
16 void gatherPart(const raft::handle_t& h,
17  float* gatheredPart,
18  std::vector<Matrix::Data<float>*>& parts,
19  Matrix::PartDescriptor& desc,
20  int partIndex,
21  int rootRank,
22  int myRank,
23  cudaStream_t stream);
24 
25 void allGatherPart(const raft::handle_t& h,
26  float* gatheredPart,
27  std::vector<Matrix::Data<float>*>& parts,
28  Matrix::PartDescriptor& desc,
29  int partIndex,
30  int myRank,
31  cudaStream_t stream);
32 void gather(const raft::handle_t& h,
33  float* gatheredMatrix,
34  std::vector<Matrix::Data<float>*>& parts,
35  Matrix::PartDescriptor& desc,
36  int rootRank,
37  int myRank,
38  cudaStream_t stream);
39 
40 void allGather(const raft::handle_t& h,
41  float* gatheredMatrix,
42  std::vector<Matrix::Data<float>*>& parts,
43  Matrix::PartDescriptor& desc,
44  int myRank,
45  cudaStream_t stream);
46 
47 void allocate(const raft::handle_t& h,
48  std::vector<Matrix::Data<float>*>& parts,
49  Matrix::PartDescriptor& desc,
50  int myRank,
51  cudaStream_t stream);
52 
53 void deallocate(const raft::handle_t& h,
54  std::vector<Matrix::Data<float>*>& parts,
55  Matrix::PartDescriptor& desc,
56  int myRank,
57  cudaStream_t stream);
58 
59 void randomize(const raft::handle_t& h,
60  raft::random::Rng& r,
61  std::vector<Matrix::Data<float>*>& parts,
62  Matrix::PartDescriptor& desc,
63  int myRank,
64  cudaStream_t stream,
65  float low = -1.0f,
66  float high = 1.0f);
67 
68 void reset(const raft::handle_t& h,
69  std::vector<Matrix::Data<float>*>& parts,
70  Matrix::PartDescriptor& desc,
71  int myRank,
72  cudaStream_t stream);
73 
74 void printRaw2D(float* buffer, int rows, int cols, bool isColMajor, cudaStream_t stream);
75 
76 void print(const raft::handle_t& h,
77  std::vector<Matrix::Data<float>*>& parts,
78  Matrix::PartDescriptor& desc,
79  const char* matrixName,
80  int myRank,
81  cudaStream_t stream);
82 
83 //------------------------------------------------------------------------------
84 
85 void gatherPart(const raft::handle_t& h,
86  double* gatheredPart,
87  std::vector<Matrix::Data<double>*>& parts,
88  Matrix::PartDescriptor& desc,
89  int partIndex,
90  int rootRank,
91  int myRank,
92  cudaStream_t stream);
93 
94 void allGatherPart(const raft::handle_t& h,
95  double* gatheredPart,
96  std::vector<Matrix::Data<double>*>& parts,
97  Matrix::PartDescriptor& desc,
98  int partIndex,
99  int myRank,
100  cudaStream_t stream);
101 
102 void gather(const raft::handle_t& h,
103  double* gatheredMatrix,
104  std::vector<Matrix::Data<double>*>& parts,
105  Matrix::PartDescriptor& desc,
106  int rootRank,
107  int myRank,
108  cudaStream_t stream);
109 
110 void allGather(const raft::handle_t& h,
111  double* gatheredMatrix,
112  std::vector<Matrix::Data<double>*>& parts,
113  Matrix::PartDescriptor& desc,
114  int myRank,
115  cudaStream_t stream);
116 
117 void allocate(const raft::handle_t& h,
118  std::vector<Matrix::Data<double>*>& parts,
119  Matrix::PartDescriptor& desc,
120  int myRank,
121  cudaStream_t stream);
122 
123 void deallocate(const raft::handle_t& h,
124  std::vector<Matrix::Data<double>*>& parts,
125  Matrix::PartDescriptor& desc,
126  int myRank,
127  cudaStream_t stream);
128 
129 void randomize(const raft::handle_t& h,
130  raft::random::Rng& r,
131  std::vector<Matrix::Data<double>*>& parts,
132  Matrix::PartDescriptor& desc,
133  int myRank,
134  cudaStream_t stream,
135  double low = -1.0,
136  double high = 1.0);
137 
138 void reset(const raft::handle_t& h,
139  std::vector<Matrix::Data<double>*>& parts,
140  Matrix::PartDescriptor& desc,
141  int myRank,
142  cudaStream_t stream);
143 
144 void printRaw2D(double* buffer, int rows, int cols, bool isColMajor, cudaStream_t stream);
145 
146 void print(const raft::handle_t& h,
147  std::vector<Matrix::Data<double>*>& parts,
148  Matrix::PartDescriptor& desc,
149  const char* matrixName,
150  int myRank,
151  cudaStream_t stream);
152 } // end namespace opg
153 } // namespace Matrix
154 } // end namespace MLCommon
MLCommon::Matrix::opg::printRaw2D
void printRaw2D(float *buffer, int rows, int cols, bool isColMajor, cudaStream_t stream)
MLCommon::Matrix::opg::allGather
void allGather(const raft::handle_t &h, float *gatheredMatrix, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::gather
void gather(const raft::handle_t &h, float *gatheredMatrix, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int rootRank, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::deallocate
void deallocate(const raft::handle_t &h, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::allGatherPart
void allGatherPart(const raft::handle_t &h, float *gatheredPart, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int partIndex, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::randomize
void randomize(const raft::handle_t &h, raft::random::Rng &r, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int myRank, cudaStream_t stream, float low=-1.0f, float high=1.0f)
MLCommon::Matrix::opg::reset
void reset(const raft::handle_t &h, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::print
void print(const raft::handle_t &h, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, const char *matrixName, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::gatherPart
void gatherPart(const raft::handle_t &h, float *gatheredPart, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int partIndex, int rootRank, int myRank, cudaStream_t stream)
MLCommon::Matrix::opg::allocate
void allocate(const raft::handle_t &h, std::vector< Matrix::Data< float > * > &parts, Matrix::PartDescriptor &desc, int myRank, cudaStream_t stream)
MLCommon
Definition: comm_utils.h:11
part_descriptor.hpp
MLCommon::Matrix::Data
This is a helper wrapper around the multi-gpu data blocks owned by a worker. It's design is NOT final...
Definition: data.hpp:18
MLCommon::Matrix::PartDescriptor
Definition: part_descriptor.hpp:40
Generated by doxygen 1.9.1