Add distributed SpGEMM

common/cuda_hip/distributed/matrix_kernels.cpp

@@ -9,6 +9,7 @@
 #include <thrust/distance.h>
 #include <thrust/execution_policy.h>
 #include <thrust/for_each.h>
+#include <thrust/iterator/permutation_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
 #include <thrust/iterator/zip_iterator.h>
 #include <thrust/sequence.h>
@@ -198,6 +199,172 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE(
     GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG);
 
 
+template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
+struct local_rows_input_type {
+    LocalIndexType row;
+    GlobalIndexType col;
+    ValueType val;
+    size_type col_range;
+
+    __forceinline__ __device__ __host__ local_rows_input_type(
+        thrust::tuple<LocalIndexType, GlobalIndexType, ValueType, size_type> t)
+        : row(thrust::get<0>(t)),
+          col(thrust::get<1>(t)),
+          val(thrust::get<2>(t)),
+          col_range(thrust::get<3>(t))
+    {}
+};
+
+
+template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
+void separate_diag_off_diag_local_rows(
+    std::shared_ptr<const DefaultExecutor> exec,
+    const array<LocalIndexType>& row_idxs,
+    const array<LocalIndexType>& col_idxs,
+    const array<GlobalIndexType>& col_map, const array<ValueType>& values,
+    const experimental::distributed::Partition<LocalIndexType, GlobalIndexType>*
+        col_partition,
+    comm_index_type local_part, array<LocalIndexType>& diag_row_idxs,
+    array<LocalIndexType>& diag_col_idxs, array<ValueType>& diag_values,
+    array<LocalIndexType>& off_diag_row_idxs,
+    array<GlobalIndexType>& off_diag_col_idxs,
+    array<ValueType>& off_diag_values)
+{
+    auto col_part_ids = col_partition->get_part_ids();
+    const auto* col_range_bounds = col_partition->get_range_bounds();
+    const auto* col_range_starting_indices =
+        col_partition->get_range_starting_indices();
+    const auto num_col_ranges = col_partition->get_num_ranges();
+    const auto num_input_elements = row_idxs.get_size();
+
+    auto policy = thrust_policy(exec);
+
+    // the input columns are compact indices into col_map; look up each entry's
+    // global column through a permutation iterator, then classify by range.
+    auto global_col_it = thrust::make_permutation_iterator(
+        col_map.get_const_data(), col_idxs.get_const_data());
+    array<size_type> col_range_ids{exec, num_input_elements};
+    thrust::upper_bound(policy, col_range_bounds + 1,
+                        col_range_bounds + num_col_ranges + 1, global_col_it,
+                        global_col_it + num_input_elements,
+                        col_range_ids.get_data());
+
+    // count number of diag<0> and off-diag<1> elements.
+    auto num_elements_pair = thrust::transform_reduce(
+        policy, col_range_ids.get_const_data(),
+        col_range_ids.get_const_data() + num_input_elements,
+        [local_part,
+         col_part_ids] __host__ __device__(const size_type col_range_id) {
+            bool is_diag_entry = col_part_ids[col_range_id] == local_part;
+            return thrust::make_tuple(
+                is_diag_entry ? size_type{1} : size_type{0},
+                is_diag_entry ? size_type{0} : size_type{1});
+        },
+        thrust::make_tuple(size_type{}, size_type{}),
+        [] __host__ __device__(const thrust::tuple<size_type, size_type>& a,
+                               const thrust::tuple<size_type, size_type>& b) {
+            return thrust::make_tuple(thrust::get<0>(a) + thrust::get<0>(b),
+                                      thrust::get<1>(a) + thrust::get<1>(b));
+        });
+    auto num_diag_elements = thrust::get<0>(num_elements_pair);
+    auto num_off_diag_elements = thrust::get<1>(num_elements_pair);
+
+    // define global-to-local map for column indices
+    auto map_to_local_col =
+        [col_range_bounds, col_range_starting_indices] __host__ __device__(
+            const GlobalIndexType col, const size_type range_id) {
+            return static_cast<LocalIndexType>(col -
+                                               col_range_bounds[range_id]) +
+                   col_range_starting_indices[range_id];
+        };
+
+    using input_type = local_rows_input_type<device_type<ValueType>,
+                                             LocalIndexType, GlobalIndexType>;
+    auto input_it = thrust::make_zip_iterator(
+        thrust::make_tuple(row_idxs.get_const_data(), global_col_it,
+                           as_device_type(values.get_const_data()),
+                           col_range_ids.get_const_data()));
+
+    // copy and transform diag entries into arrays. the row is already local
+    // and passed through unchanged, only the column is mapped to a local
+    // index
+    diag_row_idxs.resize_and_reset(num_diag_elements);
+    diag_col_idxs.resize_and_reset(num_diag_elements);
+    diag_values.resize_and_reset(num_diag_elements);
+    auto diag_it = thrust::make_transform_iterator(
+        input_it,
+        [map_to_local_col] __host__ __device__(const input_type input) {
+            auto local_col = map_to_local_col(input.col, input.col_range);
+            return thrust::make_tuple(input.row, local_col, input.val);
+        });
+    thrust::copy_if(policy, diag_it, diag_it + num_input_elements,
+                    col_range_ids.get_const_data(),
+                    thrust::make_zip_iterator(thrust::make_tuple(
+                        diag_row_idxs.get_data(), diag_col_idxs.get_data(),
+                        as_device_type(diag_values.get_data()))),
+                    [local_part, col_part_ids] __host__ __device__(
+                        const size_type col_range_id) {
+                        return col_part_ids[col_range_id] == local_part;
+                    });
+
+
+    // copy off-diag entries into arrays. this keeps global column indices
+    off_diag_row_idxs.resize_and_reset(num_off_diag_elements);
+    off_diag_col_idxs.resize_and_reset(num_off_diag_elements);
+    off_diag_values.resize_and_reset(num_off_diag_elements);
+    auto off_diag_it = thrust::make_transform_iterator(
+        input_it, [] __host__ __device__(const input_type input) {
+            return thrust::make_tuple(input.row, input.col, input.val);
+        });
+    thrust::copy_if(
+        policy, off_diag_it, off_diag_it + num_input_elements,
+        col_range_ids.get_const_data(),
+        thrust::make_zip_iterator(thrust::make_tuple(
+            off_diag_row_idxs.get_data(), off_diag_col_idxs.get_data(),
+            as_device_type(off_diag_values.get_data()))),
+        [local_part,
+         col_part_ids] __host__ __device__(const size_type col_range_id) {
+            return col_part_ids[col_range_id] != local_part;
+        });
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE(
+    GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG_LOCAL_ROWS);
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void compress_columns(std::shared_ptr<const DefaultExecutor> exec,
+                      const array<GlobalIndexType>& global_cols,
+                      array<LocalIndexType>& compact_cols,
+                      array<GlobalIndexType>& distinct_cols)
+{
+    const auto n = global_cols.get_size();
+    auto policy = thrust_policy(exec);
+
+    // distinct_cols = sorted unique of the input columns
+    array<GlobalIndexType> sorted{exec, n};
+    thrust::copy_n(policy, global_cols.get_const_data(), n, sorted.get_data());
+    thrust::sort(policy, sorted.get_data(), sorted.get_data() + n);
+    auto sorted_end =
+        thrust::unique(policy, sorted.get_data(), sorted.get_data() + n);
+    const auto num_distinct =
+        static_cast<size_type>(thrust::distance(sorted.get_data(), sorted_end));
+    distinct_cols.resize_and_reset(num_distinct);
+    thrust::copy_n(policy, sorted.get_data(), num_distinct,
+                   distinct_cols.get_data());
+
+    // compact_cols[i] = position of global_cols[i] in distinct_cols
+    compact_cols.resize_and_reset(n);
+    thrust::lower_bound(policy, distinct_cols.get_const_data(),
+                        distinct_cols.get_const_data() + num_distinct,
+                        global_cols.get_const_data(),
+                        global_cols.get_const_data() + n,
+                        compact_cols.get_data());
+}
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(GKO_DECLARE_COMPRESS_COLUMNS);
+
+
 }  // namespace distributed_matrix
 }  // namespace GKO_DEVICE_NAMESPACE
 }  // namespace kernels

