Fix distributed vector merge

CMakeLists.txt

@@ -98,6 +98,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)
 
 set(GINKGO_CIRCULAR_DEPS_FLAGS "-Wl,--no-undefined")
 
@@ -189,8 +191,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()
 
 # We keep using NVCC/HCC for consistency with previous releases even if AMD

cmake/create_test.cmake

@@ -3,11 +3,18 @@ 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)
 
-function(ginkgo_set_test_target_properties test_name test_target_name)
-    file(RELATIVE_PATH REL_BINARY_DIR
-        ${PROJECT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR})
+function(ginkgo_create_gtest_mpi_main)
+    add_library(gtest_mpi_main "")
+    target_sources(gtest_mpi_main
+        PRIVATE
+        ${PROJECT_SOURCE_DIR}/core/test/mpi/gtest/mpi_listener.cpp)
+    find_package(MPI REQUIRED)
+    target_link_libraries(gtest_mpi_main PRIVATE GTest::GTest MPI::MPI_CXX)
+endfunction(ginkgo_create_gtest_mpi_main)
+
+function(ginkgo_set_test_target_default_properties test_name test_target_name)
     set_target_properties(${test_target_name} PROPERTIES
         OUTPUT_NAME ${test_name})
     if (GINKGO_FAST_TESTS)
@@ -19,20 +26,41 @@ function(ginkgo_set_test_target_properties test_name test_target_name)
     if (GINKGO_CHECK_CIRCULAR_DEPS)
         target_link_libraries(${test_target_name} PRIVATE "${GINKGO_CIRCULAR_DEPS_FLAGS}")
     endif()
+    target_include_directories(${test_target_name} PRIVATE ${Ginkgo_BINARY_DIR})
+    target_link_libraries(${test_target_name} PRIVATE ginkgo)
+endfunction(ginkgo_set_test_target_default_properties)
+
+function(ginkgo_internal_add_test test_name test_target_name)
+    file(RELATIVE_PATH REL_BINARY_DIR
+        ${PROJECT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR})
+    ginkgo_set_test_target_default_properties(${test_name} ${test_target_name})
     add_test(NAME ${REL_BINARY_DIR}/${test_name}
         COMMAND ${test_target_name}
         WORKING_DIRECTORY "$<TARGET_FILE_DIR:ginkgo>")
-    target_include_directories(${test_target_name} PRIVATE ${Ginkgo_BINARY_DIR})
-    target_link_libraries(${test_target_name} PRIVATE ginkgo GTest::Main GTest::GTest)
-endfunction()
+    target_link_libraries(${test_target_name} PRIVATE GTest::Main GTest::GTest)
+endfunction(ginkgo_internal_add_test)
+
+function(ginkgo_internal_add_mpi_test test_name test_target_name num_mpi_procs)
+    file(RELATIVE_PATH REL_BINARY_DIR
+        ${PROJECT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR})
+    ginkgo_set_test_target_default_properties(${test_name} ${test_target_name})
+    if(NOT TARGET gtest_mpi_main)
+        ginkgo_create_gtest_mpi_main()
+    endif()
+    target_link_libraries(${test_target_name} PRIVATE gtest_mpi_main GTest::GTest MPI::MPI_CXX)
+    set(test_param ${MPIEXEC_NUMPROC_FLAG} ${num_mpi_procs} ${OPENMPI_RUN_AS_ROOT_FLAG}
+                   ${CMAKE_BINARY_DIR}/${REL_BINARY_DIR}/${test_name})
+    add_test(NAME ${REL_BINARY_DIR}/${test_name}
+        COMMAND ${MPIEXEC_EXECUTABLE} ${test_param})
+endfunction(ginkgo_internal_add_mpi_test)
 
 function(ginkgo_create_test test_name)
     ginkgo_build_test_name(${test_name} test_target_name)
     add_executable(${test_target_name} ${test_name}.cpp)
     target_compile_features(${test_target_name} PUBLIC cxx_std_14)
     target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
     target_link_libraries(${test_target_name} PRIVATE ${ARGN})
-    ginkgo_set_test_target_properties(${test_name} ${test_target_name})
+    ginkgo_internal_add_test(${test_name} ${test_target_name})
 endfunction(ginkgo_create_test)
 
 function(ginkgo_create_dpcpp_test test_name)
