unit_test/batched: Update BatchedGemm epsilon

kokkos · Jan 7, 2022 · f4c1c66 · f4c1c66
1 parent 8b97147
commit f4c1c66
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 7 deletions.
diff --git a/unit_test/batched/dense/Test_Batched_BatchedGemm.hpp b/unit_test/batched/dense/Test_Batched_BatchedGemm.hpp
@@ -135,7 +135,12 @@ void impl_test_batched_gemm_with_handle(BatchedGemmHandle* batchedGemmHandle,
   using mag_type = float;
   mag_type sum(1), diff(0);
 
-  mag_type eps = (mag_type)(1 << 1) * KOKKOSKERNELS_IMPL_FP16_EPSILON;
+  auto eps = static_cast<mag_type>(ats::epsilon());
+
+  eps *= std::is_same<ScalarType, Kokkos::Experimental::half_t>::value ||
+                 std::is_same<ScalarType, Kokkos::Experimental::bhalf_t>::value
+             ? 4
+             : 1e3;
 
   for (int k = 0; k < N; ++k) {
     for (int i = 0; i < matCdim1; ++i) {

diff --git a/unit_test/batched/dense/Test_Batched_SerialGemm.hpp b/unit_test/batched/dense/Test_Batched_SerialGemm.hpp
@@ -125,8 +125,10 @@ void impl_test_batched_gemm(const int N, const int matAdim1, const int matAdim2,
 
   mag_type eps = ats::epsilon();
 
-  eps *=
-      std::is_same<value_type, Kokkos::Experimental::half_t>::value ? 4 : 1e3;
+  eps *= std::is_same<value_type, Kokkos::Experimental::half_t>::value ||
+                 std::is_same<value_type, Kokkos::Experimental::bhalf_t>::value
+             ? 4
+             : 1e3;
 
   for (int k = 0; k < N; ++k)
     for (int i = 0; i < matCdim1; ++i)

diff --git a/unit_test/batched/dense/Test_Batched_TeamGemm.hpp b/unit_test/batched/dense/Test_Batched_TeamGemm.hpp
@@ -132,8 +132,10 @@ void impl_test_batched_teamgemm(const int N, const int matAdim1,
   mag_type sum(1), diff(0);
   mag_type eps = ats::epsilon();
 
-  eps *=
-      std::is_same<value_type, Kokkos::Experimental::half_t>::value ? 4 : 1e3;
+  eps *= std::is_same<value_type, Kokkos::Experimental::half_t>::value ||
+                 std::is_same<value_type, Kokkos::Experimental::bhalf_t>::value
+             ? 4
+             : 1e3;
 
   for (int k = 0; k < N; ++k)
     for (int i = 0; i < matCdim1; ++i)

diff --git a/unit_test/batched/dense/Test_Batched_TeamVectorGemm.hpp b/unit_test/batched/dense/Test_Batched_TeamVectorGemm.hpp
@@ -131,8 +131,10 @@ void impl_test_batched_teamvectorgemm(const int N, const int matAdim1,
 
   mag_type eps = ats::epsilon();
 
-  eps *=
-      std::is_same<value_type, Kokkos::Experimental::half_t>::value ? 4 : 1e3;
+  eps *= std::is_same<value_type, Kokkos::Experimental::half_t>::value ||
+                 std::is_same<value_type, Kokkos::Experimental::bhalf_t>::value
+             ? 4
+             : 1e3;
 
   for (int k = 0; k < N; ++k)
     for (int i = 0; i < matCdim1; ++i)