mm_aTa.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 
7 #include "../matrix/data.hpp"
8 #include "../matrix/part_descriptor.hpp"
9 
10 #include <raft/core/comms.hpp>
11 #include <raft/core/handle.hpp>
12 
13 namespace MLCommon {
14 namespace LinAlg {
15 namespace opg {
16 
26 void mm_aTa(const raft::handle_t& handle,
27  Matrix::Data<double>& out,
28  const std::vector<Matrix::Data<double>*>& in1,
29  const Matrix::PartDescriptor& in1Desc,
30  cudaStream_t* streams,
31  int n_streams);
32 
33 void mm_aTa(const raft::handle_t& handle,
34  Matrix::Data<float>& out,
35  const std::vector<Matrix::Data<float>*>& in1,
36  const Matrix::PartDescriptor& in1Desc,
37  cudaStream_t* streams,
38  int n_streams);
39 
40 } // end namespace opg
41 } // end namespace LinAlg
42 } // end namespace MLCommon
MLCommon::LinAlg::opg::mm_aTa
void mm_aTa(const raft::handle_t &handle, Matrix::Data< double > &out, const std::vector< Matrix::Data< double > * > &in1, const Matrix::PartDescriptor &in1Desc, cudaStream_t *streams, int n_streams)
performs MNMG A^T x A calculation.
MLCommon
Definition: comm_utils.h:11
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