@@ -42,7 +70,7 @@ function(ginkgo_create_dpcpp_test test_name)
     target_compile_options(${test_target_name} PRIVATE "${GINKGO_DPCPP_FLAGS}")
     target_compile_options(${test_target_name} PRIVATE "${GINKGO_COMPILER_FLAGS}")
     target_link_options(${test_target_name} PRIVATE -fsycl-device-code-split=per_kernel)
-    ginkgo_set_test_target_properties(${test_name} ${test_target_name})
+    ginkgo_internal_add_test(${test_name} ${test_target_name})
     # Note: MKL_ENV is empty on linux. Maybe need to apply MKL_ENV to all test.
     if (MKL_ENV)
        set_tests_properties(${test_target_name} PROPERTIES ENVIRONMENT "${MKL_ENV}")
@@ -57,29 +85,16 @@ function(ginkgo_create_thread_test test_name)
     target_compile_features(${test_target_name} PUBLIC cxx_std_14)
     target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
     target_link_libraries(${test_target_name} PRIVATE Threads::Threads ${ARGN})
-    ginkgo_set_test_target_properties(${test_name} ${test_target_name})
+    ginkgo_internal_add_test(${test_name} ${test_target_name})
 endfunction(ginkgo_create_thread_test)
 
 function(ginkgo_create_mpi_test test_name num_mpi_procs)
-  file(RELATIVE_PATH REL_BINARY_DIR
-    ${PROJECT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR})
-  string(REPLACE "/" "_" TEST_TARGET_NAME "${REL_BINARY_DIR}/${test_name}")
-  add_executable(${TEST_TARGET_NAME} ${test_name}.cpp)
-  set_target_properties(${TEST_TARGET_NAME} PROPERTIES
-    OUTPUT_NAME ${test_name})
-  if (GINKGO_CHECK_CIRCULAR_DEPS)
-    target_link_libraries(${TEST_TARGET_NAME} PRIVATE "${GINKGO_CIRCULAR_DEPS_FLAGS}")
-  endif()
-  if("${GINKGO_MPI_EXEC_SUFFIX}" MATCHES ".openmpi" AND MPI_RUN_AS_ROOT)
-      set(OPENMPI_RUN_AS_ROOT_FLAG "--allow-run-as-root")
-  else()
-      set(OPENMPI_RUN_AS_ROOT_FLAG "")
-  endif()
-  target_link_libraries(${TEST_TARGET_NAME} PRIVATE ginkgo gtest_mpi_main GTest::GTest ${ARGN})
-  target_link_libraries(${TEST_TARGET_NAME} PRIVATE MPI::MPI_CXX)
-  set(test_param ${MPIEXEC_NUMPROC_FLAG} ${num_mpi_procs} ${OPENMPI_RUN_AS_ROOT_FLAG} ${CMAKE_BINARY_DIR}/${REL_BINARY_DIR}/${test_name})
-  add_test(NAME ${REL_BINARY_DIR}/${test_name}
-    COMMAND ${MPIEXEC_EXECUTABLE} ${test_param})
+    ginkgo_build_test_name(${test_name} test_target_name)
+    add_executable(${test_target_name} ${test_name}.cpp)
+    target_compile_features(${test_target_name} PUBLIC cxx_std_14)
+    target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
+    target_link_libraries(${test_target_name} PRIVATE ${ARGN})
+    ginkgo_internal_add_mpi_test(${test_name} ${test_target_name} ${num_mpi_procs})
 endfunction(ginkgo_create_mpi_test)
 
 function(ginkgo_create_test_cpp_cuda_header test_name)
@@ -89,7 +104,7 @@ function(ginkgo_create_test_cpp_cuda_header test_name)
     target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
     target_include_directories(${test_target_name} PRIVATE "${CUDA_INCLUDE_DIRS}")
     target_link_libraries(${test_target_name} PRIVATE ${ARGN})
-    ginkgo_set_test_target_properties(${test_name} ${test_target_name})
+    ginkgo_internal_add_test(${test_name} ${test_target_name})
 endfunction(ginkgo_create_test_cpp_cuda_header)
 
 function(ginkgo_create_cuda_test test_name)
@@ -112,7 +127,7 @@ function(ginkgo_create_cuda_test test_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_name} ${test_target_name})
+    ginkgo_internal_add_test(${test_name} ${test_target_name})
 endfunction(ginkgo_create_cuda_test)
 
 function(ginkgo_create_hip_test test_name)
