GEMM: implicit MemberType

kokkos · Sep 20, 2022 · 2292855 · 2292855
1 parent bdff413
commit 2292855
Show file tree

Hide file tree

Showing 8 changed files with 96 additions and 124 deletions.
diff --git a/batched/dense/unit_test/Test_Batched_TeamInverseLU.hpp b/batched/dense/unit_test/Test_Batched_TeamInverseLU.hpp
@@ -52,7 +52,7 @@ struct Functor_BatchedTeamGemm {
     }
     member.team_barrier();
 
-    KokkosBlas::TeamGemm<MemberType, typename ParamTagType::transA,
+    KokkosBlas::TeamGemm<typename ParamTagType::transA,
                          typename ParamTagType::transB,
                          AlgoTagType>::invoke(member, _alpha, aa, bb, _beta,
                                               cc);

diff --git a/batched/dense/unit_test/Test_Batched_TeamSolveLU.hpp b/batched/dense/unit_test/Test_Batched_TeamSolveLU.hpp
@@ -52,7 +52,7 @@ struct Functor_BatchedTeamGemm {
     }
     member.team_barrier();
 
-    KokkosBlas::TeamGemm<MemberType, typename ParamTagType::transA,
+    KokkosBlas::TeamGemm<typename ParamTagType::transA,
                          typename ParamTagType::transB,
                          AlgoTagType>::invoke(member, _alpha, aa, bb, _beta,
                                               cc);

diff --git a/batched/dense/unit_test/Test_Batched_TeamVectorSolveUTV2.hpp b/batched/dense/unit_test/Test_Batched_TeamVectorSolveUTV2.hpp
@@ -81,8 +81,7 @@ struct Functor_TestBatchedTeamVectorSolveUTV2 {
     TeamVectorCopy<MemberType, Trans::NoTranspose>::invoke(member, aa, ac);
 
     /// bb = AA*xx
-    KokkosBlas::TeamVectorGemm<MemberType, Trans::NoTranspose,
-                               Trans::NoTranspose,
+    KokkosBlas::TeamVectorGemm<Trans::NoTranspose, Trans::NoTranspose,
                                Algo::Gemm::Unblocked>::invoke(member, one, aa,
                                                               xx, zero, bb);
     member.team_barrier();

diff --git a/blas/impl/KokkosBlas3_team_gemm_impl.hpp b/blas/impl/KokkosBlas3_team_gemm_impl.hpp
@@ -64,11 +64,10 @@ namespace KokkosBlas {
 /// NT/NT
 ///
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::NoTranspose, Trans::NoTranspose,
-                Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::NoTranspose, Trans::NoTranspose, Algo::Gemm::Unblocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -81,11 +80,10 @@ struct TeamGemm<MemberType, Trans::NoTranspose, Trans::NoTranspose,
   }
 };
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::NoTranspose, Trans::NoTranspose,
-                Algo::Gemm::Blocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::NoTranspose, Trans::NoTranspose, Algo::Gemm::Blocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -102,11 +100,10 @@ struct TeamGemm<MemberType, Trans::NoTranspose, Trans::NoTranspose,
 /// T/NT
 ///
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::Transpose, Trans::NoTranspose,
-                Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::Transpose, Trans::NoTranspose, Algo::Gemm::Unblocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -119,11 +116,10 @@ struct TeamGemm<MemberType, Trans::Transpose, Trans::NoTranspose,
   }
 };
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::Transpose, Trans::NoTranspose,
-                Algo::Gemm::Blocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::Transpose, Trans::NoTranspose, Algo::Gemm::Blocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -140,11 +136,10 @@ struct TeamGemm<MemberType, Trans::Transpose, Trans::NoTranspose,
 /// NT/T
 ///
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::NoTranspose, Trans::Transpose,
-                Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::NoTranspose, Trans::Transpose, Algo::Gemm::Unblocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -157,11 +152,10 @@ struct TeamGemm<MemberType, Trans::NoTranspose, Trans::Transpose,
   }
 };
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::NoTranspose, Trans::Transpose,
-                Algo::Gemm::Blocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::NoTranspose, Trans::Transpose, Algo::Gemm::Blocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -178,11 +172,10 @@ struct TeamGemm<MemberType, Trans::NoTranspose, Trans::Transpose,
 /// T/T
 ///
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::Transpose, Trans::Transpose,
-                Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::Transpose, Trans::Transpose, Algo::Gemm::Unblocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -195,11 +188,10 @@ struct TeamGemm<MemberType, Trans::Transpose, Trans::Transpose,
   }
 };
 