core/CMakeLists.txt

@@ -154,6 +154,7 @@ if(GINKGO_BUILD_MPI)
             distributed/partition_helpers.cpp
             distributed/preconditioner/schwarz.cpp
             distributed/row_gatherer.cpp
+            distributed/spgemm.cpp
             distributed/vector.cpp
             distributed/vector_cache.cpp
             mpi/exception.cpp

core/device_hooks/common_kernels.inc.cpp

@@ -391,6 +391,9 @@ namespace distributed_matrix {
 
 GKO_STUB_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE_BASE(
     GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG);
+GKO_STUB_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE_BASE(
+    GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG_LOCAL_ROWS);
+GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_COMPRESS_COLUMNS);
 
 
 }  // namespace distributed_matrix

core/distributed/matrix.cpp

@@ -212,6 +212,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::convert_to(
     as<Cloneable>(result->row_gatherer_)
         ->copy_from(as<Cloneable>(this->row_gatherer_));
     result->imap_ = this->imap_;
+    result->row_partition_ = this->row_partition_;
     result->set_size(this->get_size());
 }
 
@@ -229,6 +230,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::move_to(
     as<Cloneable>(result->row_gatherer_)
         ->move_from(as<Cloneable>(this->row_gatherer_));
     result->imap_ = std::move(this->imap_);
+    result->row_partition_ = std::move(this->row_partition_);
     result->set_size(this->get_size());
     this->set_size({});
 }
