example/gmres: Use Kokkos_ArithTraits for abs

kokkos · Nov 18, 2021 · 21b2519 · 21b2519
1 parent beb4e5b
commit 21b2519
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/example/gmres/gmres.hpp b/example/gmres/gmres.hpp
@@ -43,6 +43,7 @@
 */
 
 #include <math.h>
+#include "Kokkos_ArithTraits.hpp"
 #include "KokkosKernels_IOUtils.hpp"
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Random.hpp>
@@ -54,22 +55,31 @@
 ////////////////////////////////////////////////////////////////////////////////
 // libstdc++ half_t overloads
 ////////////////////////////////////////////////////////////////////////////////
+namespace Kokkos {
+namespace Experimental {
 #if !KOKKOS_HALF_T_IS_FLOAT
 Kokkos::Experimental::half_t fabs(Kokkos::Experimental::half_t arg) {
-  return arg < 0.0 ? -arg : arg;
+  using AT = Kokkos::Details::ArithTraits<half_t>;
+  return AT::abs(arg);
 }
 
-Kokkos::complex<Kokkos::Experimental::half_t> fabs(
+Kokkos::Experimental::half_t fabs(
     Kokkos::complex<Kokkos::Experimental::half_t> arg) noexcept {
-  return Kokkos::complex<Kokkos::Experimental::half_t>(
-      fabs(Kokkos::complex<double>((double)arg.real(), (double)arg.imag())));
+  using AT = Kokkos::Details::ArithTraits<Kokkos::complex<double>>;
+  return Kokkos::Experimental::half_t(
+      AT::abs(Kokkos::complex<double>((double)arg.real(), (double)arg.imag())));
 }
 #endif  // KOKKOS_HALF_T_IS_FLOAT
 
+// This fabs wrapper was added to resolve:
+// https://github.com/kokkos/kokkos-kernels/issues/1172
 template <class T>
 KOKKOS_INLINE_FUNCTION T fabs(const Kokkos::complex<T> &x) {
-  return abs(x);
+  using AT = Kokkos::Details::ArithTraits<Kokkos::complex<T>>;
+  return AT::abs(x);
 }
+}  // namespace Experimental
+}  // namespace Kokkos
 
 // This struct is returned to the user to give solver
 // statistics and convergence status.
@@ -109,6 +119,7 @@ GmresStats gmres(
     const GmresOpts<ScalarType> &opts,
     const KokkosSparse::Experimental::Preconditioner<
         ScalarType, Layout, EXSP, OrdinalType> *const M = NULL) {
+  using namespace Kokkos::Experimental;  // For 'fabs' wrappers above
   Kokkos::Profiling::pushRegion("GMRES::TotalTime:");
   typedef Kokkos::Details::ArithTraits<ScalarType> AT;
   typedef typename AT::val_type

diff --git a/src/sparse/impl/KokkosSparse_spmv_bsrmatrix_impl.hpp b/src/sparse/impl/KokkosSparse_spmv_bsrmatrix_impl.hpp
@@ -433,6 +433,7 @@ struct BsrMatrixSpMVTensorCoreFunctor {
     mbr.team_barrier();
 
     // Suppress unused var warnings
+    // TODO (@cwpeason): Should this functor only compile on device?
     (void)fx;
     (void)fa;
     (void)fy;