move TeamScale and TeamVectorScale to KokkosBlas

kokkos · Jun 23, 2022 · 3046f2a · 3046f2a
1 parent ad4da89
commit 3046f2a
Show file tree

Hide file tree

Showing 14 changed files with 139 additions and 78 deletions.
diff --git a/src/batched/dense/KokkosBatched_Scale_Decl.hpp b/src/batched/dense/KokkosBatched_Scale_Decl.hpp
@@ -3,9 +3,6 @@
 
 /// \author Kyungjoo Kim (kyukim@sandia.gov)
 
-#include "KokkosBatched_Util.hpp"
-#include "KokkosBatched_Vector.hpp"
-
 namespace KokkosBatched {
 
 ///
@@ -30,7 +27,10 @@ struct TeamScale {
   template <typename ScalarType, typename AViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member,
                                            const ScalarType alpha,
-                                           const AViewType &A);
+                                           const AViewType &A) {
+    assert(false && "Deprecated: use KokkosBlas::TeamScale");
+    return 0;
+  }
 };
 
 ///
@@ -42,11 +42,13 @@ struct TeamVectorScale {
   template <typename ScalarType, typename AViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member,
                                            const ScalarType alpha,
-                                           const AViewType &A);
+                                           const AViewType &A) {
+    // static_assert(false);
+    assert(false && "Deprecated: use KokkosBlas::TeamVectorScale");
+    return 0;
+  }
 };
 
 }  // namespace KokkosBatched
 
-#include "KokkosBatched_Scale_Impl.hpp"
-
 #endif
