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"
10using namespace KokkosBatched;
13namespace GMLS_LinearAlgebra {
15 template<
typename DeviceType,
17 typename MatrixViewType_A,
18 typename MatrixViewType_B,
19 typename MatrixViewType_X>
29 KOKKOS_INLINE_FUNCTION
34 const MatrixViewType_A &a,
35 const MatrixViewType_B &b,
36 const bool implicit_RHS)
40 template<
typename MemberType>
41 KOKKOS_INLINE_FUNCTION
44 const int k = member.league_rank();
51 _a.extent(1),
_a.extent(2));
53 _b.extent(1),
_b.extent(2));
55 _b.extent(1),
_b.extent(2));
58 if ((
size_t)
_M!=
_a.extent(1) || (
size_t)
_N!=
_a.extent(2)) {
62 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,0,
_M),[&](
const int &i) {
63 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member,0,
_N),[&](
const int &j) {
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) {
73 member.team_barrier();
77 if (std::is_same<typename MatrixViewType_B::array_layout, layout_left>::value) {
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);
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) {
101 bool do_print =
false;
103 Kokkos::single(Kokkos::PerTeam(member), [&] () {
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));
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));
128 member.team_barrier();
129 TeamVectorUTV<MemberType,AlgoTagType>
130 ::invoke(member, aa, pp, uu, vv, ww_fast, matrix_rank);
131 member.team_barrier();
134 Kokkos::single(Kokkos::PerTeam(member), [&] () {
135 printf(
"matrix_rank: %d\n", matrix_rank);
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));
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();
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() );
165 int scratch_size = scratch_matrix_right_type::shmem_size(
_N,
_N);
166 scratch_size += scratch_matrix_right_type::shmem_size(
_M,
_N );
167 scratch_size += scratch_vector_type::shmem_size(
_N*
_NRHS);
169 int l0_scratch_size = scratch_vector_type::shmem_size(
_N);
170 l0_scratch_size += scratch_vector_type::shmem_size(3*
_M);
179 Kokkos::Profiling::popRegion();
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) {
188 typedef Algo::UTV::Unblocked algo_tag_type;
189 typedef Kokkos::View<double***, A_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
191 typedef Kokkos::View<double***, B_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
193 typedef Kokkos::View<double***, X_layout, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
196 MatrixViewType_A mat_A(A, num_matrices, lda, nda);
197 MatrixViewType_B mat_B(B, num_matrices, ldb, ndb);
200 <
device_execution_space, algo_tag_type, MatrixViewType_A, MatrixViewType_B, MatrixViewType_X>(M,N,NRHS,mat_A,mat_B,implicit_RHS).run(pm);
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);
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)
void run(ParallelManager pm)
int _pm_getTeamScratchLevel_1
int _pm_getTeamScratchLevel_0
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)