Merge distributed vector

This PR adds support for a (row-wise) distributed (multi-)vector. It supports most operation of the dense class. These vector operations are supported on all devices that support the corresponding dense operation. Only the initialization through `read_distributed` is only supported on reference and openmp.

Related PR: #961
Related PR: #1030

CMakeLists.txt

@@ -88,6 +88,8 @@ option(GINKGO_DPCPP_SINGLE_MODE "Do not compile double kernels for the DPC++ bac
 option(GINKGO_INSTALL_RPATH "Set the RPATH when installing its libraries." ON)
 option(GINKGO_INSTALL_RPATH_ORIGIN "Add $ORIGIN (Linux) or @loader_path (MacOS) to the installation RPATH." ON)
 option(GINKGO_INSTALL_RPATH_DEPENDENCIES "Add dependencies to the installation RPATH." OFF)
+option(GINKGO_FORCE_GPU_AWARE_MPI "Assert that the MPI library is GPU aware. This forces Ginkgo to assume that GPU aware functionality is available (OFF (default) or ON), but may fail
+     catastrophically in case the MPI implementation is not GPU Aware, and GPU aware functionality has been forced" OFF)
 
 # load executor-specific configuration
 if(GINKGO_BUILD_CUDA)
@@ -204,8 +206,14 @@ else()
     message(STATUS "HWLOC is being forcibly switched off")
 endif()
 
+set(GINKGO_HAVE_GPU_AWARE_MPI OFF)
 if(GINKGO_BUILD_MPI)
     find_package(MPI REQUIRED)
+    if(GINKGO_FORCE_GPU_AWARE_MPI)
+        set(GINKGO_HAVE_GPU_AWARE_MPI ON)
+    else()
+        set(GINKGO_HAVE_GPU_AWARE_MPI OFF)
+    endif()
 endif()
 
 # Try to find the third party packages before using our subdirectories

cmake/create_test.cmake

@@ -7,7 +7,7 @@ function(ginkgo_build_test_name test_name target_name)
         ${PROJECT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR})
     string(REPLACE "/" "_" TEST_TARGET_NAME "${REL_BINARY_DIR}/${test_name}")
     set(${target_name} ${TEST_TARGET_NAME} PARENT_SCOPE)
-endfunction()
+endfunction(ginkgo_build_test_name)
 
 ## Set up shared target properties and handle ADDITIONAL_LIBRARIES/ADDITIONAL_INCLUDES
 ## `MPI_SIZE size` causes the tests to be run with `size` MPI processes.
@@ -80,7 +80,7 @@ function(ginkgo_create_dpcpp_test test_name)
     target_compile_features(${test_target_name} PUBLIC cxx_std_17)
     target_compile_options(${test_target_name} PRIVATE ${GINKGO_DPCPP_FLAGS})
     target_link_options(${test_target_name} PRIVATE -fsycl-device-code-split=per_kernel)
-    ginkgo_set_test_target_properties(${test_target_name} ${ARGN})
+    ginkgo_internal_add_test(${test_target_name} ${ARGN})
     ginkgo_add_test(${test_name} ${test_target_name} ${ARGN})
     # Note: MKL_ENV is empty on linux. Maybe need to apply MKL_ENV to all test.
     if (MKL_ENV)
@@ -115,7 +115,7 @@ function(ginkgo_create_cuda_test_internal test_name filename test_target_name)
     if(CMAKE_VERSION VERSION_GREATER_EQUAL 3.18)
         set_target_properties(${test_target_name} PROPERTIES CUDA_ARCHITECTURES OFF)
     endif()
-    ginkgo_set_test_target_properties(${test_target_name} ${ARGN})
+    ginkgo_internal_add_test(${test_target_name} ${ARGN})
     ginkgo_add_test(${test_name} ${test_target_name} ${ARGN})
 endfunction(ginkgo_create_cuda_test_internal)
 
@@ -205,7 +205,7 @@ function(ginkgo_create_common_test_internal test_name exec_type exec)
         target_compile_definitions(${test_target_name} PRIVATE GINKGO_COMMON_SINGLE_MODE=1)
         target_compile_definitions(${test_target_name} PRIVATE GINKGO_DPCPP_SINGLE_MODE=1)
     endif()
-    ginkgo_set_test_target_properties(${test_target_name} ${ARGN})
+    ginkgo_internal_add_test(${test_target_name} ${ARGN})
     ginkgo_add_test(${test_name}_${exec} ${test_target_name} ${ARGN})
 endfunction(ginkgo_create_common_test_internal)
 

cmake/get_info.cmake

@@ -130,7 +130,7 @@ foreach(log_type ${log_types})
         "GINKGO_BUILD_OMP;GINKGO_BUILD_MPI;GINKGO_BUILD_REFERENCE;GINKGO_BUILD_CUDA;GINKGO_BUILD_HIP;GINKGO_BUILD_DPCPP")
     ginkgo_print_module_footer(${${log_type}} "  Enabled features:")
     ginkgo_print_foreach_variable(${${log_type}}
-        "GINKGO_MIXED_PRECISION")
+        "GINKGO_MIXED_PRECISION;GINKGO_HAVE_GPU_AWARE_MPI")
     ginkgo_print_module_footer(${${log_type}} "  Tests, benchmarks and examples:")
     ginkgo_print_foreach_variable(${${log_type}}
         "GINKGO_BUILD_TESTS;GINKGO_FAST_TESTS;GINKGO_BUILD_EXAMPLES;GINKGO_EXTLIB_EXAMPLE;GINKGO_BUILD_BENCHMARKS;GINKGO_BENCHMARK_ENABLE_TUNING")