diff --git a/src/batched/dense/impl/KokkosBatched_Gemm_TeamVector_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Gemm_TeamVector_Internal.hpp
@@ -40,7 +40,8 @@ TeamVectorGemmInternal<Algo::Gemm::Unblocked, false>::invoke(
   if (beta == zero)
     TeamVectorSetInternal ::invoke(member, m, n, zero, C, cs0, cs1);
   else if (beta != one)
-    TeamVectorScaleInternal::invoke(member, m, n, beta, C, cs0, cs1);
+    KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, n, beta, C,
+                                                      cs0, cs1);
 
   if (alpha != ScalarType(0.0)) {
     if (m <= 0 || n <= 0 || k <= 0) return 0;
@@ -80,7 +81,8 @@ TeamVectorGemmInternal<Algo::Gemm::Unblocked, true>::invoke(
   if (beta == zero)
     TeamVectorSetInternal ::invoke(member, m, n, zero, C, cs0, cs1);
   else if (beta != one)
-    TeamVectorScaleInternal::invoke(member, m, n, beta, C, cs0, cs1);
+    KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, n, beta, C,
+                                                      cs0, cs1);
 
   if (alpha != ScalarType(0.0)) {
     if (m <= 0 || n <= 0 || k <= 0) return 0;

diff --git a/src/batched/dense/impl/KokkosBatched_Gemm_Team_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Gemm_Team_Internal.hpp
@@ -7,7 +7,7 @@
 #include "KokkosKernels_ExecSpaceUtils.hpp"
 
 #include "KokkosBatched_Set_Internal.hpp"
-#include "KokkosBatched_Scale_Internal.hpp"
+#include "KokkosBlas1_team_scal.hpp"
 
 #include "KokkosBatched_InnerGemmFixC_Serial_Impl.hpp"
 
@@ -43,7 +43,8 @@ KOKKOS_INLINE_FUNCTION int TeamGemmInternal<Algo::Gemm::Unblocked>::invoke(
   if (beta == zero)
     TeamSetInternal ::invoke(member, m, n, zero, C, cs0, cs1);
   else if (beta != one)
-    TeamScaleInternal::invoke(member, m, n, beta, C, cs0, cs1);
+    KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, n, beta, C, cs0,
+                                                cs1);
 
   if (alpha != ScalarType(0.0)) {
     if (m <= 0 || n <= 0 || k <= 0) return 0;
@@ -84,7 +85,8 @@ KOKKOS_INLINE_FUNCTION int TeamGemmInternal<Algo::Gemm::Blocked>::invoke(
   if (beta == zero)
     TeamSetInternal ::invoke(member, m, n, zero, C, cs0, cs1);
   else if (beta != one)
-    TeamScaleInternal::invoke(member, m, n, beta, C, cs0, cs1);
+    KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, n, beta, C, cs0,
+                                                cs1);
 
   if (alpha != ScalarType(0.0)) {
     if (m <= 0 || n <= 0 || k <= 0) return 0;

diff --git a/src/batched/dense/impl/KokkosBatched_Gemv_TeamVector_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Gemv_TeamVector_Internal.hpp
@@ -6,7 +6,7 @@
 #include "KokkosBatched_Util.hpp"
 
 #include "KokkosBatched_Set_Internal.hpp"
-#include "KokkosBatched_Scale_Internal.hpp"
+#include "KokkosBlas1_team_scal_impl.hpp"
 
 #include "KokkosBatched_InnerMultipleDotProduct_Serial_Impl.hpp"
 
@@ -60,7 +60,7 @@ TeamVectorGemvInternal<Algo::Gemv::Unblocked>::invoke(
   if (beta == zero)
     TeamVectorSetInternal ::invoke(member, m, zero, y, ys0);
   else if (beta != one)
-    TeamVectorScaleInternal::invoke(member, m, beta, y, ys0);
+    KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, beta, y, ys0);
 
   if (alpha != zero) {
     if (m <= 0 || n <= 0) return 0;

diff --git a/src/batched/dense/impl/KokkosBatched_Gemv_Team_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Gemv_Team_Internal.hpp
@@ -49,7 +49,7 @@ KOKKOS_INLINE_FUNCTION int TeamGemvInternal<Algo::Gemv::Unblocked>::invoke(
   if (beta == zero)
     TeamSetInternal ::invoke(member, m, zero, y, ys0);
   else if (beta != one)
-    TeamScaleInternal::invoke(member, m, beta, y, ys0);
+    KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, beta, y, ys0);
 
   if (alpha != zero) {
     if (m <= 0 || n <= 0) return 0;
@@ -88,7 +88,7 @@ KOKKOS_INLINE_FUNCTION int TeamGemvInternal<Algo::Gemv::Blocked>::invoke(
   if (beta == zero)
     TeamSetInternal ::invoke(member, m, zero, y, ys0);
   else if (beta != one)
-    TeamScaleInternal::invoke(member, m, beta, y, ys0);
+    KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, beta, y, ys0);
 
   if (alpha != zero) {
     if (m <= 0 || n <= 0) return 0;

diff --git a/src/batched/dense/impl/KokkosBatched_Scale_Impl.hpp b/src/batched/dense/impl/KokkosBatched_Scale_Impl.hpp
diff --git a/src/batched/dense/impl/KokkosBatched_Trsm_TeamVector_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Trsm_TeamVector_Internal.hpp
@@ -37,7 +37,8 @@ TeamVectorTrsmInternalLeftLower<Algo::Trsm::Unblocked>::invoke(
     TeamVectorSetInternal ::invoke(member, m, n, zero, B, bs0, bs1);
   else {
     if (alpha != one)
-      TeamVectorScaleInternal::invoke(member, m, n, alpha, B, bs0, bs1);
+      KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, n, alpha, B,
+                                                        bs0, bs1);
     if (m <= 0 || n <= 0) return 0;
 
     for (int p = 0; p < m; ++p) {
@@ -98,7 +99,8 @@ TeamVectorTrsmInternalLeftUpper<Algo::Trsm::Unblocked>::invoke(
     TeamVectorSetInternal ::invoke(member, m, n, zero, B, bs0, bs1);
   else {
     if (alpha != one)
-      TeamVectorScaleInternal::invoke(member, m, n, alpha, B, bs0, bs1);
+      KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, n, alpha, B,
+                                                        bs0, bs1);
     if (m <= 0 || n <= 0) return 0;
 
     ValueType *KOKKOS_RESTRICT B0 = B;

diff --git a/src/batched/dense/impl/KokkosBatched_Trsm_Team_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Trsm_Team_Internal.hpp
@@ -41,7 +41,8 @@ TeamTrsmInternalLeftLower<Algo::Trsm::Unblocked>::invoke(
     TeamSetInternal ::invoke(member, m, n, zero, B, bs0, bs1);
   else {
     if (alpha != one)
-      TeamScaleInternal::invoke(member, m, n, alpha, B, bs0, bs1);
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, n, alpha, B, bs0,
+                                                  bs1);
     if (m <= 0 || n <= 0) return 0;
 
     for (int p = 0; p < m; ++p) {
@@ -92,7 +93,8 @@ TeamTrsmInternalLeftLower<Algo::Trsm::Blocked>::invoke(
     TeamSetInternal ::invoke(member, m, n, zero, B, bs0, bs1);
   else {
     if (alpha != one)
-      TeamScaleInternal::invoke(member, m, n, alpha, B, bs0, bs1);
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, n, alpha, B, bs0,
+                                                  bs1);
     if (m <= 0 || n <= 0) return 0;
 
     ///
@@ -175,7 +177,8 @@ TeamTrsmInternalLeftUpper<Algo::Trsm::Unblocked>::invoke(
     TeamSetInternal ::invoke(member, m, n, zero, B, bs0, bs1);
   else {
     if (alpha != one)
-      TeamScaleInternal::invoke(member, m, n, alpha, B, bs0, bs1);
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, n, alpha, B, bs0,
+                                                  bs1);
     if (m <= 0 || n <= 0) return 0;
 
     ValueType *KOKKOS_RESTRICT B0 = B;
@@ -231,7 +234,8 @@ TeamTrsmInternalLeftUpper<Algo::Trsm::Blocked>::invoke(
     TeamSetInternal ::invoke(member, m, n, zero, B, bs0, bs1);
   else {
     if (alpha != one)
-      TeamScaleInternal::invoke(member, m, n, alpha, B, bs0, bs1);
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, n, alpha, B, bs0,
+                                                  bs1);
     if (m <= 0 || n <= 0) return 0;
 
     InnerTrsmLeftUpperUnitDiag<mbAlgo> trsm_u(as0, as1, bs0, bs1);

diff --git a/src/batched/dense/impl/KokkosBatched_Trsv_TeamVector_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Trsv_TeamVector_Internal.hpp
@@ -44,7 +44,9 @@ TeamVectorTrsvInternalLower<Algo::Trsv::Unblocked>::invoke(
   if (alpha == zero)
     TeamVectorSetInternal::invoke(member, m, zero, b, bs0);
   else {
-    if (alpha != one) TeamVectorScaleInternal::invoke(member, m, alpha, b, bs0);
+    if (alpha != one)
+      KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, alpha, b,
+                                                        bs0);
     if (m <= 0) return 0;
 
     for (int p = 0; p < m; ++p) {
@@ -106,7 +108,9 @@ TeamVectorTrsvInternalUpper<Algo::Trsv::Unblocked>::invoke(
   if (alpha == zero)
     TeamVectorSetInternal::invoke(member, m, zero, b, bs0);
   else {
-    if (alpha != one) TeamVectorScaleInternal::invoke(member, m, alpha, b, bs0);
+    if (alpha != one)
+      KokkosBlas::Impl::TeamVectorScaleInternal::invoke(member, m, alpha, b,
+                                                        bs0);
     if (m <= 0) return 0;
 
     ValueType *KOKKOS_RESTRICT b0 = b;

diff --git a/src/batched/dense/impl/KokkosBatched_Trsv_Team_Internal.hpp b/src/batched/dense/impl/KokkosBatched_Trsv_Team_Internal.hpp
@@ -46,7 +46,8 @@ KOKKOS_INLINE_FUNCTION int TeamTrsvInternalLower<Algo::Trsv::Unblocked>::invoke(
   if (alpha == zero)
     TeamSetInternal::invoke(member, m, zero, b, bs0);
   else {
-    if (alpha != one) TeamScaleInternal::invoke(member, m, alpha, b, bs0);
+    if (alpha != one)
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, alpha, b, bs0);
     if (m <= 0) return 0;
 
     for (int p = 0; p < m; ++p) {
@@ -91,7 +92,8 @@ KOKKOS_INLINE_FUNCTION int TeamTrsvInternalLower<Algo::Trsv::Blocked>::invoke(
   if (alpha == zero)
     TeamSetInternal::invoke(member, m, zero, b, bs0);
   else {
-    if (alpha != one) TeamScaleInternal::invoke(member, m, alpha, b, bs0);
+    if (alpha != one)
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, alpha, b, bs0);
     if (m <= 0) return 0;
 
     /// case GPU: team size is large and blocksize (mb,nb) is small
@@ -155,7 +157,8 @@ KOKKOS_INLINE_FUNCTION int TeamTrsvInternalUpper<Algo::Trsv::Unblocked>::invoke(
   if (alpha == zero)
     TeamSetInternal::invoke(member, m, zero, b, bs0);
   else {
-    if (alpha != one) TeamScaleInternal::invoke(member, m, alpha, b, bs0);
+    if (alpha != one)
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, alpha, b, bs0);
     if (m <= 0) return 0;
 
     ValueType *KOKKOS_RESTRICT b0 = b;
@@ -198,7 +201,8 @@ KOKKOS_INLINE_FUNCTION int TeamTrsvInternalUpper<Algo::Trsv::Blocked>::invoke(
   if (alpha == zero)
     TeamSetInternal::invoke(member, m, zero, b, bs0);
   else {
-    if (alpha != one) TeamScaleInternal::invoke(member, m, alpha, b, bs0);
+    if (alpha != one)
+      KokkosBlas::Impl::TeamScaleInternal::invoke(member, m, alpha, b, bs0);
     if (m <= 0) return 0;
 
     InnerTrsmLeftUpperUnitDiag<mbAlgo> trsm_u(as0, as1, bs0, 0);

diff --git a/src/blas/KokkosBlas1_team_scal.hpp b/src/blas/KokkosBlas1_team_scal.hpp
@@ -45,9 +45,46 @@
 #ifndef KOKKOSBLAS1_TEAM_SCAL_HPP_
 #define KOKKOSBLAS1_TEAM_SCAL_HPP_
 
+#include <KokkosBlas1_team_scal_impl.hpp>
+
+// TODO: deprecate/remove ?
 #include <KokkosBlas1_team_scal_spec.hpp>
 
 namespace KokkosBlas {
+
+///
+/// Team Scale
+///
+
+template <typename MemberType>
+struct TeamScale {
+  template <typename ScalarType, typename AViewType>
+  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType& member,
+                                           const ScalarType alpha,
+                                           const AViewType& A) {
+    return Impl::TeamScaleInternal::invoke(member, A.extent(0), A.extent(1),
+                                           alpha, A.data(), A.stride_0(),
+                                           A.stride_1());
+  }
+};
+
+///
+/// TeamVector Scale
+///
+
+template <typename MemberType>
+struct TeamVectorScale {
+  template <typename ScalarType, typename AViewType>
+  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType& member,
+                                           const ScalarType alpha,
+                                           const AViewType& A) {
+    return Impl::TeamVectorScaleInternal::invoke(member, A.extent(0),
+                                                 A.extent(1), alpha, A.data(),
+                                                 A.stride_0(), A.stride_1());
+  }
+};
+
+// TODO: deprecate/remove ?
 namespace Experimental {
 
 template <class TeamType, class RVector, class XVector>