Compadre 1.5.5
Loading...
Searching...
No Matches
Compadre_LinearAlgebra.cpp
Go to the documentation of this file.
2
3#include "KokkosBatched_Copy_Decl.hpp"
4#include "KokkosBatched_ApplyPivot_Decl.hpp"
5#include "KokkosBatched_Gemv_Decl.hpp"
6#include "KokkosBatched_Trsv_Decl.hpp"
7#include "KokkosBatched_UTV_Decl.hpp"
8#include "KokkosBatched_SolveUTV_Decl_Compadre.hpp"
9
10using namespace KokkosBatched;
11
12namespace Compadre{
13namespace GMLS_LinearAlgebra {
14
15 template<typename DeviceType,
16 typename AlgoTagType,
17 typename MatrixViewType_A,
18 typename MatrixViewType_B,
19 typename MatrixViewType_X>
21 MatrixViewType_A _a;
22 MatrixViewType_B _b;
23
26 int _M, _N, _NRHS;
28
29 KOKKOS_INLINE_FUNCTION
31 const int M,
32 const int N,
33 const int NRHS,
34 const MatrixViewType_A &a,
35 const MatrixViewType_B &b,
36 const bool implicit_RHS)
37 : _a(a), _b(b), _M(M), _N(N), _NRHS(NRHS), _implicit_RHS(implicit_RHS)
39
40 template<typename MemberType>
41 KOKKOS_INLINE_FUNCTION
42 void operator()(const MemberType &member) const {
43
44 const int k = member.league_rank();
45
46 // workspace vectors
47 scratch_vector_type ww_fast(member.team_scratch(_pm_getTeamScratchLevel_0), 3*_M);
48 scratch_vector_type ww_slow(member.team_scratch(_pm_getTeamScratchLevel_1), _N*_NRHS);
49
50 scratch_matrix_right_type aa(_a.data() + TO_GLOBAL(k)*TO_GLOBAL(_a.extent(1))*TO_GLOBAL(_a.extent(2)),
51 _a.extent(1), _a.extent(2));
52 scratch_matrix_right_type bb(_b.data() + TO_GLOBAL(k)*TO_GLOBAL(_b.extent(1))*TO_GLOBAL(_b.extent(2)),
53 _b.extent(1), _b.extent(2));
54 scratch_matrix_right_type xx(_b.data() + TO_GLOBAL(k)*TO_GLOBAL(_b.extent(1))*TO_GLOBAL(_b.extent(2)),
55 _b.extent(1), _b.extent(2));
56
57 // if sizes don't match extents, then copy to a view with extents matching sizes
58 if ((size_t)_M!=_a.extent(1) || (size_t)_N!=_a.extent(2)) {
59 scratch_matrix_right_type tmp(ww_slow.data(), _M, _N);
60 auto aaa = scratch_matrix_right_type(_a.data() + TO_GLOBAL(k)*TO_GLOBAL(_a.extent(1))*TO_GLOBAL(_a.extent(2)), _M, _N);
61 // copy A to W, then back to A
62 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_M),[&](const int &i) {
63 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_N),[&](const int &j) {
64 tmp(i,j) = aa(i,j);
65 });
66 });
67 member.team_barrier();
68 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_M),[&](const int &i) {
69 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_N),[&](const int &j) {
70 aaa(i,j) = tmp(i,j);
71 });
72 });
73 member.team_barrier();
74 aa = aaa;
75 }
76
77 if (std::is_same<typename MatrixViewType_B::array_layout, layout_left>::value) {
78 scratch_matrix_right_type tmp(ww_slow.data(), _N, _NRHS);
79 // coming from LU
80 // then copy B to W, then back to B
81 auto bb_left =
82 scratch_matrix_left_type(_b.data() + TO_GLOBAL(k)*TO_GLOBAL(_b.extent(1))*TO_GLOBAL(_b.extent(2)),
83 _b.extent(1), _b.extent(2));
84 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_N),[&](const int &i) {
85 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_NRHS),[&](const int &j) {
86 tmp(i,j) = bb_left(i,j);
87 });
88 });
89 member.team_barrier();
90 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_N),[&](const int &i) {
91 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_NRHS),[&](const int &j) {
92 bb(i,j) = tmp(i,j);
93 });
94 });
95 }
96
97 scratch_matrix_right_type uu(member.team_scratch(_pm_getTeamScratchLevel_1), _M, _N /* only N columns of U are filled, maximum */);
100
101 bool do_print = false;
102 if (do_print) {
103 Kokkos::single(Kokkos::PerTeam(member), [&] () {
104 //print a
105 printf("a=zeros(%lu,%lu);\n", aa.extent(0), aa.extent(1));
106 for (size_t i=0; i<aa.extent(0); ++i) {
107 for (size_t j=0; j<aa.extent(1); ++j) {
108 printf("a(%lu,%lu)= %f;\n", i+1,j+1, aa(i,j));
109 }
110 }
111 //print b
112 printf("b=zeros(%lu,%lu);\n", bb.extent(0), bb.extent(1));
113 for (size_t i=0; i<bb.extent(0); ++i) {
114 for (size_t j=0; j<bb.extent(1); ++j) {
115 printf("b(%lu,%lu)= %f;\n", i+1,j+1, bb(i,j));
116 }
117 }
118 });
119 }
120 do_print = false;
121
122 /// Solving Ax = b using UTV transformation
123 /// A P^T P x = b
124 /// UTV P x = b;
125
126 /// UTV = A P^T
127 int matrix_rank(0);
128 member.team_barrier();
129 TeamVectorUTV<MemberType,AlgoTagType>
130 ::invoke(member, aa, pp, uu, vv, ww_fast, matrix_rank);
131 member.team_barrier();
132
133 if (do_print) {
134 Kokkos::single(Kokkos::PerTeam(member), [&] () {
135 printf("matrix_rank: %d\n", matrix_rank);
136 //print u
137 printf("u=zeros(%lu,%lu);\n", uu.extent(0), uu.extent(1));
138 for (size_t i=0; i<uu.extent(0); ++i) {
139 for (size_t j=0; j<uu.extent(1); ++j) {
140 printf("u(%lu,%lu)= %f;\n", i+1,j+1, uu(i,j));
141 }
142 }
143 });
144 }
145 TeamVectorSolveUTVCompadre<MemberType,AlgoTagType>
146 ::invoke(member, matrix_rank, _M, _N, _NRHS, uu, aa, vv, pp, bb, xx, ww_slow, ww_fast, _implicit_RHS);
147 member.team_barrier();
148
149 }
150
151 inline
153 typedef typename MatrixViewType_A::non_const_value_type value_type;
154 std::string name_region("KokkosBatched::Test::TeamVectorSolveUTVCompadre");
155 std::string name_value_type = ( std::is_same<value_type,float>::value ? "::Float" :
156 std::is_same<value_type,double>::value ? "::Double" :
157 std::is_same<value_type,Kokkos::complex<float> >::value ? "::ComplexFloat" :
158 std::is_same<value_type,Kokkos::complex<double> >::value ? "::ComplexDouble" : "::UnknownValueType" );
159 std::string name = name_region + name_value_type;
160 Kokkos::Profiling::pushRegion( name.c_str() );
161
164
165 int scratch_size = scratch_matrix_right_type::shmem_size(_N, _N); // V
166 scratch_size += scratch_matrix_right_type::shmem_size(_M, _N /* only N columns of U are filled, maximum */); // U
167 scratch_size += scratch_vector_type::shmem_size(_N*_NRHS); // W (for SolveUTV)
168
169 int l0_scratch_size = scratch_vector_type::shmem_size(_N); // P (temporary)
170 l0_scratch_size += scratch_vector_type::shmem_size(3*_M); // W (for UTV)
171
173 pm.setTeamScratchSize(0, l0_scratch_size);
174 pm.setTeamScratchSize(1, scratch_size);
175
176 pm.CallFunctorWithTeamThreadsAndVectors(*this, _a.extent(0));
177 Kokkos::fence();
178
179 Kokkos::Profiling::popRegion();
180 }
181 };
182
183
184
185template <typename A_layout, typename B_layout, typename X_layout>
186void batchQRPivotingSolve(ParallelManager pm, double *A, int lda, int nda, double *B, int ldb, int ndb, int M, int N, int NRHS, const int num_matrices, const bool implicit_RHS) {
187
188 typedef Algo::UTV::Unblocked algo_tag_type;
189 typedef Kokkos::View<double***, A_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
190 MatrixViewType_A;
191 typedef Kokkos::View<double***, B_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
192 MatrixViewType_B;
193 typedef Kokkos::View<double***, X_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
194 MatrixViewType_X;
195
196 MatrixViewType_A mat_A(A, num_matrices, lda, nda);
197 MatrixViewType_B mat_B(B, num_matrices, ldb, ndb);
198
200 <device_execution_space, algo_tag_type, MatrixViewType_A, MatrixViewType_B, MatrixViewType_X>(M,N,NRHS,mat_A,mat_B,implicit_RHS).run(pm);
201
202}
203
204template void batchQRPivotingSolve<layout_right, layout_right, layout_right>(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
205template void batchQRPivotingSolve<layout_right, layout_right, layout_left >(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
206template void batchQRPivotingSolve<layout_right, layout_left , layout_right>(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
207template void batchQRPivotingSolve<layout_right, layout_left , layout_left >(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
208template void batchQRPivotingSolve<layout_left , layout_right, layout_right>(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
209template void batchQRPivotingSolve<layout_left , layout_right, layout_left >(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
210template void batchQRPivotingSolve<layout_left , layout_left , layout_right>(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
211template void batchQRPivotingSolve<layout_left , layout_left , layout_left >(ParallelManager,double*,int,int,double*,int,int,int,int,int,const int,const bool);
212
213} // GMLS_LinearAlgebra
214} // Compadre
Kokkos::DefaultExecutionSpace device_execution_space
Kokkos::View< double **, layout_left, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_matrix_left_type
Kokkos::View< double *, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_vector_type
#define TO_GLOBAL(variable)
Kokkos::View< int *, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_local_index_type
Kokkos::View< double **, layout_right, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_matrix_right_type
void setTeamScratchSize(const int level, const int value)
void CallFunctorWithTeamThreadsAndVectors(C functor, const global_index_type batch_size, const int threads_per_team=-1, const int vector_lanes_per_thread=-1) const
Calls a parallel_for parallel_for will break out over loops over teams with each vector lane executin...
KOKKOS_INLINE_FUNCTION int getTeamScratchLevel(const int level) const
template void batchQRPivotingSolve< layout_left, layout_right, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
template void batchQRPivotingSolve< layout_left, layout_left, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
void batchQRPivotingSolve(ParallelManager pm, double *A, int lda, int nda, double *B, int ldb, int ndb, int M, int N, int NRHS, const int num_matrices, const bool implicit_RHS)
Solves a batch of problems with QR+Pivoting.
template void batchQRPivotingSolve< layout_left, layout_right, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
template void batchQRPivotingSolve< layout_left, layout_left, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
template void batchQRPivotingSolve< layout_right, layout_left, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
template void batchQRPivotingSolve< layout_right, layout_right, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
template void batchQRPivotingSolve< layout_right, layout_left, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
template void batchQRPivotingSolve< layout_right, layout_right, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int, const bool)
KOKKOS_INLINE_FUNCTION void operator()(const MemberType &member) const
KOKKOS_INLINE_FUNCTION Functor_TestBatchedTeamVectorSolveUTV(const int M, const int N, const int NRHS, const MatrixViewType_A &a, const MatrixViewType_B &b, const bool implicit_RHS)