@@ -159,11 +174,12 @@ ginkgo_build_test_name(${test_name} test_target_name)
         ${HIPSPARSE_INCLUDE_DIRS}
         )
     target_link_libraries(${test_target_name} PRIVATE ${ARGN})
-    ginkgo_set_test_target_properties(${test_name} ${test_target_name})
+    ginkgo_internal_add_test(${test_name} ${test_target_name})
 endfunction(ginkgo_create_hip_test)
 
-function(ginkgo_create_common_test test_name)
-    cmake_parse_arguments(PARSE_ARGV 1 common_test "" "" "DISABLE_EXECUTORS;ADDITIONAL_LIBRARIES")
+function(ginkgo_internal_create_common_test_template test_name)
+    cmake_parse_arguments(PARSE_ARGV 1 common_test "" "TEST_TYPE" "DISABLE_EXECUTORS;ADDITIONAL_LIBRARIES;ADDITIONAL_TEST_PARAMETERS")
+    string(TOLOWER ${common_test_TEST_TYPE} test_type)
     set(executors)
     if(GINKGO_BUILD_OMP)
         list(APPEND executors omp)
@@ -200,18 +216,45 @@ function(ginkgo_create_common_test test_name)
             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_name}_${exec} ${test_target_name})
+        if(${test_type} STREQUAL default)
+            ginkgo_internal_add_test(${test_name}_${exec} ${test_target_name})
+        elseif(${test_type} STREQUAL mpi)
+            ginkgo_internal_add_mpi_test(${test_name}_${exec} ${test_target_name} ${common_test_ADDITIONAL_TEST_PARAMETERS})
+        else()
+            message(FATAL_ERROR "Encountered unrecognized test type ${test_type} during common test creation.")
+        endif()
     endforeach()
+endfunction(ginkgo_internal_create_common_test_template)
+
+function(ginkgo_create_common_test test_name)
+    ginkgo_internal_create_common_test_template(${test_name} TEST_TYPE default ${ARGN})
 endfunction(ginkgo_create_common_test)
 
+function(ginkgo_create_common_mpi_test test_name num_mpi_procs)
+    ginkgo_internal_create_common_test_template(${test_name} TEST_TYPE mpi ADDITIONAL_TEST_PARAMETERS ${num_mpi_procs} ${ARGN})
+endfunction(ginkgo_create_common_mpi_test)
+
 function(ginkgo_create_common_and_reference_test test_name)
     ginkgo_create_common_test(${test_name})
     ginkgo_build_test_name(${test_name} test_target_name)
     set(test_target_name ${test_target_name}_reference)
     add_executable(${test_target_name} ${test_name}.cpp)
     target_compile_features(${test_target_name} PUBLIC cxx_std_14)
     target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
-    target_compile_definitions(${test_target_name} PRIVATE EXEC_TYPE=ReferenceExecutor EXEC_NAMESPACE=reference)
+    target_compile_definitions(${test_target_name} PRIVATE EXEC_TYPE=ReferenceExecutor EXEC_NAMESPACE=reference GKO_COMPILING_REFERENCE)
+    target_link_libraries(${test_target_name} PRIVATE ${ARGN})
+    ginkgo_internal_add_test(${test_name}_reference ${test_target_name})
+endfunction()
+
+
+function(ginkgo_create_common_and_reference_mpi_test test_name num_mpi_procs)
+    ginkgo_create_common_mpi_test(${test_name} ${num_mpi_procs})
+    ginkgo_build_test_name(${test_name} test_target_name)
+    set(test_target_name ${test_target_name}_reference)
+    add_executable(${test_target_name} ${test_name}.cpp)
+    target_compile_features(${test_target_name} PUBLIC cxx_std_14)
+    target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
+    target_compile_definitions(${test_target_name} PRIVATE EXEC_TYPE=ReferenceExecutor EXEC_NAMESPACE=reference GKO_COMPILING_REFERENCE)
     target_link_libraries(${test_target_name} PRIVATE ${ARGN})
-    ginkgo_set_test_target_properties(${test_name}_reference ${test_target_name})
+    ginkgo_internal_add_mpi_test(${test_name}_reference ${test_target_name} ${num_mpi_procs})
 endfunction()

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
@@ -65,7 +66,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