Merge pull request #91 from zacharyjbaker/Zack

Added set_values methods to kokkos_types.h
lanl · Sep 6, 2024 · 93923f4 · 93923f4
2 parents 1438a6b + 382ed88
commit 93923f4
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 2 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,4 +1,5 @@
 install-*
 benchmark/benchmark*
 build-matar-*
-install/*
+install/*
+heffte/
diff --git a/src/include/kokkos_types.h b/src/include/kokkos_types.h
@@ -128,7 +128,7 @@ namespace mtr
 template <typename T, typename Layout = DefaultLayout, typename ExecSpace = DefaultExecSpace, typename MemoryTraits = void>
 class FArrayKokkos {
 
-    using TArray1D = Kokkos::View<T*, Layout, ExecSpace, MemoryTraits>;
+    using TArray1D = Kokkos::View<T*, Layout, ExecSpace, MemoryTraits>; 
 
 private:
     size_t dims_[7];
@@ -229,6 +229,10 @@ class FArrayKokkos {
 
     KOKKOS_INLINE_FUNCTION
     T* pointer() const;
+
+    // set values
+    KOKKOS_INLINE_FUNCTION
+    void set_values(T val);
 
     //return kokkos view
     KOKKOS_INLINE_FUNCTION
@@ -525,6 +529,15 @@ Kokkos::View<T*, Layout, ExecSpace, MemoryTraits> FArrayKokkos<T,Layout,ExecSpac
     return this_array_;
 }
 
+// set values of array
+template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
+KOKKOS_INLINE_FUNCTION
+void FArrayKokkos<T,Layout,ExecSpace,MemoryTraits>::set_values(T val) {
+    Kokkos::parallel_for("SetValues_FArrayKokkos", length_, KOKKOS_CLASS_LAMBDA(const int i) {
+        this_array_(i) = val;
+    });
+}
+
 // Destructor
 template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
 KOKKOS_INLINE_FUNCTION
@@ -950,6 +963,10 @@ class FMatrixKokkos {
 
     KOKKOS_INLINE_FUNCTION
     T* pointer() const;
+
+    // set values
+    KOKKOS_INLINE_FUNCTION
+    void set_values(T val);
 
     //return kokkos view
     KOKKOS_INLINE_FUNCTION
@@ -1229,6 +1246,15 @@ Kokkos::View<T*, Layout, ExecSpace, MemoryTraits> FMatrixKokkos<T,Layout,ExecSpa
     return this_matrix_;
 }
 
+// set values of array
+template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
+KOKKOS_INLINE_FUNCTION
+void FMatrixKokkos<T, Layout, ExecSpace, MemoryTraits>::set_values(T val) {
+    Kokkos::parallel_for("SetValues_FMatrixKokkos", length_, KOKKOS_CLASS_LAMBDA(const int i) {
+        this_matrix_(i) = val;
+    });
+}
+
 template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
 KOKKOS_INLINE_FUNCTION
 FMatrixKokkos<T,Layout,ExecSpace,MemoryTraits>::~FMatrixKokkos() {}
@@ -3334,6 +3360,10 @@ class CArrayKokkos {
     // Methods returns the raw pointer (most likely GPU) of the Kokkos View
     KOKKOS_INLINE_FUNCTION
     T* pointer() const;
+
+    // set values
+    KOKKOS_INLINE_FUNCTION
+    void set_values(T val);
 
     //return the view
     KOKKOS_INLINE_FUNCTION
@@ -3610,6 +3640,15 @@ Kokkos::View<T*, Layout, ExecSpace, MemoryTraits> CArrayKokkos<T,Layout,ExecSpac
     return this_array_;
 }
 
+// set values of array
+template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
+KOKKOS_INLINE_FUNCTION
+void CArrayKokkos<T,Layout,ExecSpace,MemoryTraits>::set_values(T val) {
+    Kokkos::parallel_for("SetValues_CArrayKokkos", length_, KOKKOS_CLASS_LAMBDA(const int i) {
+        this_array_(i) = val;
+    });
+}
+
 template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
 KOKKOS_INLINE_FUNCTION
 CArrayKokkos<T,Layout,ExecSpace,MemoryTraits>::~CArrayKokkos() {}
@@ -4032,6 +4071,10 @@ class CMatrixKokkos {
     KOKKOS_INLINE_FUNCTION
     T* pointer() const;
 
+    // set values
+    KOKKOS_INLINE_FUNCTION
+    void set_values(T val);
+
     //return the view
     KOKKOS_INLINE_FUNCTION
     TArray1D get_kokkos_view() const;
@@ -4313,6 +4356,15 @@ Kokkos::View<T*, Layout, ExecSpace, MemoryTraits> CMatrixKokkos<T,Layout,ExecSpa
     return this_matrix_;
 }
 
+// set values of array
+template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
+KOKKOS_INLINE_FUNCTION
+void CMatrixKokkos<T, Layout, ExecSpace, MemoryTraits>::set_values(T val) {
+    Kokkos::parallel_for("SetValues_CMatrixKokkos", length_, KOKKOS_CLASS_LAMBDA(const int i) {
+        this_matrix_(i) = val;
+    });
+}
+
 // Deconstructor
 template <typename T, typename Layout, typename ExecSpace, typename MemoryTraits>
 KOKKOS_INLINE_FUNCTION