4 #include "KokkosBatched_Copy_Decl.hpp" 5 #include "KokkosBatched_ApplyPivot_Decl.hpp" 6 #include "KokkosBatched_Gemv_Decl.hpp" 7 #include "KokkosBatched_Trsv_Decl.hpp" 8 #include "KokkosBatched_UTV_Decl.hpp" 9 #include "KokkosBatched_SolveUTV_Decl_Compadre.hpp" 14 namespace GMLS_LinearAlgebra {
16 template<
typename DeviceType,
18 typename MatrixViewType_A,
19 typename MatrixViewType_B,
20 typename MatrixViewType_X>
29 KOKKOS_INLINE_FUNCTION
34 const MatrixViewType_A &a,
35 const MatrixViewType_B &b)
36 : _a(a), _b(b), _M(M), _N(N), _NRHS(NRHS) { _pm_getTeamScratchLevel_0 = 0; _pm_getTeamScratchLevel_1 = 0; }
38 template<
typename MemberType>
39 KOKKOS_INLINE_FUNCTION
42 const int k = member.league_rank();
49 _a.extent(1), _a.extent(2));
51 _b.extent(1), _b.extent(2));
53 _b.extent(1), _b.extent(2));
56 if ((
size_t)_M!=_a.extent(1) || (size_t)_N!=_a.extent(2)) {
60 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_M),[&](
const int &i) {
61 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_N),[&](
const int &j) {
65 member.team_barrier();
66 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_M),[&](
const int &i) {
67 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_N),[&](
const int &j) {
71 member.team_barrier();
75 if (std::is_same<typename MatrixViewType_B::array_layout, layout_left>::value) {
81 _b.extent(1), _b.extent(2));
82 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_N),[&](
const int &i) {
83 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_NRHS),[&](
const int &j) {
84 tmp(i,j) = bb_left(i,j);
87 member.team_barrier();
88 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,_N),[&](
const int &i) {
89 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,_NRHS),[&](
const int &j) {
99 bool do_print =
false;
101 Kokkos::single(Kokkos::PerTeam(member), [&] () {
103 printf(
"a=zeros(%lu,%lu);\n", aa.extent(0), aa.extent(1));
104 for (
size_t j=0; j<aa.extent(0); ++j) {
105 for (
size_t k=0; k<aa.extent(1); ++k) {
106 printf(
"a(%lu,%lu)= %f;\n", j+1,k+1, aa(j,k));
110 printf(
"b=zeros(%lu,%lu);\n", bb.extent(0), bb.extent(1));
111 for (
size_t j=0; j<bb.extent(0); ++j) {
112 for (
size_t k=0; k<bb.extent(1); ++k) {
113 printf(
"b(%lu,%lu)= %f;\n", j+1,k+1, bb(j,k));
126 member.team_barrier();
127 TeamVectorUTV<MemberType,AlgoTagType>
128 ::invoke(member, aa, pp, uu, vv, ww_fast, matrix_rank);
129 member.team_barrier();
132 Kokkos::single(Kokkos::PerTeam(member), [&] () {
133 printf(
"matrix_rank: %d\n", matrix_rank);
135 printf(
"u=zeros(%lu,%lu);\n", uu.extent(0), uu.extent(1));
136 for (
size_t j=0; j<uu.extent(0); ++j) {
137 for (
size_t k=0; k<uu.extent(1); ++k) {
138 printf(
"u(%lu,%lu)= %f;\n", j+1,k+1, uu(j,k));
143 TeamVectorSolveUTVCompadre<MemberType,AlgoTagType>
144 ::invoke(member, matrix_rank, _M, _N, _NRHS, uu, aa, vv, pp, bb, xx, ww_slow, ww_fast);
145 member.team_barrier();
151 typedef typename MatrixViewType_A::non_const_value_type value_type;
152 std::string name_region(
"KokkosBatched::Test::TeamVectorSolveUTVCompadre");
153 std::string name_value_type = ( std::is_same<value_type,float>::value ?
"::Float" :
154 std::is_same<value_type,double>::value ?
"::Double" :
155 std::is_same<value_type,Kokkos::complex<float> >::value ?
"::ComplexFloat" :
156 std::is_same<value_type,Kokkos::complex<double> >::value ?
"::ComplexDouble" :
"::UnknownValueType" );
157 std::string name = name_region + name_value_type;
158 Kokkos::Profiling::pushRegion( name.c_str() );
163 int scratch_size = scratch_matrix_right_type::shmem_size(_N, _N);
164 scratch_size += scratch_matrix_right_type::shmem_size(_M, _N );
165 scratch_size += scratch_vector_type::shmem_size(_N*_NRHS);
167 int l0_scratch_size = scratch_vector_type::shmem_size(_N);
168 l0_scratch_size += scratch_vector_type::shmem_size(3*_M);
177 Kokkos::Profiling::popRegion();
183 template <
typename A_layout,
typename B_layout,
typename X_layout>
184 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) {
186 typedef Algo::UTV::Unblocked algo_tag_type;
187 typedef Kokkos::View<double***, A_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
189 typedef Kokkos::View<double***, B_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
191 typedef Kokkos::View<double***, X_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
194 MatrixViewType_A mat_A(A, num_matrices, lda, nda);
195 MatrixViewType_B mat_B(B, num_matrices, ldb, ndb);
198 <
device_execution_space, algo_tag_type, MatrixViewType_A, MatrixViewType_B, MatrixViewType_X>(M,N,NRHS,mat_A,mat_B).run(pm);
202 template void batchQRPivotingSolve<layout_right, layout_right, layout_right>(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
203 template void batchQRPivotingSolve<layout_right, layout_right, layout_left >(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
204 template void batchQRPivotingSolve<layout_right, layout_left , layout_right>(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
205 template void batchQRPivotingSolve<layout_right, layout_left , layout_left >(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
206 template void batchQRPivotingSolve<layout_left , layout_right, layout_right>(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
207 template void batchQRPivotingSolve<layout_left , layout_right, layout_left >(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
208 template void batchQRPivotingSolve<layout_left , layout_left , layout_right>(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
209 template void batchQRPivotingSolve<layout_left , layout_left , layout_left >(
ParallelManager,
double*,int,int,
double*,int,int,int,int,int,
const int);
template void batchQRPivotingSolve< layout_right, layout_left, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
template void batchQRPivotingSolve< layout_left, layout_left, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
KOKKOS_INLINE_FUNCTION int getTeamScratchLevel(const int level) const
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...
int _pm_getTeamScratchLevel_0
template void batchQRPivotingSolve< layout_right, layout_right, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
int _pm_getTeamScratchLevel_1
template void batchQRPivotingSolve< layout_right, layout_right, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
template void batchQRPivotingSolve< layout_left, layout_right, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
KOKKOS_INLINE_FUNCTION void operator()(const MemberType &member) const
Kokkos::DefaultExecutionSpace device_execution_space
#define TO_GLOBAL(variable)
void run(ParallelManager pm)
template void batchQRPivotingSolve< layout_left, layout_right, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
Kokkos::View< double **, layout_right, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_matrix_right_type
Kokkos::View< double *, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_vector_type
KOKKOS_INLINE_FUNCTION Functor_TestBatchedTeamVectorSolveUTV(const int M, const int N, const int NRHS, const MatrixViewType_A &a, const MatrixViewType_B &b)
Kokkos::View< double **, layout_left, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_matrix_left_type
template void batchQRPivotingSolve< layout_left, layout_left, layout_right >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
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)
Solves a batch of problems with QR+Pivoting.
void setTeamScratchSize(const int level, const int value)
template void batchQRPivotingSolve< layout_right, layout_left, layout_left >(ParallelManager, double *, int, int, double *, int, int, int, int, int, const int)
Kokkos::View< int *, Kokkos::MemoryTraits< Kokkos::Unmanaged > > scratch_local_index_type