Add a serial populate subsets reference kernel.

+ Update kernels to take Arrays.
ginkgo-project · Dec 15, 2020 · f84fb2c · f84fb2c
1 parent 655d9c7
commit f84fb2c
Show file tree

Hide file tree

Showing 10 changed files with 356 additions and 70 deletions.
diff --git a/core/base/index_set.cpp b/core/base/index_set.cpp
@@ -51,6 +51,8 @@ namespace index_set {
 
 
 GKO_REGISTER_OPERATION(populate_subsets, index_set::populate_subsets);
+GKO_REGISTER_OPERATION(global_to_local, index_set::global_to_local);
+GKO_REGISTER_OPERATION(local_to_global, index_set::local_to_global);
 
 
 }  // namespace index_set
@@ -60,32 +62,42 @@ template <typename IndexType>
 void IndexSet<IndexType>::populate_subsets(const gko::Array<IndexType> &indices)
 {
     auto exec = this->get_executor();
-
-    auto num_indices = static_cast<IndexType>(indices.get_num_elems());
-
+    this->num_stored_indices_ = indices.get_num_elems();
     exec->run(index_set::make_populate_subsets(
-        this->index_space_size_, this->num_stored_indices_,
-        indices.get_const_data(), num_indices, this->subsets_begin_.get_data(),
-        this->subsets_end_.get_data(),
-        this->superset_cumulative_indices_.get_data()));
+        this->index_space_size_, &indices, &this->subsets_begin_,
+        &this->subsets_end_, &this->superset_cumulative_indices_));
 }
 
 
 template <typename IndexType>
-bool IndexSet<IndexType>::is_element(const IndexType index) const
-{}
+Array<IndexType> IndexSet<IndexType>::get_global_indices_from_local(
+    const Array<IndexType> &local_indices) const
+{
+    auto exec = this->get_executor();
+    auto global_indices =
+        gko::Array<IndexType>(exec, local_indices.get_num_elems());
+
+    GKO_ASSERT(this->get_num_subsets() >= 1);
+    exec->run(index_set::make_global_to_local(
+        this->index_space_size_, &this->subsets_begin_, &this->subsets_end_,
+        &this->superset_cumulative_indices_, &local_indices, &global_indices));
+    return std::move(global_indices);
+}
 
 
 template <typename IndexType>
-void IndexSet<IndexType>::get_global_index(const IndexType &local_index,
-                                           IndexType &global_index) const
-{}
-
+Array<IndexType> IndexSet<IndexType>::get_local_indices_from_global(
+    const Array<IndexType> &global_indices) const
+{
+    auto exec = this->get_executor();
+    auto local_indices = gko::Array<IndexType>(exec);
 
-template <typename IndexType>
-void IndexSet<IndexType>::get_local_index(const IndexType &global_index,
-                                          IndexType &local_index) const
-{}
+    GKO_ASSERT(this->get_num_subsets() >= 1);
+    exec->run(index_set::make_local_to_global(
+        this->index_space_size_, &this->subsets_begin_, &this->subsets_end_,
+        &this->superset_cumulative_indices_, &global_indices, &local_indices));
+    return std::move(local_indices);
+}
 
 
 #define GKO_DECLARE_INDEX_SET(_type) class IndexSet<_type>

diff --git a/core/base/index_set_kernels.hpp b/core/base/index_set_kernels.hpp
@@ -45,16 +45,38 @@ namespace kernels {
 
 
 #define GKO_DECLARE_INDEX_SET_POPULATE_KERNEL(IndexType)                   \
-    void populate_subsets(std::shared_ptr<const DefaultExecutor> exec,     \
-                          const IndexType &index_space_size,               \
-                          IndexType &num_elems, const IndexType *indices,  \
-                          IndexType &num_indices, IndexType *subset_begin, \
-                          IndexType *subset_end, IndexType *superset_indices)
-
-
-#define GKO_DECLARE_ALL_AS_TEMPLATES \
-    template <typename IndexType>    \
-    GKO_DECLARE_INDEX_SET_POPULATE_KERNEL(IndexType)
+    void populate_subsets(                                                 \
+        std::shared_ptr<const DefaultExecutor> exec,                       \
+        const IndexType index_space_size, const Array<IndexType> *indices, \
+        Array<IndexType> *subset_begin, Array<IndexType> *subset_end,      \
+        Array<IndexType> *superset_indices)
+
+#define GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL(IndexType)       \
+    void global_to_local(std::shared_ptr<const DefaultExecutor> exec, \
+                         const IndexType index_space_size,            \
+                         const Array<IndexType> *subset_begin,        \
+                         const Array<IndexType> *subset_end,          \
+                         const Array<IndexType> *superset_indices,    \
+                         const Array<IndexType> *global_indices,      \
+                         Array<IndexType> *local_indices)
+
+#define GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL(IndexType)       \
+    void local_to_global(std::shared_ptr<const DefaultExecutor> exec, \
+                         const IndexType index_space_size,            \
+                         const Array<IndexType> *subset_begin,        \
+                         const Array<IndexType> *subset_end,          \
+                         const Array<IndexType> *superset_indices,    \
+                         const Array<IndexType> *local_indices,       \
+                         Array<IndexType> *global_indices)
+
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                         \
+    template <typename IndexType>                            \
+    GKO_DECLARE_INDEX_SET_POPULATE_KERNEL(IndexType);        \
+    template <typename IndexType>                            \
+    GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL(IndexType); \
+    template <typename IndexType>                            \
+    GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL(IndexType)
 
 
 namespace omp {

diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp
@@ -120,6 +120,18 @@ GKO_DECLARE_INDEX_SET_POPULATE_KERNEL(IndexType)
 GKO_NOT_COMPILED(GKO_HOOK_MODULE);
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL);
 
+template <typename IndexType>
+GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL(IndexType)
+GKO_NOT_COMPILED(GKO_HOOK_MODULE);
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL);
+
+template <typename IndexType>
+GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL(IndexType)
+GKO_NOT_COMPILED(GKO_HOOK_MODULE);
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL);
+
 
 }  // namespace index_set
 