@@ -249,6 +251,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::convert_to(
     as<Cloneable>(result->row_gatherer_)
         ->copy_from(as<Cloneable>(this->row_gatherer_));
     result->imap_ = this->imap_;
+    result->row_partition_ = this->row_partition_;
     result->set_size(this->get_size());
 }
 
@@ -266,6 +269,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::move_to(
         ->move_from(as<Cloneable>(this->off_diag_mtx_.get()));
     result->row_gatherer_->move_from(this->row_gatherer_);
     result->imap_ = std::move(this->imap_);
+    result->row_partition_ = std::move(this->row_partition_);
     result->set_size(this->get_size());
     this->set_size({});
 }
@@ -286,6 +290,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::convert_to(
         ->copy_from(as<Cloneable>(this->off_diag_mtx_.get()));
     result->row_gatherer_->copy_from(this->row_gatherer_);
     result->imap_ = this->imap_;
+    result->row_partition_ = this->row_partition_;
     result->set_size(this->get_size());
 }
 
@@ -303,6 +308,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::move_to(
         ->move_from(as<Cloneable>(this->off_diag_mtx_.get()));
     result->row_gatherer_->move_from(this->row_gatherer_);
     result->imap_ = std::move(this->imap_);
+    result->row_partition_ = std::move(this->row_partition_);
     result->set_size(this->get_size());
     this->set_size({});
 }
@@ -364,6 +370,7 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::read_distributed(
 
     imap_ = index_map<local_index_type, global_index_type>(
         exec, col_partition, comm.rank(), global_off_diag_col_idxs);
+    row_partition_ = row_partition;
 
     auto off_diag_col_idxs =
         imap_.map_to_local(global_off_diag_col_idxs, index_space::non_local);
@@ -722,6 +729,7 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::operator=(
         as<Cloneable>(row_gatherer_)
             ->copy_from(as<Cloneable>(other.row_gatherer_));
         imap_ = other.imap_;
+        row_partition_ = other.row_partition_;
     }
     return *this;
 }
@@ -742,6 +750,7 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::operator=(Matrix&& other)
         as<Cloneable>(row_gatherer_)
             ->move_from(as<Cloneable>(other.row_gatherer_));
         imap_ = std::move(other.imap_);
+        row_partition_ = std::move(other.row_partition_);
     }
     return *this;
 }

core/distributed/matrix_kernels.hpp

@@ -35,12 +35,41 @@ namespace kernels {
         array<ValueType>& off_diag_values)
 
 
-#define GKO_DECLARE_ALL_AS_TEMPLATES                                    \
-    using comm_index_type = experimental::distributed::comm_index_type; \
-    template <typename ValueType, typename LocalIndexType,              \
-              typename GlobalIndexType>                                 \
-    GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG(ValueType, LocalIndexType,       \
-                                       GlobalIndexType)
+#define GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG_LOCAL_ROWS(                         \
+    ValueType, LocalIndexType, GlobalIndexType)                                \
+    void separate_diag_off_diag_local_rows(                                    \
+        std::shared_ptr<const DefaultExecutor> exec,                           \
+        const array<LocalIndexType>& row_idxs,                                 \
+        const array<LocalIndexType>& col_idxs,                                 \
+        const array<GlobalIndexType>& col_map, const array<ValueType>& values, \
+        const experimental::distributed::Partition<                            \
+            LocalIndexType, GlobalIndexType>* col_partition,                   \
+        comm_index_type local_part, array<LocalIndexType>& diag_row_idxs,      \
+        array<LocalIndexType>& diag_col_idxs, array<ValueType>& diag_values,   \
+        array<LocalIndexType>& off_diag_row_idxs,                              \
+        array<GlobalIndexType>& off_diag_col_idxs,                             \
+        array<ValueType>& off_diag_values)
+
+
+#define GKO_DECLARE_COMPRESS_COLUMNS(LocalIndexType, GlobalIndexType)  \
+    void compress_columns(std::shared_ptr<const DefaultExecutor> exec, \
+                          const array<GlobalIndexType>& global_cols,   \
+                          array<LocalIndexType>& compact_cols,         \
+                          array<GlobalIndexType>& distinct_cols)
+
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                                         \
+    using comm_index_type = experimental::distributed::comm_index_type;      \
+    template <typename ValueType, typename LocalIndexType,                   \
+              typename GlobalIndexType>                                      \
+    GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG(ValueType, LocalIndexType,            \
+                                       GlobalIndexType);                     \
+    template <typename ValueType, typename LocalIndexType,                   \
+              typename GlobalIndexType>                                      \
+    GKO_DECLARE_SEPARATE_DIAG_OFF_DIAG_LOCAL_ROWS(ValueType, LocalIndexType, \
+                                                  GlobalIndexType);          \
+    template <typename LocalIndexType, typename GlobalIndexType>             \
+    GKO_DECLARE_COMPRESS_COLUMNS(LocalIndexType, GlobalIndexType)
 
 
 GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(distributed_matrix,