Skip to content

Commit 3c84a52

Browse files
ndellingwoodndellingwood
committed
ifpack2: update ArithTraits namespace for KOKKOS_VERSION > 40799
Compatibility update corresponding to kokkos/kokkos-kernels#2771 Signed-off-by: Nathan Ellingwood <ndellin@sandia.gov>
1 parent f9a986f commit 3c84a52

28 files changed

+529
-1
lines changed

‎packages/ifpack2/example/RelaxationWithEquilibration.cpp‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1091,7 +1091,11 @@ template <class ViewType1,
10911091
const bool takeSquareRootsOfScalingFactors,
10921092
const bool takeAbsoluteValueOfScalingFactors =
10931093
!std::is_same<
1094+
#if KOKKOS_VERSION > 40799
1095+
typename KokkosKernels::ArithTraits<
1096+
#else
10941097
typename Kokkos::ArithTraits<
1098+
#endif
10951099
typename ViewType1::non_const_value_type>::mag_type,
10961100
typename ViewType2::non_const_value_type>::value,
10971101
const int rank = ViewType1::rank>
@@ -1120,9 +1124,17 @@ class ElementWiseMultiply<ViewType1,
11201124

11211125
KOKKOS_INLINE_FUNCTION void operator()(const IndexType i) const {
11221126
using val_type = typename ViewType2::non_const_value_type;
1127+
#if KOKKOS_VERSION > 40799
1128+
using KAT = KokkosKernels::ArithTraits<val_type>;
1129+
#else
11231130
using KAT = Kokkos::ArithTraits<val_type>;
1131+
#endif
11241132
using mag_type = typename KAT::mag_type;
1133+
#if KOKKOS_VERSION > 40799
1134+
using KAM = KokkosKernels::ArithTraits<mag_type>;
1135+
#else
11251136
using KAM = Kokkos::ArithTraits<mag_type>;
1137+
#endif
11261138

11271139
if (takeAbsoluteValueOfScalingFactors) {
11281140
const mag_type scalFactAbs = KAT::abs(scalingFactors_(i));
@@ -1163,9 +1175,17 @@ class ElementWiseMultiply<ViewType1,
11631175

11641176
KOKKOS_INLINE_FUNCTION void operator()(const IndexType i) const {
11651177
using val_type = typename ViewType2::non_const_value_type;
1178+
#if KOKKOS_VERSION > 40799
1179+
using KAT = KokkosKernels::ArithTraits<val_type>;
1180+
#else
11661181
using KAT = Kokkos::ArithTraits<val_type>;
1182+
#endif
11671183
using mag_type = typename KAT::mag_type;
1184+
#if KOKKOS_VERSION > 40799
1185+
using KAM = KokkosKernels::ArithTraits<mag_type>;
1186+
#else
11681187
using KAM = Kokkos::ArithTraits<mag_type>;
1188+
#endif
11691189

11701190
for (IndexType j = 0; j < static_cast<IndexType>(X_.extent(1)); ++j) {
11711191
if (takeAbsoluteValueOfScalingFactors) {
@@ -1194,7 +1214,11 @@ void elementWiseMultiply(const MultiVectorViewType& X,
11941214
const bool takeSquareRootsOfScalingFactors,
11951215
const bool takeAbsoluteValueOfScalingFactors =
11961216
!std::is_same<
1217+
#if KOKKOS_VERSION > 40799
1218+
typename KokkosKernels::ArithTraits<
1219+
#else
11971220
typename Kokkos::ArithTraits<
1221+
#endif
11981222
typename MultiVectorViewType::non_const_value_type>::mag_type,
11991223
typename ScalingFactorsViewType::non_const_value_type>::value) {
12001224
using execution_space = typename MultiVectorViewType::device_type::execution_space;
@@ -1243,7 +1267,11 @@ void elementWiseMultiplyMultiVector(MultiVectorType& X,
12431267
const bool takeSquareRootsOfScalingFactors,
12441268
const bool takeAbsoluteValueOfScalingFactors =
12451269
!std::is_same<
1270+
#if KOKKOS_VERSION > 40799
1271+
typename KokkosKernels::ArithTraits<
1272+
#else
12461273
typename Kokkos::ArithTraits<
1274+
#endif
12471275
typename MultiVectorType::scalar_type>::mag_type,
12481276
typename ScalingFactorsViewType::non_const_value_type>::value) {
12491277
using index_type = typename MultiVectorType::local_ordinal_type;
@@ -1269,7 +1297,11 @@ template <class ViewType1,
12691297
const bool takeSquareRootsOfScalingFactors,
12701298
const bool takeAbsoluteValueOfScalingFactors =
12711299
!std::is_same<
1300+
#if KOKKOS_VERSION > 40799
1301+
typename KokkosKernels::ArithTraits<
1302+
#else
12721303
typename Kokkos::ArithTraits<
1304+
#endif
12731305
typename ViewType1::non_const_value_type>::mag_type,
12741306
typename ViewType2::non_const_value_type>::value,
12751307
const int rank = ViewType1::rank>
@@ -1298,9 +1330,17 @@ class ElementWiseDivide<ViewType1,
12981330

12991331
KOKKOS_INLINE_FUNCTION void operator()(const IndexType i) const {
13001332
using val_type = typename ViewType2::non_const_value_type;
1333+
#if KOKKOS_VERSION > 40799
1334+
using KAT = KokkosKernels::ArithTraits<val_type>;
1335+
#else
13011336
using KAT = Kokkos::ArithTraits<val_type>;
1337+
#endif
13021338
using mag_type = typename KAT::mag_type;
1339+
#if KOKKOS_VERSION > 40799
1340+
using KAM = KokkosKernels::ArithTraits<mag_type>;
1341+
#else
13031342
using KAM = Kokkos::ArithTraits<mag_type>;
1343+
#endif
13041344

13051345
if (takeAbsoluteValueOfScalingFactors) {
13061346
const mag_type scalFactAbs = KAT::abs(scalingFactors_(i));
@@ -1341,9 +1381,17 @@ class ElementWiseDivide<ViewType1,
13411381

13421382
KOKKOS_INLINE_FUNCTION void operator()(const IndexType i) const {
13431383
using val_type = typename ViewType2::non_const_value_type;
1384+
#if KOKKOS_VERSION > 40799
1385+
using KAT = KokkosKernels::ArithTraits<val_type>;
1386+
#else
13441387
using KAT = Kokkos::ArithTraits<val_type>;
1388+
#endif
13451389
using mag_type = typename KAT::mag_type;
1390+
#if KOKKOS_VERSION > 40799
1391+
using KAM = KokkosKernels::ArithTraits<mag_type>;
1392+
#else
13461393
using KAM = Kokkos::ArithTraits<mag_type>;
1394+
#endif
13471395

13481396
for (IndexType j = 0; j < static_cast<IndexType>(X_.extent(1)); ++j) {
13491397
if (takeAbsoluteValueOfScalingFactors) {
@@ -1372,7 +1420,11 @@ void elementWiseDivide(const MultiVectorViewType& X,
13721420
const bool takeSquareRootsOfScalingFactors,
13731421
const bool takeAbsoluteValueOfScalingFactors =
13741422
!std::is_same<
1423+
#if KOKKOS_VERSION > 40799
1424+
typename KokkosKernels::ArithTraits<
1425+
#else
13751426
typename Kokkos::ArithTraits<
1427+
#endif
13761428
typename MultiVectorViewType::non_const_value_type>::mag_type,
13771429
typename ScalingFactorsViewType::non_const_value_type>::value) {
13781430
using execution_space = typename MultiVectorViewType::device_type::execution_space;
@@ -1421,7 +1473,11 @@ void elementWiseDivideMultiVector(MultiVectorType& X,
14211473
const bool takeSquareRootsOfScalingFactors,
14221474
const bool takeAbsoluteValueOfScalingFactors =
14231475
!std::is_same<
1476+
#if KOKKOS_VERSION > 40799
1477+
typename KokkosKernels::ArithTraits<
1478+
#else
14241479
typename Kokkos::ArithTraits<
1480+
#endif
14251481
typename MultiVectorType::scalar_type>::mag_type,
14261482
typename ScalingFactorsViewType::non_const_value_type>::value) {
14271483
using index_type = typename MultiVectorType::local_ordinal_type;
@@ -1670,7 +1726,11 @@ class BelosIfpack2Solver {
16701726
equibResult_ = computeRowAndColumnOneNorms(*A_, assumeSymmetric_);
16711727
if (useDiagonalToEquilibrate_) {
16721728
using device_type = typename node_type::device_type;
1729+
#if KOKKOS_VERSION > 40799
1730+
using mag_type = typename KokkosKernels::ArithTraits<scalar_type>::mag_type;
1731+
#else
16731732
using mag_type = typename Kokkos::ArithTraits<scalar_type>::mag_type;
1733+
#endif
16741734
using view_type = Kokkos::View<mag_type*, device_type>;
16751735

16761736
view_type rowDiagAbsVals("rowDiagAbsVals",
@@ -1968,7 +2028,11 @@ void solveAndReport(BelosIfpack2Solver<CrsMatrixType>& solver,
19682028
using mag_type = typename MultiVectorType::mag_type;
19692029
Teuchos::Array<mag_type> norms(R.getNumVectors());
19702030
R.norm2(norms());
2031+
#if KOKKOS_VERSION > 40799
2032+
mag_type B_norm2_max = KokkosKernels::ArithTraits<mag_type>::zero();
2033+
#else
19712034
mag_type B_norm2_max = Kokkos::ArithTraits<mag_type>::zero();
2035+
#endif
19722036
for (std::size_t j = 0; j < B.getNumVectors(); ++j) {
19732037
// Any NaN will persist (since the first test will fail);
19742038
// this is what we want
@@ -1984,15 +2048,23 @@ void solveAndReport(BelosIfpack2Solver<CrsMatrixType>& solver,
19842048
A_original.apply(X, R, Teuchos::NO_TRANS, -ONE, ONE); // R := -A*X + B
19852049
R.norm2(norms());
19862050

2051+
#if KOKKOS_VERSION > 40799
2052+
mag_type R_norm2_max = KokkosKernels::ArithTraits<mag_type>::zero();
2053+
#else
19872054
mag_type R_norm2_max = Kokkos::ArithTraits<mag_type>::zero();
2055+
#endif
19882056
for (std::size_t j = 0; j < R.getNumVectors(); ++j) {
19892057
// Any NaN will persist (since the first test will fail);
19902058
// this is what we want
19912059
R_norm2_max = norms[j] < R_norm2_max ? R_norm2_max : norms[j];
19922060
}
19932061

19942062
X.norm2(norms());
2063+
#if KOKKOS_VERSION > 40799
2064+
mag_type X_norm2_max = KokkosKernels::ArithTraits<mag_type>::zero();
2065+
#else
19952066
mag_type X_norm2_max = Kokkos::ArithTraits<mag_type>::zero();
2067+
#endif
19962068
for (std::size_t j = 0; j < R.getNumVectors(); ++j) {
19972069
// Any NaN will persist (since the first test will fail);
19982070
// this is what we want
@@ -2020,7 +2092,11 @@ void solveAndReport(BelosIfpack2Solver<CrsMatrixType>& solver,
20202092
<< " ||B-A*X||_2: " << R_norm2_max << endl
20212093
<< " ||B||_2: " << B_norm2_max << endl
20222094
<< " ||X||_2: " << X_norm2_max << endl;
2095+
#if KOKKOS_VERSION > 40799
2096+
if (B_norm2_max != KokkosKernels::ArithTraits<mag_type>::zero()) {
2097+
#else
20232098
if (B_norm2_max != Kokkos::ArithTraits<mag_type>::zero()) {
2099+
#endif
20242100
cout << " ||B-A*X||_2 / ||B||_2: " << (R_norm2_max / B_norm2_max)
20252101
<< endl;
20262102
}

‎packages/ifpack2/src/Ifpack2_BlockComputeResidualAndSolve.hpp‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,11 @@ struct ComputeResidualAndSolve_1Pass {
149149
Kokkos::TeamThreadRange(member, 0, numEntries), [&](const int k) {
150150
int64_t A_offset = A_x_offsets(rowidx, 0, k);
151151
int64_t x_offset = A_x_offsets(rowidx, 1, k);
152+
#if KOKKOS_VERSION > 40799
153+
if (A_offset != KokkosKernels::ArithTraits<int64_t>::min()) {
154+
#else
152155
if (A_offset != Kokkos::ArithTraits<int64_t>::min()) {
156+
#endif
153157
A_block_cst.assign_data(tpetra_values.data() + A_offset);
154158
// Pull x into local memory
155159
int64_t remote_cutoff = blocksize * num_local_rows;
@@ -195,9 +199,17 @@ struct ComputeResidualAndSolve_1Pass {
195199
// entry.
196200
impl_scalar_type old_y = x(row + k, col);
197201
impl_scalar_type y_update = local_Dinv_residual[k] - old_y;
202+
#if KOKKOS_VERSION > 40799
203+
if constexpr (KokkosKernels::ArithTraits<impl_scalar_type>::is_complex) {
204+
#else
198205
if constexpr (Kokkos::ArithTraits<impl_scalar_type>::is_complex) {
206+
#endif
199207
magnitude_type ydiff =
208+
#if KOKKOS_VERSION > 40799
209+
KokkosKernels::ArithTraits<impl_scalar_type>::abs(y_update);
210+
#else
200211
Kokkos::ArithTraits<impl_scalar_type>::abs(y_update);
212+
#endif
201213
update += ydiff * ydiff;
202214
} else {
203215
update += y_update * y_update;
@@ -354,7 +366,11 @@ struct ComputeResidualAndSolve_2Pass {
354366
Kokkos::TeamThreadRange(member, 0, numEntries), [&](const int k) {
355367
int64_t A_offset = A_x_offsets(rowidx, 0, k);
356368
int64_t x_offset = A_x_offsets(rowidx, 1, k);
369+
#if KOKKOS_VERSION > 40799
370+
if (A_offset != KokkosKernels::ArithTraits<int64_t>::min()) {
371+
#else
357372
if (A_offset != Kokkos::ArithTraits<int64_t>::min()) {
373+
#endif
358374
A_block_cst.assign_data(tpetra_values.data() + A_offset);
359375
// Pull x into local memory
360376
Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, blocksize),
@@ -420,7 +436,11 @@ struct ComputeResidualAndSolve_2Pass {
420436
Kokkos::TeamThreadRange(member, 0, numEntries), [&](const int k) {
421437
int64_t A_offset = A_x_offsets_remote(rowidx, 0, k);
422438
int64_t x_offset = A_x_offsets_remote(rowidx, 1, k);
439+
#if KOKKOS_VERSION > 40799
440+
if (A_offset != KokkosKernels::ArithTraits<int64_t>::min()) {
441+
#else
423442
if (A_offset != Kokkos::ArithTraits<int64_t>::min()) {
443+
#endif
424444
A_block_cst.assign_data(tpetra_values.data() + A_offset);
425445
// Pull x into local memory
426446
Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, blocksize),
@@ -461,9 +481,17 @@ struct ComputeResidualAndSolve_2Pass {
461481
// entry.
462482
impl_scalar_type old_y = x(row + k, col);
463483
impl_scalar_type y_update = local_Dinv_residual[k] - old_y;
484+
#if KOKKOS_VERSION > 40799
485+
if constexpr (KokkosKernels::ArithTraits<impl_scalar_type>::is_complex) {
486+
#else
464487
if constexpr (Kokkos::ArithTraits<impl_scalar_type>::is_complex) {
488+
#endif
465489
magnitude_type ydiff =
490+
#if KOKKOS_VERSION > 40799
491+
KokkosKernels::ArithTraits<impl_scalar_type>::abs(y_update);
492+
#else
466493
Kokkos::ArithTraits<impl_scalar_type>::abs(y_update);
494+
#endif
467495
update += ydiff * ydiff;
468496
} else {
469497
update += y_update * y_update;
@@ -646,9 +674,17 @@ struct ComputeResidualAndSolve_SolveOnly {
646674
// Compute the change in y (assuming damping_factor == 1) for this
647675
// entry.
648676
impl_scalar_type y_update = local_Dinv_residual[k];
677+
#if KOKKOS_VERSION > 40799
678+
if constexpr (KokkosKernels::ArithTraits<impl_scalar_type>::is_complex) {
679+
#else
649680
if constexpr (Kokkos::ArithTraits<impl_scalar_type>::is_complex) {
681+
#endif
650682
magnitude_type ydiff =
683+
#if KOKKOS_VERSION > 40799
684+
KokkosKernels::ArithTraits<impl_scalar_type>::abs(y_update);
685+
#else
651686
Kokkos::ArithTraits<impl_scalar_type>::abs(y_update);
687+
#endif
652688
update += ydiff * ydiff;
653689
} else {
654690
update += y_update * y_update;

‎packages/ifpack2/src/Ifpack2_BlockComputeResidualVector.hpp‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,8 +191,16 @@ void precompute_A_x_offsets(
191191
A_x_offsets(i, 0, entry) = int64_t(j) * blocksize_square;
192192
A_x_offsets(i, 1, entry) = int64_t(loc) * blocksize;
193193
} else {
194+
#if KOKKOS_VERSION > 40799
195+
A_x_offsets(i, 0, entry) = KokkosKernels::ArithTraits<int64_t>::min();
196+
#else
194197
A_x_offsets(i, 0, entry) = Kokkos::ArithTraits<int64_t>::min();
198+
#endif
199+
#if KOKKOS_VERSION > 40799
200+
A_x_offsets(i, 1, entry) = KokkosKernels::ArithTraits<int64_t>::min();
201+
#else
195202
A_x_offsets(i, 1, entry) = Kokkos::ArithTraits<int64_t>::min();
203+
#endif
196204
}
197205
}
198206
// Nonowned entries
@@ -207,8 +215,16 @@ void precompute_A_x_offsets(
207215
A_x_offsets_remote(i, 0, entry) = int64_t(j) * blocksize_square;
208216
A_x_offsets_remote(i, 1, entry) = int64_t(loc) * blocksize;
209217
} else {
218+
#if KOKKOS_VERSION > 40799
219+
A_x_offsets_remote(i, 0, entry) = KokkosKernels::ArithTraits<int64_t>::min();
220+
#else
210221
A_x_offsets_remote(i, 0, entry) = Kokkos::ArithTraits<int64_t>::min();
222+
#endif
223+
#if KOKKOS_VERSION > 40799
224+
A_x_offsets_remote(i, 1, entry) = KokkosKernels::ArithTraits<int64_t>::min();
225+
#else
211226
A_x_offsets_remote(i, 1, entry) = Kokkos::ArithTraits<int64_t>::min();
227+
#endif
212228
}
213229
}
214230
}
@@ -249,8 +265,16 @@ void precompute_A_x_offsets(
249265
A_x_offsets(i, 1, entry) = int64_t(A_colind_at_j) * blocksize;
250266
}
251267
} else {
268+
#if KOKKOS_VERSION > 40799
269+
A_x_offsets(i, 0, entry) = KokkosKernels::ArithTraits<int64_t>::min();
270+
#else
252271
A_x_offsets(i, 0, entry) = Kokkos::ArithTraits<int64_t>::min();
272+
#endif
273+
#if KOKKOS_VERSION > 40799
274+
A_x_offsets(i, 1, entry) = KokkosKernels::ArithTraits<int64_t>::min();
275+
#else
253276
A_x_offsets(i, 1, entry) = Kokkos::ArithTraits<int64_t>::min();
277+
#endif
254278
}
255279
}
256280
});
@@ -752,7 +776,11 @@ struct ComputeResidualVector {
752776
[&](const int k) {
753777
int64_t A_offset = overlap ? A_x_offsets_remote(rowidx, 0, k) : A_x_offsets(rowidx, 0, k);
754778
int64_t x_offset = overlap ? A_x_offsets_remote(rowidx, 1, k) : A_x_offsets(rowidx, 1, k);
779+
#if KOKKOS_VERSION > 40799
780+
if (A_offset != KokkosKernels::ArithTraits<int64_t>::min()) {
781+
#else
755782
if (A_offset != Kokkos::ArithTraits<int64_t>::min()) {
783+
#endif
756784
A_block_cst.assign_data(tpetra_values.data() + A_offset);
757785
// Pull x into local memory
758786
if constexpr (async) {

‎packages/ifpack2/src/Ifpack2_BlockHelper.hpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,11 +280,23 @@ struct ImplType {
280280
///
281281
/// kokkos arithmetic traits of scalar_type
282282
///
283+
#if KOKKOS_VERSION > 40799
284+
typedef typename KokkosKernels::ArithTraits<scalar_type>::val_type impl_scalar_type;
285+
#else
283286
typedef typename Kokkos::ArithTraits<scalar_type>::val_type impl_scalar_type;
287+
#endif
288+
#if KOKKOS_VERSION > 40799
289+
typedef typename KokkosKernels::ArithTraits<impl_scalar_type>::mag_type magnitude_type;
290+
#else
284291
typedef typename Kokkos::ArithTraits<impl_scalar_type>::mag_type magnitude_type;
292+
#endif
285293

286294
typedef typename BlockTridiagScalarType<impl_scalar_type>::type btdm_scalar_type;
295+
#if KOKKOS_VERSION > 40799
296+
typedef typename KokkosKernels::ArithTraits<btdm_scalar_type>::mag_type btdm_magnitude_type;
297+
#else
287298
typedef typename Kokkos::ArithTraits<btdm_scalar_type>::mag_type btdm_magnitude_type;
299+
#endif
288300

289301
///
290302
/// default host execution space
@@ -517,7 +529,11 @@ void reduceVector(const ConstUnmanaged<typename BlockHelperDetails::ImplType<Mat
517529
},
518530
norm2);
519531
#endif
532+
#if KOKKOS_VERSION > 40799
533+
vals[0] = KokkosKernels::ArithTraits<impl_scalar_type>::abs(norm2);
534+
#else
520535
vals[0] = Kokkos::ArithTraits<impl_scalar_type>::abs(norm2);
536+
#endif
521537

522538
IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
523539
IFPACK2_BLOCKHELPER_TIMER_FENCE(typename ImplType<MatrixType>::execution_space)

0 commit comments

Comments
 (0)