-template <typename MemberType>
-struct TeamGemm<MemberType, Trans::Transpose, Trans::Transpose,
-                Algo::Gemm::Blocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <>
+struct TeamGemm<Trans::Transpose, Trans::Transpose, Algo::Gemm::Blocked> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -224,11 +216,11 @@ struct TeamGemm<MemberType, Trans::Transpose, Trans::Transpose,
 /// NT/NT
 ///
 
-template <typename MemberType>
-struct TeamVectorGemm<MemberType, Trans::NoTranspose, Trans::NoTranspose,
+template <>
+struct TeamVectorGemm<Trans::NoTranspose, Trans::NoTranspose,
                       Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -245,11 +237,11 @@ struct TeamVectorGemm<MemberType, Trans::NoTranspose, Trans::NoTranspose,
 /// T/NT
 ///
 
-template <typename MemberType>
-struct TeamVectorGemm<MemberType, Trans::Transpose, Trans::NoTranspose,
+template <>
+struct TeamVectorGemm<Trans::Transpose, Trans::NoTranspose,
                       Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -266,11 +258,11 @@ struct TeamVectorGemm<MemberType, Trans::Transpose, Trans::NoTranspose,
 /// NT/T
 ///
 
-template <typename MemberType>
-struct TeamVectorGemm<MemberType, Trans::NoTranspose, Trans::Transpose,
+template <>
+struct TeamVectorGemm<Trans::NoTranspose, Trans::Transpose,
                       Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {
@@ -287,11 +279,11 @@ struct TeamVectorGemm<MemberType, Trans::NoTranspose, Trans::Transpose,
 /// T/T
 ///
 
-template <typename MemberType>
-struct TeamVectorGemm<MemberType, Trans::Transpose, Trans::Transpose,
+template <>
+struct TeamVectorGemm<Trans::Transpose, Trans::Transpose,
                       Algo::Gemm::Unblocked> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType &member, const ScalarType alpha, const AViewType &A,
       const BViewType &B, const ScalarType beta, const CViewType &C) {

diff --git a/blas/src/KokkosBlas3_gemm.hpp b/blas/src/KokkosBlas3_gemm.hpp
@@ -277,11 +277,10 @@ struct SerialGemm {
 /// Team Impl
 /// =========
 
-template <typename MemberType, typename ArgTransA, typename ArgTransB,
-          typename ArgAlgo>
+template <typename ArgTransA, typename ArgTransB, typename ArgAlgo>
 struct TeamGemm {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType& member, const ScalarType alpha, const AViewType& A,
       const BViewType& B, const ScalarType beta, const CViewType& C);
@@ -291,11 +290,10 @@ struct TeamGemm {
 /// TeamVector Impl
 /// =========
 
-template <typename MemberType, typename ArgTransA, typename ArgTransB,
-          typename ArgAlgo>
+template <typename ArgTransA, typename ArgTransB, typename ArgAlgo>
 struct TeamVectorGemm {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(
       const MemberType& member, const ScalarType alpha, const AViewType& A,
       const BViewType& B, const ScalarType beta, const CViewType& C);
@@ -304,20 +302,19 @@ struct TeamVectorGemm {
 ///
 /// Selective Interface
 ///
-template <typename MemberType, typename ArgTransA, typename ArgTransB,
-          typename ArgMode, typename ArgAlgo = Algo::Gemm::Default>
+template <typename ArgTransA, typename ArgTransB, typename ArgMode,
+          typename ArgAlgo = Algo::Gemm::Default>
 struct Gemm {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_FORCEINLINE_FUNCTION static int invoke(
       const MemberType& member, const ScalarType alpha, const AViewType& A,
       const BViewType& B, const ScalarType beta, const CViewType& C);
 };
 
-template <typename MemberType, typename ArgTransA, typename ArgTransB,
-          typename ArgAlgo>
-struct Gemm<MemberType, ArgTransA, ArgTransB, Mode::Serial, ArgAlgo> {
-  template <typename ScalarType, typename AViewType,
+template <typename ArgTransA, typename ArgTransB, typename ArgAlgo>
+struct Gemm<ArgTransA, ArgTransB, Mode::Serial, ArgAlgo> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
             typename BViewType, typename CViewType>
   KOKKOS_FORCEINLINE_FUNCTION static int invoke(const MemberType& /* member */,
                                                 const ScalarType alpha,
@@ -330,29 +327,27 @@ struct Gemm<MemberType, ArgTransA, ArgTransB, Mode::Serial, ArgAlgo> {
   }
 };
 
-template <typename MemberType, typename ArgTransA, typename ArgTransB,
-          typename ArgAlgo>
-struct Gemm<MemberType, ArgTransA, ArgTransB, Mode::Team, ArgAlgo> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <typename ArgTransA, typename ArgTransB, typename ArgAlgo>
+struct Gemm<ArgTransA, ArgTransB, Mode::Team, ArgAlgo> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_FORCEINLINE_FUNCTION static int invoke(
       const MemberType& member, const ScalarType alpha, const AViewType& A,
       const BViewType& B, const ScalarType beta, const CViewType& C) {
-    return TeamGemm<MemberType, ArgTransA, ArgTransB, ArgAlgo>::invoke(
-        member, alpha, A, B, beta, C);
+    return TeamGemm<ArgTransA, ArgTransB, ArgAlgo>::invoke(member, alpha, A, B,
+                                                           beta, C);
   }
 };
 
-template <typename MemberType, typename ArgTransA, typename ArgTransB,
-          typename ArgAlgo>
-struct Gemm<MemberType, ArgTransA, ArgTransB, Mode::TeamVector, ArgAlgo> {
-  template <typename ScalarType, typename AViewType, typename BViewType,
-            typename CViewType>
+template <typename ArgTransA, typename ArgTransB, typename ArgAlgo>
+struct Gemm<ArgTransA, ArgTransB, Mode::TeamVector, ArgAlgo> {
+  template <typename MemberType, typename ScalarType, typename AViewType,
+            typename BViewType, typename CViewType>
   KOKKOS_FORCEINLINE_FUNCTION static int invoke(
       const MemberType& member, const ScalarType alpha, const AViewType& A,
       const BViewType& B, const ScalarType beta, const CViewType& C) {
-    return TeamVectorGemm<MemberType, ArgTransA, ArgTransB, ArgAlgo>::invoke(
-        member, alpha, A, B, beta, C);
+    return TeamVectorGemm<ArgTransA, ArgTransB, ArgAlgo>::invoke(member, alpha,
+                                                                 A, B, beta, C);
   }
 };
 

diff --git a/blas/unit_test/Test_Blas3_team_gemm.hpp b/blas/unit_test/Test_Blas3_team_gemm.hpp
@@ -42,7 +42,7 @@ struct Functor_TestBatchedTeamGemm {
     auto bb = Kokkos::subview(_b, k, Kokkos::ALL(), Kokkos::ALL());
     auto cc = Kokkos::subview(_c, k, Kokkos::ALL(), Kokkos::ALL());
 
-    KokkosBlas::TeamGemm<MemberType, typename ParamTagType::transA,
+    KokkosBlas::TeamGemm<typename ParamTagType::transA,
                          typename ParamTagType::transB,
                          AlgoTagType>::invoke(member, _alpha, aa, bb, _beta,
                                               cc);

diff --git a/blas/unit_test/Test_Blas3_teamvector_gemm.hpp b/blas/unit_test/Test_Blas3_teamvector_gemm.hpp
@@ -39,7 +39,7 @@ struct Functor_TestBatchedTeamVector {
     auto bb = Kokkos::subview(_b, k, Kokkos::ALL(), Kokkos::ALL());
     auto cc = Kokkos::subview(_c, k, Kokkos::ALL(), Kokkos::ALL());
 
-    KokkosBlas::TeamVectorGemm<MemberType, typename ParamTagType::transA,
+    KokkosBlas::TeamVectorGemm<typename ParamTagType::transA,
                                typename ParamTagType::transB,
                                AlgoTagType>::invoke(member, _alpha, aa, bb,
                                                     _beta, cc);