common/unified/matrix/dense_kernels.cpp

@@ -380,6 +380,37 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
     GKO_DECLARE_DENSE_COUNT_NONZEROS_PER_ROW_KERNEL_SIZE_T);
 
 
+template <typename ValueType>
+void compute_squared_norm2(std::shared_ptr<const DefaultExecutor> exec,
+                           const matrix::Dense<ValueType>* x,
+                           matrix::Dense<remove_complex<ValueType>>* result)
+{
+    run_kernel_col_reduction(
+        exec,
+        [] GKO_KERNEL(auto i, auto j, auto x) { return squared_norm(x(i, j)); },
+        GKO_KERNEL_REDUCE_SUM(remove_complex<ValueType>), result->get_values(),
+        x->get_size(), x);
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
+    GKO_DECLARE_DENSE_COMPUTE_SQUARED_NORM2_KERNEL);
+
+
+template <typename ValueType>
+void compute_sqrt(std::shared_ptr<const DefaultExecutor> exec,
+                  matrix::Dense<ValueType>* x)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto row, auto col, auto x) {
+            x(row, col) = sqrt(x(row, col));
+        },
+        x->get_size(), x);
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_SQRT_KERNEL);
+
+
 template <typename ValueType, typename IndexType>
 void symm_permute(std::shared_ptr<const DefaultExecutor> exec,
                   const array<IndexType>* permutation_indices,

core/CMakeLists.txt

@@ -6,6 +6,7 @@ target_sources(ginkgo
     base/array.cpp
     base/combination.cpp
     base/composition.cpp
+    base/dense_cache.cpp
     base/device_matrix_data.cpp
     base/executor.cpp
     base/index_set.cpp
@@ -67,7 +68,9 @@ endif()
 
 if (GINKGO_BUILD_MPI)
     target_sources(ginkgo
-        PRIVATE mpi/exception.cpp)
+        PRIVATE
+        mpi/exception.cpp
+        distributed/vector.cpp)
 endif()
 
 ginkgo_compile_features(ginkgo)

core/base/dense_cache.cpp

@@ -0,0 +1,69 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2022, the Ginkgo authors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+******************************<GINKGO LICENSE>*******************************/
+
+#include <ginkgo/core/base/dense_cache.hpp>
+
+
+#include <ginkgo/core/matrix/dense.hpp>
+
+
+namespace gko {
+namespace detail {
+
+
+template <typename ValueType>
+void DenseCache<ValueType>::init(std::shared_ptr<const Executor> exec,
+                                 dim<2> size) const
+{
+    if (!vec || vec->get_size() != size || vec->get_executor() != exec) {
+        vec = matrix::Dense<ValueType>::create(exec, size);
+    }
+}
+
+
+template <typename ValueType>
+void DenseCache<ValueType>::init_from(
+    const matrix::Dense<ValueType>* template_vec) const
+{
+    if (!vec || vec->get_size() != template_vec->get_size() ||
+        vec->get_executor() != template_vec->get_executor()) {
+        vec = matrix::Dense<ValueType>::create_with_config_of(template_vec);
+    }
+}
+
+
+#define GKO_DECLARE_DENSE_CACHE(_type) class DenseCache<_type>
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_CACHE);
+
+
+}  // namespace detail
+}  // namespace gko

core/device_hooks/common_kernels.inc.cpp

@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "core/components/prefix_sum_kernels.hpp"
 #include "core/components/reduce_array_kernels.hpp"
 #include "core/distributed/partition_kernels.hpp"
+#include "core/distributed/vector_kernels.hpp"
 #include "core/factorization/cholesky_kernels.hpp"
 #include "core/factorization/factorization_kernels.hpp"
 #include "core/factorization/ic_kernels.hpp"
@@ -144,6 +145,13 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
         GKO_NOT_COMPILED(GKO_HOOK_MODULE);                       \
     GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE_2(_macro)
 
+#define GKO_STUB_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE(_macro) \
+    template <typename ValueType, typename LocalIndexType, \
+              typename GlobalIndexType>                    \
+    _macro(ValueType, LocalIndexType, GlobalIndexType)     \
+        GKO_NOT_COMPILED(GKO_HOOK_MODULE);                 \
+    GKO_INSTANTIATE_FOR_EACH_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE(_macro)
+
 #define GKO_STUB_TEMPLATE_TYPE(_macro)                   \
     template <typename IndexType>                        \
     _macro(IndexType) GKO_NOT_COMPILED(GKO_HOOK_MODULE); \
@@ -243,6 +251,16 @@ GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_PARTITION_IS_ORDERED);
 }  // namespace partition
 
 
+namespace distributed_vector {
+
+
+GKO_STUB_VALUE_AND_LOCAL_GLOBAL_INDEX_TYPE(
+    GKO_DECLARE_DISTRIBUTED_VECTOR_BUILD_LOCAL);
+
+
+}
+
+
 namespace dense {
 
 
@@ -264,6 +282,8 @@ GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_CONJ_DOT_DISPATCH_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM2_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM2_DISPATCH_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM1_KERNEL);
+GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_SQUARED_NORM2_KERNEL);
+GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_SQRT_KERNEL);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DENSE_FILL_IN_MATRIX_DATA_KERNEL);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DENSE_CONVERT_TO_COO_KERNEL);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DENSE_CONVERT_TO_CSR_KERNEL);