diff --git a/core/test/base/index_set.cpp b/core/test/base/index_set.cpp
@@ -78,6 +78,15 @@ class IndexSet : public ::testing::Test {
         }
     }
 
+    static void assert_equal_arrays(const T num_elems, const T *a, const T *b)
+    {
+        if (num_elems > 0) {
+            for (auto i = 0; i < num_elems; ++i) {
+                ASSERT_EQ(a[i], b[i]);
+            }
+        }
+    }
+
 
     std::shared_ptr<const gko::Executor> exec;
 };
@@ -140,4 +149,48 @@ TYPED_TEST(IndexSet, KnowsItsSize)
 }
 
 
+TYPED_TEST(IndexSet, CanBeConstructedFromIndices)
+{
+    auto idx_arr = gko::Array<TypeParam>{this->exec, {0, 1, 2, 4, 6, 7, 8, 9}};
+    auto begin_comp = gko::Array<TypeParam>{this->exec, {0, 4, 6}};
+    auto end_comp = gko::Array<TypeParam>{this->exec, {3, 5, 10}};
+    auto superset_comp = gko::Array<TypeParam>{this->exec, {0, 3, 4, 8}};
+
+    auto idx_set = gko::IndexSet<TypeParam>{this->exec, 10, idx_arr};
+
+    ASSERT_EQ(idx_set.get_size(), 10);
+    ASSERT_EQ(idx_set.get_num_subsets(), 3);
+    ASSERT_EQ(idx_set.get_num_subsets(), begin_comp.get_num_elems());
+    auto num_elems = idx_set.get_num_subsets();
+    this->assert_equal_arrays(num_elems, idx_set.get_subsets_begin(),
+                              begin_comp.get_data());
+    this->assert_equal_arrays(num_elems, idx_set.get_subsets_end(),
+                              end_comp.get_data());
+    this->assert_equal_arrays(num_elems, idx_set.get_superset_indices(),
+                              superset_comp.get_data());
+}
+
+
+TYPED_TEST(IndexSet, CanBeConstructedFromNonSortedIndices)
+{
+    auto idx_arr = gko::Array<TypeParam>{this->exec, {9, 1, 4, 2, 6, 8, 0, 7}};
+    auto begin_comp = gko::Array<TypeParam>{this->exec, {0, 4, 6}};
+    auto end_comp = gko::Array<TypeParam>{this->exec, {3, 5, 10}};
+    auto superset_comp = gko::Array<TypeParam>{this->exec, {0, 3, 4, 8}};
+
+    auto idx_set = gko::IndexSet<TypeParam>{this->exec, 10, idx_arr};
+
+    ASSERT_EQ(idx_set.get_size(), 10);
+    ASSERT_EQ(idx_set.get_num_subsets(), 3);
+    ASSERT_EQ(idx_set.get_num_subsets(), begin_comp.get_num_elems());
+    auto num_elems = idx_set.get_num_subsets();
+    this->assert_equal_arrays(num_elems, idx_set.get_subsets_begin(),
+                              begin_comp.get_data());
+    this->assert_equal_arrays(num_elems, idx_set.get_subsets_end(),
+                              end_comp.get_data());
+    this->assert_equal_arrays(num_elems, idx_set.get_superset_indices(),
+                              superset_comp.get_data());
+}
+
+
 }  // namespace
diff --git a/cuda/base/index_set_kernels.cpp b/cuda/base/index_set_kernels.cpp
@@ -30,9 +30,6 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************<GINKGO LICENSE>*******************************/
 
-#include "core/base/index_set_kernels.hpp"
-
-
 #include <algorithm>
 #include <iostream>
 #include <mutex>
