Rename Impl::alignPtr to Impl::alignPtrTo, allow it to infer argu…

…ment type (#2261) * KokkosKernels::Impl::alignPtr infers argument type * Rename KokkosKernels::Impl::alignPtr -> alignPtrTo
kokkos · Jun 28, 2024 · 41954e2 · 41954e2
1 parent bbfc3ff
commit 41954e2
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 27 deletions.
diff --git a/common/src/KokkosKernels_Utils.hpp b/common/src/KokkosKernels_Utils.hpp
@@ -1527,8 +1527,8 @@ struct array_sum_reduce {
   }
 };
 
-template <typename InPtr, typename T>
-KOKKOS_INLINE_FUNCTION T *alignPtr(InPtr p) {
+template <typename T, typename InPtr>
+KOKKOS_INLINE_FUNCTION T *alignPtrTo(InPtr p) {
   // ugly but computationally free and the "right" way to do this in C++
   std::uintptr_t ptrVal = reinterpret_cast<std::uintptr_t>(p);
   // ptrVal + (align - 1) lands inside the next valid aligned scalar_t,

diff --git a/sparse/impl/KokkosSparse_bspgemm_impl_kkmem.hpp b/sparse/impl/KokkosSparse_bspgemm_impl_kkmem.hpp
@@ -270,8 +270,7 @@ struct KokkosBSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *hash_ids = (nnz_lno_t *)(tmp);
     tmp += pow2_hash_size;
 
-    scalar_t *hash_values =
-        KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(tmp);
+    scalar_t *hash_values = KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp);
 
     BlockAccumulator hm(block_dim, pow2_hash_size, pow2_hash_func, nullptr,
                         nullptr, hash_ids, hash_values);
@@ -414,7 +413,7 @@ struct KokkosBSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     all_shared_memory += sizeof(nnz_lno_t) * thread_shmem_key_size;
     // remainder of shmem allocation for vals
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     BlockAccumulator hm(block_dim, thread_shmem_key_size,
                         thread_shared_memory_hash_func, begins, nexts, keys,
@@ -554,7 +553,7 @@ struct KokkosBSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * team_cuckoo_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     int thread_rank = teamMember.team_rank();
 
@@ -601,8 +600,7 @@ struct KokkosBSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
           }
           global_acc_row_keys = (nnz_lno_t *)(tmp);
           global_acc_row_vals =
-              KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(
-                  tmp + pow2_hash_size);
+              KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp + pow2_hash_size);
         }
         // initialize begins.
         {
@@ -885,7 +883,7 @@ struct KokkosBSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * team_cuckoo_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     int thread_rank = teamMember.team_rank();
 

diff --git a/sparse/impl/KokkosSparse_bspgemm_impl_speed.hpp b/sparse/impl/KokkosSparse_bspgemm_impl_speed.hpp
@@ -325,7 +325,7 @@ struct KokkosBSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * shmem_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     KokkosKernels::Experimental::BlockHashmapAccumulator<
         nnz_lno_t, nnz_lno_t, scalar_t,

diff --git a/sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp b/sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp
@@ -261,8 +261,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *hash_ids = (nnz_lno_t *)(tmp);
     tmp += pow2_hash_size;
 
-    scalar_t *hash_values =
-        KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(tmp);
+    scalar_t *hash_values = KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp);
 
     Kokkos::parallel_for(
         Kokkos::TeamThreadRange(teamMember, team_row_begin, team_row_end),
@@ -409,8 +408,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
 
     hm2.keys = (nnz_lno_t *)(tmp);
     tmp += max_nnz;
-    hm2.values =
-        KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(tmp);
+    hm2.values = KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp);
 
     Kokkos::parallel_for(
         Kokkos::TeamThreadRange(teamMember, team_row_begin, team_row_end),
@@ -498,7 +496,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     all_shared_memory += sizeof(nnz_lno_t) * thread_shmem_key_size;
     // remainder of shmem allocation for vals
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     KokkosKernels::Experimental::HashmapAccumulator<
         nnz_lno_t, nnz_lno_t, scalar_t,
@@ -639,7 +637,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * team_cuckoo_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     int thread_rank = teamMember.team_rank();
 
@@ -686,8 +684,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
           }
           global_acc_row_keys = (nnz_lno_t *)(tmp);
           global_acc_row_vals =
-              KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(
-                  tmp + pow2_hash_size);
+              KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp + pow2_hash_size);
         }
         // initialize begins.
         {
@@ -970,7 +967,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * team_cuckoo_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     int thread_rank = teamMember.team_rank();
 

diff --git a/sparse/impl/KokkosSparse_spgemm_impl_speed.hpp b/sparse/impl/KokkosSparse_spgemm_impl_speed.hpp
@@ -304,7 +304,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * shmem_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     KokkosKernels::Experimental::HashmapAccumulator<
         nnz_lno_t, nnz_lno_t, scalar_t,

diff --git a/sparse/impl/KokkosSparse_spgemm_jacobi_sparseacc_impl.hpp b/sparse/impl/KokkosSparse_spgemm_jacobi_sparseacc_impl.hpp
@@ -260,8 +260,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     tmp += max_nnz;
     nnz_lno_t *hash_ids = (nnz_lno_t *)(tmp);
     tmp += pow2_hash_size;
-    scalar_t *hash_values =
-        KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(tmp);
+    scalar_t *hash_values = KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp);
 
     Kokkos::parallel_for(
         Kokkos::TeamThreadRange(teamMember, team_row_begin, team_row_end),
@@ -452,7 +451,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
 
     // Remainder of shmem allocation for vals
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     // Create the hashmaps
     KokkosKernels::Experimental::HashmapAccumulator<
@@ -610,7 +609,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * team_cuckoo_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     int thread_rank = teamMember.team_rank();
     int vector_rank = 0;
@@ -826,7 +825,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     nnz_lno_t *keys = (nnz_lno_t *)(all_shared_memory);
     all_shared_memory += sizeof(nnz_lno_t) * team_cuckoo_key_size;
     scalar_t *vals =
-        KokkosKernels::Impl::alignPtr<char *, scalar_t>(all_shared_memory);
+        KokkosKernels::Impl::alignPtrTo<scalar_t>(all_shared_memory);
 
     int thread_rank = teamMember.team_rank();
     int vector_rank = 0;
@@ -871,8 +870,7 @@ struct KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
         }
         global_acc_row_keys = (nnz_lno_t *)(tmp);
         global_acc_row_vals =
-            KokkosKernels::Impl::alignPtr<volatile nnz_lno_t *, scalar_t>(
-                tmp + pow2_hash_size);
+            KokkosKernels::Impl::alignPtrTo<scalar_t>(tmp + pow2_hash_size);
 
         nnz_lno_t num_threads = pow2_hash_size / vector_size;
         Kokkos::parallel_for(