diff --git a/core/base/index_set.cpp b/core/base/index_set.cpp index 6aa1a2ef574..eb383882f49 100644 --- 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 void IndexSet::populate_subsets(const gko::Array &indices) { auto exec = this->get_executor(); - - auto num_indices = static_cast(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 -bool IndexSet::is_element(const IndexType index) const -{} +Array IndexSet::get_global_indices_from_local( + const Array &local_indices) const +{ + auto exec = this->get_executor(); + auto global_indices = + gko::Array(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 -void IndexSet::get_global_index(const IndexType &local_index, - IndexType &global_index) const -{} - +Array IndexSet::get_local_indices_from_global( + const Array &global_indices) const +{ + auto exec = this->get_executor(); + auto local_indices = gko::Array(exec); -template -void IndexSet::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 index 951d2365b61..5ce0c3cb35f 100644 --- 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 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 \ - GKO_DECLARE_INDEX_SET_POPULATE_KERNEL(IndexType) + void populate_subsets( \ + std::shared_ptr exec, \ + const IndexType index_space_size, const Array *indices, \ + Array *subset_begin, Array *subset_end, \ + Array *superset_indices) + +#define GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL(IndexType) \ + void global_to_local(std::shared_ptr exec, \ + const IndexType index_space_size, \ + const Array *subset_begin, \ + const Array *subset_end, \ + const Array *superset_indices, \ + const Array *global_indices, \ + Array *local_indices) + +#define GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL(IndexType) \ + void local_to_global(std::shared_ptr exec, \ + const IndexType index_space_size, \ + const Array *subset_begin, \ + const Array *subset_end, \ + const Array *superset_indices, \ + const Array *local_indices, \ + Array *global_indices) + + +#define GKO_DECLARE_ALL_AS_TEMPLATES \ + template \ + GKO_DECLARE_INDEX_SET_POPULATE_KERNEL(IndexType); \ + template \ + GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL(IndexType); \ + template \ + 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 index 038f2779611..0640d07463a 100644 --- 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 +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 +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 index 34317f7440f..8a2945064b5 100644 --- 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 exec; }; @@ -140,4 +149,48 @@ TYPED_TEST(IndexSet, KnowsItsSize) } +TYPED_TEST(IndexSet, CanBeConstructedFromIndices) +{ + auto idx_arr = gko::Array{this->exec, {0, 1, 2, 4, 6, 7, 8, 9}}; + auto begin_comp = gko::Array{this->exec, {0, 4, 6}}; + auto end_comp = gko::Array{this->exec, {3, 5, 10}}; + auto superset_comp = gko::Array{this->exec, {0, 3, 4, 8}}; + + auto idx_set = gko::IndexSet{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{this->exec, {9, 1, 4, 2, 6, 8, 0, 7}}; + auto begin_comp = gko::Array{this->exec, {0, 4, 6}}; + auto end_comp = gko::Array{this->exec, {3, 5, 10}}; + auto superset_comp = gko::Array{this->exec, {0, 3, 4, 8}}; + + auto idx_set = gko::IndexSet{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 index 8d6a3f9834f..302f9202395 100644 --- 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. *************************************************************/ -#include "core/base/index_set_kernels.hpp" - - #include #include #include @@ -43,6 +40,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include +#include "core/base/index_set_kernels.hpp" + + namespace gko { namespace kernels { /** @@ -61,15 +61,44 @@ namespace index_set { template void populate_subsets(std::shared_ptr 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 *indices, + Array *subset_begin, + Array *subset_end, + Array *superset_indices) {} GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL); +template +void global_to_local(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *global_indices, + Array *local_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL); + + +template +void local_to_global(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *local_indices, + Array *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 index d28bfcf1be3..ce8c5f78290 100644 --- a/dpcpp/base/index_set_kernels.dp.cpp +++ b/dpcpp/base/index_set_kernels.dp.cpp @@ -61,15 +61,44 @@ namespace index_set { template void populate_subsets(std::shared_ptr 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 *indices, + Array *subset_begin, + Array *subset_end, + Array *superset_indices) {} GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL); +template +void global_to_local(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *global_indices, + Array *local_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL); + + +template +void local_to_global(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *local_indices, + Array *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 index 1317258dfa0..eed6f615e0c 100644 --- 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. *************************************************************/ -#include "core/base/index_set_kernels.hpp" - - #include #include #include @@ -43,6 +40,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include +#include "core/base/index_set_kernels.hpp" + + namespace gko { namespace kernels { /** @@ -61,15 +61,44 @@ namespace index_set { template void populate_subsets(std::shared_ptr 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 *indices, + Array *subset_begin, + Array *subset_end, + Array *superset_indices) {} GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL); +template +void global_to_local(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *global_indices, + Array *local_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL); + + +template +void local_to_global(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *local_indices, + Array *global_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL); + + } // namespace index_set } // namespace hip } // namespace kernels diff --git a/include/ginkgo/core/base/index_set.hpp b/include/ginkgo/core/base/index_set.hpp index 5c10c76f630..9c0ea29885b 100644 --- a/include/ginkgo/core/base/index_set.hpp +++ b/include/ginkgo/core/base/index_set.hpp @@ -97,9 +97,7 @@ class IndexSet { subsets_end_(exec_), superset_cumulative_indices_(exec_) { - auto tmp_arr = gko::Array(exec_); - tmp_arr = std::move(indices); - this->populate_subsets(tmp_arr); + this->populate_subsets(indices); } /** @@ -112,20 +110,17 @@ class IndexSet { return exec_; } - index_type get_size() const { return this->index_space_size_; } - bool is_element(const index_type index) const; - bool is_contiguous() const { return (this->get_num_subsets() <= 1); } index_type get_num_elems() const { return this->num_stored_indices_; }; - void get_global_index(const index_type &local_index, - index_type &global_index) const; + Array get_global_indices_from_local( + const Array &local_indices) const; - void get_local_index(const index_type &global_index, - index_type &local_index) const; + Array get_local_indices_from_global( + const Array &global_indices) const; index_type get_num_subsets() const { diff --git a/omp/base/index_set_kernels.cpp b/omp/base/index_set_kernels.cpp index 0d0fb2b52ab..ef70b949fed 100644 --- a/omp/base/index_set_kernels.cpp +++ b/omp/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. *************************************************************/ -#include "core/base/index_set_kernels.hpp" - - #include #include #include @@ -43,6 +40,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include +#include "core/base/index_set_kernels.hpp" + + namespace gko { namespace kernels { /** @@ -61,15 +61,44 @@ namespace index_set { template void populate_subsets(std::shared_ptr 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 *indices, + Array *subset_begin, + Array *subset_end, + Array *superset_indices) {} GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL); +template +void global_to_local(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *global_indices, + Array *local_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL); + + +template +void local_to_global(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *local_indices, + Array *global_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL); + + } // namespace index_set } // namespace omp } // namespace kernels diff --git a/reference/base/index_set_kernels.cpp b/reference/base/index_set_kernels.cpp index 5463efa052f..24842c65b9a 100644 --- a/reference/base/index_set_kernels.cpp +++ b/reference/base/index_set_kernels.cpp @@ -30,19 +30,21 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. *************************************************************/ -#include "core/base/index_set_kernels.hpp" - - #include #include #include #include +#include +#include #include #include +#include "core/base/index_set_kernels.hpp" + + namespace gko { namespace kernels { /** @@ -61,15 +63,89 @@ namespace index_set { template void populate_subsets(std::shared_ptr 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 *indices, + Array *subset_begin, + Array *subset_end, + Array *superset_indices) +{ + auto num_indices = indices->get_num_elems(); + auto tmp_indices = gko::Array(*indices); + GKO_ASSERT(*std::max_element(indices->get_const_data(), + indices->get_const_data() + num_indices) <= + index_space_size); + // Sort the indices. + std::sort(tmp_indices.get_data(), tmp_indices.get_data() + num_indices); + + // Detect subsets. + auto tmp_subset_begin = gko::vector(exec); + auto tmp_subset_end = gko::vector(exec); + auto tmp_subset_superset_index = gko::vector(exec); + tmp_subset_begin.push_back(tmp_indices.get_data()[0]); + tmp_subset_superset_index.push_back(0); + for (auto i = 1; i < num_indices; ++i) { + if (tmp_indices.get_data()[i] == (tmp_indices.get_data()[i - 1]) + 1) { + continue; + } else { + tmp_subset_end.push_back(tmp_indices.get_data()[i - 1] + 1); + tmp_subset_superset_index.push_back( + tmp_subset_superset_index.back() + tmp_subset_end.back() - + tmp_subset_begin.back()); + if (i + 1 < num_indices) { + tmp_subset_begin.push_back(tmp_indices.get_data()[i]); + } + } + } + tmp_subset_end.push_back(tmp_indices.get_data()[num_indices - 1] + 1); + tmp_subset_superset_index.push_back(tmp_subset_superset_index.back() + + tmp_subset_end.back() - + tmp_subset_begin.back()); + + GKO_ASSERT(tmp_subset_begin.size() == tmp_subset_end.size()); + GKO_ASSERT((tmp_subset_begin.size() + 1) == + tmp_subset_superset_index.size()); + *subset_begin = std::move(gko::Array( + exec, tmp_subset_begin.data(), + tmp_subset_begin.data() + tmp_subset_begin.size())); + *subset_end = std::move( + gko::Array(exec, tmp_subset_end.data(), + tmp_subset_end.data() + tmp_subset_end.size())); + *superset_indices = std::move(gko::Array( + exec, tmp_subset_superset_index.data(), + tmp_subset_superset_index.data() + tmp_subset_superset_index.size())); +} GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_INDEX_SET_POPULATE_KERNEL); +template +void global_to_local(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *global_indices, + Array *local_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_GLOBAL_TO_LOCAL_KERNEL); + + +template +void local_to_global(std::shared_ptr exec, + const IndexType index_space_size, + const Array *subset_begin, + const Array *subset_end, + const Array *superset_indices, + const Array *local_indices, + Array *global_indices) +{} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_INDEX_SET_LOCAL_TO_GLOBAL_KERNEL); + + } // namespace index_set } // namespace reference } // namespace kernels