@@ -43,6 +40,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/types.hpp>
 
 
+#include "core/base/index_set_kernels.hpp"
+
+
 namespace gko {
 namespace kernels {
 /**
@@ -61,15 +61,44 @@ namespace index_set {
 
 template <typename IndexType>
 void populate_subsets(std::shared_ptr<const DefaultExecutor> exec,
-                      const IndexType &index_space_size, IndexType &num_elems,
-                      const IndexType *indices, IndexType &num_indices,
-                      IndexType *subset_begin, IndexType *subset_end,
-                      IndexType *superset_indices)
+                      const IndexType index_space_size,
+                      const Array<IndexType> *indices,
+                      Array<IndexType> *subset_begin,
+                      Array<IndexType> *subset_end,
+                      Array<IndexType> *superset_indices)
 {}
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL);
 
 
+template <typename IndexType>
+void global_to_local(std::shared_ptr<const DefaultExecutor> exec,
+                     const IndexType index_space_size,
+                     const Array<IndexType> *subset_begin,
+                     const Array<IndexType> *subset_end,
+                     const Array<IndexType> *superset_indices,
+                     const Array<IndexType> *global_indices,
+                     Array<IndexType> *local_indices)
+{}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL);
+
+
+template <typename IndexType>
+void local_to_global(std::shared_ptr<const DefaultExecutor> exec,
+                     const IndexType index_space_size,
+                     const Array<IndexType> *subset_begin,
+                     const Array<IndexType> *subset_end,
+                     const Array<IndexType> *superset_indices,
+                     const Array<IndexType> *local_indices,
+                     Array<IndexType> *global_indices)
+{}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL);
+
+
 }  // namespace index_set
 }  // namespace cuda
 }  // namespace kernels

diff --git a/dpcpp/base/index_set_kernels.dp.cpp b/dpcpp/base/index_set_kernels.dp.cpp
@@ -61,15 +61,44 @@ namespace index_set {
 
 template <typename IndexType>
 void populate_subsets(std::shared_ptr<const DefaultExecutor> exec,
-                      const IndexType &index_space_size, IndexType &num_elems,
-                      const IndexType *indices, IndexType &num_indices,
-                      IndexType *subset_begin, IndexType *subset_end,
-                      IndexType *superset_indices)
+                      const IndexType index_space_size,
+                      const Array<IndexType> *indices,
+                      Array<IndexType> *subset_begin,
+                      Array<IndexType> *subset_end,
+                      Array<IndexType> *superset_indices)
 {}
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL);
 
 
+template <typename IndexType>
+void global_to_local(std::shared_ptr<const DefaultExecutor> exec,
+                     const IndexType index_space_size,
+                     const Array<IndexType> *subset_begin,
+                     const Array<IndexType> *subset_end,
+                     const Array<IndexType> *superset_indices,
+                     const Array<IndexType> *global_indices,
+                     Array<IndexType> *local_indices)
+{}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL);
+
+
+template <typename IndexType>
+void local_to_global(std::shared_ptr<const DefaultExecutor> exec,
+                     const IndexType index_space_size,
+                     const Array<IndexType> *subset_begin,
+                     const Array<IndexType> *subset_end,
+                     const Array<IndexType> *superset_indices,
+                     const Array<IndexType> *local_indices,
+                     Array<IndexType> *global_indices)
+{}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL);
+
+
 }  // namespace index_set
 }  // namespace dpcpp
 }  // namespace kernels

diff --git a/hip/base/index_set_kernels.hip.cpp b/hip/base/index_set_kernels.hip.cpp
@@ -30,9 +30,6 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************<GINKGO LICENSE>*******************************/
 
-#include "core/base/index_set_kernels.hpp"
-
-
 #include <algorithm>
 #include <iostream>
 #include <mutex>
@@ -43,6 +40,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/types.hpp>
 
 
+#include "core/base/index_set_kernels.hpp"
+
+
 namespace gko {
 namespace kernels {
 /**
@@ -61,15 +61,44 @@ namespace index_set {
 
 template <typename IndexType>
 void populate_subsets(std::shared_ptr<const DefaultExecutor> exec,
-                      const IndexType &index_space_size, IndexType &num_elems,
-                      const IndexType *indices, IndexType &num_indices,
-                      IndexType *subset_begin, IndexType *subset_end,
-                      IndexType *superset_indices)
+                      const IndexType index_space_size,
+                      const Array<IndexType> *indices,
+                      Array<IndexType> *subset_begin,
+                      Array<IndexType> *subset_end,
+                      Array<IndexType> *superset_indices)
 {}
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL);
 
 
+template <typename IndexType>
+void global_to_local(std::shared_ptr<const DefaultExecutor> exec,
+                     const IndexType index_space_size,
+                     const Array<IndexType> *subset_begin,
+                     const Array<IndexType> *subset_end,
+                     const Array<IndexType> *superset_indices,
+                     const Array<IndexType> *global_indices,
+                     Array<IndexType> *local_indices)
+{}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL);
+
+
+template <typename IndexType>
+void local_to_global(std::shared_ptr<const DefaultExecutor> exec,
+                     const IndexType index_space_size,
+                     const Array<IndexType> *subset_begin,
+                     const Array<IndexType> *subset_end,
+                     const Array<IndexType> *superset_indices,
+                     const Array<IndexType> *local_indices,
+                     Array<IndexType> *global_indices)
+{}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL);
+
+
 }  // namespace index_set
 }  // namespace hip
 }  // namespace kernels