diff --git a/.github/workflows/multi_gpu_cpp_test.yaml b/.github/workflows/multi_gpu_cpp_test.yaml index 515be4a330..328bd97353 100644 --- a/.github/workflows/multi_gpu_cpp_test.yaml +++ b/.github/workflows/multi_gpu_cpp_test.yaml @@ -39,6 +39,13 @@ on: script: type: string default: "ci/test_cpp_multi_gpu.sh" + container-options: + description: | + Extra 'docker run' options. NCCL allocates its inter-GPU buffers in + /dev/shm; the container default of 64 MB is too small and makes NCCL + setup fail with "No space left on device", so bump it with --shm-size. + type: string + default: "--shm-size=1g" continue-on-error: description: "Treat failures as non-blocking (useful while the MG suite is maturing)." type: boolean @@ -63,5 +70,6 @@ jobs: node_type: ${{ inputs.node_type }} arch: ${{ inputs.arch }} container_image: ${{ inputs.container_image }} + container-options: ${{ inputs.container-options }} script: ${{ inputs.script }} continue-on-error: ${{ inputs.continue-on-error }} diff --git a/ci/test_cpp_multi_gpu.sh b/ci/test_cpp_multi_gpu.sh index 2213f56dab..dc54637474 100755 --- a/ci/test_cpp_multi_gpu.sh +++ b/ci/test_cpp_multi_gpu.sh @@ -64,6 +64,10 @@ rapids-print-env rapids-logger "Check GPU usage" nvidia-smi +# Dump the GPU topology matrix +rapids-logger "Check GPU topology" +nvidia-smi topo -m + # Multi-GPU tests are meaningless on a single device — fail loudly rather than # passing a run that never exercised NCCL. GPU_COUNT=$(nvidia-smi -L | wc -l) @@ -74,6 +78,17 @@ if [ "${GPU_COUNT}" -lt 2 ]; then exit 1 fi +# Distributed PDLP parity tests use a small set of git-ignored MPS instances. +# Download only those direct-MPS fixtures here: downloading the full PDLP suite +# falls back to netlib conversion when S3 credentials are not available, and that +# path needs gcc (not present in this test environment). +rapids-logger "Download datasets" +python benchmarks/linear_programming/utils/get_datasets.py -datasets graph40-40 +python benchmarks/linear_programming/utils/get_datasets.py -datasets ex10 + +RAPIDS_DATASET_ROOT_DIR="$(realpath datasets)" +export RAPIDS_DATASET_ROOT_DIR + # Locate the installed gtest binaries (same search order as ci/run_ctests.sh). installed_test_location="${INSTALL_PREFIX:-${CONDA_PREFIX:-/usr}}/bin/gtests/libcuopt/" devcontainers_test_location="$(dirname "$(realpath "${BASH_SOURCE[0]}")")/../cpp/build/latest/gtests/libcuopt/" @@ -97,6 +112,11 @@ if [ "${#mg_tests[@]}" -eq 0 ]; then exit 0 fi +# PHB topology on the 2-GPU runner: disable direct GPU peer access (see topo above). +# NCCL then uses the SHM transport, which needs a larger /dev/shm than the 64 MB +# container default — the workflow launches the container with --shm-size for this. +export NCCL_P2P_DISABLE=1 + EXITCODE=0 for gt in "${mg_tests[@]}"; do test_name=$(basename "${gt}") diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index a603c69098..a5ea2a2a48 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -22,11 +22,11 @@ PYDISTCHECK_ARGS=( if [[ "${package_dir}" == "python/libcuopt" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '690Mi' + --max-allowed-size-compressed '960Mi' ) else PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '550Mi' + --max-allowed-size-compressed '800Mi' ) fi elif [[ "${package_dir}" != "python/cuopt" ]] && \ diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 6975b7a552..82ad4010c5 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -42,6 +42,7 @@ dependencies: - msgpack-numpy==0.4.8 - msgpack-python==1.2.1 - myst-parser +- nccl >=2.19 - ninja - notebook - numba-cuda>=0.22.1 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 2785e5e66c..b293e2d8c0 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -42,6 +42,7 @@ dependencies: - msgpack-numpy==0.4.8 - msgpack-python==1.2.1 - myst-parser +- nccl >=2.19 - ninja - notebook - numba-cuda>=0.22.1 diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 25ac9c4048..1cb8efbe75 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -42,6 +42,7 @@ dependencies: - msgpack-numpy==0.4.8 - msgpack-python==1.2.1 - myst-parser +- nccl >=2.19 - ninja - notebook - numba-cuda>=0.22.1 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index ad24412add..c49b656860 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -42,6 +42,7 @@ dependencies: - msgpack-numpy==0.4.8 - msgpack-python==1.2.1 - myst-parser +- nccl >=2.19 - ninja - notebook - numba-cuda>=0.22.1 diff --git a/conda/recipes/libcuopt/recipe.yaml b/conda/recipes/libcuopt/recipe.yaml index 77c957eade..f7c8c244db 100644 --- a/conda/recipes/libcuopt/recipe.yaml +++ b/conda/recipes/libcuopt/recipe.yaml @@ -89,6 +89,7 @@ cache: - libcusparse-dev - libnvjitlink-dev - libboost-devel + - nccl >=2.19 - tbb-devel - zlib - bzip2 @@ -127,6 +128,7 @@ outputs: - libcudss-dev >=0.7,<0.8 - libcusparse-dev - libnvjitlink-dev + - nccl >=2.19 - openssl - c-ares - libuuid @@ -141,6 +143,7 @@ outputs: - librmm =${{ minor_version }} - cuda-nvrtc - libcudss >=0.7,<0.8 + - nccl >=2.19 - openssl - c-ares - libuuid @@ -189,6 +192,7 @@ outputs: - libcublas - libcudss-dev >=0.7,<0.8 - libcusparse-dev + - nccl >=2.19 - openssl - c-ares - libgrpc >=1.78.0,<1.80.0a0 diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 67412162c2..6a287d72f7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -311,6 +311,44 @@ create_logger_macros(CUOPT "cuopt::default_logger()" include/cuopt) find_package(CUDSS REQUIRED) +# ################################################################################################## +# - NCCL (multi-GPU distributed PDLP) ------------------------------------------------------------- +# NCCL is shipped via the conda env; no canonical CMake config target, so look it +# up by name in the standard lib paths (plus CONDA_PREFIX as a hint). +set(NCCL_HINT_PREFIXES "") +if (DEFINED ENV{CONDA_PREFIX} AND NOT "$ENV{CONDA_PREFIX}" STREQUAL "") + list(APPEND NCCL_HINT_PREFIXES "$ENV{CONDA_PREFIX}") +endif () +find_path(NCCL_INCLUDE_DIR + NAMES nccl.h + HINTS ${NCCL_HINT_PREFIXES} + PATH_SUFFIXES include +) +find_library(NCCL_LIBRARY + NAMES nccl + HINTS ${NCCL_HINT_PREFIXES} + PATH_SUFFIXES lib lib64 +) +if (NOT NCCL_INCLUDE_DIR OR NOT NCCL_LIBRARY) + message(FATAL_ERROR "NCCL not found. Looked in ${NCCL_HINT_PREFIXES}. Install nccl-dev / libnccl-dev in the active env.") +endif () +add_library(nccl_external UNKNOWN IMPORTED GLOBAL) +set_target_properties(nccl_external PROPERTIES + IMPORTED_LOCATION "${NCCL_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${NCCL_INCLUDE_DIR}" +) +message(STATUS "Using NCCL: ${NCCL_LIBRARY}") + +# ################################################################################################## +# - KaMinPar (multi-threaded partitioning for distributed PDLP) ------------------------------------ +# Brought in the RAPIDS way (rapids_cpm_find): uses an installed KaMinPar (deb/rpm/conda, +# discovered via its CMake config) if present, otherwise builds the pinned source via CPM. +include(cmake/thirdparty/get_kaminpar.cmake) +if (NOT TARGET KaMinPar::KaMinPar) + message(FATAL_ERROR "KaMinPar::KaMinPar was not made available by get_kaminpar.cmake") +endif () +message(STATUS "Using KaMinPar (distributed PDLP graph partitioner)") + # ################################################################################################## # - gRPC and Protobuf setup ----------------------------------------------------------------------- @@ -555,6 +593,32 @@ target_include_directories(cuopt target_link_libraries(cuopt PRIVATE $) add_dependencies(cuopt PSLP) +# Link KaMinPar by file to avoid export dependency tracking (mirrors PSLP above). +# KaMinPar is a from-source static library fully embedded into libcuopt.so; it is never +# installed (INSTALL_KAMINPAR OFF) and consumers of cuopt::cuopt never use it, so it must +# not leak into cuopt's exported link interface (otherwise rapids_export fails with +# "target KaMinPar is not in any export set"). libKaMinPar.a is self-contained (the +# kaminpar-common OBJECT lib is archived into it); we only need its public headers +# (, which pulls in stdlib + TBB) at compile time. +target_include_directories(cuopt SYSTEM PRIVATE + $) +# partitioner.cpp includes . Because KaMinPar is linked by file, cuopt +# does NOT inherit KaMinPar's PUBLIC compile definitions, so we must replicate the ones that +# affect what compiles to: +# * TBB_PREVIEW_GLOBAL_CONTROL - includes , which older +# TBB gates behind this macro. +# * KAMINPAR_64BIT_EDGE_IDS - selects EdgeID = uint64 in . MUST stay in sync +# with the KAMINPAR_64BIT_* options in get_kaminpar.cmake; +# otherwise the public API types (e.g. copy_graph's span<> +# widths) mismatch the prebuilt archive => undefined refs. +target_compile_definitions(cuopt PRIVATE TBB_PREVIEW_GLOBAL_CONTROL KAMINPAR_64BIT_EDGE_IDS) +target_link_libraries(cuopt PRIVATE $) +# Real (non-imported) target only exists when KaMinPar is built from source via CPM; +# force it to build before cuopt since it is EXCLUDE_FROM_ALL. +if (TARGET KaMinPar) + add_dependencies(cuopt KaMinPar) +endif () + # ################################################################################################## # - link libraries -------------------------------------------------------------------------------- @@ -605,6 +669,7 @@ target_link_libraries(cuopt ${CUDSS_LIB_FILE} PRIVATE ${CUOPT_PRIVATE_CUDA_LIBS} + nccl_external $<$:protobuf::libprotobuf> $<$:gRPC::grpc++> ) diff --git a/cpp/cmake/thirdparty/get_kaminpar.cmake b/cpp/cmake/thirdparty/get_kaminpar.cmake new file mode 100644 index 0000000000..0ea104b360 --- /dev/null +++ b/cpp/cmake/thirdparty/get_kaminpar.cmake @@ -0,0 +1,95 @@ +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on + +# Multi-threaded graph partitioner for distributed PDLP. +# Uses rapids_cpm_find so a system / conda / .deb install of KaMinPar (which ships a +# CMake config package exporting KaMinPar::KaMinPar) is used when available, and +# otherwise the pinned source is cloned and built via CPM. KaMinPar depends on TBB, +# which cuOpt already requires (see find_package(TBB) for papilo). +function(find_and_configure_kaminpar) + set(oneValueArgs VERSION PINNED_TAG) + cmake_parse_arguments(PKG "" "${oneValueArgs}" "" ${ARGN}) + + # Silence the (many) warnings from KaMinPar and its bundled third-party + # dependencies (KaHyPar, KaGen, growt, ...). They are compiled from pinned + # upstream source we do not maintain, so their diagnostics only drown out + # cuOpt's own build output. Appending -w here is function-scoped: the CPM + # add_subdirectory below inherits these flags, but cuOpt's targets (compiled + # in the parent scope with -Wall -Werror) are unaffected. A system/conda + # install of KaMinPar builds nothing, so the flag is simply unused there. + string(APPEND CMAKE_C_FLAGS " -w") + string(APPEND CMAKE_CXX_FLAGS " -w") + + # NOTE: KaMinPar is intentionally NOT added to cuopt's BUILD/INSTALL export sets. + # It is a from-source static dependency that is fully embedded into libcuopt.so and + # never installed (INSTALL_KAMINPAR OFF below). Registering it in cuopt-exports would + # both break export generation ("target KaMinPar is not in any export set") and emit a + # bogus find_dependency(KaMinPar) into the installed cuopt config. It is linked by file + # in cpp/CMakeLists.txt (mirroring PSLP) so it stays out of cuopt's export interface. + rapids_cpm_find(KaMinPar ${PKG_VERSION} + GLOBAL_TARGETS KaMinPar::KaMinPar + CPM_ARGS + GIT_REPOSITORY https://github.com/KaHIP/KaMinPar.git + GIT_TAG ${PKG_PINNED_TAG} + EXCLUDE_FROM_ALL + OPTIONS + "KAMINPAR_BUILD_APPS OFF" + "KAMINPAR_BUILD_TOOLS OFF" + "KAMINPAR_BUILD_TESTS OFF" + "KAMINPAR_BUILD_BENCHMARKS OFF" + "KAMINPAR_BUILD_EXAMPLES OFF" + "KAMINPAR_BUILD_DISTRIBUTED OFF" + # KaMinPar defaults this to ON, which injects `-march=native -mtune=native` + # (and propagates the same flag to its bundled KaHyPar / KaGen via + # KAHYPAR_ENABLE_ARCH_COMPILE_OPTIMIZATIONS / KAGEN_USE_MARCH_NATIVE). + # That bakes in whatever ISA the build host happens to support + # (e.g. AVX-512 on Sapphire Rapids), which then SIGILLs on any test + # host with a leaner CPU. Force generic codegen instead. + "KAMINPAR_BUILD_WITH_MTUNE_NATIVE OFF" + # Timers use global state and force single-threaded use of the library + # interface; disable so cuOpt can call the partitioner freely. + "KAMINPAR_ENABLE_TIMERS OFF" + # Avoid an extra hard dependency on Google Sparsehash. + "KAMINPAR_BUILD_WITH_SPARSEHASH OFF" + # cuOpt's TBB is discovered via a legacy find that only exposes TBB::tbb + # (no TBB::tbbmalloc target); disable KaMinPar's optional tbbmalloc use. + "KAMINPAR_ENABLE_TBB_MALLOC OFF" + # Large LP constraint graphs can exceed 2^31 directed edges. + "KAMINPAR_64BIT_EDGE_IDS ON" + "INSTALL_KAMINPAR OFF" + # Build KaMinPar as a STATIC library that is embedded into libcuopt.so (linked + # by file in cpp/CMakeLists.txt). The wheel build configures with + # BUILD_SHARED_LIBS=ON; without this override KaMinPar would build a separate + # libKaMinPar.so that is neither embedded nor shipped in the wheel. Forcing PIC + # is required so the static objects can be linked into the shared libcuopt.so + # (KaMinPar's KaMinParCommon OBJECT lib otherwise lacks -fPIC). + "BUILD_SHARED_LIBS OFF" + "CMAKE_POSITION_INDEPENDENT_CODE ON" + ) + + if(KaMinPar_ADDED) + message(VERBOSE "CUOPT: Using KaMinPar located in ${KaMinPar_SOURCE_DIR}") + # KaMinPar's public header pulls in . On older TBB releases + # that header is gated behind TBB_PREVIEW_GLOBAL_CONTROL (KaMinPar upstream assumes a + # newer oneTBB and never defines it). Define it on KaMinParCommon PUBLIC so it + # propagates to all KaMinPar translation units (KaMinPar links KaMinParCommon PUBLIC). + # Harmless on newer oneTBB where global_control is no longer a preview feature. + # Also force PIC on every KaMinPar target (the KaMinParCommon OBJECT library does not + # reliably inherit CMAKE_POSITION_INDEPENDENT_CODE) so the static archive can be + # embedded into the shared libcuopt.so. + foreach(_kaminpar_tgt KaMinParCommon KaMinPar KaMinParIO) + if(TARGET ${_kaminpar_tgt}) + set_target_properties(${_kaminpar_tgt} PROPERTIES POSITION_INDEPENDENT_CODE ON) + endif() + endforeach() + if(TARGET KaMinParCommon) + target_compile_definitions(KaMinParCommon PUBLIC TBB_PREVIEW_GLOBAL_CONTROL) + endif() + else() + message(VERBOSE "CUOPT: Using KaMinPar located in ${KaMinPar_DIR}") + endif() +endfunction() + +find_and_configure_kaminpar(VERSION 3.7.3 PINNED_TAG v3.7.3) diff --git a/cpp/cuopt_cli.cpp b/cpp/cuopt_cli.cpp index bf16e9e9da..e02c12fade 100644 --- a/cpp/cuopt_cli.cpp +++ b/cpp/cuopt_cli.cpp @@ -5,6 +5,7 @@ */ /* clang-format on */ +#include #include #include #include @@ -138,6 +139,32 @@ int run_single_file(const std::string& file_path, std::make_unique>(); } + // Distributed PDLP is used for large problems that don't fit on a single GPU. + // We need to debranch before the problem_interface is created and tries to materialize the + // problem in device memory. + if (settings.get_pdlp_settings().use_distributed_pdlp) { + cuopt::cuopt_expects( + handle_ptr != nullptr, + cuopt::error_type_t::ValidationError, + "Distributed PDLP requires the GPU memory backend; no GPU handle is available for the " + "selected memory backend."); + cuopt::cuopt_expects(!solve_relaxation, + cuopt::error_type_t::ValidationError, + "Solving the LP relaxation is not allowed for distributed PDLP."); + cuopt::cuopt_expects(initial_solution_file.empty(), + cuopt::error_type_t::ValidationError, + "Initial solution file is not allowed for distributed PDLP."); + const auto& var_type_chars = mps_data_model.get_variable_types(); + cuopt::cuopt_expects( + std::none_of( + var_type_chars.begin(), var_type_chars.end(), [](char t) { return t == 'I' || t == 'B'; }), + cuopt::error_type_t::ValidationError, + "Distributed PDLP does not support mixed integer problems."); + auto solution = cuopt::mathematical_optimization::solve_lp( + handle_ptr.get(), mps_data_model, settings.get_pdlp_settings()); + return 0; + } + cuopt::mathematical_optimization::populate_from_mps_data_model(problem_interface.get(), mps_data_model); @@ -177,6 +204,8 @@ int run_single_file(const std::string& file_path, auto solution = cuopt::mathematical_optimization::solve_mip(problem_interface.get(), mip_settings); } else { + // Distributed PDLP was handled by the early-exit branch above; this + // path is always single-GPU LP going through problem_interface. auto& lp_settings = settings.get_pdlp_settings(); auto solution = cuopt::mathematical_optimization::solve_lp(problem_interface.get(), lp_settings); @@ -424,15 +453,24 @@ int main(int argc, char* argv[]) std::vector memory_resources; if (memory_backend == cuopt::mathematical_optimization::memory_backend_t::GPU) { - const int num_gpus = settings.get_parameter(CUOPT_NUM_GPUS); + // Get the right number of GPUs + // Distributed PDLP uses its own knob: distributed_pdlp_num_gpus + // Everything else uses num_gpus which is capped at 2 + const bool use_distributed_pdlp = settings.get_parameter(CUOPT_USE_DISTRIBUTED_PDLP); + int requested_gpus = use_distributed_pdlp + ? settings.get_parameter(CUOPT_DISTRIBUTED_PDLP_NUM_GPUS) + : settings.get_parameter(CUOPT_NUM_GPUS); + if (use_distributed_pdlp && requested_gpus == -1) { + requested_gpus = raft::device_setter::get_device_count(); + } + const int provisioned_gpus = std::min(raft::device_setter::get_device_count(), requested_gpus); - memory_resources.reserve(std::min(raft::device_setter::get_device_count(), num_gpus)); - for (int i = 0; i < std::min(raft::device_setter::get_device_count(), num_gpus); ++i) { - RAFT_CUDA_TRY(cudaSetDevice(i)); + memory_resources.reserve(provisioned_gpus); + for (int i = 0; i < provisioned_gpus; ++i) { + raft::device_setter guard(i); memory_resources.emplace_back(); rmm::mr::set_per_device_resource(rmm::cuda_device_id{i}, memory_resources.back()); } - RAFT_CUDA_TRY(cudaSetDevice(0)); } return run_single_file(file_name, initial_solution_file, solve_relaxation, settings); diff --git a/cpp/include/cuopt/mathematical_optimization/constants.h b/cpp/include/cuopt/mathematical_optimization/constants.h index 5e720a645e..c6bb26e872 100644 --- a/cpp/include/cuopt/mathematical_optimization/constants.h +++ b/cpp/include/cuopt/mathematical_optimization/constants.h @@ -83,14 +83,17 @@ #define CUOPT_MIP_STRONG_BRANCHING_SIMPLEX_ITERATION_LIMIT \ "mip_strong_branching_simplex_iteration_limit" -#define CUOPT_SOLUTION_FILE "solution_file" -#define CUOPT_NUM_CPU_THREADS "num_cpu_threads" -#define CUOPT_NUM_GPUS "num_gpus" -#define CUOPT_USER_PROBLEM_FILE "user_problem_file" -#define CUOPT_PRESOLVE_FILE "presolve_file" -#define CUOPT_RANDOM_SEED "random_seed" -#define CUOPT_PDLP_PRECISION "pdlp_precision" -#define CUOPT_MIP_SEMICONTINUOUS_BIG_M "mip_semi_continuous_big_m" +#define CUOPT_SOLUTION_FILE "solution_file" +#define CUOPT_NUM_CPU_THREADS "num_cpu_threads" +#define CUOPT_NUM_GPUS "num_gpus" +#define CUOPT_DISTRIBUTED_PDLP_NUM_GPUS "distributed_pdlp_num_gpus" +#define CUOPT_DISTRIBUTED_PDLP_PARTITIONER "distributed_pdlp_partitioner" +#define CUOPT_USE_DISTRIBUTED_PDLP "use_distributed_pdlp" +#define CUOPT_USER_PROBLEM_FILE "user_problem_file" +#define CUOPT_PRESOLVE_FILE "presolve_file" +#define CUOPT_RANDOM_SEED "random_seed" +#define CUOPT_PDLP_PRECISION "pdlp_precision" +#define CUOPT_MIP_SEMICONTINUOUS_BIG_M "mip_semi_continuous_big_m" #define CUOPT_MIP_HYPER_HEURISTIC_POPULATION_SIZE "mip_hyper_heuristic_population_size" #define CUOPT_MIP_HYPER_HEURISTIC_NUM_CPUFJ_THREADS "mip_hyper_heuristic_num_cpufj_threads" @@ -206,6 +209,14 @@ #define CUOPT_PRESOLVE_PAPILO 1 #define CUOPT_PRESOLVE_PSLP 2 +/* @brief distributed_pdlp_partitioner values. + * Auto: pick automatically (Dummy on 1 GPU, KaMinPar otherwise). + * KaMinPar: multi-threaded KaMinPar graph partitioner. + * Dummy: round-robin, no graph. */ +#define CUOPT_DISTRIBUTED_PDLP_PARTITIONER_AUTO 0 +#define CUOPT_DISTRIBUTED_PDLP_PARTITIONER_KAMINPAR 1 +#define CUOPT_DISTRIBUTED_PDLP_PARTITIONER_DUMMY 2 + /* @brief MIP scaling mode constants */ #define CUOPT_MIP_SCALING_OFF 0 #define CUOPT_MIP_SCALING_ON 1 diff --git a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp index 96f548ec32..5162ae8c4b 100644 --- a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp +++ b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp @@ -96,6 +96,19 @@ enum pdlp_precision_t : int { MixedPrecision = CUOPT_PDLP_MIXED_PRECISION }; +/** + * @brief Which graph partitioner distributed PDLP uses. + * + * Auto: pick automatically (Dummy on 1 GPU, KaMinPar otherwise). + * KaMinPar: multi-threaded KaMinPar graph partitioner. + * Dummy: round-robin, no graph (trivial). + */ +enum distributed_pdlp_partitioner_t : int { + Auto = CUOPT_DISTRIBUTED_PDLP_PARTITIONER_AUTO, + KaMinPar = CUOPT_DISTRIBUTED_PDLP_PARTITIONER_KAMINPAR, + Dummy = CUOPT_DISTRIBUTED_PDLP_PARTITIONER_DUMMY, +}; + template class pdlp_solver_settings_t { public: @@ -307,6 +320,14 @@ class pdlp_solver_settings_t { presolver_t presolver{presolver_t::Default}; bool dual_postsolve{true}; int num_gpus{1}; + // Dispatch the LP to the multi-GPU distributed PDLP engine + bool use_distributed_pdlp{false}; + // Number of GPUs to use specifically for distributed PDLP (use_distributed_pdlp=true). + // -1 means auto-detect + int distributed_pdlp_num_gpus{-1}; + // Which graph partitioner distributed PDLP uses. See + // distributed_pdlp_partitioner_t for the meaning of each value. + distributed_pdlp_partitioner_t distributed_pdlp_partitioner{distributed_pdlp_partitioner_t::Auto}; method_t method{method_t::Concurrent}; bool inside_mip{false}; // For concurrent termination diff --git a/cpp/include/cuopt/mathematical_optimization/solve.hpp b/cpp/include/cuopt/mathematical_optimization/solve.hpp index 961fcd29b6..f7f138ed53 100644 --- a/cpp/include/cuopt/mathematical_optimization/solve.hpp +++ b/cpp/include/cuopt/mathematical_optimization/solve.hpp @@ -148,6 +148,15 @@ optimization_problem_t mps_data_model_to_optimization_problem( raft::handle_t const* handle_ptr, const cuopt::mathematical_optimization::io::mps_data_model_t& data_model); +// Device->host projection of an optimization_problem_t onto an +// mps_data_model_t. The inverse of mps_data_model_to_optimization_problem +// above. Copies every vector buffer (CSR, var/row bounds, types, optional +// row_types + RHS) and the relevant metadata (objective offset/scale, names, +// sense). One stream sync at the end. +template +cuopt::mathematical_optimization::io::mps_data_model_t op_problem_to_mps_data_model( + const optimization_problem_t& op_problem); + // ============================================================================ // CPU problem overloads (convert to GPU, solve, convert solution back) // ============================================================================ diff --git a/cpp/src/math_optimization/solver_settings.cu b/cpp/src/math_optimization/solver_settings.cu index 8bfd7cd8ba..dd94aa2f39 100644 --- a/cpp/src/math_optimization/solver_settings.cu +++ b/cpp/src/math_optimization/solver_settings.cu @@ -145,11 +145,13 @@ solver_settings_t::solver_settings_t() : pdlp_settings(), mip_settings {CUOPT_MIP_OBJECTIVE_STEP, &mip_settings.objective_step, 0, 1, 1}, {CUOPT_NUM_GPUS, &pdlp_settings.num_gpus, 1, 2, 1}, {CUOPT_NUM_GPUS, &mip_settings.num_gpus, 1, 2, 1}, + {CUOPT_DISTRIBUTED_PDLP_NUM_GPUS, &pdlp_settings.distributed_pdlp_num_gpus, -1, 72, -1}, {CUOPT_MIP_BATCH_PDLP_STRONG_BRANCHING, &mip_settings.mip_batch_pdlp_strong_branching, 0, 2, 0}, {CUOPT_MIP_BATCH_PDLP_RELIABILITY_BRANCHING, &mip_settings.mip_batch_pdlp_reliability_branching, 0, 2, 0}, {CUOPT_MIP_STRONG_BRANCHING_SIMPLEX_ITERATION_LIMIT, &mip_settings.strong_branching_simplex_iteration_limit, -1,std::numeric_limits::max(), -1}, {CUOPT_PRESOLVE, reinterpret_cast(&pdlp_settings.presolver), CUOPT_PRESOLVE_DEFAULT, CUOPT_PRESOLVE_PSLP, CUOPT_PRESOLVE_DEFAULT}, {CUOPT_PRESOLVE, reinterpret_cast(&mip_settings.presolver), CUOPT_PRESOLVE_DEFAULT, CUOPT_PRESOLVE_PSLP, CUOPT_PRESOLVE_DEFAULT}, + {CUOPT_DISTRIBUTED_PDLP_PARTITIONER, reinterpret_cast(&pdlp_settings.distributed_pdlp_partitioner), CUOPT_DISTRIBUTED_PDLP_PARTITIONER_AUTO, CUOPT_DISTRIBUTED_PDLP_PARTITIONER_DUMMY, CUOPT_DISTRIBUTED_PDLP_PARTITIONER_AUTO}, {CUOPT_MIP_DETERMINISM_MODE, &mip_settings.determinism_mode, CUOPT_MODE_OPPORTUNISTIC, CUOPT_MODE_DETERMINISTIC, CUOPT_MODE_OPPORTUNISTIC}, {CUOPT_RANDOM_SEED, &mip_settings.seed, -1, std::numeric_limits::max(), -1}, {CUOPT_MIP_RELIABILITY_BRANCHING, &mip_settings.reliability_branching, -1, std::numeric_limits::max(), -1}, @@ -192,6 +194,7 @@ solver_settings_t::solver_settings_t() : pdlp_settings(), mip_settings {CUOPT_DUAL_POSTSOLVE, &pdlp_settings.dual_postsolve, true}, {CUOPT_BARRIER_ITERATIVE_REFINEMENT, &pdlp_settings.barrier_iterative_refinement, true}, {CUOPT_MIP_PROBING, &mip_settings.probing, true}, + {CUOPT_USE_DISTRIBUTED_PDLP, &pdlp_settings.use_distributed_pdlp, false}, // Diving heuristic hyper-parameters (hidden from default --help: name contains "hyper_") {CUOPT_MIP_HYPER_DIVING_SHOW_TYPE, &mip_settings.diving_params.show_type, false, "log diving heuristic type when it finds a new incumbent"}, }; @@ -204,7 +207,7 @@ solver_settings_t::solver_settings_t() : pdlp_settings(), mip_settings {CUOPT_USER_PROBLEM_FILE, &mip_settings.user_problem_file, ""}, {CUOPT_USER_PROBLEM_FILE, &pdlp_settings.user_problem_file, ""}, {CUOPT_PRESOLVE_FILE, &mip_settings.presolve_file, ""}, - {CUOPT_PRESOLVE_FILE, &pdlp_settings.presolve_file, ""} + {CUOPT_PRESOLVE_FILE, &pdlp_settings.presolve_file, ""}, }; // clang-format on } diff --git a/cpp/src/mip_heuristics/presolve/third_party_presolve.cpp b/cpp/src/mip_heuristics/presolve/third_party_presolve.cpp index 64a25ff556..0e544b57f3 100644 --- a/cpp/src/mip_heuristics/presolve/third_party_presolve.cpp +++ b/cpp/src/mip_heuristics/presolve/third_party_presolve.cpp @@ -35,6 +35,9 @@ #else #pragma GCC diagnostic pop #endif + +#include +#include #include #include #include @@ -44,145 +47,139 @@ #include +#include +#include +#include #include namespace cuopt::mathematical_optimization::mip { +// Backend-agnostic normalisation of the mutable presolve fields: +// * sign-flip `obj_coeffs` / `objective_offset` when maximise, +// * materialise ranged `constr_lb` / `constr_ub` from `row_types` + +// `constraint_bounds` when the ranged pair is absent from the mps. +// +// Precondition: the caller has already copy-initialised each vector from with mps data template -papilo::Problem build_papilo_problem(const optimization_problem_t& op_problem, - problem_category_t category, - bool maximize) +void normalize_for_presolve(io::mps_data_model_t const& mps, + bool maximize, + std::vector& obj_coeffs, + f_t& objective_offset, + std::vector& var_lb, + std::vector& var_ub, + std::vector& constr_lb, + std::vector& constr_ub) { - raft::common::nvtx::range fun_scope("Build papilo problem"); - // Build papilo problem from optimization problem - papilo::ProblemBuilder builder; - - // Get problem dimensions - const i_t num_cols = op_problem.get_n_variables(); - const i_t num_rows = op_problem.get_n_constraints(); - const i_t nnz = op_problem.get_nnz(); - - builder.reserve(nnz, num_rows, num_cols); - - // Get problem data from optimization problem - const auto& coefficients = op_problem.get_constraint_matrix_values(); - const auto& offsets = op_problem.get_constraint_matrix_offsets(); - const auto& variables = op_problem.get_constraint_matrix_indices(); - const auto& obj_coeffs = op_problem.get_objective_coefficients(); - const auto& var_lb = op_problem.get_variable_lower_bounds(); - const auto& var_ub = op_problem.get_variable_upper_bounds(); - const auto& bounds = op_problem.get_constraint_bounds(); - const auto& row_types = op_problem.get_row_types(); - const auto& constr_lb = op_problem.get_constraint_lower_bounds(); - const auto& constr_ub = op_problem.get_constraint_upper_bounds(); - const auto& var_types = op_problem.get_variable_types(); - - // Copy data to host - std::vector h_coefficients(coefficients.size()); - auto stream_view = op_problem.get_handle_ptr()->get_stream(); - raft::copy(h_coefficients.data(), coefficients.data(), coefficients.size(), stream_view); - std::vector h_offsets(offsets.size()); - raft::copy(h_offsets.data(), offsets.data(), offsets.size(), stream_view); - std::vector h_variables(variables.size()); - raft::copy(h_variables.data(), variables.data(), variables.size(), stream_view); - std::vector h_obj_coeffs(obj_coeffs.size()); - raft::copy(h_obj_coeffs.data(), obj_coeffs.data(), obj_coeffs.size(), stream_view); - std::vector h_var_lb(var_lb.size()); - raft::copy(h_var_lb.data(), var_lb.data(), var_lb.size(), stream_view); - std::vector h_var_ub(var_ub.size()); - raft::copy(h_var_ub.data(), var_ub.data(), var_ub.size(), stream_view); - std::vector h_bounds(bounds.size()); - raft::copy(h_bounds.data(), bounds.data(), bounds.size(), stream_view); - std::vector h_row_types(row_types.size()); - raft::copy(h_row_types.data(), row_types.data(), row_types.size(), stream_view); - std::vector h_constr_lb(constr_lb.size()); - raft::copy(h_constr_lb.data(), constr_lb.data(), constr_lb.size(), stream_view); - std::vector h_constr_ub(constr_ub.size()); - raft::copy(h_constr_ub.data(), constr_ub.data(), constr_ub.size(), stream_view); - std::vector h_var_types(var_types.size()); - raft::copy(h_var_types.data(), var_types.data(), var_types.size(), stream_view); - if (maximize) { - for (size_t i = 0; i < h_obj_coeffs.size(); ++i) { - h_obj_coeffs[i] = -h_obj_coeffs[i]; + for (auto& c : obj_coeffs) { + c = -c; } + objective_offset = -objective_offset; } - auto constr_bounds_empty = h_constr_lb.empty() && h_constr_ub.empty(); - if (constr_bounds_empty) { - for (size_t i = 0; i < h_row_types.size(); ++i) { - if (h_row_types[i] == 'L') { - h_constr_lb.push_back(-std::numeric_limits::infinity()); - h_constr_ub.push_back(h_bounds[i]); - } else if (h_row_types[i] == 'G') { - h_constr_lb.push_back(h_bounds[i]); - h_constr_ub.push_back(std::numeric_limits::infinity()); - } else if (h_row_types[i] == 'E') { - h_constr_lb.push_back(h_bounds[i]); - h_constr_ub.push_back(h_bounds[i]); + if (constr_lb.empty() && constr_ub.empty()) { + const auto& row_types = mps.get_row_types(); + const auto& constraint_bounds = mps.get_constraint_bounds(); + for (size_t i = 0; i < row_types.size(); ++i) { + if (row_types[i] == 'L') { + constr_lb.push_back(-std::numeric_limits::infinity()); + constr_ub.push_back(constraint_bounds[i]); + } else if (row_types[i] == 'G') { + constr_lb.push_back(constraint_bounds[i]); + constr_ub.push_back(std::numeric_limits::infinity()); + } else if (row_types[i] == 'E') { + constr_lb.push_back(constraint_bounds[i]); + constr_ub.push_back(constraint_bounds[i]); } } } +} - builder.setNumCols(num_cols); - builder.setNumRows(num_rows); - - builder.setObjAll(h_obj_coeffs); - builder.setObjOffset(maximize ? -op_problem.get_objective_offset() - : op_problem.get_objective_offset()); +// Build a papilo::Problem +template +papilo::Problem build_papilo_problem(io::mps_data_model_t const& mps, + bool maximize, + problem_category_t category) +{ + raft::common::nvtx::range fun_scope("Build papilo::Problem from mps_data_model"); + + const i_t n_cols = mps.get_n_variables(); + const i_t n_rows = mps.get_n_constraints(); + const i_t nnz = mps.get_nnz(); + + std::vector obj_coeffs(mps.get_objective_coefficients()); + std::vector var_lb(mps.get_variable_lower_bounds()); + std::vector var_ub(mps.get_variable_upper_bounds()); + std::vector constr_lb(mps.get_constraint_lower_bounds()); + std::vector constr_ub(mps.get_constraint_upper_bounds()); + f_t objective_offset = mps.get_objective_offset(); + normalize_for_presolve( + mps, maximize, obj_coeffs, objective_offset, var_lb, var_ub, constr_lb, constr_ub); + + const auto& coefficients = mps.get_constraint_matrix_values(); + const auto& indices = mps.get_constraint_matrix_indices(); + const auto& offsets = mps.get_constraint_matrix_offsets(); + const auto& variable_names = mps.get_variable_names(); + const auto& mps_var_types = mps.get_variable_types(); - if (!h_var_lb.empty() && !h_var_ub.empty()) { - builder.setColLbAll(h_var_lb); - builder.setColUbAll(h_var_ub); - if (op_problem.get_variable_names().size() == h_var_lb.size()) { - builder.setColNameAll(op_problem.get_variable_names()); + papilo::ProblemBuilder builder; + builder.reserve(nnz, n_rows, n_cols); + builder.setNumCols(n_cols); + builder.setNumRows(n_rows); + builder.setObjAll(obj_coeffs); + builder.setObjOffset(objective_offset); + + if (!var_lb.empty() && !var_ub.empty()) { + builder.setColLbAll(var_lb); + builder.setColUbAll(var_ub); + if (variable_names.size() == static_cast(n_cols)) { + builder.setColNameAll(variable_names); } } if (category == problem_category_t::MIP) { - for (size_t i = 0; i < h_var_types.size(); ++i) { - builder.setColIntegral(i, h_var_types[i] == var_t::INTEGER); + for (size_t i = 0; i < mps_var_types.size(); ++i) { + builder.setColIntegral(i, char_to_var_type(mps_var_types[i]) == var_t::INTEGER); } } - if (!h_constr_lb.empty() && !h_constr_ub.empty()) { - builder.setRowLhsAll(h_constr_lb); - builder.setRowRhsAll(h_constr_ub); + if (!constr_lb.empty() && !constr_ub.empty()) { + builder.setRowLhsAll(constr_lb); + builder.setRowRhsAll(constr_ub); } - std::vector h_row_flags(h_constr_lb.size()); + std::vector h_row_flags(constr_lb.size()); std::vector> h_entries; - // Add constraints row by row - for (size_t i = 0; i < h_constr_lb.size(); ++i) { - // Get row entries - i_t row_start = h_offsets[i]; - i_t row_end = h_offsets[i + 1]; - i_t num_entries = row_end - row_start; + for (size_t i = 0; i < constr_lb.size(); ++i) { + const i_t row_start = offsets[i]; + const i_t row_end = offsets[i + 1]; + const i_t num_entries = row_end - row_start; for (size_t j = 0; j < num_entries; ++j) { - h_entries.push_back( - std::make_tuple(i, h_variables[row_start + j], h_coefficients[row_start + j])); + h_entries.push_back(std::make_tuple(i, indices[row_start + j], coefficients[row_start + j])); } - if (h_constr_lb[i] == -std::numeric_limits::infinity()) { + if (constr_lb[i] == -std::numeric_limits::infinity()) { h_row_flags[i].set(papilo::RowFlag::kLhsInf); } else { h_row_flags[i].unset(papilo::RowFlag::kLhsInf); } - if (h_constr_ub[i] == std::numeric_limits::infinity()) { + if (constr_ub[i] == std::numeric_limits::infinity()) { h_row_flags[i].set(papilo::RowFlag::kRhsInf); } else { h_row_flags[i].unset(papilo::RowFlag::kRhsInf); } - if (h_constr_lb[i] == -std::numeric_limits::infinity()) { h_constr_lb[i] = 0; } - if (h_constr_ub[i] == std::numeric_limits::infinity()) { h_constr_ub[i] = 0; } + // Papilo stores finite dummies in place of ±inf; the flags above are the + // source of truth for infinity. Zero out the dummies before setConstraintMatrix. + if (constr_lb[i] == -std::numeric_limits::infinity()) { constr_lb[i] = 0; } + if (constr_ub[i] == std::numeric_limits::infinity()) { constr_ub[i] = 0; } } - for (size_t i = 0; i < h_var_lb.size(); ++i) { - builder.setColLbInf(i, h_var_lb[i] == -std::numeric_limits::infinity()); - builder.setColUbInf(i, h_var_ub[i] == std::numeric_limits::infinity()); - if (h_var_lb[i] == -std::numeric_limits::infinity()) { builder.setColLb(i, 0); } - if (h_var_ub[i] == std::numeric_limits::infinity()) { builder.setColUb(i, 0); } + for (size_t i = 0; i < var_lb.size(); ++i) { + builder.setColLbInf(i, var_lb[i] == -std::numeric_limits::infinity()); + builder.setColUbInf(i, var_ub[i] == std::numeric_limits::infinity()); + if (var_lb[i] == -std::numeric_limits::infinity()) { builder.setColLb(i, 0); } + if (var_ub[i] == std::numeric_limits::infinity()) { builder.setColUb(i, 0); } } auto problem = builder.build(); @@ -193,8 +190,8 @@ papilo::Problem build_papilo_problem(const optimization_problem_t const double spare_ratio = category == problem_category_t::MIP ? 4.0 : 2.0; const int min_inter_row_space = category == problem_category_t::MIP ? 30 : 4; auto csr_storage = papilo::SparseStorage( - h_entries, num_rows, num_cols, sorted_entries, spare_ratio, min_inter_row_space); - problem.setConstraintMatrix(csr_storage, h_constr_lb, h_constr_ub, h_row_flags); + h_entries, n_rows, n_cols, sorted_entries, spare_ratio, min_inter_row_space); + problem.setConstraintMatrix(csr_storage, constr_lb, constr_ub, h_row_flags); papilo::ConstraintMatrix& matrix = problem.getConstraintMatrix(); for (int i = 0; i < problem.getNRows(); ++i) { @@ -208,222 +205,93 @@ papilo::Problem build_papilo_problem(const optimization_problem_t return problem; } -struct PSLPContext { - Presolver* presolver = nullptr; - Settings* settings = nullptr; - PresolveStatus status = PresolveStatus_::UNCHANGED; -}; - +// Presolved mps_data_model builder from PSLP template -PSLPContext build_and_run_pslp_presolver(const optimization_problem_t& op_problem, - bool maximize, - const double time_limit) +io::mps_data_model_t build_reduced_mps_from_pslp(Presolver* pslp_presolver, + bool maximize, + f_t original_obj_offset) { - PSLPContext ctx; - raft::common::nvtx::range fun_scope("Build and run PSLP presolver"); - - // Get problem dimensions - const i_t num_cols = op_problem.get_n_variables(); - const i_t num_rows = op_problem.get_n_constraints(); - const i_t nnz = op_problem.get_nnz(); - - // Get problem data from optimization problem - const auto& coefficients = op_problem.get_constraint_matrix_values(); - const auto& offsets = op_problem.get_constraint_matrix_offsets(); - const auto& variables = op_problem.get_constraint_matrix_indices(); - const auto& obj_coeffs = op_problem.get_objective_coefficients(); - const auto& var_lb = op_problem.get_variable_lower_bounds(); - const auto& var_ub = op_problem.get_variable_upper_bounds(); - const auto& bounds = op_problem.get_constraint_bounds(); - const auto& row_types = op_problem.get_row_types(); - const auto& constr_lb = op_problem.get_constraint_lower_bounds(); - const auto& constr_ub = op_problem.get_constraint_upper_bounds(); - const auto& var_types = op_problem.get_variable_types(); - - // Copy data to host - std::vector h_coefficients(coefficients.size()); - auto stream_view = op_problem.get_handle_ptr()->get_stream(); - raft::copy(h_coefficients.data(), coefficients.data(), coefficients.size(), stream_view); - std::vector h_offsets(offsets.size()); - raft::copy(h_offsets.data(), offsets.data(), offsets.size(), stream_view); - std::vector h_variables(variables.size()); - raft::copy(h_variables.data(), variables.data(), variables.size(), stream_view); - std::vector h_obj_coeffs(obj_coeffs.size()); - raft::copy(h_obj_coeffs.data(), obj_coeffs.data(), obj_coeffs.size(), stream_view); - std::vector h_var_lb(var_lb.size()); - raft::copy(h_var_lb.data(), var_lb.data(), var_lb.size(), stream_view); - std::vector h_var_ub(var_ub.size()); - raft::copy(h_var_ub.data(), var_ub.data(), var_ub.size(), stream_view); - std::vector h_bounds(bounds.size()); - raft::copy(h_bounds.data(), bounds.data(), bounds.size(), stream_view); - std::vector h_row_types(row_types.size()); - raft::copy(h_row_types.data(), row_types.data(), row_types.size(), stream_view); - std::vector h_constr_lb(constr_lb.size()); - raft::copy(h_constr_lb.data(), constr_lb.data(), constr_lb.size(), stream_view); - std::vector h_constr_ub(constr_ub.size()); - raft::copy(h_constr_ub.data(), constr_ub.data(), constr_ub.size(), stream_view); - std::vector h_var_types(var_types.size()); - raft::copy(h_var_types.data(), var_types.data(), var_types.size(), stream_view); + raft::common::nvtx::range fun_scope("Build mps_data_model from PSLP"); + io::mps_data_model_t mps; - stream_view.synchronize(); - if (maximize) { - for (size_t i = 0; i < h_obj_coeffs.size(); ++i) { - h_obj_coeffs[i] = -h_obj_coeffs[i]; + if constexpr (std::is_same_v) { + auto* reduced = pslp_presolver->reduced_prob; + const i_t n_rows = static_cast(reduced->m); + const i_t n_cols = static_cast(reduced->n); + const i_t nnz = static_cast(reduced->nnz); + // PSLP folds the sign flip into obj_offset for maximise problems, and does + // not track the original mps's objective offset — put both back. + const f_t obj_offset = + (maximize ? -reduced->obj_offset : reduced->obj_offset) + original_obj_offset; + mps.set_maximize(maximize); + mps.set_objective_offset(obj_offset); + + if (n_cols == 0 && n_rows == 0) { + std::vector empty_offsets{0}; + mps.set_csr_constraint_matrix( + {}, {}, std::span(empty_offsets.data(), empty_offsets.size())); + return mps; } - } - auto constr_bounds_empty = h_constr_lb.empty() && h_constr_ub.empty(); - if (constr_bounds_empty) { - for (size_t i = 0; i < h_row_types.size(); ++i) { - if (h_row_types[i] == 'L') { - h_constr_lb.push_back(-std::numeric_limits::infinity()); - h_constr_ub.push_back(h_bounds[i]); - } else if (h_row_types[i] == 'G') { - h_constr_lb.push_back(h_bounds[i]); - h_constr_ub.push_back(std::numeric_limits::infinity()); - } else if (h_row_types[i] == 'E') { - h_constr_lb.push_back(h_bounds[i]); - h_constr_ub.push_back(h_bounds[i]); - } - } - } + mps.set_csr_constraint_matrix( + std::span(reduced->Ax, static_cast(nnz)), + std::span(reduced->Ai, static_cast(nnz)), + std::span(reduced->Ap, static_cast(n_rows + 1))); - // handle empty variable bounds - if (h_var_lb.empty()) { - h_var_lb = std::vector(num_cols, -std::numeric_limits::infinity()); - } - if (h_var_ub.empty()) { - h_var_ub = std::vector(num_cols, std::numeric_limits::infinity()); - } - - // Call PSLP presolver - ctx.settings = default_settings(); - ctx.settings->verbose = false; - ctx.settings->max_time = time_limit; - auto start_time = std::chrono::high_resolution_clock::now(); - ctx.presolver = new_presolver(h_coefficients.data(), - h_variables.data(), - h_offsets.data(), - num_rows, - num_cols, - nnz, - h_constr_lb.data(), - h_constr_ub.data(), - h_var_lb.data(), - h_var_ub.data(), - h_obj_coeffs.data(), - ctx.settings); - - assert(ctx.presolver != nullptr && "Presolver initialization failed"); - ctx.status = run_presolver(ctx.presolver); - auto end_time = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end_time - start_time); - CUOPT_LOG_DEBUG("PSLP presolver time: %d milliseconds", duration.count()); - CUOPT_LOG_INFO("PSLP Presolved problem: %d constraints, %d variables, %d non-zeros", - ctx.presolver->stats->n_rows_reduced, - ctx.presolver->stats->n_cols_reduced, - ctx.presolver->stats->nnz_reduced); - - return ctx; -} - -template -optimization_problem_t build_optimization_problem_from_pslp( - Presolver* pslp_presolver, - raft::handle_t const* handle_ptr, - bool maximize, - f_t original_obj_offset) -{ - raft::common::nvtx::range fun_scope("Build optimization problem from PSLP"); - cuopt_expects(pslp_presolver != nullptr && pslp_presolver->reduced_prob != nullptr, - error_type_t::RuntimeError, - "PSLP presolver is not initialized"); - auto reduced_prob = pslp_presolver->reduced_prob; - int n_rows = reduced_prob->m; - int n_cols = reduced_prob->n; - int nnz = reduced_prob->nnz; - double obj_offset = reduced_prob->obj_offset; - - optimization_problem_t op_problem(handle_ptr); - - obj_offset = maximize ? -obj_offset : obj_offset; - - // PSLP does not allow setting the objective offset, so we add the original objective offset to - // the reduced objective offset - obj_offset += original_obj_offset; - op_problem.set_objective_offset(obj_offset); - op_problem.set_maximize(maximize); - op_problem.set_problem_category(problem_category_t::LP); - - // Handle empty reduced problem (presolve completely solved it) - if (n_cols == 0 && n_rows == 0) { - // Set empty constraint matrix with proper offsets - std::vector empty_offsets = {0}; - op_problem.set_csr_constraint_matrix(nullptr, 0, nullptr, 0, empty_offsets.data(), 1); - return op_problem; - } - - op_problem.set_csr_constraint_matrix( - reduced_prob->Ax, nnz, reduced_prob->Ai, nnz, reduced_prob->Ap, n_rows + 1); - - std::vector h_obj_coeffs(n_cols); - std::copy(reduced_prob->c, reduced_prob->c + n_cols, h_obj_coeffs.begin()); - if (maximize) { - for (size_t i = 0; i < n_cols; ++i) { - h_obj_coeffs[i] = -h_obj_coeffs[i]; + if (maximize) { + std::vector h_obj_coeffs(reduced->c, reduced->c + n_cols); + for (auto& c : h_obj_coeffs) { + c = -c; + } + mps.set_objective_coefficients( + std::span(h_obj_coeffs.data(), h_obj_coeffs.size())); + } else { + mps.set_objective_coefficients(std::span(reduced->c, static_cast(n_cols))); } + mps.set_constraint_lower_bounds( + std::span(reduced->lhs, static_cast(n_rows))); + mps.set_constraint_upper_bounds( + std::span(reduced->rhs, static_cast(n_rows))); + mps.set_variable_lower_bounds(std::span(reduced->lbs, static_cast(n_cols))); + mps.set_variable_upper_bounds(std::span(reduced->ubs, static_cast(n_cols))); + } else { + cuopt_expects(false, error_type_t::ValidationError, "PSLP only supports double precision"); } - op_problem.set_objective_coefficients(h_obj_coeffs.data(), n_cols); - op_problem.set_constraint_lower_bounds(reduced_prob->lhs, n_rows); - op_problem.set_constraint_upper_bounds(reduced_prob->rhs, n_rows); - op_problem.set_variable_lower_bounds(reduced_prob->lbs, n_cols); - op_problem.set_variable_upper_bounds(reduced_prob->ubs, n_cols); - return op_problem; + return mps; } template -optimization_problem_t build_optimization_problem( - papilo::Problem const& papilo_problem, - raft::handle_t const* handle_ptr, - problem_category_t category, - bool maximize) +io::mps_data_model_t build_reduced_mps_from_papilo( + papilo::Problem const& papilo_problem, bool maximize) { - raft::common::nvtx::range fun_scope("Build optimization problem"); - optimization_problem_t op_problem(handle_ptr); + raft::common::nvtx::range fun_scope("Reduced mps <- Papilo"); + io::mps_data_model_t mps; auto obj = papilo_problem.getObjective(); - op_problem.set_objective_offset(maximize ? -obj.offset : obj.offset); - op_problem.set_maximize(maximize); - op_problem.set_problem_category(category); + mps.set_maximize(maximize); + mps.set_objective_offset(maximize ? -obj.offset : obj.offset); if (papilo_problem.getNRows() == 0 && papilo_problem.getNCols() == 0) { - // FIXME: Shouldn't need to set offsets - std::vector h_offsets{0}; - std::vector h_indices{}; - std::vector h_values{}; - op_problem.set_csr_constraint_matrix(h_values.data(), - h_values.size(), - h_indices.data(), - h_indices.size(), - h_offsets.data(), - h_offsets.size()); - - return op_problem; + std::vector empty_offsets{0}; + mps.set_csr_constraint_matrix( + {}, {}, std::span(empty_offsets.data(), empty_offsets.size())); + return mps; } + if (maximize) { - for (size_t i = 0; i < obj.coefficients.size(); ++i) { - obj.coefficients[i] = -obj.coefficients[i]; + for (auto& c : obj.coefficients) { + c = -c; } } - op_problem.set_objective_coefficients(obj.coefficients.data(), obj.coefficients.size()); + mps.set_objective_coefficients( + std::span(obj.coefficients.data(), obj.coefficients.size())); auto& constraint_matrix = papilo_problem.getConstraintMatrix(); - auto row_lower = constraint_matrix.getLeftHandSides(); - auto row_upper = constraint_matrix.getRightHandSides(); - auto col_lower = papilo_problem.getLowerBounds(); - auto col_upper = papilo_problem.getUpperBounds(); + // Row bounds: copy out (papilo returns by value) then substitute ±inf per flag. + auto row_lower = constraint_matrix.getLeftHandSides(); + auto row_upper = constraint_matrix.getRightHandSides(); auto row_flags = constraint_matrix.getRowFlags(); for (size_t i = 0; i < row_flags.size(); i++) { if (row_flags[i].test(papilo::RowFlag::kLhsInf)) { @@ -433,40 +301,33 @@ optimization_problem_t build_optimization_problem( row_upper[i] = std::numeric_limits::infinity(); } } + mps.set_constraint_lower_bounds(std::span(row_lower.data(), row_lower.size())); + mps.set_constraint_upper_bounds(std::span(row_upper.data(), row_upper.size())); - op_problem.set_constraint_lower_bounds(row_lower.data(), row_lower.size()); - op_problem.set_constraint_upper_bounds(row_upper.data(), row_upper.size()); - + // CSR offsets have to be synthesised from papilo's RangeInfo (non-contiguous + // in general); values and column indices are contiguous, so span in-place. auto [index_range, nrows] = constraint_matrix.getRangeInfo(); - std::vector offsets(nrows + 1); - // papilo indices do not start from 0 after presolve - size_t start = index_range[0].start; + const size_t start = index_range[0].start; for (i_t i = 0; i < nrows; i++) { - offsets[i] = index_range[i].start - start; + offsets[i] = static_cast(index_range[i].start - start); } - offsets[nrows] = index_range[nrows - 1].end - start; - - i_t nnz = constraint_matrix.getNnz(); + offsets[nrows] = static_cast(index_range[nrows - 1].end - start); + const i_t nnz = static_cast(constraint_matrix.getNnz()); assert(offsets[nrows] == nnz); - const int* cols = constraint_matrix.getConstraintMatrix().getColumns(); const f_t* coeffs = constraint_matrix.getConstraintMatrix().getValues(); + mps.set_csr_constraint_matrix(std::span(&coeffs[start], static_cast(nnz)), + std::span(&cols[start], static_cast(nnz)), + std::span(offsets.data(), offsets.size())); - i_t ncols = papilo_problem.getNCols(); - cuopt_assert( - std::all_of( - cols + start, cols + start + nnz, [ncols](i_t col) { return col >= 0 && col < ncols; }), - "Papilo produced invalid column indices in presolved matrix"); - - op_problem.set_csr_constraint_matrix( - &(coeffs[start]), nnz, &(cols[start]), nnz, offsets.data(), nrows + 1); - + // Col bounds + var_types: same copy-then-fixup pattern. + auto col_lower = papilo_problem.getLowerBounds(); + auto col_upper = papilo_problem.getUpperBounds(); auto col_flags = papilo_problem.getColFlags(); - std::vector var_types(col_flags.size()); + std::vector var_types(col_flags.size()); for (size_t i = 0; i < col_flags.size(); i++) { - var_types[i] = - col_flags[i].test(papilo::ColFlag::kIntegral) ? var_t::INTEGER : var_t::CONTINUOUS; + var_types[i] = col_flags[i].test(papilo::ColFlag::kIntegral) ? 'I' : 'C'; if (col_flags[i].test(papilo::ColFlag::kLbInf)) { col_lower[i] = -std::numeric_limits::infinity(); } @@ -474,12 +335,11 @@ optimization_problem_t build_optimization_problem( col_upper[i] = std::numeric_limits::infinity(); } } + mps.set_variable_lower_bounds(std::span(col_lower.data(), col_lower.size())); + mps.set_variable_upper_bounds(std::span(col_upper.data(), col_upper.size())); + mps.set_variable_types(var_types); - op_problem.set_variable_lower_bounds(col_lower.data(), col_lower.size()); - op_problem.set_variable_upper_bounds(col_upper.data(), col_upper.size()); - op_problem.set_variable_types(var_types.data(), var_types.size()); - - return op_problem; + return mps; } void check_presolve_status(const papilo::PresolveStatus& status) @@ -630,71 +490,97 @@ void set_presolve_parameters(papilo::Presolve& presolver, } template -third_party_presolve_result_t third_party_presolve_t::apply_pslp( - optimization_problem_t const& op_problem, const double time_limit) +third_party_presolve_status_t third_party_presolve_t::apply_pslp( + io::mps_data_model_t const& mps, double time_limit) { - if constexpr (std::is_same_v) { - double original_obj_offset = op_problem.get_objective_offset(); - auto ctx = build_and_run_pslp_presolver(op_problem, maximize_, time_limit); + raft::common::nvtx::range fun_scope("Apply PSLP presolver on host"); - // Free previously allocated presolver and settings if they exist + if constexpr (std::is_same_v) { + const i_t n_cols = mps.get_n_variables(); + const i_t n_rows = mps.get_n_constraints(); + const i_t nnz = mps.get_nnz(); + + // Local owned copies of the fields that need mutation (sign flip on + // maximise / ±inf fill for empty bounds / row_types materialisation); + // matrix arrays are read straight from mps below (const&). + std::vector obj_coeffs(mps.get_objective_coefficients()); + std::vector var_lb(mps.get_variable_lower_bounds()); + std::vector var_ub(mps.get_variable_upper_bounds()); + std::vector constr_lb(mps.get_constraint_lower_bounds()); + std::vector constr_ub(mps.get_constraint_upper_bounds()); + f_t objective_offset = mps.get_objective_offset(); + normalize_for_presolve( + mps, maximize_, obj_coeffs, objective_offset, var_lb, var_ub, constr_lb, constr_ub); + if (var_lb.empty()) { var_lb.assign(n_cols, -std::numeric_limits::infinity()); } + if (var_ub.empty()) { var_ub.assign(n_cols, std::numeric_limits::infinity()); } + const auto& coefficients = mps.get_constraint_matrix_values(); + const auto& indices = mps.get_constraint_matrix_indices(); + const auto& offsets = mps.get_constraint_matrix_offsets(); + + Settings* settings = default_settings(); + settings->verbose = false; + settings->max_time = time_limit; + + auto start_time = std::chrono::high_resolution_clock::now(); + Presolver* presolver = new_presolver(coefficients.data(), + indices.data(), + offsets.data(), + n_rows, + n_cols, + nnz, + constr_lb.data(), + constr_ub.data(), + var_lb.data(), + var_ub.data(), + obj_coeffs.data(), + settings); + assert(presolver != nullptr && "Presolver initialization failed"); + const PresolveStatus pslp_status = run_presolver(presolver); + auto end_time = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(end_time - start_time); + CUOPT_LOG_DEBUG("PSLP presolver time: %d milliseconds", duration.count()); + CUOPT_LOG_INFO("PSLP Presolved problem: %d constraints, %d variables, %d non-zeros", + presolver->stats->n_rows_reduced, + presolver->stats->n_cols_reduced, + presolver->stats->nnz_reduced); + + // Free previously allocated presolver and settings (if any) and stash the + // new ones so undo_pslp / build_reduced_mps_from_pslp can find them later. if (pslp_presolver_ != nullptr) { free_presolver(pslp_presolver_); } if (pslp_stgs_ != nullptr) { free_settings(pslp_stgs_); } + pslp_presolver_ = presolver; + pslp_stgs_ = settings; - pslp_presolver_ = ctx.presolver; - pslp_stgs_ = ctx.settings; - - auto status = convert_pslp_presolve_status_to_third_party_presolve_status(ctx.status); - if (ctx.status == PresolveStatus_::INFEASIBLE || ctx.status == PresolveStatus_::UNBNDORINFEAS) { - optimization_problem_t empty_problem(op_problem.get_handle_ptr()); - return third_party_presolve_result_t{status, std::move(empty_problem), {}, {}, {}}; - } - - auto opt_problem = build_optimization_problem_from_pslp( - pslp_presolver_, op_problem.get_handle_ptr(), maximize_, original_obj_offset); - return third_party_presolve_result_t{status, std::move(opt_problem), {}, {}, {}}; + return convert_pslp_presolve_status_to_third_party_presolve_status(pslp_status); } else { cuopt_expects( false, error_type_t::ValidationError, "PSLP presolver only supports double precision"); - return third_party_presolve_result_t{ - third_party_presolve_status_t::UNCHANGED, - optimization_problem_t(op_problem.get_handle_ptr()), - {}, - {}, - {}}; + return third_party_presolve_status_t::UNCHANGED; // unreachable } } template -third_party_presolve_result_t third_party_presolve_t::apply( - optimization_problem_t const& op_problem, +third_party_presolve_status_t third_party_presolve_t::apply_papilo( + papilo::Problem& papilo_problem, problem_category_t category, - cuopt::mathematical_optimization::presolver_t presolver, bool dual_postsolve, f_t absolute_tolerance, f_t relative_tolerance, double time_limit, i_t num_cpu_threads) { - presolver_ = presolver; - maximize_ = op_problem.get_sense(); - if (category == problem_category_t::MIP && - presolver == cuopt::mathematical_optimization::presolver_t::PSLP) { - cuopt_expects( - false, error_type_t::RuntimeError, "PSLP presolver is not supported for MIP problems"); - } - - if (presolver == cuopt::mathematical_optimization::presolver_t::PSLP) { - return apply_pslp(op_problem, time_limit); - } + raft::common::nvtx::range fun_scope("Apply Papilo presolve on host"); - papilo::Problem papilo_problem = build_papilo_problem(op_problem, category, maximize_); + // Capture original dimensions before papilo.apply() mutates papilo_problem + // in place into its reduced form. + const i_t original_n_vars = static_cast(papilo_problem.getNCols()); + const i_t original_n_cons = static_cast(papilo_problem.getNRows()); + const i_t original_nnz = static_cast(papilo_problem.getConstraintMatrix().getNnz()); CUOPT_LOG_DEBUG("Original problem: %d constraints, %d variables, %d nonzeros", - papilo_problem.getNRows(), - papilo_problem.getNCols(), - papilo_problem.getConstraintMatrix().getNnz()); - + original_n_cons, + original_n_vars, + original_nnz); CUOPT_LOG_INFO("\nRunning Papilo presolve (git hash %s)", PAPILO_GITHASH); if (category == problem_category_t::MIP) { dual_postsolve = false; } papilo::Presolve papilo_presolver; @@ -706,10 +592,7 @@ third_party_presolve_result_t third_party_presolve_t::apply( time_limit, dual_postsolve, num_cpu_threads); - set_presolve_parameters( - papilo_presolver, category, op_problem.get_n_constraints(), op_problem.get_n_variables()); - - // Disable papilo logs + set_presolve_parameters(papilo_presolver, category, original_n_cons, original_n_vars); papilo_presolver.setVerbosityLevel(papilo::VerbosityLevel::kQuiet); auto result = papilo_presolver.apply(papilo_problem); @@ -718,19 +601,18 @@ third_party_presolve_result_t third_party_presolve_t::apply( if (result.status == papilo::PresolveStatus::kInfeasible || result.status == papilo::PresolveStatus::kUnbndOrInfeas || result.status == papilo::PresolveStatus::kUnbounded) { - optimization_problem_t empty_problem(op_problem.get_handle_ptr()); - return third_party_presolve_result_t{status, std::move(empty_problem), {}, {}, {}}; + return status; } papilo_post_solve_storage_.reset(new papilo::PostsolveStorage(result.postsolve)); CUOPT_LOG_INFO("Presolve removed: %d constraints, %d variables, %d nonzeros", - op_problem.get_n_constraints() - papilo_problem.getNRows(), - op_problem.get_n_variables() - papilo_problem.getNCols(), - op_problem.get_nnz() - papilo_problem.getConstraintMatrix().getNnz()); + original_n_cons - papilo_problem.getNRows(), + original_n_vars - papilo_problem.getNCols(), + original_nnz - papilo_problem.getConstraintMatrix().getNnz()); i_t n_integer = 0; { auto col_flags = papilo_problem.getColFlags(); - for (size_t i = 0; i < col_flags.size(); i++) { + for (size_t i = 0; i < col_flags.size(); ++i) { if (col_flags[i].test(papilo::ColFlag::kIntegral)) n_integer++; } } @@ -745,62 +627,224 @@ third_party_presolve_result_t third_party_presolve_t::apply( status = third_party_presolve_status_t::OPTIMAL; } - auto opt_problem = build_optimization_problem( - papilo_problem, op_problem.get_handle_ptr(), category, maximize_); - // metadata from original optimization problem that is not filled - opt_problem.set_problem_name(op_problem.get_problem_name()); - opt_problem.set_objective_scaling_factor(op_problem.get_objective_scaling_factor()); - // when an objective offset outside (e.g. from mps file), handle accordingly - auto col_flags = papilo_problem.getColFlags(); - std::vector implied_integer_indices; - for (size_t i = 0; i < col_flags.size(); i++) { - if (col_flags[i].test(papilo::ColFlag::kImplInt)) implied_integer_indices.push_back(i); - } - auto const& col_map = result.postsolve.origcol_mapping; reduced_to_original_map_.assign(col_map.begin(), col_map.end()); - original_to_reduced_map_.assign(op_problem.get_n_variables(), -1); + original_to_reduced_map_.assign(original_n_vars, -1); for (size_t i = 0; i < reduced_to_original_map_.size(); ++i) { auto original_idx = reduced_to_original_map_[i]; if (original_idx >= 0 && static_cast(original_idx) < original_to_reduced_map_.size()) { original_to_reduced_map_[original_idx] = static_cast(i); } } + return status; +} + +// Wrapper around apply_presolve_from_mps_data +// Generates an mps_data, presolve it and turn it back into a reduced op_problem +template +third_party_presolve_device_result_t +third_party_presolve_t::apply_presolve_from_op_problem( + optimization_problem_t const& op_problem, + problem_category_t category, + cuopt::mathematical_optimization::presolver_t presolver, + bool dual_postsolve, + f_t absolute_tolerance, + f_t relative_tolerance, + double time_limit, + i_t num_cpu_threads) +{ + auto* handle = op_problem.get_handle_ptr(); + + cuopt_expects(!op_problem.has_quadratic_objective(), + error_type_t::ValidationError, + "Presolve does not support optimization_problem with a quadratic objective"); + cuopt_expects(!op_problem.has_quadratic_constraints(), + error_type_t::ValidationError, + "Presolve does not support optimization_problem with quadratic constraints"); + + auto mps = ::cuopt::mathematical_optimization::op_problem_to_mps_data_model(op_problem); + + auto host_res = apply_presolve_from_mps_data(mps, + category, + presolver, + dual_postsolve, + absolute_tolerance, + relative_tolerance, + time_limit, + num_cpu_threads); + + // On terminal statuses the mps entry returns an empty reduced problem; + // mirror that shape on the device side without going through H->D. + if (host_res.status == third_party_presolve_status_t::INFEASIBLE || + host_res.status == third_party_presolve_status_t::UNBOUNDED || + host_res.status == third_party_presolve_status_t::UNBNDORINFEAS) { + return third_party_presolve_device_result_t{ + host_res.status, + optimization_problem_t(handle), + std::move(host_res.implied_integer_indices), + std::move(host_res.reduced_to_original_map), + std::move(host_res.original_to_reduced_map)}; + } - return third_party_presolve_result_t{status, - std::move(opt_problem), - implied_integer_indices, - reduced_to_original_map_, - original_to_reduced_map_}; + auto reduced_opt = + ::cuopt::mathematical_optimization::mps_data_model_to_optimization_problem( + handle, host_res.reduced_problem); + // mps_data_model doesn't carry problem_category; plumb it here so the + // reduced op_problem matches the input's category + reduced_opt.set_problem_category(category); + + return third_party_presolve_device_result_t{ + host_res.status, + std::move(reduced_opt), + std::move(host_res.implied_integer_indices), + std::move(host_res.reduced_to_original_map), + std::move(host_res.original_to_reduced_map)}; } template -void third_party_presolve_t::undo(rmm::device_uvector& primal_solution, - rmm::device_uvector& dual_solution, - rmm::device_uvector& reduced_costs, - problem_category_t category, - bool status_to_skip, - bool dual_postsolve, - rmm::cuda_stream_view stream_view) +third_party_presolve_host_result_t +third_party_presolve_t::apply_presolve_from_mps_data( + io::mps_data_model_t const& mps, + problem_category_t category, + cuopt::mathematical_optimization::presolver_t presolver, + bool dual_postsolve, + f_t absolute_tolerance, + f_t relative_tolerance, + double time_limit, + i_t num_cpu_threads) { - if (presolver_ == cuopt::mathematical_optimization::presolver_t::PSLP) { - undo_pslp(primal_solution, dual_solution, reduced_costs, stream_view); - return; + presolver_ = presolver; + maximize_ = mps.get_sense(); + + cuopt_expects(!(category == problem_category_t::MIP && + presolver == cuopt::mathematical_optimization::presolver_t::PSLP), + error_type_t::RuntimeError, + "PSLP presolver is not supported for MIP problems"); + + // Neither PSLP nor Papilo handle quadratic objective / constraints. + cuopt_expects(!mps.has_quadratic_objective(), + error_type_t::ValidationError, + "Presolve does not support mps_data_models with a quadratic objective"); + cuopt_expects(!mps.has_quadratic_constraints(), + error_type_t::ValidationError, + "Presolve does not support mps_data_models with quadratic constraints"); + + // PSLP branch: apply_pslp -> reduced mps. + if (presolver == cuopt::mathematical_optimization::presolver_t::PSLP) { + const f_t original_obj_offset = mps.get_objective_offset(); + auto status = apply_pslp(mps, time_limit); + + if (status == third_party_presolve_status_t::INFEASIBLE || + status == third_party_presolve_status_t::UNBNDORINFEAS) { + return third_party_presolve_host_result_t{ + status, io::mps_data_model_t{}, {}, {}, {}}; + } + + auto reduced_mps = + build_reduced_mps_from_pslp(pslp_presolver_, maximize_, original_obj_offset); + reduced_mps.set_problem_name(mps.get_problem_name()); + reduced_mps.set_objective_scaling_factor(mps.get_objective_scaling_factor()); + return third_party_presolve_host_result_t{status, std::move(reduced_mps), {}, {}, {}}; + } else { + // Papilo branch: build papilo::Problem -> + // apply_papilo -> reduced mps. + auto papilo_problem = build_papilo_problem(mps, maximize_, category); + auto status = apply_papilo(papilo_problem, + category, + dual_postsolve, + absolute_tolerance, + relative_tolerance, + time_limit, + num_cpu_threads); + + if (status == third_party_presolve_status_t::INFEASIBLE || + status == third_party_presolve_status_t::UNBOUNDED || + status == third_party_presolve_status_t::UNBNDORINFEAS) { + return third_party_presolve_host_result_t{ + status, io::mps_data_model_t{}, {}, {}, {}}; + } + + auto reduced_mps = build_reduced_mps_from_papilo(papilo_problem, maximize_); + reduced_mps.set_problem_name(mps.get_problem_name()); + reduced_mps.set_objective_scaling_factor(mps.get_objective_scaling_factor()); + + auto col_flags = papilo_problem.getColFlags(); + std::vector implied_integer_indices; + for (size_t i = 0; i < col_flags.size(); ++i) { + if (col_flags[i].test(papilo::ColFlag::kImplInt)) { + implied_integer_indices.push_back(static_cast(i)); + } + } + + return third_party_presolve_host_result_t{status, + std::move(reduced_mps), + std::move(implied_integer_indices), + reduced_to_original_map_, + original_to_reduced_map_}; } +} - if (status_to_skip) { return; } +template +void third_party_presolve_t::undo_from_device(rmm::device_uvector& primal_solution, + rmm::device_uvector& dual_solution, + rmm::device_uvector& reduced_costs, + problem_category_t category, + bool status_to_skip, + bool dual_postsolve, + rmm::cuda_stream_view stream_view) +{ + std::vector h_primal(primal_solution.size()); + std::vector h_dual(dual_solution.size()); + std::vector h_rc(reduced_costs.size()); + raft::copy(h_primal.data(), primal_solution.data(), primal_solution.size(), stream_view); + raft::copy(h_dual.data(), dual_solution.data(), dual_solution.size(), stream_view); + raft::copy(h_rc.data(), reduced_costs.data(), reduced_costs.size(), stream_view); + stream_view.synchronize(); - std::vector primal_sol_vec_h(primal_solution.size()); - raft::copy(primal_sol_vec_h.data(), primal_solution.data(), primal_solution.size(), stream_view); - std::vector dual_sol_vec_h(dual_solution.size()); - raft::copy(dual_sol_vec_h.data(), dual_solution.data(), dual_solution.size(), stream_view); - std::vector reduced_costs_vec_h(reduced_costs.size()); - raft::copy(reduced_costs_vec_h.data(), reduced_costs.data(), reduced_costs.size(), stream_view); + undo(h_primal, h_dual, h_rc, category, status_to_skip, dual_postsolve); - papilo::Solution reduced_sol(primal_sol_vec_h); + primal_solution.resize(h_primal.size(), stream_view); + dual_solution.resize(h_dual.size(), stream_view); + reduced_costs.resize(h_rc.size(), stream_view); + raft::copy(primal_solution.data(), h_primal.data(), h_primal.size(), stream_view); + raft::copy(dual_solution.data(), h_dual.data(), h_dual.size(), stream_view); + raft::copy(reduced_costs.data(), h_rc.data(), h_rc.size(), stream_view); + stream_view.synchronize(); +} + +template +void third_party_presolve_t::undo_pslp(std::vector& primal_solution, + std::vector& dual_solution, + std::vector& reduced_costs) +{ + if constexpr (std::is_same_v) { + // PSLP postsolve reads from the passed-in host buffers and writes the + // uncrushed solution into pslp_presolver_->sol->{x, y, z}. + postsolve(pslp_presolver_, primal_solution.data(), dual_solution.data(), reduced_costs.data()); + + auto uncrushed_sol = pslp_presolver_->sol; + const int n_cols = uncrushed_sol->dim_x; + const int n_rows = uncrushed_sol->dim_y; + + primal_solution.assign(uncrushed_sol->x, uncrushed_sol->x + n_cols); + dual_solution.assign(uncrushed_sol->y, uncrushed_sol->y + n_rows); + reduced_costs.assign(uncrushed_sol->z, uncrushed_sol->z + n_cols); + } else { + cuopt_expects( + false, error_type_t::ValidationError, "PSLP postsolve only supports double precision"); + } +} + +template +void third_party_presolve_t::undo_papilo(std::vector& primal_solution, + std::vector& dual_solution, + std::vector& reduced_costs, + bool dual_postsolve) +{ + papilo::Solution reduced_sol(primal_solution); if (dual_postsolve) { - reduced_sol.dual = dual_sol_vec_h; - reduced_sol.reducedCosts = reduced_costs_vec_h; + reduced_sol.dual = dual_solution; + reduced_sol.reducedCosts = reduced_costs; reduced_sol.type = papilo::SolutionType::kPrimalDual; } papilo::Solution full_sol; @@ -813,51 +857,26 @@ void third_party_presolve_t::undo(rmm::device_uvector& primal_sol auto status = post_solver.undo(reduced_sol, full_sol, *papilo_post_solve_storage_, is_optimal); check_postsolve_status(status); - primal_solution.resize(full_sol.primal.size(), stream_view); - dual_solution.resize(full_sol.dual.size(), stream_view); - reduced_costs.resize(full_sol.reducedCosts.size(), stream_view); - raft::copy(primal_solution.data(), full_sol.primal.data(), full_sol.primal.size(), stream_view); - raft::copy(dual_solution.data(), full_sol.dual.data(), full_sol.dual.size(), stream_view); - raft::copy( - reduced_costs.data(), full_sol.reducedCosts.data(), full_sol.reducedCosts.size(), stream_view); + primal_solution = std::move(full_sol.primal); + dual_solution = std::move(full_sol.dual); + reduced_costs = std::move(full_sol.reducedCosts); } template -void third_party_presolve_t::undo_pslp(rmm::device_uvector& primal_solution, - rmm::device_uvector& dual_solution, - rmm::device_uvector& reduced_costs, - rmm::cuda_stream_view stream_view) +void third_party_presolve_t::undo(std::vector& primal_solution, + std::vector& dual_solution, + std::vector& reduced_costs, + problem_category_t /*category*/, + bool status_to_skip, + bool dual_postsolve) { - if constexpr (std::is_same_v) { - // PSLP uses double internally, so we can use the data directly - std::vector h_primal_solution(primal_solution.size()); - std::vector h_dual_solution(dual_solution.size()); - std::vector h_reduced_costs(reduced_costs.size()); - raft::copy( - h_primal_solution.data(), primal_solution.data(), primal_solution.size(), stream_view); - raft::copy(h_dual_solution.data(), dual_solution.data(), dual_solution.size(), stream_view); - raft::copy(h_reduced_costs.data(), reduced_costs.data(), reduced_costs.size(), stream_view); - stream_view.synchronize(); - - postsolve( - pslp_presolver_, h_primal_solution.data(), h_dual_solution.data(), h_reduced_costs.data()); - - auto uncrushed_sol = pslp_presolver_->sol; - int n_cols = uncrushed_sol->dim_x; - int n_rows = uncrushed_sol->dim_y; - - primal_solution.resize(n_cols, stream_view); - dual_solution.resize(n_rows, stream_view); - reduced_costs.resize(n_cols, stream_view); - raft::copy(primal_solution.data(), uncrushed_sol->x, n_cols, stream_view); - raft::copy(dual_solution.data(), uncrushed_sol->y, n_rows, stream_view); - raft::copy(reduced_costs.data(), uncrushed_sol->z, n_cols, stream_view); - } else { - cuopt_expects( - false, error_type_t::ValidationError, "PSLP postsolve only supports double precision"); + if (presolver_ == cuopt::mathematical_optimization::presolver_t::PSLP) { + undo_pslp(primal_solution, dual_solution, reduced_costs); + return; + } else { // Papilo branch + if (status_to_skip) { return; } + undo_papilo(primal_solution, dual_solution, reduced_costs, dual_postsolve); } - - stream_view.synchronize(); } template diff --git a/cpp/src/mip_heuristics/presolve/third_party_presolve.hpp b/cpp/src/mip_heuristics/presolve/third_party_presolve.hpp index a5fa089d71..776a12b50e 100644 --- a/cpp/src/mip_heuristics/presolve/third_party_presolve.hpp +++ b/cpp/src/mip_heuristics/presolve/third_party_presolve.hpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -18,7 +19,10 @@ namespace papilo { template class PostsolveStorage; -} + +template +class Problem; +} // namespace papilo namespace cuopt::mathematical_optimization::mip { @@ -36,16 +40,24 @@ enum class third_party_presolve_status_t { UNCHANGED, }; -template +template struct third_party_presolve_result_t { third_party_presolve_status_t status; - optimization_problem_t reduced_problem; + ProblemT reduced_problem; std::vector implied_integer_indices; std::vector reduced_to_original_map; std::vector original_to_reduced_map; // clique info, etc... }; +template +using third_party_presolve_device_result_t = + third_party_presolve_result_t>; + +template +using third_party_presolve_host_result_t = + third_party_presolve_result_t>; + template class third_party_presolve_t { public: @@ -58,7 +70,9 @@ class third_party_presolve_t { third_party_presolve_t(third_party_presolve_t&&) = delete; third_party_presolve_t& operator=(third_party_presolve_t&&) = delete; - third_party_presolve_result_t apply( + // Device entry: takes an optimization_problem_t and returns a device-side + // reduced optimization_problem_t. This is a wrapper around apply_presolve_from_mps_data. + third_party_presolve_device_result_t apply_presolve_from_op_problem( optimization_problem_t const& op_problem, problem_category_t category, cuopt::mathematical_optimization::presolver_t presolver, @@ -68,13 +82,34 @@ class third_party_presolve_t { double time_limit, i_t num_cpu_threads = 0); - void undo(rmm::device_uvector& primal_solution, - rmm::device_uvector& dual_solution, - rmm::device_uvector& reduced_costs, + // Host entry: takes an mps_data_model_t and returns a host-side reduced + // mps_data_model_t. Pure-host throughout + third_party_presolve_host_result_t apply_presolve_from_mps_data( + io::mps_data_model_t const& mps_problem, + problem_category_t category, + cuopt::mathematical_optimization::presolver_t presolver, + bool dual_postsolve, + f_t absolute_tolerance, + f_t relative_tolerance, + double time_limit, + i_t num_cpu_threads = 0); + + void undo_from_device(rmm::device_uvector& primal_solution, + rmm::device_uvector& dual_solution, + rmm::device_uvector& reduced_costs, + problem_category_t category, + bool status_to_skip, + bool dual_postsolve, + rmm::cuda_stream_view stream_view); + + // Host-only postsolve. Resizes the vectors to original-problem dimensions. + // The device-side `undo_from_device` above is a thin wrapper around this method. + void undo(std::vector& primal_solution, + std::vector& dual_solution, + std::vector& reduced_costs, problem_category_t category, bool status_to_skip, - bool dual_postsolve, - rmm::cuda_stream_view stream_view); + bool dual_postsolve); void uncrush_primal_solution(const std::vector& reduced_primal, std::vector& full_primal) const; @@ -97,13 +132,27 @@ class third_party_presolve_t { ~third_party_presolve_t(); private: - third_party_presolve_result_t apply_pslp( - optimization_problem_t const& op_problem, const double time_limit); - - void undo_pslp(rmm::device_uvector& primal_solution, - rmm::device_uvector& dual_solution, - rmm::device_uvector& reduced_costs, - rmm::cuda_stream_view stream_view); + third_party_presolve_status_t apply_pslp(io::mps_data_model_t const& mps, + double time_limit); + + third_party_presolve_status_t apply_papilo(papilo::Problem& papilo_problem, + problem_category_t category, + bool dual_postsolve, + f_t absolute_tolerance, + f_t relative_tolerance, + double time_limit, + i_t num_cpu_threads); + + // Host-only per-backend postsolve helpers. Both resize their vector args + // to original-problem dimensions. + void undo_pslp(std::vector& primal_solution, + std::vector& dual_solution, + std::vector& reduced_costs); + + void undo_papilo(std::vector& primal_solution, + std::vector& dual_solution, + std::vector& reduced_costs, + bool dual_postsolve); bool maximize_ = false; cuopt::mathematical_optimization::presolver_t presolver_ = diff --git a/cpp/src/mip_heuristics/solve.cu b/cpp/src/mip_heuristics/solve.cu index 59b582188f..95786f3c5a 100644 --- a/cpp/src/mip_heuristics/solve.cu +++ b/cpp/src/mip_heuristics/solve.cu @@ -448,7 +448,7 @@ mip_solution_t solve_mip_helper(optimization_problem_t& op_p } double presolve_time = 0.0; std::unique_ptr> presolver; - std::optional> presolve_result_opt; + std::optional> presolve_result_opt; mip::problem_t problem( op_problem, settings.get_tolerances(), settings.determinism_mode == CUOPT_MODE_DETERMINISTIC); @@ -568,14 +568,15 @@ mip_solution_t solve_mip_helper(optimization_problem_t& op_p presolve_time_limit = std::numeric_limits::infinity(); } presolver = std::make_unique>(); - auto result = presolver->apply(op_problem, - cuopt::mathematical_optimization::problem_category_t::MIP, - settings.presolver, - dual_postsolve, - settings.tolerances.absolute_tolerance, - settings.tolerances.relative_tolerance, - presolve_time_limit, - settings.num_cpu_threads); + auto result = presolver->apply_presolve_from_op_problem( + op_problem, + cuopt::mathematical_optimization::problem_category_t::MIP, + settings.presolver, + dual_postsolve, + settings.tolerances.absolute_tolerance, + settings.tolerances.relative_tolerance, + presolve_time_limit, + settings.num_cpu_threads); if (result.status == mip::third_party_presolve_status_t::INFEASIBLE) { return mip_solution_t(mip_termination_status_t::Infeasible, @@ -671,13 +672,13 @@ mip_solution_t solve_mip_helper(optimization_problem_t& op_p cuopt::device_copy(sol.get_solution(), op_problem.get_handle_ptr()->get_stream()); rmm::device_uvector dual_solution(0, op_problem.get_handle_ptr()->get_stream()); rmm::device_uvector reduced_costs(0, op_problem.get_handle_ptr()->get_stream()); - presolver->undo(primal_solution, - dual_solution, - reduced_costs, - cuopt::mathematical_optimization::problem_category_t::MIP, - status_to_skip, - dual_postsolve, - op_problem.get_handle_ptr()->get_stream()); + presolver->undo_from_device(primal_solution, + dual_solution, + reduced_costs, + cuopt::mathematical_optimization::problem_category_t::MIP, + status_to_skip, + dual_postsolve, + op_problem.get_handle_ptr()->get_stream()); if (!status_to_skip) { thrust::fill(rmm::exec_policy(op_problem.get_handle_ptr()->get_stream()), dual_solution.data(), diff --git a/cpp/src/pdlp/CMakeLists.txt b/cpp/src/pdlp/CMakeLists.txt index f5f26837b6..2f90f94872 100644 --- a/cpp/src/pdlp/CMakeLists.txt +++ b/cpp/src/pdlp/CMakeLists.txt @@ -29,6 +29,11 @@ set(LP_CORE_FILES ${CMAKE_CURRENT_SOURCE_DIR}/termination_strategy/convergence_information.cu ${CMAKE_CURRENT_SOURCE_DIR}/optimal_batch_size_handler/optimal_batch_size_handler.cu ${CMAKE_CURRENT_SOURCE_DIR}/utilities/ping_pong_graph.cu + ${CMAKE_CURRENT_SOURCE_DIR}/distributed_pdlp/shard.cu + ${CMAKE_CURRENT_SOURCE_DIR}/distributed_pdlp/multi_gpu_engine.cu + ${CMAKE_CURRENT_SOURCE_DIR}/distributed_pdlp/distributed_algorithms.cu + ${CMAKE_CURRENT_SOURCE_DIR}/distributed_pdlp/distributed_utils.cu + ${CMAKE_CURRENT_SOURCE_DIR}/distributed_pdlp/partitioner.cpp ) # C and Python adapter files diff --git a/cpp/src/pdlp/cusparse_view.cu b/cpp/src/pdlp/cusparse_view.cu index 0974a49c82..08db620986 100644 --- a/cpp/src/pdlp/cusparse_view.cu +++ b/cpp/src/pdlp/cusparse_view.cu @@ -566,14 +566,17 @@ cusparse_view_t::cusparse_view_t( // setup cusparse view A.create(op_problem_scaled.n_constraints, op_problem_scaled.n_variables, - op_problem_scaled.nnz, + static_cast(A_.size()), const_cast(op_problem_scaled.offsets.data()), const_cast(op_problem_scaled.variables.data()), const_cast(op_problem_scaled.coefficients.data())); + // A_T can have a different nnz than A in multi-GPU shards + // A is just what is needed to compute A_x for owned constraints + // A_T is just what is needed to compute A_T_y for owned variables A_T.create(op_problem_scaled.n_variables, op_problem_scaled.n_constraints, - op_problem_scaled.nnz, + static_cast(A_T_.size()), const_cast(A_T_offsets_.data()), const_cast(A_T_indices_.data()), const_cast(A_T_.data())); @@ -610,7 +613,9 @@ cusparse_view_t::cusparse_view_t( op_problem_scaled.n_variables, current_saddle_point_state.get_next_AtY().data(), CUSPARSE_ORDER_COL); - cuopt_assert(_reflected_primal_solution.size() > 0, "Reflected primal solution empty"); + cuopt_assert(_reflected_primal_solution.size() >= + static_cast(op_problem_scaled.n_variables) * climber_strategies.size(), + "Reflected primal solution undersized"); batch_reflected_primal_solutions.create(op_problem_scaled.n_variables, climber_strategies.size(), climber_strategies.size(), @@ -661,7 +666,9 @@ cusparse_view_t::cusparse_view_t( tmp_primal.create(op_problem_scaled.n_variables, _tmp_primal.data()); tmp_dual.create(op_problem_scaled.n_constraints, _tmp_dual.data()); if (hyper_params.use_reflected_primal_dual) { - cuopt_assert(_reflected_primal_solution.size() > 0, "Reflected primal solution empty"); + cuopt_assert( + _reflected_primal_solution.size() >= static_cast(op_problem_scaled.n_variables), + "Reflected primal solution undersized"); reflected_primal_solution.create(op_problem_scaled.n_variables, _reflected_primal_solution.data()); } @@ -982,14 +989,14 @@ cusparse_view_t::cusparse_view_t( // setup cusparse view A.create(op_problem.n_constraints, op_problem.n_variables, - op_problem.nnz, + static_cast(A_.size()), const_cast(op_problem.offsets.data()), const_cast(op_problem.variables.data()), const_cast(op_problem.coefficients.data())); A_T.create(op_problem.n_variables, op_problem.n_constraints, - op_problem.nnz, + static_cast(A_T_.size()), const_cast(A_T_offsets_.data()), const_cast(A_T_indices_.data()), const_cast(A_T_.data())); @@ -1197,16 +1204,18 @@ cusparse_view_t::cusparse_view_t( // Copying them from the existing cuSparse view is a bad practice and creates segfault post // CUDA 12.4 Using the saved pointer of the existing cusparse view to make sure we capture the // correct pointer + // See comment in the PDHG cusparse_view_t ctor: bind the descriptor nnz to + // the actual value-buffer length so A and A_T stay symmetric and shard-safe. A.create(op_problem.n_constraints, op_problem.n_variables, - op_problem.nnz, + static_cast(A_.size()), const_cast(A_offsets_.data()), const_cast(A_indices_.data()), const_cast(A_.data())); A_T.create(op_problem.n_variables, op_problem.n_constraints, - op_problem.nnz, + static_cast(existing_cusparse_view.A_T_.size()), const_cast(existing_cusparse_view.A_T_offsets_.data()), const_cast(existing_cusparse_view.A_T_indices_.data()), const_cast(existing_cusparse_view.A_T_.data())); diff --git a/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu b/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu new file mode 100644 index 0000000000..af9bef5571 --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu @@ -0,0 +1,428 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Out-of-line definitions of multi_gpu_engine_t's high-level algorithm methods +// used by the pdlp solver. +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// -------- Solution gather (shards -> master) ------------------------------ +// Gather the potential next primal/dual solutions and the reduced cost from shards to master. +template +void multi_gpu_engine_t::gather_potential_next_solutions_to_master() +{ + cuopt_assert(master_pdlp_ != nullptr, + "gather_potential_next_solutions_to_master requires set_master(...)"); + + gather_owned_var_to_master([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_potential_next_primal_solution(); + }); + + gather_owned_cstr_to_master([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_potential_next_dual_solution(); + }); + + gather_owned_var_to_master([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.get_current_termination_strategy().get_convergence_information().get_reduced_cost(); + }); +} + +// -------- Distributed bound / objective rescaling ------------------------- +// compute and apply_bound_objective_rescaling_to_problem, unfused because we need a +// raw squared-sum on device to reduce and the base version comptues +// tranform->reduce->transform in one cub call for efficiency +template +void multi_gpu_engine_t::distributed_bound_objective_rescaling(f_t c_scaling_weight) +{ + raft::common::nvtx::range scope("distributed_bound_objective_rescaling"); + + // 1) + 2) Local transform-reduce on each shard, accumulate the global + // squared L2 norms on host as we go. + f_t global_bound_sq = f_t(0); + f_t global_obj_sq = f_t(0); + for_each_shard([&](auto& s) { + const auto& scaled = s.sub_pdlp->get_initial_scaling_strategy().get_scaled_op_problem(); + auto policy = rmm::exec_policy(s.stream.view()); + auto bounds_begin = thrust::make_zip_iterator(scaled.constraint_lower_bounds.data(), + scaled.constraint_upper_bounds.data()); + global_bound_sq += thrust::transform_reduce(policy, + bounds_begin, + bounds_begin + s.rank_data.owned_cstr_size, + rhs_sum_of_squares_t{}, + f_t(0), + thrust::plus{}); + global_obj_sq += + thrust::transform_reduce(policy, + scaled.objective_coefficients.data(), + scaled.objective_coefficients.data() + s.rank_data.owned_var_size, + weighted_square_op{c_scaling_weight}, + f_t(0), + thrust::plus{}); + }); + + // 3) Host-side derivation of the (identical on every shard) scaling scalars. + const f_t bound_rescaling = rescaling_from_squared_norm_op{}(global_bound_sq); + const f_t obj_rescaling = rescaling_from_squared_norm_op{}(global_obj_sq); + + // 4) Publish + apply on every shard via the shared helpers. + for_each_shard([&](auto& s) { + auto& scaling = s.sub_pdlp->get_initial_scaling_strategy(); + scaling.set_h_bound_rescaling(bound_rescaling); + scaling.set_h_objective_rescaling(obj_rescaling); + scaling.apply_bound_objective_rescaling_to_problem(); + }); + + synchronize_shards(); +} + +// -------- Refresh halo of cumulative scalings ----------------------------- +// Refreshes the halo copies of the cumulative variable + constraint scalings on +// every shard. Called before and after each matrix-scaling pass in ruiz and pock-chambolle +template +void multi_gpu_engine_t::refresh_halo_cummulative_scalings() +{ + halo_exchange_var([](pdlp_solver_t& p) -> auto& { + return p.get_initial_scaling_strategy().get_cummulative_variable_scaling(); + }); + halo_exchange_cstr([](pdlp_solver_t& p) -> auto& { + return p.get_initial_scaling_strategy().get_cummulative_constraint_matrix_scaling(); + }); +} + +// -------- Distributed Ruiz inf-scaling ------------------------------------ +// Each shard owns its rows AND its columns and stores both complete (h_A = +// owned rows, h_A_t = owned columns) +template +void multi_gpu_engine_t::distributed_ruiz_inf_scaling(int num_iter, i_t n_global_vars) +{ + if (num_iter <= 0 || n_global_vars <= 0) return; + raft::common::nvtx::range scope("distributed_ruiz_inf_scaling"); + + for (int it = 0; it < num_iter; ++it) { + refresh_halo_cummulative_scalings(); + + // Shard-local Ruiz iteration + // rows: inf norm only over OWNED (full) rows from A + // cols: inf norm only over OWNED (full) cols from A_T + // Then fold into cumulative on owned entries (halo entries get refreshed by + // the next iteration's halo update) + for_each_shard( + [](auto& shard) { shard.sub_pdlp->get_initial_scaling_strategy().ruiz_iter_local(); }); + } + + // Final refresh after last iteration + refresh_halo_cummulative_scalings(); + + synchronize_shards(); +} + +// -------- Distributed Pock-Chambolle scaling ------------------------------ +// Distributed Pock-Chambolle: one pass, mirroring single-GPU +// pock_chambolle_scaling. Row sum-of-powers come from the row-major matrix +// (owned rows) and column sum-of-powers from A_T (owned columns). +template +void multi_gpu_engine_t::distributed_pock_chambolle_scaling(f_t alpha, i_t n_global_vars) +{ + if (n_global_vars <= 0) return; + raft::common::nvtx::range scope("distributed_pock_chambolle_scaling"); + + refresh_halo_cummulative_scalings(); + + for_each_shard([alpha](auto& shard) { + shard.sub_pdlp->get_initial_scaling_strategy().pock_chambolle_scaling(alpha); + }); + + // Final refresh for downstream consumers. + refresh_halo_cummulative_scalings(); + + synchronize_shards(); +} + +// -------- Distributed scaling orchestration ------------------------------ +// Mirrors single GPU scaling +template +void multi_gpu_engine_t::distributed_scaling(pdlp_hyper_params_t const& hyper_params, + i_t n_global_vars, + bool inside_mip) +{ + raft::common::nvtx::range scope("distributed_scaling"); + + // 1) Matrix scaling passes populate the cumulative row/col scalings on + // every shard. Each pass keeps the halo copies refreshed internally. + if (hyper_params.do_ruiz_scaling) { + distributed_ruiz_inf_scaling(hyper_params.default_l_inf_ruiz_iterations, n_global_vars); + } + if (hyper_params.do_pock_chambolle_scaling) { + distributed_pock_chambolle_scaling( + static_cast(hyper_params.default_alpha_pock_chambolle_rescaling), n_global_vars); + } + + // 2) Per-shard apply of the accumulated scaling to A, c, variable and + // constraint bounds. This is scale_problem() minus its local + // bound/objective rescaling; the equivalent global step happens in (3). + for_each_shard([](auto& shard) { + shard.sub_pdlp->get_initial_scaling_strategy().apply_cummulative_scaling_to_problem(); + }); + synchronize_shards(); + + // 3) Global bound/objective rescaling (all shards get the identical scalar). + if (hyper_params.bound_objective_rescaling) { + distributed_bound_objective_rescaling( + static_cast(hyper_params.initial_primal_weight_c_scaling)); + } +} + +// -------- Distributed sigma_max(A)^2 via power iteration ------------------ +// Owns per-shard scratch (q / z / atq / scalar reductions) and drives the +// iteration; every cross-shard operation goes through multi_gpu_engine_t's +// *_bufs helpers (halo_exchange_{cstr,var}_bufs, distributed_l2_norm_bufs, +// distributed_dot_bufs), so this function contains no NCCL calls directly. +template +f_t multi_gpu_engine_t::distributed_max_singular_value_squared(i_t n_global_cstrs, + int max_iterations, + f_t tolerance) +{ + raft::common::nvtx::range scope("distributed_max_singular_value_squared"); + + // ┌──────────────────────────────────────────────────────────────┐ + // │ Setup │ + // └──────────────────────────────────────────────────────────────┘ + + const int nb = static_cast(shards.size()); + // Generate the GLOBAL z[] sequence in cstr-index order. + // Scatter it to the shards according to the partition. + std::vector h_global_z = + make_singular_value_probe(static_cast(n_global_cstrs)); + + // Per-shard scratch lives on each shard's device. + std::vector> q; + std::vector> z; + std::vector> atq; + std::vector> sigma_sq; + std::vector> norm_q; + std::vector> residual_norm; + + std::vector> q_dn(nb); + std::vector> z_dn(nb); + std::vector> atq_dn(nb); + + // Per-shard owned-slice spans consumed by the engine's *_bufs helpers. + std::vector> q_owned, z_owned; + for (auto* v : {&q, &z, &atq}) + v->reserve(nb); + for (auto* v : {&sigma_sq, &norm_q, &residual_norm}) + v->reserve(nb); + for (auto* v : {&q_owned, &z_owned}) + v->reserve(nb); + + // Allocate per-shard scratch, scatter z according to partition, and build + // the *_bufs views for the power iteration below. + for_each_shard([&](auto& s, int r) { + const i_t cstr_total = s.rank_data.total_cstr_size; + const i_t var_total = s.rank_data.total_var_size; + const i_t n_owned = s.rank_data.owned_cstr_size; + + q.emplace_back(static_cast(cstr_total), s.stream.view()); + z.emplace_back(static_cast(cstr_total), s.stream.view()); + atq.emplace_back(static_cast(var_total), s.stream.view()); + sigma_sq.emplace_back(s.stream.view()); + norm_q.emplace_back(s.stream.view()); + residual_norm.emplace_back(s.stream.view()); + q_dn[r].create(static_cast(cstr_total), q.back().data()); + z_dn[r].create(static_cast(cstr_total), z.back().data()); + atq_dn[r].create(static_cast(var_total), atq.back().data()); + + q_owned.emplace_back(q.back().data(), static_cast(n_owned)); + z_owned.emplace_back(z.back().data(), static_cast(n_owned)); + + // Scatter z according to partition + std::vector h_owned_z(static_cast(n_owned)); + thrust::gather(thrust::host, + s.rank_data.local_to_global_cstr.begin(), + s.rank_data.local_to_global_cstr.begin() + n_owned, + h_global_z.begin(), + h_owned_z.begin()); + raft::copy( + z.back().data(), h_owned_z.data(), static_cast(n_owned), s.stream.view()); + thrust::fill(rmm::exec_policy_nosync(s.stream.view()), + z.back().data() + n_owned, + z.back().data() + cstr_total, + f_t(0)); + + // Sync to ensure h_owned_z stays valid through the H2D copy (it goes + // out of scope at end of this iteration of the per-shard loop). + s.stream.synchronize(); + }); + + // ┌──────────────────────────────────────────────────────────────┐ + // │ Power iteration │ + // └──────────────────────────────────────────────────────────────┘ + + // Mirrors single-GPU compute_initial_step_size. + // copy -> l2 norm -> transform -> SpMV -> SpMV -> dot -> transform -> norm -> convergence check + for (int it = 0; it < max_iterations; ++it) { + // q := z on the owned slice (the carried iterate). + for_each_shard([&](auto& s, int r) { + const i_t n_owned = s.rank_data.owned_cstr_size; + raft::copy(q[r].data(), z[r].data(), n_owned, s.stream.view()); + }); + + // ||q||₂ over the global OWNED cstr slice (one allreduce-sum + sqrt). + distributed_l2_norm_bufs(q_owned, norm_q); + + // q /= ||q||₂ on owned slice (halo gets refreshed by next exchange). + // Kept as per-shard cub launch: the divisor is a per-shard scalar. + for_each_shard([&](auto& s, int r) { + const i_t n_owned = s.rank_data.owned_cstr_size; + cub::DeviceTransform::Transform(q[r].data(), + q[r].data(), + n_owned, + divide_by_device_scalar_t{norm_q[r].data()}, + s.stream.view().value()); + }); + + // atq = A^T q (fused halo-refresh of q + per-shard local SpMV). + distributed_spmv_At(q, q_dn, atq_dn); + + // z = A atq (fused halo-refresh of atq + per-shard local SpMV). + distributed_spmv_A(atq, atq_dn, z_dn); + + // σ² = q · z over the global OWNED cstr slice (= q^T A A^T q = σ_max² + // when q is the dominant left-singular vector). + distributed_dot_bufs(q_owned, z_owned, sigma_sq); + + // q := -σ² q + z (owned slice) — residual of the eigen-equation. + for_each_shard([&](auto& s, int r) { + const i_t n_owned = s.rank_data.owned_cstr_size; + cub::DeviceTransform::Transform(cuda::std::make_tuple(q[r].data(), z[r].data()), + q[r].data(), + n_owned, + residual_fma_neg_scalar_t{sigma_sq[r].data()}, + s.stream.view().value()); + }); + + // Convergence check via global residual norm. + distributed_l2_norm_bufs(q_owned, residual_norm); + auto& s0 = *shards[0]; + raft::device_setter guard0(s0.device_id); + f_t h_res{}; + raft::copy(&h_res, residual_norm[0].data(), 1, s0.stream.view()); + s0.stream.synchronize(); + if (h_res < tolerance) break; + } + + // σ_max² is the same on every shard after the last allreduce. + auto& s0 = *shards[0]; + raft::device_setter guard0(s0.device_id); + f_t sigma_sq_h{}; + raft::copy(&sigma_sq_h, sigma_sq[0].data(), 1, s0.stream.view()); + s0.stream.synchronize(); + + return sigma_sq_h; +} + +// -------- Distributed initial step size --------------------------------- +// Sigma_max(A) via the shared power-iteration primitive. +// +// This function mirrors single-GPU's compute_initial_step_size exactly +// and broadcasts the result to every shard +template +void multi_gpu_engine_t::distributed_compute_initial_step_size( + pdlp_hyper_params_t const& hyper_params, + i_t n_global_cstrs, + f_t scaling_factor, + int max_iterations, + f_t tolerance) +{ + raft::common::nvtx::range scope("distributed_compute_initial_step_size"); + cuopt_assert(master_pdlp_ != nullptr, + "distributed_compute_initial_step_size requires set_master(...)"); + cuopt_expects(hyper_params.initial_step_size_max_singular_value, + error_type_t::ValidationError, + "distributed_compute_initial_step_size requires " + "initial_step_size_max_singular_value = true; the max-abs-value " + "of A fallback is single-GPU only. This should have been rejected " + "earlier in solve_lp_distributed_from_mps."); + + const f_t sigma_max_sq = + distributed_max_singular_value_squared(n_global_cstrs, max_iterations, tolerance); + + const f_t h_step_size = scaling_factor / std::sqrt(sigma_max_sq); + + set_scalar_on_master_and_shards(h_step_size, [](auto& sp) { return sp.get_step_size().data(); }); +} + +// -------- Distributed initial primal weight ------------------------------ +// Distributed PDLP is currently restricted to the Stable3-shaped hyper-param +// profile. Single-GPU compute_initial_primal_weight +// short-circuits to primal_weight = 1 without touching the norms (see +// pdlp.cu: +// !initial_primal_weight_combined_bounds && bound_objective_rescaling +// -> uninitialized_fill(primal_weight_ / best_primal_weight_, 1); return +// This function also fills the shards and masters primal_weight / step_size buffers +template +void multi_gpu_engine_t::distributed_compute_initial_primal_weight( + pdlp_hyper_params_t const& hyper_params) +{ + raft::common::nvtx::range scope("distributed_compute_initial_primal_weight"); + cuopt_assert(master_pdlp_ != nullptr, + "distributed_compute_initial_primal_weight requires set_master(...)"); + cuopt_expects( + !hyper_params.initial_primal_weight_combined_bounds && hyper_params.bound_objective_rescaling, + error_type_t::ValidationError, + "distributed_compute_initial_primal_weight: only the Stable3-shaped " + "short-circuit is supported (initial_primal_weight_combined_bounds=false " + "and bound_objective_rescaling=true). This should have been rejected " + "earlier in solve_lp_distributed_from_mps."); + const f_t h_primal_weight = f_t(1); + + set_scalar_on_master_and_shards(h_primal_weight, + [](auto& sp) { return sp.get_primal_weight().data(); }); + set_scalar_on_master_and_shards(h_primal_weight, + [](auto& sp) { return sp.get_best_primal_weight().data(); }); +} + +// ----- Explicit instantiations (member-by-member) -------------------------- +// The class template is instantiated in multi_gpu_engine.cu; here we only +// explicit-instantiate the out-of-line members defined in this TU. +#define INSTANTIATE(F_TYPE) \ + template void multi_gpu_engine_t::gather_potential_next_solutions_to_master(); \ + template void multi_gpu_engine_t::refresh_halo_cummulative_scalings(); \ + template void multi_gpu_engine_t::distributed_bound_objective_rescaling(F_TYPE); \ + template void multi_gpu_engine_t::distributed_ruiz_inf_scaling(int, int); \ + template void multi_gpu_engine_t::distributed_pock_chambolle_scaling(F_TYPE, int); \ + template void multi_gpu_engine_t::distributed_scaling( \ + pdlp_hyper_params_t const&, int, bool); \ + template F_TYPE multi_gpu_engine_t::distributed_max_singular_value_squared( \ + int, int, F_TYPE); \ + template void multi_gpu_engine_t::distributed_compute_initial_step_size( \ + pdlp_hyper_params_t const&, int, F_TYPE, int, F_TYPE); \ + template void multi_gpu_engine_t::distributed_compute_initial_primal_weight( \ + pdlp_hyper_params_t const&); + +INSTANTIATE(double) +INSTANTIATE(float) + +#undef INSTANTIATE + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/distributed_utils.cu b/cpp/src/pdlp/distributed_pdlp/distributed_utils.cu new file mode 100644 index 0000000000..c7c746449f --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/distributed_utils.cu @@ -0,0 +1,263 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +template +std::vector> create_rank_data_from_parts( + const std::vector& parts, + const std::vector& A_row_offsets, + const std::vector& A_col_indices, + const std::vector& A_values, + const std::vector& A_t_row_offsets, + const std::vector& A_t_col_indices, + const std::vector& A_t_values, + i_t nb_parts, + i_t nb_cstr, + i_t nb_vars, + i_t nnz) +{ + std::vector> rank_data(nb_parts, rank_data_t(nb_parts)); + cuopt_expects(static_cast(parts.size()) == nb_cstr + nb_vars, + error_type_t::ValidationError, + "parts size mismatch: expected nb_cstr + nb_vars"); + + // Same as two vectors but faster + auto cstr_owner = [&](i_t c) { return parts[c]; }; + auto var_owner = [&](i_t v) { return parts[nb_cstr + v]; }; + + std::vector owned_cstr_counts(nb_parts, 0); + std::vector owned_var_counts(nb_parts, 0); + std::vector owned_A_nnz(nb_parts, 0); + std::vector owned_A_t_nnz(nb_parts, 0); + + // Pre-count ownership and nnz to reserve exact capacities and avoid + // repeated growth/reallocation across huge vectors. + for (i_t i = 0; i < nb_cstr; ++i) { + const i_t owner = cstr_owner(i); + ++owned_cstr_counts[owner]; + owned_A_nnz[owner] += (A_row_offsets[i + 1] - A_row_offsets[i]); + } + for (i_t j = 0; j < nb_vars; ++j) { + const i_t owner = var_owner(j); + ++owned_var_counts[owner]; + owned_A_t_nnz[owner] += (A_t_row_offsets[j + 1] - A_t_row_offsets[j]); + } + // Reserve exact capacities + for (i_t rank = 0; rank < nb_parts; ++rank) { + rank_data[rank].owned_cstr_indices.reserve(static_cast(owned_cstr_counts[rank])); + rank_data[rank].owned_var_indices.reserve(static_cast(owned_var_counts[rank])); + } + + // 1. Compute ownership + for (i_t i = 0; i < nb_cstr; i++) { + rank_data[cstr_owner(i)].owned_cstr_indices.push_back(i); + } + for (i_t i = 0; i < nb_vars; i++) { + rank_data[var_owner(i)].owned_var_indices.push_back(i); + } + + // 2. Compute local matrices and rank_data +#pragma omp parallel for + for (i_t rank = 0; rank < nb_parts; rank++) { + auto& rd = rank_data[rank]; + rd.owned_var_size = rd.owned_var_indices.size(); + rd.owned_cstr_size = rd.owned_cstr_indices.size(); + // ---- A side ---- + std::vector local_A_row_offsets; + std::vector local_A_col_indices; + std::vector local_A_values; + local_A_row_offsets.reserve(static_cast(rd.owned_cstr_size) + 1); + local_A_col_indices.reserve(static_cast(owned_A_nnz[rank])); + local_A_values.reserve(static_cast(owned_A_nnz[rank])); + + i_t local_A_nnz = 0; + local_A_row_offsets.push_back(local_A_nnz); + + // For each owned constraint, add associated matrix row to local A + // Keep the same indices here but re-order later + for (auto owned_cstr : rd.owned_cstr_indices) { + i_t cstr_len = A_row_offsets[owned_cstr + 1] - A_row_offsets[owned_cstr]; + i_t row_start = A_row_offsets[owned_cstr]; + local_A_col_indices.insert(local_A_col_indices.end(), + A_col_indices.begin() + row_start, + A_col_indices.begin() + row_start + cstr_len); + local_A_values.insert(local_A_values.end(), + A_values.begin() + row_start, + A_values.begin() + row_start + cstr_len); + local_A_nnz += cstr_len; + local_A_row_offsets.push_back(local_A_nnz); + } + + // Compute halo + std::vector> needed_var_from_peer(nb_parts); + std::unordered_set seen_needed_vars; + + // size / 2 + 1 is a heuristic to avoid overestimating and resizing + seen_needed_vars.reserve(local_A_col_indices.size() / 2 + 1); + for (auto indice : local_A_col_indices) { + const i_t owner = var_owner(indice); + if (owner != rank && seen_needed_vars.insert(indice).second) { + needed_var_from_peer[owner].push_back(indice); + } + } + + // Compute counts and offsets of halo data to stack them all at the end of the vector + for (i_t peer = 0; peer < nb_parts; peer++) { + i_t nb_recv_from_peer = needed_var_from_peer[peer].size(); + rd.var_recv_counts[peer] = nb_recv_from_peer; + rd.var_recv_offsets[peer] = + peer == 0 ? 0 : rd.var_recv_offsets[peer - 1] + rd.var_recv_counts[peer - 1]; + rank_data[peer].var_send_per_peer[rank] = std::move(needed_var_from_peer[peer]); + } + + rd.h_A_row_offsets = std::move(local_A_row_offsets); + rd.h_A_col_indices = std::move(local_A_col_indices); + rd.h_A_values = std::move(local_A_values); + + // ---- A_t side ---- + // conceptually same as A side + std::vector local_A_t_row_offsets; + std::vector local_A_t_col_indices; + std::vector local_A_t_values; + local_A_t_row_offsets.reserve(static_cast(rd.owned_var_size) + 1); + local_A_t_col_indices.reserve(static_cast(owned_A_t_nnz[rank])); + local_A_t_values.reserve(static_cast(owned_A_t_nnz[rank])); + i_t local_A_t_nnz = 0; + local_A_t_row_offsets.push_back(local_A_t_nnz); + + for (auto owned_var : rd.owned_var_indices) { + i_t var_len = A_t_row_offsets[owned_var + 1] - A_t_row_offsets[owned_var]; + i_t row_start = A_t_row_offsets[owned_var]; + local_A_t_col_indices.insert(local_A_t_col_indices.end(), + A_t_col_indices.begin() + row_start, + A_t_col_indices.begin() + row_start + var_len); + local_A_t_values.insert(local_A_t_values.end(), + A_t_values.begin() + row_start, + A_t_values.begin() + row_start + var_len); + local_A_t_nnz += var_len; + local_A_t_row_offsets.push_back(local_A_t_nnz); + } + + std::vector> needed_cstr_from_peer(nb_parts); + std::unordered_set seen_needed_cstrs; + seen_needed_cstrs.reserve(local_A_t_col_indices.size() / 2 + 1); + for (auto indice : local_A_t_col_indices) { + const i_t owner = cstr_owner(indice); + if (owner != rank && seen_needed_cstrs.insert(indice).second) { + needed_cstr_from_peer[owner].push_back(indice); + } + } + + for (i_t peer = 0; peer < nb_parts; peer++) { + i_t nb_recv_from_peer = needed_cstr_from_peer[peer].size(); + rd.cstr_recv_counts[peer] = nb_recv_from_peer; + rd.cstr_recv_offsets[peer] = + peer == 0 ? 0 : rd.cstr_recv_offsets[peer - 1] + rd.cstr_recv_counts[peer - 1]; + rank_data[peer].cstr_send_per_peer[rank] = std::move(needed_cstr_from_peer[peer]); + } + + rd.h_A_t_row_offsets = std::move(local_A_t_row_offsets); + rd.h_A_t_col_indices = std::move(local_A_t_col_indices); + rd.h_A_t_values = std::move(local_A_t_values); + + rd.total_var_size = rd.owned_var_size + static_cast(seen_needed_vars.size()); + rd.total_cstr_size = rd.owned_cstr_size + static_cast(seen_needed_cstrs.size()); + + // Pad row-offset arrays so cuSPARSE sees the local matrices as + // (total_cstr x total_var) for A and (total_var x total_cstr) for A_T + const i_t a_last_nnz = rd.h_A_row_offsets.empty() ? i_t{0} : rd.h_A_row_offsets.back(); + rd.h_A_row_offsets.resize(rd.total_cstr_size + 1, a_last_nnz); + + const i_t at_last_nnz = rd.h_A_t_row_offsets.empty() ? i_t{0} : rd.h_A_t_row_offsets.back(); + rd.h_A_t_row_offsets.resize(rd.total_var_size + 1, at_last_nnz); + } + + // 3. Generate local indices for contiguous [[self], [peer1], ..., [peer_k]] + // Build scatter_gather_maps +#pragma omp parallel for + for (i_t rank = 0; rank < nb_parts; rank++) { + auto& rd = rank_data[rank]; + rd.global_to_local_cstr.reserve(static_cast(rd.total_cstr_size)); + rd.global_to_local_var.reserve(static_cast(rd.total_var_size)); + rd.local_to_global_cstr.reserve(static_cast(rd.total_cstr_size)); + rd.local_to_global_var.reserve(static_cast(rd.total_var_size)); + + i_t curr_id = 0; + for (auto owned_cstr : rd.owned_cstr_indices) { + rd.global_to_local_cstr[owned_cstr] = curr_id; + rd.local_to_global_cstr.push_back(owned_cstr); + curr_id++; + } + for (i_t peer = 0; peer < nb_parts; peer++) { + if (peer == rank) continue; + for (auto recv_cstr : rank_data[peer].cstr_send_per_peer[rank]) { + rd.global_to_local_cstr[recv_cstr] = curr_id; + rd.local_to_global_cstr.push_back(recv_cstr); + curr_id++; + } + } + + curr_id = 0; + for (auto owned_var : rd.owned_var_indices) { + rd.global_to_local_var[owned_var] = curr_id; + rd.local_to_global_var.push_back(owned_var); + curr_id++; + } + for (i_t peer = 0; peer < nb_parts; peer++) { + if (peer == rank) continue; + for (auto recv_var : rank_data[peer].var_send_per_peer[rank]) { + rd.global_to_local_var[recv_var] = curr_id; + rd.local_to_global_var.push_back(recv_var); + curr_id++; + } + } + } + + // 4. Remap global -> local everywhere + // Including in A_local and At_local +#pragma omp parallel for + for (i_t rank = 0; rank < nb_parts; rank++) { + auto& rd = rank_data[rank]; + + for (auto& send_vec : rd.var_send_per_peer) { + for (auto& v : send_vec) + v = rd.global_to_local_var.at(v); + } + for (auto& send_vec : rd.cstr_send_per_peer) { + for (auto& v : send_vec) + v = rd.global_to_local_cstr.at(v); + } + + for (auto& v : rd.h_A_col_indices) + v = rd.global_to_local_var.at(v); + for (auto& v : rd.h_A_t_col_indices) + v = rd.global_to_local_cstr.at(v); + } + + return rank_data; +} + +template std::vector> create_rank_data_from_parts( + const std::vector& parts, + const std::vector& A_row_offsets, + const std::vector& A_col_indices, + const std::vector& A_values, + const std::vector& A_t_row_offsets, + const std::vector& A_t_col_indices, + const std::vector& A_t_values, + int nb_parts, + int nb_cstr, + int nb_vars, + int nnz); + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/distributed_utils.hpp b/cpp/src/pdlp/distributed_pdlp/distributed_utils.hpp new file mode 100644 index 0000000000..a96a5178f9 --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/distributed_utils.hpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// Slices the data to prepare a split from graph partitioning with halo communication. +template +std::vector> create_rank_data_from_parts( + const std::vector& parts, + const std::vector& A_row_offsets, + const std::vector& A_col_indices, + const std::vector& A_values, + const std::vector& A_t_row_offsets, + const std::vector& A_t_col_indices, + const std::vector& A_t_values, + i_t nb_parts, + i_t nb_cstr, + i_t nb_vars, + i_t nnz); + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.cu b/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.cu new file mode 100644 index 0000000000..f651c1cea9 --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.cu @@ -0,0 +1,426 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include + +#include + +#include + +#include + +namespace cuopt::mathematical_optimization::pdlp { + +template +multi_gpu_engine_t::multi_gpu_engine_t( + std::vector>&& rank_data, + io::mps_data_model_t const& mps, + pdlp_solver_settings_t const& sub_solver_settings) + : stream() +{ + const int nb_parts = static_cast(rank_data.size()); + cuopt_expects( + nb_parts > 0, error_type_t::ValidationError, "multi_gpu_engine_t: rank_data must be non-empty"); + + shards.reserve(nb_parts); + std::vector devices(nb_parts); + std::iota(devices.begin(), devices.end(), 0); + + // Create NCCL Comms, then immediately wrap each in a RAII owner so they are + // destroyed on any exception (e.g. a shard ctor throwing) before being + // handed off to a shard. + std::vector comms; + comms.reserve(nb_parts); + std::vector raw_comms(nb_parts, nullptr); + cuopt_expects(ncclCommInitAll(raw_comms.data(), nb_parts, devices.data()) == ncclSuccess, + error_type_t::RuntimeError, + "ncclCommInitAll failed"); + + for (int r = 0; r < nb_parts; ++r) { + comms.emplace_back(raw_comms[r], nccl_comm_deleter_t{devices[r]}); + } + + // 3. Construct one shard per rank, pinned to its device. Ownership of each + // communicator moves into its shard. + for (int r = 0; r < nb_parts; ++r) { + raft::device_setter guard(devices[r]); // shard ctor needs device set + shards.emplace_back(std::make_unique>( + devices[r], std::move(rank_data[r]), std::move(comms[r]), mps, sub_solver_settings)); + } + + // Two different events + // capture_*_event_ are used inside graph capture + // ext_*_event_ are used when sync is needed outside of graph + graph_master_ready_event_ = std::make_unique(); + sync_master_ready_event_ = std::make_unique(); + graph_shard_ready_events_.reserve(nb_parts); + sync_shard_ready_events_.reserve(nb_parts); + for_each_shard([&](auto&) { + graph_shard_ready_events_.emplace_back(std::make_unique()); + sync_shard_ready_events_.emplace_back(std::make_unique()); + }); + + // Cache per-shard partition metadata for gather_owned_*_to_master_bufs. + local_to_global_vars_.reserve(nb_parts); + local_to_global_cstrs_.reserve(nb_parts); + for (int r = 0; r < nb_parts; ++r) { + auto const& rd = shards[r]->rank_data; + local_to_global_vars_.push_back(rd.local_to_global_var); + local_to_global_cstrs_.push_back(rd.local_to_global_cstr); + } +} + +// -------- High-level: A @ x and A_T @ y ----------------------------------- +template +void multi_gpu_engine_t::distributed_compute_A_x() +{ + halo_exchange_var([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_reflected_primal(); + }); + for_each_shard([](auto& shard) { shard.sub_pdlp->pdhg_solver_.spmvop_A_x(); }); +} + +template +void multi_gpu_engine_t::distributed_compute_At_y() +{ + halo_exchange_cstr([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_dual_solution(); + }); + for_each_shard([](auto& shard) { shard.sub_pdlp->pdhg_solver_.spmvop_At_y(); }); +} + +// -------- Cross-stream fork / join / sync --------------------------------- +template +void multi_gpu_engine_t::graph_capture_fork_to_shards(rmm::cuda_stream_view master_stream) +{ + graph_master_ready_event_->record(master_stream); + for_each_shard([&](auto& s) { graph_master_ready_event_->stream_wait(s.stream.view()); }); +} + +template +void multi_gpu_engine_t::graph_capture_join_from_shards( + rmm::cuda_stream_view master_stream) +{ + for_each_shard([&](auto& s, int r) { graph_shard_ready_events_[r]->record(s.stream.view()); }); + for (auto& e : graph_shard_ready_events_) { + e->stream_wait(master_stream); + } +} + +template +void multi_gpu_engine_t::sync_await_master(rmm::cuda_stream_view master_stream) +{ + sync_master_ready_event_->record(master_stream); + for_each_shard([&](auto& s) { sync_master_ready_event_->stream_wait(s.stream.view()); }); +} + +template +void multi_gpu_engine_t::sync_await_shards(rmm::cuda_stream_view master_stream) +{ + for_each_shard([&](auto& s, int r) { sync_shard_ready_events_[r]->record(s.stream.view()); }); + for (auto& e : sync_shard_ready_events_) { + e->stream_wait(master_stream); + } +} + +// -------- Halo exchange ------------ +// typename pdlp_shard_t::halo_axis_t is the unified view of the halo exchange metadata +// for both variables and constraints. It contains the send and receive buffers for each peer, the +// owned size, and the receive offsets and counts for each peer. There is one for variables and one +// for constraints, on each shard. This allows to avoid duplicating the logic for variables and +// constraints. +template +void multi_gpu_engine_t::halo_exchange_bufs_impl( + std::vector> const& bufs, + std::vector::halo_axis_t> const& axes) +{ + const int nb = static_cast(shards.size()); + cuopt_expects(static_cast(bufs.size()) == nb && static_cast(axes.size()) == nb, + error_type_t::RuntimeError, + "halo_exchange_bufs_impl: bufs / axes must have size == shards.size()"); + + // Step 1: gather owned values that each peer needs into per-peer staging. + for_each_shard([&](auto& s, int r) { + auto const& ax = axes[r]; + auto x = bufs[r]; + for (int peer = 0; peer < nb; ++peer) { + if (peer == r) continue; + if (ax.send_indices[peer].size() == 0) continue; + thrust::gather(rmm::exec_policy_nosync(s.stream.view()), + ax.send_indices[peer].begin(), + ax.send_indices[peer].end(), + x.data(), + ax.send_buf[peer].begin()); + } + }); + + // Step 2: matched send / recv across the whole topology in one NCCL group. + CUOPT_NCCL_TRY(ncclGroupStart()); + for_each_shard([&](auto& s, int r) { + auto const& ax = axes[r]; + for (int peer = 0; peer < nb; ++peer) { + if (peer == r) continue; + CUOPT_NCCL_TRY(ncclSend(ax.send_buf[peer].data(), + ax.send_buf[peer].size(), + nccl_data_type(), + peer, + s.comm.get(), + s.stream.view().value())); + } + }); + for_each_shard([&](auto& s, int r) { + auto const& ax = axes[r]; + auto x = bufs[r]; + for (int peer = 0; peer < nb; ++peer) { + if (peer == r) continue; + f_t* recv_ptr = x.data() + ax.owned_size + ax.recv_offsets[peer]; + CUOPT_NCCL_TRY(ncclRecv(recv_ptr, + static_cast(ax.recv_counts[peer]), + nccl_data_type(), + peer, + s.comm.get(), + s.stream.view().value())); + } + }); + CUOPT_NCCL_TRY(ncclGroupEnd()); +} + +template +void multi_gpu_engine_t::halo_exchange_var_bufs( + std::vector> const& bufs) +{ + std::vector::halo_axis_t> axes; + axes.reserve(shards.size()); + for (auto& s : shards) + axes.push_back(s->var_halo_axis()); + halo_exchange_bufs_impl(bufs, axes); +} + +template +void multi_gpu_engine_t::halo_exchange_var_bufs( + std::vector>& bufs) +{ + std::vector> spans; + spans.reserve(bufs.size()); + for (auto& b : bufs) + spans.emplace_back(b.data(), b.size()); + halo_exchange_var_bufs(spans); +} + +template +void multi_gpu_engine_t::halo_exchange_cstr_bufs( + std::vector> const& bufs) +{ + std::vector::halo_axis_t> axes; + axes.reserve(shards.size()); + for (auto& s : shards) + axes.push_back(s->cstr_halo_axis()); + halo_exchange_bufs_impl(bufs, axes); +} + +template +void multi_gpu_engine_t::halo_exchange_cstr_bufs( + std::vector>& bufs) +{ + std::vector> spans; + spans.reserve(bufs.size()); + for (auto& b : bufs) + spans.emplace_back(b.data(), b.size()); + halo_exchange_cstr_bufs(spans); +} + +// -------- Gather owned slices to master ------------------------------------- +template +void multi_gpu_engine_t::gather_owned_to_master_bufs_impl( + std::vector> const& shard_owned, + raft::device_span master_buf, + std::vector> const& local_to_globals) +{ + cuopt_assert(master_pdlp_ != nullptr, + "gather_owned_to_master_bufs_impl requires set_master(...)"); + const int nb = static_cast(shards.size()); + cuopt_expects( + static_cast(shard_owned.size()) == nb && static_cast(local_to_globals.size()) == nb, + error_type_t::RuntimeError, + "gather_owned_to_master_bufs_impl: shard_owned / local_to_globals " + "must have size == shards.size()"); + + // Assemble on host in global-index order. + std::vector h_master(master_buf.size()); + for_each_shard([&](auto& s, int r) { + const std::size_t n_owned = shard_owned[r].size(); + if (n_owned == 0) return; + std::vector tmp(n_owned); + raft::copy(tmp.data(), shard_owned[r].data(), n_owned, s.stream.view()); + // Sync so tmp is populated before the host scatter (and stays valid). + s.stream.synchronize(); + thrust::scatter( + thrust::host, tmp.begin(), tmp.end(), local_to_globals[r].begin(), h_master.begin()); + }); + + // Single H->D onto master's device (`stream` lives on the master device). + raft::copy(master_buf.data(), h_master.data(), master_buf.size(), stream.view()); + stream.synchronize(); +} + +template +void multi_gpu_engine_t::gather_owned_var_to_master_bufs( + std::vector> const& shard_owned, raft::device_span master_buf) +{ + gather_owned_to_master_bufs_impl(shard_owned, master_buf, local_to_global_vars_); +} + +template +void multi_gpu_engine_t::gather_owned_cstr_to_master_bufs( + std::vector> const& shard_owned, raft::device_span master_buf) +{ + gather_owned_to_master_bufs_impl(shard_owned, master_buf, local_to_global_cstrs_); +} + +// -------- NCCL allreduce (sum, in place) ------------------------------------ +template +void multi_gpu_engine_t::allreduce_sum_inplace_bufs( + std::vector> const& scalars) +{ + const int nb = static_cast(shards.size()); + cuopt_expects(static_cast(scalars.size()) == nb, + error_type_t::RuntimeError, + "allreduce_sum_inplace_bufs: scalars.size() must equal shards.size()"); + if (nb == 0) return; + + CUOPT_NCCL_TRY(ncclGroupStart()); + for_each_shard([&](auto& s, int r) { + f_t* p = scalars[r].data_handle(); + CUOPT_NCCL_TRY(ncclAllReduce(p, + p, + /*count=*/1, + nccl_data_type(), + ncclSum, + s.comm.get(), + s.stream.view().value())); + }); + CUOPT_NCCL_TRY(ncclGroupEnd()); +} + +template +void multi_gpu_engine_t::allreduce_sum_inplace_to_master_buf( + std::vector> const& shard_scalars, + raft::device_scalar_view master_dst) +{ + cuopt_assert(master_pdlp_ != nullptr, + "allreduce_sum_inplace_to_master_buf requires set_master(...)"); + allreduce_sum_inplace_bufs(shard_scalars); + if (shards.empty()) return; + auto master_stream = master_pdlp_->get_handle_ptr()->get_stream(); + sync_await_shards(master_stream); + auto& s0 = *shards[0]; + raft::device_setter guard(s0.device_id); + raft::copy(master_dst.data_handle(), shard_scalars[0].data_handle(), 1, master_stream); +} + +// -------- Distributed dot / L2 norm ----------------------------------------- +template +void multi_gpu_engine_t::distributed_dot_bufs( + std::vector> const& a_bufs, + std::vector> const& b_bufs, + std::vector> const& out_scalars) +{ + const int nb = static_cast(shards.size()); + cuopt_expects(static_cast(a_bufs.size()) == nb && static_cast(b_bufs.size()) == nb && + static_cast(out_scalars.size()) == nb, + error_type_t::RuntimeError, + "distributed_dot_bufs: a_bufs / b_bufs / out_scalars must " + "all have size == shards.size()"); + + for_each_shard([&](auto& s, int r) { + cuopt_expects(a_bufs[r].size() == b_bufs[r].size(), + error_type_t::RuntimeError, + "distributed_dot_bufs: a_bufs[r] and b_bufs[r] must have equal size"); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(s.handle.get_cublas_handle(), + static_cast(a_bufs[r].size()), + a_bufs[r].data(), + 1, + b_bufs[r].data(), + 1, + out_scalars[r].data_handle(), + s.stream.view().value())); + }); + + allreduce_sum_inplace_bufs(out_scalars); +} + +template +void multi_gpu_engine_t::distributed_dot_bufs( + std::vector> const& a_bufs, + std::vector> const& b_bufs, + std::vector>& out_scalars) +{ + std::vector> views; + views.reserve(out_scalars.size()); + for (auto& s : out_scalars) + views.emplace_back(raft::make_device_scalar_view(s.data())); + distributed_dot_bufs(a_bufs, b_bufs, views); +} + +template +void multi_gpu_engine_t::distributed_l2_norm_bufs( + std::vector> const& in_bufs, + std::vector> const& out_scalars) +{ + distributed_dot_bufs(in_bufs, in_bufs, out_scalars); + for_each_shard([&](pdlp_shard_t& s, int r) { + cub::DeviceTransform::Transform( + out_scalars[r].data_handle(), + out_scalars[r].data_handle(), + 1, + [] __device__(f_t x) { return cuda::std::sqrt(x); }, + s.stream.view().value()); + }); +} + +template +void multi_gpu_engine_t::distributed_l2_norm_bufs( + std::vector> const& in_bufs, + std::vector>& out_scalars) +{ + std::vector> views; + views.reserve(out_scalars.size()); + for (auto& s : out_scalars) + views.emplace_back(raft::make_device_scalar_view(s.data())); + distributed_l2_norm_bufs(in_bufs, views); +} + +// -------- Fused halo-exchange + SpMV ---------------------------------------- +template +void multi_gpu_engine_t::distributed_spmv_At( + std::vector>& in_bufs, + std::vector>& in_descs, + std::vector>& out_descs) +{ + halo_exchange_cstr_bufs(in_bufs); + for_each_shard( + [&](auto& s, int r) { s.sub_pdlp->pdhg_solver_.spmv_At_into(in_descs[r], out_descs[r]); }); +} + +template +void multi_gpu_engine_t::distributed_spmv_A( + std::vector>& in_bufs, + std::vector>& in_descs, + std::vector>& out_descs) +{ + halo_exchange_var_bufs(in_bufs); + for_each_shard( + [&](auto& s, int r) { s.sub_pdlp->pdhg_solver_.spmv_A_into(in_descs[r], out_descs[r]); }); +} + +template struct multi_gpu_engine_t; +template struct multi_gpu_engine_t; + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp b/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp new file mode 100644 index 0000000000..5637576eee --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp @@ -0,0 +1,498 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// Maps the solver floating-point type to the matching NCCL datatype so that +// halo exchanges / all-reduces transfer the correct element size for both +// double and float instantiations. +template +constexpr ncclDataType_t nccl_data_type() +{ + static_assert(std::is_same_v || std::is_same_v, + "Unsupported floating-point type for NCCL transfers"); + if constexpr (std::is_same_v) { + return ncclFloat64; + } else { + return ncclFloat32; + } +} + +template +struct multi_gpu_engine_t { + // Constructs shards from rank_data. The global (unpartitioned) problem is + // read straight from `mps`; each shard slices out the entries it owns. + multi_gpu_engine_t(std::vector>&& rank_data, + io::mps_data_model_t const& mps, + pdlp_solver_settings_t const& sub_solver_settings); + + multi_gpu_engine_t(const multi_gpu_engine_t&) = delete; + multi_gpu_engine_t& operator=(const multi_gpu_engine_t&) = delete; + + // Invokes `fn` on every shard with the shard's device pre-set. `fn` may be + // (pdlp_shard_t&) or + // (pdlp_shard_t&, int r) — the second overload also gets the shard's rank. + template + void for_each_shard(Fn&& fn) + { + for (int r = 0; r < static_cast(shards.size()); ++r) { + auto& s = *shards[r]; + raft::device_setter guard(s.device_id); + // If the function is invocable with a pdlp_shard_t& and an int, call it with the + // shard and the rank. + if constexpr (std::is_invocable_v&, int>) { + fn(s, r); + // If the function is invocable only with a pdlp_shard_t&, call it with the shard. + } else if constexpr (std::is_invocable_v&>) { + fn(s); + } else { + // Otherwise, the function has an invalid signature. + cuopt_expects( + false, error_type_t::RuntimeError, "for_each_shard: invalid function signature"); + } + } + } + + // Host-blocking barrier: waits until every shard stream has drained. + void synchronize_shards() + { + for_each_shard([](auto& s) { s.stream.synchronize(); }); + } + + // Core: launches cub::DeviceTransform on every shard using per-shard + // pre-resolved inputs / outputs / sizes. + // - in_tuples[r] is the tuple passed as cub input for shard r (any + // iterator-shaped types cub accepts: raw pointers, thrust iterators, ...) + // - outs[r] is the output iterator for shard r + // - sizes[r] is the element count for shard r + // All three must have size == shards.size(). + template + void distributed_transform_bufs(std::vector const& in_tuples, + std::vector const& outs, + std::vector const& sizes, + Op op) + { + const int nb = static_cast(shards.size()); + cuopt_expects(static_cast(in_tuples.size()) == nb && static_cast(outs.size()) == nb && + static_cast(sizes.size()) == nb, + error_type_t::RuntimeError, + "distributed_transform_bufs: in_tuples / outs / sizes must " + "all have size == shards.size()"); + for_each_shard([&](auto& s, int r) { + cub::DeviceTransform::Transform(in_tuples[r], outs[r], sizes[r], op, s.stream.view()); + }); + } + + // Wrapper: accessor form. Resolves each shard's cub input_tuple / output / + // size via the provided accessors, then delegates to + // distributed_transform_bufs. + template + void distributed_transform(std::tuple in_accessors, + OutAccess out, + SizeAccess sz, + Op op) + { + cuopt_expects( + !shards.empty(), error_type_t::RuntimeError, "distributed_transform: engine has no shards"); + + // Deduce per-shard tuple / output types from the accessors themselves so + // the runtime vector doesn't complain + auto& sample_sub = *shards[0]->sub_pdlp; + using in_tuple_t = decltype(std::apply( + [&sample_sub](auto&... acc) { return cuda::std::make_tuple(acc(sample_sub)...); }, + in_accessors)); + using out_iter_t = decltype(out(sample_sub)); + + std::vector in_tuples; + std::vector outs; + std::vector sizes; + in_tuples.reserve(shards.size()); + outs.reserve(shards.size()); + sizes.reserve(shards.size()); + + for_each_shard([&](auto& s) { + auto& sub = *s.sub_pdlp; + // apply() = Turns a tuple of accessors into a tuple of values. + in_tuples.emplace_back(std::apply( + [&sub](auto&... acc) { return cuda::std::make_tuple(acc(sub)...); }, in_accessors)); + outs.emplace_back(out(sub)); + sizes.emplace_back(sz(sub)); + }); + distributed_transform_bufs(in_tuples, outs, sizes, op); + } + + // --- 2) convenience: single input accessor (delegates) --- + // Allows to use distributed_transform on single input without having to do a std::make_tuple(in) + template + void distributed_transform(InAccess in, OutAccess out, SizeAccess sz, Op op) + { + distributed_transform(std::make_tuple(in), out, sz, op); + } + + // -------- Halo exchange (owner -> halo) --------------------------------- + // {var/cstr}-agnostic core: + // Step 1: thrust::gather per-peer outgoing values into staging buffers. + // Step 2: one NCCL group with matched ncclSend / ncclRecv across all + // (rank, peer) pairs, receiving into each shard's halo tail. + void halo_exchange_bufs_impl( + std::vector> const& bufs, + std::vector::halo_axis_t> const& axes); + + // -------- Halo exchange (variables / x) --------------------------------- + void halo_exchange_var_bufs(std::vector> const& bufs); + + // Overload: accept the owning device_uvector directly (rmm doesn't provide + // an implicit conversion to raft::device_span) + void halo_exchange_var_bufs(std::vector>& bufs); + + // Wrapper: pdlp_solver_t accessor. Resolves one uvector per shard into a + // vector of spans, then delegates to halo_exchange_var_bufs. + // buf_access : pdlp_solver_t& -> rmm::device_uvector& + template + void halo_exchange_var(BufAccess&& buf_access) + { + std::vector> bufs; + bufs.reserve(shards.size()); + for_each_shard([&](pdlp_shard_t& s) { + auto& x = buf_access(*s.sub_pdlp); + bufs.emplace_back(x.data(), x.size()); + }); + halo_exchange_var_bufs(bufs); + } + + // -------- Halo exchange (constraints / y) ------------------------------- + void halo_exchange_cstr_bufs(std::vector> const& bufs); + + // Overload: same rationale as halo_exchange_var_bufs above. + void halo_exchange_cstr_bufs(std::vector>& bufs); + + // Wrapper: pdlp_solver_t accessor. Resolves one uvector per shard into a + // vector of spans, then delegates to halo_exchange_cstr_bufs. + // buf_access : pdlp_solver_t& -> rmm::device_uvector& + template + void halo_exchange_cstr(BufAccess&& buf_access) + { + std::vector> bufs; + bufs.reserve(shards.size()); + for_each_shard([&](pdlp_shard_t& s) { + auto& y = buf_access(*s.sub_pdlp); + bufs.emplace_back(y.data(), y.size()); + }); + halo_exchange_cstr_bufs(bufs); + } + + // -------- Gather owned slices to master --------------------------------- + // {var/cstr}-agnostic core: scatters each shard's owned slice into + // master_buf using local_to_globals[r] as the destination index list. + // local_to_globals[r] is + // the axis-specific rank_data.local_to_global_{var/cstr} for shard r. + void gather_owned_to_master_bufs_impl( + std::vector> const& shard_owned, + raft::device_span master_buf, + std::vector> const& local_to_globals); + + // -------- Gather (variables / x) ---------------------------------------- + void gather_owned_var_to_master_bufs(std::vector> const& shard_owned, + raft::device_span master_buf); + + // Wrapper: pdlp_solver_t accessor. Slices each shard's owned prefix and + // then delegates to gather_owned_var_to_master_bufs. + // buf_access : pdlp_solver_t& -> rmm::device_uvector& + template + void gather_owned_var_to_master(BufAccess&& buf_access) + { + cuopt_assert(master_pdlp_ != nullptr, "gather_owned_var_to_master requires set_master(...)"); + std::vector> shard_bufs; + shard_bufs.reserve(shards.size()); + for_each_shard([&](pdlp_shard_t& s) { + auto& x = buf_access(*s.sub_pdlp); + shard_bufs.emplace_back(x.data(), static_cast(s.rank_data.owned_var_size)); + }); + auto& m = buf_access(*master_pdlp_); + gather_owned_var_to_master_bufs(shard_bufs, raft::device_span{m.data(), m.size()}); + } + + // -------- Gather (constraints / y) -------------------------------------- + void gather_owned_cstr_to_master_bufs( + std::vector> const& shard_owned, + raft::device_span master_buf); + + // Wrapper: same rationale as gather_owned_var_to_master. + template + void gather_owned_cstr_to_master(BufAccess&& buf_access) + { + cuopt_assert(master_pdlp_ != nullptr, "gather_owned_cstr_to_master requires set_master(...)"); + std::vector> shard_bufs; + shard_bufs.reserve(shards.size()); + for_each_shard([&](pdlp_shard_t& s) { + auto& x = buf_access(*s.sub_pdlp); + shard_bufs.emplace_back(x.data(), static_cast(s.rank_data.owned_cstr_size)); + }); + auto& m = buf_access(*master_pdlp_); + gather_owned_cstr_to_master_bufs(shard_bufs, raft::device_span{m.data(), m.size()}); + } + + // -------- NCCL allreduce (sum, in place) -------------------------------- + // Core: per-shard in-place sum-allreduce on a single f_t scalar viewed by + // scalars[r], wrapped in one NCCL group so it executes as a single + // collective. After this returns, every shard's scalar holds the global sum. + void allreduce_sum_inplace_bufs(std::vector> const& scalars); + + // Wrapper: pdlp_solver_t accessor for a single per-shard scalar. + // ptr_access : pdlp_solver_t& -> f_t* (pointer to the scalar + // to reduce; one per shard) + template + void allreduce_sum_inplace(PtrAccess&& ptr_access) + { + std::vector> scalars; + scalars.reserve(shards.size()); + for_each_shard([&](auto& s) { + scalars.emplace_back(raft::make_device_scalar_view(ptr_access(*s.sub_pdlp))); + }); + allreduce_sum_inplace_bufs(scalars); + } + + // Core: same as allreduce_sum_inplace_bufs, plus after the collective the + // value is D2D-copied from shard 0 into master_dst. On master's stream. + void allreduce_sum_inplace_to_master_buf( + std::vector> const& shard_scalars, + raft::device_scalar_view master_dst); + + // Wrapper: applies the ptr_access lambda to each shard's sub_pdlp to build + // the per-shard scalar views and to master_pdlp_ to obtain the master + // destination, then delegates to allreduce_sum_inplace_to_master_buf. + template + void allreduce_sum_inplace_to_master(PtrAccess&& ptr_access) + { + cuopt_assert(master_pdlp_ != nullptr, + "allreduce_sum_inplace_to_master requires set_master(...) to have been called"); + std::vector> shard_scalars; + shard_scalars.reserve(shards.size()); + for_each_shard([&](auto& s) { + shard_scalars.emplace_back(raft::make_device_scalar_view(ptr_access(*s.sub_pdlp))); + }); + allreduce_sum_inplace_to_master_buf( + shard_scalars, raft::make_device_scalar_view(ptr_access(*master_pdlp_))); + } + + // -------- Set a host scalar on master AND every shard ------------------- + // Writes a host-side value to master and every shard through `ptr_access`. + template + void set_scalar_on_master_and_shards(f_t value, PtrAccess&& ptr_access) + { + cuopt_assert(master_pdlp_ != nullptr, + "set_scalar_on_master_and_shards requires set_master(...) to have been called"); + auto master_stream = master_pdlp_->get_handle_ptr()->get_stream(); + raft::copy(ptr_access(*master_pdlp_), &value, 1, master_stream); + for_each_shard([&](auto& shard) { + raft::copy(ptr_access(*shard.sub_pdlp), &value, 1, shard.stream.view()); + }); + master_pdlp_->get_handle_ptr()->sync_stream(master_stream); + synchronize_shards(); + } + + // -------- Distributed dot / L2 norm ------------------------------------- + // Computes the dot product of two vectors for each shard. Returns the global result in + // out_scalars. + void distributed_dot_bufs(std::vector> const& a_bufs, + std::vector> const& b_bufs, + std::vector> const& out_scalars); + + // Overload: accept owning device_scalar outputs directly. Wraps each into a + // scalar_view and delegates to the span-based core. Convenience for the typical + // case where per-shard outputs are rmm::device_scalar. + void distributed_dot_bufs(std::vector> const& a_bufs, + std::vector> const& b_bufs, + std::vector>& out_scalars); + + // Core L2 norm: writes sqrt(sum_r ||in_bufs[r]||_2^2) into every + // *out_scalars[r].data_handle(). Delegates to distributed_dot_bufs(in, in, + // out) then does a per-shard in-place sqrt on the resulting scalar. + void distributed_l2_norm_bufs(std::vector> const& in_bufs, + std::vector> const& out_scalars); + + // Overload: same rationale as distributed_dot_bufs above. Allows to use rmm::device_scalar + // directly. + void distributed_l2_norm_bufs(std::vector> const& in_bufs, + std::vector>& out_scalars); + + // Wrapper: accessor form. Resolves per-shard input / output / owned-length + // then delegates to distributed_l2_norm_bufs. + // BufAccess : pdlp_solver_t& -> rmm::device_uvector& + // OutAccess : pdlp_solver_t& -> f_t* (single scalar) + // SizeAccess : pdlp_shard_t& -> i_t (owned slice length) + template + void distributed_l2_norm(BufAccess&& buf_access, OutAccess&& out_access, SizeAccess&& size_access) + { + std::vector> in_bufs; + std::vector> out_scalars; + in_bufs.reserve(shards.size()); + out_scalars.reserve(shards.size()); + for_each_shard([&](auto& s) { + auto& sub = *s.sub_pdlp; + auto& buf = buf_access(sub); + const i_t n = size_access(s); + in_bufs.emplace_back(buf.data(), static_cast(n)); + out_scalars.emplace_back(raft::make_device_scalar_view(out_access(sub))); + }); + distributed_l2_norm_bufs(in_bufs, out_scalars); + } + + // -------- High-level: A @ x and A_T @ y --------------------------------- + // Distributed counterpart to pdhg_solver_t::compute_A_x() / compute_At_y(). + void distributed_compute_A_x(); + void distributed_compute_At_y(); + + // Distributed A^T @ in on caller-owned scratch. Refreshes the halo of `in_bufs` + // (cstr axis, since the input is cstr-shaped), then dispatches each shard's + // local spmv_At_into that reads from in_descs[r] and writes into out_descs[r]. + void distributed_spmv_At(std::vector>& in_bufs, + std::vector>& in_descs, + std::vector>& out_descs); + + // Distributed A @ in on caller-owned scratch. Refreshes the halo of `in_bufs` + // (var axis, since the input is var-shaped), then dispatches each shard's + // local spmv_A_into. Caller owns / sizes the descriptor vectors as above + // (in_descs to var_total, out_descs to cstr_total). + void distributed_spmv_A(std::vector>& in_bufs, + std::vector>& in_descs, + std::vector>& out_descs); + + // -------- High-level algorithms (defined in distributed_algorithms.cu) --- + // Refreshes the halo copies of the cumulative variable + constraint scalings on + // every shard. Used by the matrix-scaling passes (Ruiz, Pock-Chambolle) + void refresh_halo_cummulative_scalings(); + + // Global bound/objective rescaling: allreduce the owned partial squared norms + // of the constraint bounds and (weighted) objective, then apply the identical + // scalar on every shard. + void distributed_bound_objective_rescaling(f_t c_scaling_weight); + + // Distributed Ruiz inf-scaling (num_iter passes). Each shard computes both its + // owned-row and owned-column inf-norms locally then broadcasts the cumulative scalings to all + // shards. + void distributed_ruiz_inf_scaling(int num_iter, i_t n_global_vars); + + // Distributed Pock-Chambolle scaling (one pass), mirroring the single-GPU + // pock_chambolle_scaling. + void distributed_pock_chambolle_scaling(f_t alpha, i_t n_global_vars); + + // Full distributed scaling entry point. Mirrors what scale_problem() does in + // single-GPU by orchestrating: + // - Ruiz inf-scaling -> populates cumulative row/col scalings + // - Pock-Chambolle scaling -> same + // - per-shard apply_cummulative_scaling_to_problem() + // - global bound/objective rescaling via distributed_bound_objective_rescaling + void distributed_scaling(pdlp_hyper_params_t const& hyper_params, + i_t n_global_vars, + bool inside_mip); + + // Distributed sigma_max(A)^2 via power iteration (used to seed the initial + // step size). Returns the square of the largest singular value of the scaled + // constraint matrix. + f_t distributed_max_singular_value_squared(i_t n_global_cstrs, + int max_iterations = 5000, + f_t tolerance = 1e-4); + + // Distributed counterpart of pdlp_solver_t::compute_initial_step_size. + void distributed_compute_initial_step_size(pdlp_hyper_params_t const& hyper_params, + i_t n_global_cstrs, + f_t scaling_factor, + int max_iterations, + f_t tolerance); + + // Distributed counterpart of pdlp_solver_t::compute_initial_primal_weight. + // Writes primal_weight = best_primal_weight = 1 onto master + every shard, + // mirroring the Stable3-shaped short-circuit + // (!initial_primal_weight_combined_bounds && bound_objective_rescaling). + void distributed_compute_initial_primal_weight(pdlp_hyper_params_t const& hyper_params); + + // Gather the global potential_next primal/dual solutions and the reduced cost + // onto the master from the owned slices distributed across shards. + void gather_potential_next_solutions_to_master(); + + // Engine-level stream for fork/join orchestration (master side). + rmm::cuda_stream stream; + + // Shards stored by unique_ptr because pdlp_shard_t is immovable + // (owns device-affine resources: handle, NCCL comm, RMM buffers). + std::vector>> shards; + + // Cached per-shard partition metadata, populated once at construction. + // Consumed by gather_owned_*_to_master_bufs; caching avoids copying the + // sizeable local_to_global_* host vectors on every termination check. + // local_to_global_{vars,cstrs}_[r] == shards[r]->rank_data.local_to_global_{var,cstr} + std::vector> local_to_global_vars_; + std::vector> local_to_global_cstrs_; + + // Non-owning back-pointer to the master pdlp_solver_t. + pdlp_solver_t* master_pdlp_ = nullptr; + void set_master(pdlp_solver_t* m) { master_pdlp_ = m; } + + // ===== Cross-stream synchronization events ===== + // two different events + // graph_*_event_ are used inside graph capture + // sync_*_event_ are used when sync is needed outside of graph + std::unique_ptr graph_master_ready_event_; + std::vector> graph_shard_ready_events_; + std::unique_ptr sync_master_ready_event_; + std::vector> sync_shard_ready_events_; + + // Forks master stream to shards, so that the captured graph can see the work on the shards + void graph_capture_fork_to_shards(rmm::cuda_stream_view master_stream); + + // Joins shards back to master stream for correct graph capture + void graph_capture_join_from_shards(rmm::cuda_stream_view master_stream); + + // Functionnaly same as graph_capture_fork_to_shards but on a different event to avoid race + // conditions Can be used as a way to sync shards with master stream + void sync_await_master(rmm::cuda_stream_view master_stream); + + // Same as sync_await_master + // Can be used as a way to sync master stream with shards + void sync_await_shards(rmm::cuda_stream_view master_stream); +}; + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/nccl_helpers.hpp b/cpp/src/pdlp/distributed_pdlp/nccl_helpers.hpp new file mode 100644 index 0000000000..48d7bd8b66 --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/nccl_helpers.hpp @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// Wraps a NCCL call and throws cuopt::logic_error (RuntimeError) on any non- +// ncclSuccess return code. +#define CUOPT_NCCL_TRY(call) \ + do { \ + ::ncclResult_t const _cuopt_nccl_status = (call); \ + ::cuopt::cuopt_expects(_cuopt_nccl_status == ncclSuccess, \ + ::cuopt::error_type_t::RuntimeError, \ + "NCCL error at %s:%d: %s", \ + __FILE__, \ + __LINE__, \ + ::ncclGetErrorString(_cuopt_nccl_status)); \ + } while (0) + +// Non-throwing variant: logs at ERROR level on failure and swallows the error. +// Intended for noexcept contexts (destructors, unique_ptr deleters, teardown) +// where throwing would call std::terminate. Mirrors RAFT_CUDA_TRY_NO_THROW. +#define CUOPT_NCCL_TRY_NO_THROW(call) \ + do { \ + ::ncclResult_t const _cuopt_nccl_status = (call); \ + if (_cuopt_nccl_status != ncclSuccess) { \ + CUOPT_LOG_ERROR( \ + "NCCL error at %s:%d: %s", __FILE__, __LINE__, ::ncclGetErrorString(_cuopt_nccl_status)); \ + } \ + } while (0) + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/partitioner.cpp b/cpp/src/pdlp/distributed_pdlp/partitioner.cpp new file mode 100644 index 0000000000..cf5f4f7aca --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/partitioner.cpp @@ -0,0 +1,224 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Plain C++ translation unit (not .cu): this file contains no device code, and +// KaMinPar's public header () is C++20 host code that pulls in TBB. + +#include + +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +template +std::vector dummy_partitioner_t::partition( + partitioner_input_t const& input) const +{ + cuopt_expects(input.nb_parts > 0, + error_type_t::ValidationError, + "dummy_partitioner: nb_parts must be positive"); + cuopt_expects(input.nb_cstr >= 0 && input.nb_vars >= 0, + error_type_t::ValidationError, + "dummy_partitioner: invalid problem dimensions"); + + const std::size_t nvtx = + static_cast(input.nb_cstr) + static_cast(input.nb_vars); + std::vector parts(nvtx); + for (std::size_t i = 0; i < nvtx; ++i) { + parts[i] = static_cast(i % static_cast(input.nb_parts)); + } + validate_partition(parts, + static_cast(input.nb_cstr), + static_cast(input.nb_vars), + static_cast(input.nb_parts), + "dummy_partitioner"); + return parts; +} + +// Builds the bipartite constraint/variable graph induced by A and runs the +// multi-threaded KaMinPar k-way kernel. +// * nodes [0, nb_cstr) : constraint nodes +// * nodes [nb_cstr, nb_cstr+nb_vars): variable nodes +// * each edge is a nnz between a constraint and a variable +template +std::vector kaminpar_partitioner_t::partition( + partitioner_input_t const& input) const +{ + cuopt_expects(input.nb_parts >= 1, + error_type_t::ValidationError, + "kaminpar_partitioner: nb_parts must be >= 1"); + cuopt_expects(input.nb_cstr >= 0 && input.nb_vars >= 0, + error_type_t::ValidationError, + "kaminpar_partitioner: invalid problem dimensions"); + + // return trivial partition if only one part + if (input.nb_parts == 1) { + CUOPT_LOG_INFO("KaMinPar: nb_parts == 1, returning trivial single-block partition"); + return std::vector(static_cast(input.nb_cstr + input.nb_vars), i_t{0}); + } + cuopt_expects(!input.A.row_offsets.empty() && !input.A.col_indices.empty(), + error_type_t::ValidationError, + "kaminpar_partitioner: A.row_offsets and A.col_indices are required"); + cuopt_expects(!input.A_t.row_offsets.empty() && !input.A_t.col_indices.empty(), + error_type_t::ValidationError, + "kaminpar_partitioner: A_t.row_offsets and A_t.col_indices are required"); + + auto A_offsets = input.A.row_offsets; + auto A_cols = input.A.col_indices; + auto A_t_offsets = input.A_t.row_offsets; + auto A_t_cols = input.A_t.col_indices; + + cuopt_expects(static_cast(A_offsets.size()) == input.nb_cstr + 1, + error_type_t::ValidationError, + "kaminpar_partitioner: A.row_offsets size mismatch (expected nb_cstr+1)"); + cuopt_expects(static_cast(A_t_offsets.size()) == input.nb_vars + 1, + error_type_t::ValidationError, + "kaminpar_partitioner: A_t.row_offsets size mismatch (expected nb_vars+1)"); + cuopt_expects(A_cols.size() == A_t_cols.size(), + error_type_t::ValidationError, + "kaminpar_partitioner: A and A_t nnz mismatch"); + + const i_t nb_cstr = input.nb_cstr; + const i_t nb_vars = input.nb_vars; + const i_t nnz = static_cast(A_cols.size()); + const i_t nvtx = nb_cstr + nb_vars; + + // > 0: use the specified number of threads, + // <= 0: use all hardware threads (1 as a last resort). + int nthreads = input.nb_threads > 0 ? static_cast(input.nb_threads) : 0; + if (nthreads <= 0) { + nthreads = static_cast(std::thread::hardware_concurrency()); + if (nthreads <= 0) { nthreads = 1; } + } + + // Bipartite CSR using KaMinPar index types (EdgeID for offsets, NodeID for neighbours). + std::vector xadj(static_cast(nvtx) + 1); + std::vector adjncy(2 * static_cast(nnz)); + + // CSR already represents an adjency list of cstr -> variables. + // Adding the transpose to represent the var -> cstr edges. + // Casting the types to KaMinPar friendly types + // Put A in top right corner of adjency matrix + // Put A_t in bottom left corner of adjency matrix + for (i_t i = 0; i <= nb_cstr; ++i) { + xadj[i] = static_cast(A_offsets[i]); + } + for (i_t i = 0; i <= nb_vars; ++i) { + // A_t edges live in adjncy[nnz .. 2*nnz), so their CSR offsets start at nnz + // (NOT nb_cstr). Corrupting this made xadj[nvtx] = nnz + nb_cstr instead of + // 2*nnz, causing KaMinPar to under-allocate and stomp the heap. + xadj[nb_cstr + i] = static_cast(A_t_offsets[i]) + + static_cast(nnz); + } + // cstr node/row has value in index J <=> link current cstr node with var node J + for (i_t k = 0; k < nnz; ++k) { + adjncy[k] = + static_cast(A_cols[k]) + static_cast(nb_cstr); + } + // same as right above but reversed + for (i_t k = 0; k < nnz; ++k) { + adjncy[nnz + k] = static_cast(A_t_cols[k]); + } + + std::vector block_of(static_cast(nvtx)); + + kaminpar::KaMinPar engine(nthreads, kaminpar::shm::create_default_context()); + engine.set_output_level(kaminpar::OutputLevel::QUIET); + engine.copy_graph(std::span(xadj), + std::span(adjncy)); + engine.set_k(static_cast(input.nb_parts)); + + // Allow up to 3% number of nodes imbalance. Could be knobbed for maximum performance + engine.set_uniform_max_block_weights(0.03); + + // The actual partition computation + auto t0 = std::chrono::high_resolution_clock::now(); + + const kaminpar::shm::EdgeWeight edge_cut = + engine.compute_partition(std::span(block_of)); + + auto t1 = std::chrono::high_resolution_clock::now(); + const double dt = std::chrono::duration(t1 - t0).count(); + + CUOPT_LOG_INFO( + "KaMinPar partitioned bipartite graph: nvtx=%d nnz=%d nb_parts=%d nthreads=%d edge_cut=%lld " + "in %.3fs", + static_cast(nvtx), + static_cast(nnz), + static_cast(input.nb_parts), + nthreads, + static_cast(edge_cut), + dt); + + std::vector parts(static_cast(nvtx)); + for (i_t i = 0; i < nvtx; ++i) { + parts[i] = static_cast(block_of[i]); + } + + validate_partition(parts, + static_cast(nb_cstr), + static_cast(nb_vars), + static_cast(input.nb_parts), + "kaminpar_partitioner"); + return parts; +} + +void validate_partition( + std::vector const& parts, int nb_cstr, int nb_vars, int nb_parts, char const* context) +{ + const std::size_t expected = + static_cast(nb_cstr) + static_cast(nb_vars); + cuopt_expects(parts.size() == expected, + error_type_t::ValidationError, + "%s: expected %zu part entries (cstrs + vars), got %zu", + context, + expected, + parts.size()); + cuopt_expects( + nb_parts > 0, error_type_t::ValidationError, "%s: nb_parts must be positive", context); + if (parts.empty()) { return; } + const auto [min_it, max_it] = std::minmax_element(parts.begin(), parts.end()); + cuopt_expects(*min_it >= 0, + error_type_t::ValidationError, + "%s: partition ids must be non-negative (min=%d)", + context, + static_cast(*min_it)); + cuopt_expects(*max_it < nb_parts, + error_type_t::ValidationError, + "%s: partition ids must be in [0, %d) (max=%d)", + context, + static_cast(nb_parts), + static_cast(*max_it)); +} + +template +std::unique_ptr> make_partitioner(partitioner_kind_t kind) +{ + switch (kind) { + case partitioner_kind_t::Dummy: return std::make_unique>(); + case partitioner_kind_t::KaMinPar: return std::make_unique>(); + } + cuopt_expects( + false, error_type_t::RuntimeError, "make_partitioner: unsupported partitioner kind"); + return nullptr; +} + +template class dummy_partitioner_t; +template class kaminpar_partitioner_t; +template std::unique_ptr> make_partitioner( + partitioner_kind_t); + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/partitioner.hpp b/cpp/src/pdlp/distributed_pdlp/partitioner.hpp new file mode 100644 index 0000000000..93b7065137 --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/partitioner.hpp @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// Non-owning view of a host CSR matrix (A or A_t). +template +struct csr_host_view_t { + std::span row_offsets{}; + std::span col_indices{}; + std::span values{}; // optional; unused by topology-only partitioners + i_t num_rows{0}; + i_t num_cols{0}; +}; + +// Inputs shared by all distributed-PDLP partitioners. +// Returns a flat vector of length (nb_cstr + nb_vars): constraint part-ids first, +// then variable part-ids, each in [0, nb_parts). +template +struct partitioner_input_t { + i_t nb_cstr{0}; + i_t nb_vars{0}; + i_t nb_parts{0}; + // Number of CPU threads the partitioner may use. Only honored by the + // multi-threaded KaMinPar backend. <= 0 means "auto" + i_t nb_threads{0}; + // Constraint matrix A (rows = constraints, cols = variables). + csr_host_view_t A{}; + // Transpose A_t (rows = variables, cols = constraints) + csr_host_view_t A_t{}; +}; + +// Dummy: round-robin, no graph (single-shard / debugging). +// KaMinPar: multi-threaded KaMinPar (preferred for multi-shard partitioning). +enum class partitioner_kind_t { Dummy, KaMinPar }; + +template +class partitioner_i { + public: + virtual ~partitioner_i() = default; + virtual std::vector partition(partitioner_input_t const& input) const = 0; +}; + +template +class dummy_partitioner_t : public partitioner_i { + public: + std::vector partition(partitioner_input_t const& input) const override; +}; + +// Multi-threaded k-way partitioner backed by KaMinPar. Builds a +// constraint/variable bipartite graph and runs the shared-memory parallel +// KaMinPar kernel so partitioning scales across all CPU cores of a node. +template +class kaminpar_partitioner_t : public partitioner_i { + public: + std::vector partition(partitioner_input_t const& input) const override; +}; + +void validate_partition(std::vector const& parts, + int nb_cstr, + int nb_vars, + int nb_parts, + char const* context = "partition"); + +template +std::unique_ptr> make_partitioner(partitioner_kind_t kind); + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/rank_data.hpp b/cpp/src/pdlp/distributed_pdlp/rank_data.hpp new file mode 100644 index 0000000000..ccc0430aa8 --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/rank_data.hpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { +// Pure data class representing most of the distributed data needed for mGPU operatiosn +template +struct rank_data_t { + rank_data_t(std::size_t nb_parts) + : var_send_per_peer(nb_parts), + cstr_send_per_peer(nb_parts), + var_recv_counts(nb_parts, 0), + var_recv_offsets(nb_parts, 0), + cstr_recv_counts(nb_parts, 0), + cstr_recv_offsets(nb_parts, 0) + { + } + + i_t owned_var_size{0}; + i_t total_var_size{0}; + i_t owned_cstr_size{0}; + i_t total_cstr_size{0}; + + // === Variable and Constraint indices owned by this shard, in global problem indices === + std::vector owned_var_indices; + std::vector owned_cstr_indices; + + // === Send plan: each element is a vector of indices to send to associated peer === + std::vector> var_send_per_peer; + std::vector> cstr_send_per_peer; + + // === Recv plan: per peer, contiguous slot in halo region === + std::vector var_recv_counts; + std::vector var_recv_offsets; + std::vector cstr_recv_counts; + std::vector cstr_recv_offsets; + + // === Mappings === + // global_to_local_* : full global problem indices to local shard problem indices + std::unordered_map global_to_local_var; + std::unordered_map global_to_local_cstr; + // local_to_global_* : local shard problem indices to full global problem indices + std::vector local_to_global_var; + std::vector local_to_global_cstr; + + // === Local host CSR matrices === + // A + std::vector h_A_row_offsets; + std::vector h_A_col_indices; + std::vector h_A_values; + // A_t + std::vector h_A_t_row_offsets; + std::vector h_A_t_col_indices; + std::vector h_A_t_values; +}; +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/shard.cu b/cpp/src/pdlp/distributed_pdlp/shard.cu new file mode 100644 index 0000000000..2bb4c5da6e --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/shard.cu @@ -0,0 +1,166 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include + +#include +#include + +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// This must be done in .cu file because the pdlp_solver_t is not already complete in the hpp file +// This is caused by the problematic cyclic include of pdlp_solver_t +template +pdlp_shard_t::~pdlp_shard_t() = default; + +template +pdlp_shard_t::pdlp_shard_t(int device_id, + rank_data_t&& rd, + nccl_comm_unique_ptr_t&& comm, + io::mps_data_model_t const& mps, + pdlp_solver_settings_t const& settings) + : device_id(device_id), + stream(), + handle(stream.view()), + comm(std::move(comm)), + rank_data(std::move(rd)), + opt_problem(&handle), + sub_problem(std::nullopt), + sub_pdlp(nullptr) +{ + assert(raft::device_setter::get_current_device() == device_id && + "Right device must be set before building the shard"); + + // ---- 0. Problem-level scalars, taken straight from the global mps. ---- + // Objective coefficients / offset / scaling factor are passed through + // unchanged; the max -> min conversion (negating all three) happens once, + // in mip::problem_t's constructor via convert_to_maximization_problem + const bool maximize = mps.get_sense(); + const f_t objective_offset = mps.get_objective_offset(); + const f_t objective_scaling_factor = mps.get_objective_scaling_factor(); + + // Global (unpartitioned) host arrays, indexed by global var / cstr id. + const std::vector& g_obj = mps.get_objective_coefficients(); + const std::vector& g_var_lower = mps.get_variable_lower_bounds(); + const std::vector& g_var_upper = mps.get_variable_upper_bounds(); + const std::vector& g_cstr_lower = mps.get_constraint_lower_bounds(); + const std::vector& g_cstr_upper = mps.get_constraint_upper_bounds(); + + // ---- 1. Gather per-shard host slices using rank_data's index maps. ---- + // All vectors are sized to TOTAL (owned + halo). Owned slots get real + // values; halo slots keep defaults because they should not be accessed before getting filled. + std::vector h_obj(rank_data.total_var_size, f_t{0}); + std::vector h_var_lower(rank_data.total_var_size, -std::numeric_limits::infinity()); + std::vector h_var_upper(rank_data.total_var_size, std::numeric_limits::infinity()); + std::vector h_cstr_lower(rank_data.total_cstr_size, -std::numeric_limits::infinity()); + std::vector h_cstr_upper(rank_data.total_cstr_size, std::numeric_limits::infinity()); + + for (i_t i = 0; i < rank_data.owned_var_size; ++i) { + const auto g = rank_data.local_to_global_var[i]; + h_obj[i] = g_obj[g]; + h_var_lower[i] = g_var_lower[g]; + h_var_upper[i] = g_var_upper[g]; + } + for (i_t i = 0; i < rank_data.owned_cstr_size; ++i) { + const auto g = rank_data.local_to_global_cstr[i]; + h_cstr_lower[i] = g_cstr_lower[g]; + h_cstr_upper[i] = g_cstr_upper[g]; + } + + // ---- 2. Populate opt_problem (constructed in init list) on this shard's device. ---- + opt_problem.set_csr_constraint_matrix(rank_data.h_A_values.data(), + static_cast(rank_data.h_A_values.size()), + rank_data.h_A_col_indices.data(), + static_cast(rank_data.h_A_col_indices.size()), + rank_data.h_A_row_offsets.data(), + static_cast(rank_data.h_A_row_offsets.size())); + + // Primal axis: TOTAL (owned + halo) + opt_problem.set_objective_coefficients(h_obj.data(), rank_data.total_var_size); + opt_problem.set_variable_lower_bounds(h_var_lower.data(), rank_data.total_var_size); + opt_problem.set_variable_upper_bounds(h_var_upper.data(), rank_data.total_var_size); + + // Dual axis: TOTAL (owned + halo) + opt_problem.set_constraint_lower_bounds(h_cstr_lower.data(), rank_data.total_cstr_size); + opt_problem.set_constraint_upper_bounds(h_cstr_upper.data(), rank_data.total_cstr_size); + + opt_problem.set_maximize(maximize); + opt_problem.set_objective_offset(objective_offset); + opt_problem.set_objective_scaling_factor(objective_scaling_factor); + opt_problem.set_problem_category(problem_category_t::LP); + + // ---- 3. Build problem_t from opt_problem (UNSCALED). ---- + sub_problem.emplace(opt_problem); + + // ---- 4. Override reverse_* with the real local A_T from rank_data. ---- + // problem_t's ctor computes the transpose of the LOCAL A, which is wrong + // in multi-GPU: A_local is owned_cstr x total_var, and A_t_local is the + // pre-sliced owned_var x total_cstr matrix we built during partitioning. + auto stream_view = handle.get_stream(); + sub_problem->reverse_offsets.resize(rank_data.h_A_t_row_offsets.size(), stream_view); + sub_problem->reverse_constraints.resize(rank_data.h_A_t_col_indices.size(), stream_view); + sub_problem->reverse_coefficients.resize(rank_data.h_A_t_values.size(), stream_view); + raft::copy(sub_problem->reverse_offsets.data(), + rank_data.h_A_t_row_offsets.data(), + rank_data.h_A_t_row_offsets.size(), + stream_view); + raft::copy(sub_problem->reverse_constraints.data(), + rank_data.h_A_t_col_indices.data(), + rank_data.h_A_t_col_indices.size(), + stream_view); + raft::copy(sub_problem->reverse_coefficients.data(), + rank_data.h_A_t_values.data(), + rank_data.h_A_t_values.size(), + stream_view); + handle.sync_stream(stream_view); + + // ---- 5. Build sub_pdlp (single-GPU mode). ---- + // is_distributed_sub_pdlp=true has two effects in pdlp_solver_t's ctor: + // * skip the CSR/CSC transpose validity check -- A and A_T here are two + // independent local slices, not transposes (A has all owned rows and + // A_T has all owned columns). + // * skip local Ruiz / Pock-Chambolle inside initial_scaling_strategy_'s + // ctor -- distributed scaling (multi_gpu_engine_t::distributed_scaling) + // runs a cross-shard-coherent scaling later. Local per-shard scaling + // would be incoherent across shards. + sub_pdlp = std::make_unique>( + *sub_problem, settings, /*is_legacy_batch_mode=*/false, /*is_distributed_sub_pdlp=*/true); + + // ---- 6. Build per-peer halo-exchange plans ---- + // For each peer p, we precompute: + // send_indices_d[p] : local indices to gather (uploaded from host send plan) + // send_buf_d[p] : f_t staging buffer sized to match + // Self-peer slot is present but empty (size 0). + auto build_send_plan = [&](std::vector> const& send_per_peer, + std::vector>& indices_d, + std::vector>& buf_d) { + const std::size_t n_peers = send_per_peer.size(); + indices_d.reserve(n_peers); + buf_d.reserve(n_peers); + for (auto const& send_to_peer : send_per_peer) { + indices_d.emplace_back(send_to_peer.size(), stream_view); + buf_d.emplace_back(send_to_peer.size(), stream_view); + if (!send_to_peer.empty()) { + raft::copy(indices_d.back().data(), send_to_peer.data(), send_to_peer.size(), stream_view); + } + } + }; + build_send_plan(rank_data.var_send_per_peer, var_send_indices_d, var_send_buf_d); + build_send_plan(rank_data.cstr_send_per_peer, cstr_send_indices_d, cstr_send_buf_d); + + handle.sync_stream(stream_view); +} + +template struct pdlp_shard_t; +template struct pdlp_shard_t; + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/distributed_pdlp/shard.hpp b/cpp/src/pdlp/distributed_pdlp/shard.hpp new file mode 100644 index 0000000000..dcca723c3e --- /dev/null +++ b/cpp/src/pdlp/distributed_pdlp/shard.hpp @@ -0,0 +1,108 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace cuopt::mathematical_optimization::pdlp { + +// Forward-declare to break the cyclic include with pdlp.cuh +// (pdlp.cuh -> multi_gpu_engine.hpp -> shard.hpp -> pdlp.cuh). +// Definitions of out-of-line members live in shard.cu, which includes pdlp.cuh. +template +class pdlp_solver_t; + +// RAII deleter for ncclComm_t; sets the right device before destroy. +struct nccl_comm_deleter_t { + int device_id{-1}; + void operator()(ncclComm* comm) const noexcept + { + if (comm == nullptr) return; + cuopt_assert(device_id >= 0, "nccl_comm_deleter_t: device_id not set"); + raft::device_setter guard(device_id); + CUOPT_NCCL_TRY_NO_THROW(ncclCommDestroy(comm)); + } +}; +using nccl_comm_unique_ptr_t = std::unique_ptr; + +template +struct pdlp_shard_t { + // Out-of-line (in shard.cu) because pdlp_solver_t is incomplete here. + ~pdlp_shard_t(); + + // sub worker for distributed pdlp. Owns its own view on scaled problem and unscaled problem + // Owns necessary multi-gpu data (rank_data, device_id, nccl_comm) + pdlp_shard_t(int device_id, + rank_data_t&& rd, + nccl_comm_unique_ptr_t&& comm, + io::mps_data_model_t const& mps, + pdlp_solver_settings_t const& settings); + + pdlp_shard_t(const pdlp_shard_t&) = delete; + pdlp_shard_t& operator=(const pdlp_shard_t&) = delete; + + int device_id{-1}; + rmm::cuda_stream stream; + raft::handle_t handle; + nccl_comm_unique_ptr_t comm; + rank_data_t rank_data; + optimization_problem_t opt_problem; + std::optional> sub_problem; + std::unique_ptr> sub_pdlp; + + // var_send_indices_d[peer] : local indices into primal vector to gather and ncclSend + // var_send_buf_d [peer] : staging buffer for outgoing variable values + // cstr_send_indices_d/cstr_send_buf_d : same, for dual vector + std::vector> var_send_indices_d; + std::vector> var_send_buf_d; + std::vector> cstr_send_indices_d; + std::vector> cstr_send_buf_d; + + // Non-owning bundle of per-axis halo-exchange metadata, indexed by peer. + // Consumed by multi_gpu_engine_t::halo_exchange_bufs_impl + struct halo_axis_t { + std::vector>& send_indices; // [peer] + std::vector>& send_buf; // [peer] + i_t owned_size; + std::vector const& recv_offsets; // [peer] + std::vector const& recv_counts; // [peer] + }; + halo_axis_t var_halo_axis() + { + return {var_send_indices_d, + var_send_buf_d, + rank_data.owned_var_size, + rank_data.var_recv_offsets, + rank_data.var_recv_counts}; + } + halo_axis_t cstr_halo_axis() + { + return {cstr_send_indices_d, + cstr_send_buf_d, + rank_data.owned_cstr_size, + rank_data.cstr_recv_offsets, + rank_data.cstr_recv_counts}; + } +}; + +} // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu index 551c23bd3a..51ca2e305a 100644 --- a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu +++ b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu @@ -31,17 +31,6 @@ namespace cuopt::mathematical_optimization::pdlp { -template -struct weighted_square_op { - f_t weight; - HDI f_t operator()(f_t v) { return v * v * weight; } -}; - -template -struct rescaling_from_squared_norm_op { - HDI f_t operator()(f_t sum) { return f_t(1.0) / (raft::sqrt(sum) + f_t(1.0)); } -}; - template struct inverse_rescaling_op { HDI f_t operator()(f_t v) @@ -79,7 +68,8 @@ pdlp_initial_scaling_strategy_t::pdlp_initial_scaling_strategy_t( pdhg_solver_t* pdhg_solver_ptr, const pdlp::pdlp_hyper_params_t& hyper_params, i_t original_batch_size, - bool running_mip) + bool running_mip, + bool skip_initial_scaling) : handle_ptr_(handle_ptr), stream_view_(handle_ptr_->get_stream()), primal_size_h_(op_problem_scaled.n_variables), @@ -130,10 +120,10 @@ pdlp_initial_scaling_strategy_t::pdlp_initial_scaling_strategy_t( objective_rescaling_.end(), f_t(1)); - compute_scaling_vectors(number_of_ruiz_iterations, alpha); - - iteration_constraint_matrix_scaling_.resize(0, stream_view_); - iteration_variable_scaling_.resize(0, stream_view_); + // Distributed PDLP shards defer scaling to multi_gpu_engine_t::distributed_scaling, + // which runs a cross-shard-coherent Ruiz. Local per-shard Ruiz would be incoherent + // across shards, so skip it here. + if (!skip_initial_scaling) { compute_scaling_vectors(number_of_ruiz_iterations, alpha); } } template @@ -142,6 +132,10 @@ void pdlp_initial_scaling_strategy_t::compute_scaling_vectors( { raft::common::nvtx::range fun_scope("compute_scaling_vectors"); + // Skip scaling entirely for a shape-0 problem (distributed PDLP builds the + // master pdlp_solver_t from a shape-0 placeholder) + if (primal_size_h_ == 0 || dual_size_h_ == 0) return; + if (hyper_params_.do_ruiz_scaling) { ruiz_inf_scaling(number_of_ruiz_iterations); } if (hyper_params_.do_pock_chambolle_scaling) { pock_chambolle_scaling(alpha); } } @@ -183,8 +177,9 @@ void pdlp_initial_scaling_strategy_t::bound_objective_rescaling() h_objective_rescaling_ = cuopt::host_copy(objective_rescaling_, stream_view_); } +// Row inf-norm of the scaled matrix, over the row-major matrix template -__global__ void inf_norm_row_and_col_kernel( +__global__ void inf_norm_row_kernel( const typename mip::problem_t::view_t op_problem, typename pdlp_initial_scaling_strategy_t::view_t initial_scaling_view) { @@ -198,14 +193,31 @@ __global__ void inf_norm_row_and_col_kernel( f_t scaled_val = (op_problem.coefficients[row_offset + j] * constraint_scale_factor) * variable_scale_factor; f_t abs_val = raft::abs(scaled_val); - - // row part if (abs_val > initial_scaling_view.iteration_constraint_matrix_scaling[row]) { raft::myAtomicMax(&initial_scaling_view.iteration_constraint_matrix_scaling[row], abs_val); } + } + } +} - // col part - // Add max with abs val in objective_matrix here for QP for cols +// Column inf-norm of the scaled matrix, over the column major matrix +template +__global__ void inf_norm_col_kernel( + const typename mip::problem_t::view_t op_problem, + typename pdlp_initial_scaling_strategy_t::view_t initial_scaling_view, + const f_t* A_T, + const i_t* A_T_offsets, + const i_t* A_T_indices) +{ + for (int col = blockIdx.x; col < op_problem.n_variables; col += gridDim.x) { + i_t col_offset = A_T_offsets[col]; + i_t nnz_in_col = A_T_offsets[col + 1] - col_offset; + f_t variable_scale_factor = initial_scaling_view.cummulative_variable_scaling[col]; + for (int j = threadIdx.x; j < nnz_in_col; j += blockDim.x) { + i_t row = A_T_indices[col_offset + j]; + f_t constraint_scale_factor = initial_scaling_view.cummulative_constraint_matrix_scaling[row]; + f_t scaled_val = (A_T[col_offset + j] * constraint_scale_factor) * variable_scale_factor; + f_t abs_val = raft::abs(scaled_val); if (abs_val > initial_scaling_view.iteration_variable_scaling[col]) { raft::myAtomicMax(&initial_scaling_view.iteration_variable_scaling[col], abs_val); } @@ -213,6 +225,54 @@ __global__ void inf_norm_row_and_col_kernel( } } +// One iteration of Ruiz inf-norm scaling. +// Distributed PDLP calls this per outer iteration between halo broadcasts; +template +void pdlp_initial_scaling_strategy_t::ruiz_iter_local() +{ + // Reset the iteration_scaling vectors to all 0 + RAFT_CUDA_TRY(cudaMemsetAsync( + iteration_constraint_matrix_scaling_.data(), 0, sizeof(f_t) * dual_size_h_, stream_view_)); + RAFT_CUDA_TRY(cudaMemsetAsync( + iteration_variable_scaling_.data(), 0, sizeof(f_t) * primal_size_h_, stream_view_)); + + // Inf-norm over rows and columns. Split into two kernels so the distributed path can + // touch only owned entries. + // Reading cols data from A_t allows for better cache locality on the AtomicAdd + // than it would by reading cols data from A as it is csr-represented => scattered cols + i_t number_of_blocks = op_problem_scaled_.n_constraints / block_size; + if (op_problem_scaled_.n_constraints % block_size) number_of_blocks++; + i_t number_of_threads = std::min(op_problem_scaled_.n_variables, (i_t)block_size); + inf_norm_row_kernel<<>>( + op_problem_scaled_.view(), this->view()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + i_t number_of_blocks_col = op_problem_scaled_.n_variables / block_size; + if (op_problem_scaled_.n_variables % block_size) number_of_blocks_col++; + i_t number_of_threads_col = std::min(op_problem_scaled_.n_constraints, (i_t)block_size); + inf_norm_col_kernel<<>>( + op_problem_scaled_.view(), this->view(), A_T_.data(), A_T_offsets_.data(), A_T_indices_.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + if (running_mip_) { reset_integer_variables(); } + + // Fold this iteration's inf-norms into the cumulative scalings: + // cumulative /= sqrt(iteration). + raft::linalg::binaryOp(cummulative_constraint_matrix_scaling_.data(), + cummulative_constraint_matrix_scaling_.data(), + iteration_constraint_matrix_scaling_.data(), + dual_size_h_, + a_divides_sqrt_b_bounded(), + stream_view_); + + raft::linalg::binaryOp(cummulative_variable_scaling_.data(), + cummulative_variable_scaling_.data(), + iteration_variable_scaling_.data(), + primal_size_h_, + a_divides_sqrt_b_bounded(), + stream_view_); +} + template void pdlp_initial_scaling_strategy_t::ruiz_inf_scaling(i_t number_of_ruiz_iterations) { @@ -221,36 +281,7 @@ void pdlp_initial_scaling_strategy_t::ruiz_inf_scaling(i_t number_of_r std::cout << "Doing ruiz_inf_scaling" << std::endl; #endif for (int i = 0; i < number_of_ruiz_iterations; i++) { - // find inf norm over rows and columns of the scaled matrix in given iteration (matrix is not - // actually updated, but the scaled value is computed and evaluated) - i_t number_of_blocks = op_problem_scaled_.n_constraints / block_size; - if (op_problem_scaled_.n_constraints % block_size) number_of_blocks++; - i_t number_of_threads = std::min(op_problem_scaled_.n_variables, (i_t)block_size); - inf_norm_row_and_col_kernel<<>>( - op_problem_scaled_.view(), this->view()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - - if (running_mip_) { reset_integer_variables(); } - - raft::linalg::binaryOp(cummulative_constraint_matrix_scaling_.data(), - cummulative_constraint_matrix_scaling_.data(), - iteration_constraint_matrix_scaling_.data(), - dual_size_h_, - a_divides_sqrt_b_bounded(), - stream_view_); - - raft::linalg::binaryOp(cummulative_variable_scaling_.data(), - cummulative_variable_scaling_.data(), - iteration_variable_scaling_.data(), - primal_size_h_, - a_divides_sqrt_b_bounded(), - stream_view_); - - // Reset the iteration_scaling vectors to all 0 - RAFT_CUDA_TRY(cudaMemsetAsync( - iteration_constraint_matrix_scaling_.data(), 0.0, sizeof(f_t) * dual_size_h_, stream_view_)); - RAFT_CUDA_TRY(cudaMemsetAsync( - iteration_variable_scaling_.data(), 0.0, sizeof(f_t) * primal_size_h_, stream_view_)); + ruiz_iter_local(); } } @@ -514,9 +545,9 @@ void pdlp_initial_scaling_strategy_t::resize_context(i_t new_size) } template -void pdlp_initial_scaling_strategy_t::scale_problem() +void pdlp_initial_scaling_strategy_t::apply_cummulative_scaling_to_problem() { - raft::common::nvtx::range fun_scope("scale_problem"); + raft::common::nvtx::range fun_scope("apply_cummulative_scaling_to_problem"); // scale A i_t number_of_blocks = op_problem_scaled_.n_constraints / block_size; @@ -586,56 +617,6 @@ void pdlp_initial_scaling_strategy_t::scale_problem() cuda::std::multiplies{}, stream_view_); - if (hyper_params_.bound_objective_rescaling && !running_mip_) { - // Coefficients are computed on the already scaled values - bound_objective_rescaling(); - -#ifdef CUPDLP_DEBUG_MODE - print("bound_rescaling", bound_rescaling_); - print("objective_rescaling", objective_rescaling_); -#endif - - cub::DeviceTransform::Transform( - cuda::std::make_tuple(op_problem_scaled_.constraint_lower_bounds.data(), - op_problem_scaled_.constraint_upper_bounds.data(), - batch_wrapped_container(bound_rescaling_, dual_size_h_)), - thrust::make_zip_iterator(op_problem_scaled_.constraint_lower_bounds.data(), - op_problem_scaled_.constraint_upper_bounds.data()), - op_problem_scaled_.constraint_upper_bounds.size(), - [] __device__(f_t constraint_lower_bound, - f_t constraint_upper_bound, - f_t bound_rescaling) -> thrust::tuple { - return {constraint_lower_bound * bound_rescaling, constraint_upper_bound * bound_rescaling}; - }, - stream_view_.value()); - - // In batch mode we don't scale the variable bounds (here) because they are shared across - // climbers. While the variable bounds are the same across climbers, there can be different - // bound rescaling factors for each climber. One solution would be to have per climber variable - // bounds but its costly from a memory perspective and from a memory bandwidth perspective. - // Since the variable bounds are the same across climbers but only the scaling factor changes, - // we pass the scaling factor to PDHG later. In PDHG we act the (almost fully) scaled variable - // bounds and add this missing scaling factor. - if (original_batch_size_ == 1) { - cub::DeviceTransform::Transform( - op_problem_scaled_.variable_bounds.data(), - op_problem_scaled_.variable_bounds.data(), - op_problem_scaled_.variable_bounds.size(), - [bound_rescaling = bound_rescaling_.data()] __device__(f_t2 variable_bounds) -> f_t2 { - return {variable_bounds.x * *bound_rescaling, variable_bounds.y * *bound_rescaling}; - }, - stream_view_); - } - - cub::DeviceTransform::Transform( - cuda::std::make_tuple(op_problem_scaled_.objective_coefficients.data(), - batch_wrapped_container(objective_rescaling_, primal_size_h_)), - op_problem_scaled_.objective_coefficients.data(), - op_problem_scaled_.objective_coefficients.size(), - cuda::std::multiplies{}, - stream_view_.value()); - } - #ifdef CUPDLP_DEBUG_MODE print("constraint_lower_bound", op_problem_scaled_.constraint_lower_bounds); print("constraint_upper_bound", op_problem_scaled_.constraint_upper_bounds); @@ -661,6 +642,79 @@ void pdlp_initial_scaling_strategy_t::scale_problem() } } +// Apply the already-published bound_rescaling_ / objective_rescaling_ device +// vectors to the scaled problem's constraint bounds, variable bounds, and +// objective. Used in both distributed and non-distributed PDLP. +template +void pdlp_initial_scaling_strategy_t::apply_bound_objective_rescaling_to_problem() +{ + using f_t2 = typename type_2::type; + + cub::DeviceTransform::Transform( + cuda::std::make_tuple(op_problem_scaled_.constraint_lower_bounds.data(), + op_problem_scaled_.constraint_upper_bounds.data(), + batch_wrapped_container(bound_rescaling_, dual_size_h_)), + thrust::make_zip_iterator(op_problem_scaled_.constraint_lower_bounds.data(), + op_problem_scaled_.constraint_upper_bounds.data()), + op_problem_scaled_.constraint_upper_bounds.size(), + [] __device__(f_t constraint_lower_bound, + f_t constraint_upper_bound, + f_t bound_rescaling) -> thrust::tuple { + return {constraint_lower_bound * bound_rescaling, constraint_upper_bound * bound_rescaling}; + }, + stream_view_.value()); + + // In batch mode we don't scale the variable bounds (here) because they are shared across + // climbers. While the variable bounds are the same across climbers, there can be different + // bound rescaling factors for each climber. One solution would be to have per climber variable + // bounds but its costly from a memory perspective and from a memory bandwidth perspective. + // Since the variable bounds are the same across climbers but only the scaling factor changes, + // we pass the scaling factor to PDHG later. In PDHG we act the (almost fully) scaled variable + // bounds and add this missing scaling factor. + if (original_batch_size_ == 1) { + cub::DeviceTransform::Transform( + op_problem_scaled_.variable_bounds.data(), + op_problem_scaled_.variable_bounds.data(), + op_problem_scaled_.variable_bounds.size(), + [bound_rescaling = bound_rescaling_.data()] __device__(f_t2 variable_bounds) -> f_t2 { + return {variable_bounds.x * *bound_rescaling, variable_bounds.y * *bound_rescaling}; + }, + stream_view_); + } + + cub::DeviceTransform::Transform( + cuda::std::make_tuple(op_problem_scaled_.objective_coefficients.data(), + batch_wrapped_container(objective_rescaling_, primal_size_h_)), + op_problem_scaled_.objective_coefficients.data(), + op_problem_scaled_.objective_coefficients.size(), + cuda::std::multiplies{}, + stream_view_.value()); +} + +template +void pdlp_initial_scaling_strategy_t::scale_problem() +{ + raft::common::nvtx::range fun_scope("scale_problem"); + + apply_cummulative_scaling_to_problem(); + + // Local bound/objective rescaling. Distributed PDLP intentionally does NOT + // reach this code path - it calls apply_cummulative_scaling_to_problem() + // directly and then applies the GLOBAL (allreduced) bound/objective factors + // via distributed_bound_objective_rescaling() instead. + if (hyper_params_.bound_objective_rescaling && !running_mip_) { + // Coefficients are computed on the already scaled values + bound_objective_rescaling(); + +#ifdef CUPDLP_DEBUG_MODE + print("bound_rescaling", bound_rescaling_); + print("objective_rescaling", objective_rescaling_); +#endif + + apply_bound_objective_rescaling_to_problem(); + } +} + template void pdlp_initial_scaling_strategy_t::scale_solutions( rmm::device_uvector& primal_solution, @@ -913,6 +967,24 @@ pdlp_initial_scaling_strategy_t::get_variable_scaling_vector() const return cummulative_variable_scaling_; } +template +void pdlp_initial_scaling_strategy_t::set_h_bound_rescaling(f_t value) +{ + std::fill(h_bound_rescaling_.begin(), h_bound_rescaling_.end(), value); + thrust::fill( + handle_ptr_->get_thrust_policy(), bound_rescaling_.begin(), bound_rescaling_.end(), value); +} + +template +void pdlp_initial_scaling_strategy_t::set_h_objective_rescaling(f_t value) +{ + std::fill(h_objective_rescaling_.begin(), h_objective_rescaling_.end(), value); + thrust::fill(handle_ptr_->get_thrust_policy(), + objective_rescaling_.begin(), + objective_rescaling_.end(), + value); +} + template typename pdlp_initial_scaling_strategy_t::view_t pdlp_initial_scaling_strategy_t::view() @@ -937,10 +1009,17 @@ pdlp_initial_scaling_strategy_t::view() #define INSTANTIATE(F_TYPE) \ template class pdlp_initial_scaling_strategy_t; \ \ - template __global__ void inf_norm_row_and_col_kernel( \ + template __global__ void inf_norm_row_kernel( \ const typename mip::problem_t::view_t op_problem, \ typename pdlp_initial_scaling_strategy_t::view_t initial_scaling_view); \ \ + template __global__ void inf_norm_col_kernel( \ + const typename mip::problem_t::view_t op_problem, \ + typename pdlp_initial_scaling_strategy_t::view_t initial_scaling_view, \ + const F_TYPE* A_T, \ + const int* A_T_offsets, \ + const int* A_T_indices); \ + \ template __global__ void pock_chambolle_scaling_kernel_col( \ const typename mip::problem_t::view_t op_problem, \ F_TYPE alpha, \ diff --git a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh index aee91f895e..bcda2c6f93 100644 --- a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh +++ b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh @@ -46,6 +46,10 @@ class pdlp_initial_scaling_strategy_t { raft::device_span cummulative_variable_scaling; }; // struct view_t + // skip_initial_scaling: when true, the ctor performs identity-initialization + // of the scaling vectors but does NOT run local Ruiz / Pock-Chambolle. Used + // by distributed PDLP shards, where cross-shard-coherent scaling is applied + // later by multi_gpu_engine_t::distributed_scaling. pdlp_initial_scaling_strategy_t(raft::handle_t const* handle_ptr, mip::problem_t& op_problem_scaled, i_t number_of_ruiz_iterations, @@ -56,7 +60,8 @@ class pdlp_initial_scaling_strategy_t { pdhg_solver_t* pdhg_solver_ptr, const pdlp::pdlp_hyper_params_t& hyper_params, i_t original_batch_size, - bool running_mip = false); + bool running_mip = false, + bool skip_initial_scaling = false); void scale_problem(); @@ -75,6 +80,17 @@ class pdlp_initial_scaling_strategy_t { rmm::device_uvector& dual_slack) const; void unscale_solutions(mip::solution_t& solution) const; const rmm::device_uvector& get_constraint_matrix_scaling_vector() const; + // Mutable access needed by distributed PDLP to broadcast owned constraint + rmm::device_uvector& get_cummulative_constraint_matrix_scaling() + { + return cummulative_constraint_matrix_scaling_; + } + // Mutable access needed by distributed PDLP to broadcast owned variable + // (column) scaling into the halo copies between scaling iterations. + rmm::device_uvector& get_cummulative_variable_scaling() + { + return cummulative_variable_scaling_; + } const rmm::device_uvector& get_variable_scaling_vector() const; const mip::problem_t& get_scaled_op_problem(); @@ -85,8 +101,45 @@ class pdlp_initial_scaling_strategy_t { void swap_context(const thrust::universal_host_pinned_vector>& swap_pairs); void resize_context(i_t new_size); + void set_h_bound_rescaling(f_t value); + void set_h_objective_rescaling(f_t value); + void bound_objective_rescaling(); + // Apply the already-populated bound_rescaling_ / objective_rescaling_ + // device vectors to op_problem_scaled_ (constraint bounds, variable bounds, + // objective). Extracted from scale_problem() into a shared helper so + // distributed PDLP can apply its globally-reduced scalars via the same + // three multiplies. + void apply_bound_objective_rescaling_to_problem(); + + // Public for distributed PDLP + void compute_scaling_vectors(i_t number_of_ruiz_iterations, f_t alpha); + + // ----- Distributed-PDLP hooks ----- + + // Apply the cumulative row/column scalings that Ruiz/Pock-Chambolle + // accumulated to A, A_T, c, variable bounds and constraint bounds, mark + // the problem as scaled and scale the seed primal/dual solutions. + // scale_problem() = apply_cummulative_scaling_to_problem() + local + // bound/objective rescaling + void apply_cummulative_scaling_to_problem(); + + // One Ruiz iteration (compute iteration vectors + fold into + // cumulative). Exposed for distributed PDLP so the outer loop with halo + // broadcasts lives at the distributed level + void ruiz_iter_local(); + // Shard-local end-to-end Pock-Chambolle pass. Exposed for distributed PDLP: + void pock_chambolle_scaling(f_t alpha); + // Iteration_* scratch buffers used by ruiz_iter_local / + // pock_chambolle_scaling. Exposed mutably so distributed PDLP can grow + // them back to full size after the ctor's release (see distributed_scaling). + rmm::device_uvector& get_iteration_variable_scaling() { return iteration_variable_scaling_; } + rmm::device_uvector& get_iteration_constraint_matrix_scaling() + { + return iteration_constraint_matrix_scaling_; + } + /** * @brief Gets the device-side view (with raw pointers), for ease of access * inside cuda kernels @@ -94,9 +147,7 @@ class pdlp_initial_scaling_strategy_t { view_t view(); private: - void compute_scaling_vectors(i_t number_of_ruiz_iterations, f_t alpha); void ruiz_inf_scaling(i_t number_of_ruiz_iterations); - void pock_chambolle_scaling(f_t alpha); void reset_integer_variables(); raft::handle_t const* handle_ptr_{nullptr}; diff --git a/cpp/src/pdlp/pdhg.cu b/cpp/src/pdlp/pdhg.cu index 89fc01faeb..965d6d737f 100644 --- a/cpp/src/pdlp/pdhg.cu +++ b/cpp/src/pdlp/pdhg.cu @@ -4,7 +4,9 @@ * SPDX-License-Identifier: Apache-2.0 */ /* clang-format on */ +#include #include +#include #include #include #include @@ -506,6 +508,13 @@ void pdhg_solver_t::compute_At_y() { // A_t @ y + // Multi-GPU dispatch: when the master pdhg has an engine, drive halo + // exchange + per-shard SpMV via the engine. + if (is_distributed_master()) { + mgpu_engine_->distributed_compute_At_y(); + return; + } + if (!batch_mode_) { if constexpr (std::is_same_v) { if (cusparse_view_.mixed_precision_enabled_) { @@ -555,6 +564,15 @@ template void pdhg_solver_t::compute_A_x() { // A @ x + + // Multi-GPU dispatch: see compute_At_y. The engine halo-updates the + // reflected_primal vector (the buffer this SpMV reads) and then drives + // per-shard local cusparse SpMV. + if (is_distributed_master()) { + mgpu_engine_->distributed_compute_A_x(); + return; + } + if (!batch_mode_) { if constexpr (std::is_same_v) { if (cusparse_view_.mixed_precision_enabled_) { @@ -600,6 +618,43 @@ void pdhg_solver_t::compute_A_x() } } +// out_desc = A^T @ in_desc, on this shard's local matrix. Both descriptors are +// caller-owned and can point at arbitrary scratch buffers. Used by +// multi_gpu_engine_t::distributed_spmv_At. +template +void pdhg_solver_t::spmv_At_into(cusparseDnVecDescr_t in_desc, + cusparseDnVecDescr_t out_desc) +{ + RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsespmv(handle_ptr_->get_cusparse_handle(), + CUSPARSE_OPERATION_NON_TRANSPOSE, + reusable_device_scalar_value_1_.data(), + cusparse_view_.A_T, + in_desc, + reusable_device_scalar_value_0_.data(), + out_desc, + CUSPARSE_SPMV_CSR_ALG2, + (f_t*)cusparse_view_.buffer_transpose.data(), + stream_view_)); +} + +// out_desc = A @ in_desc, the counterpart of spmv_At_into on this shard's local A. +template +void pdhg_solver_t::spmv_A_into(cusparseDnVecDescr_t in_desc, + cusparseDnVecDescr_t out_desc) +{ + RAFT_CUSPARSE_TRY( + raft::sparse::detail::cusparsespmv(handle_ptr_->get_cusparse_handle(), + CUSPARSE_OPERATION_NON_TRANSPOSE, + reusable_device_scalar_value_1_.data(), + cusparse_view_.A, + in_desc, + reusable_device_scalar_value_0_.data(), + out_desc, + CUSPARSE_SPMV_CSR_ALG2, + (f_t*)cusparse_view_.buffer_non_transpose.data(), + stream_view_)); +} + template void pdhg_solver_t::compute_primal_projection_with_gradient( rmm::device_uvector& primal_step_size) @@ -694,6 +749,22 @@ struct primal_reflected_major_projection { const f_t* scalar_; }; +template +void pdhg_solver_t::primal_reflected_major_projection_transform( + rmm::device_uvector& primal_step_size) +{ + cub::DeviceTransform::Transform( + cuda::std::make_tuple(current_saddle_point_state_.get_primal_solution().data(), + problem_ptr->objective_coefficients.data(), + current_saddle_point_state_.get_current_AtY().data(), + problem_ptr->variable_bounds.data()), + thrust::make_zip_iterator( + potential_next_primal_solution_.data(), dual_slack_.data(), reflected_primal_.data()), + primal_size_h_, + primal_reflected_major_projection(primal_step_size.data()), + stream_view_.value()); +} + template struct primal_reflected_major_projection_batch { using f_t2 = typename type_2::type; @@ -722,6 +793,21 @@ struct primal_reflected_projection { const f_t* scalar_; }; +template +void pdhg_solver_t::primal_reflected_projection_transform( + rmm::device_uvector& primal_step_size) +{ + cub::DeviceTransform::Transform( + cuda::std::make_tuple(current_saddle_point_state_.get_primal_solution().data(), + problem_ptr->objective_coefficients.data(), + current_saddle_point_state_.get_current_AtY().data(), + problem_ptr->variable_bounds.data()), + reflected_primal_.data(), + primal_size_h_, + primal_reflected_projection(primal_step_size.data()), + stream_view_.value()); +} + template struct primal_reflected_projection_batch { using f_t2 = typename type_2::type; @@ -751,6 +837,21 @@ struct dual_reflected_major_projection { const f_t* scalar_; }; +template +void pdhg_solver_t::dual_reflected_major_projection_transform( + rmm::device_uvector& dual_step_size) +{ + cub::DeviceTransform::Transform( + cuda::std::make_tuple(current_saddle_point_state_.get_dual_solution().data(), + current_saddle_point_state_.get_dual_gradient().data(), + problem_ptr->constraint_lower_bounds.data(), + problem_ptr->constraint_upper_bounds.data()), + thrust::make_zip_iterator(potential_next_dual_solution_.data(), reflected_dual_.data()), + dual_size_h_, + dual_reflected_major_projection(dual_step_size.data()), + stream_view_.value()); +} + template struct dual_reflected_major_projection_batch { HDI thrust::tuple operator()( @@ -779,6 +880,21 @@ struct dual_reflected_projection { const f_t* scalar_; }; +template +void pdhg_solver_t::dual_reflected_projection_transform( + rmm::device_uvector& dual_step_size) +{ + cub::DeviceTransform::Transform( + cuda::std::make_tuple(current_saddle_point_state_.get_dual_solution().data(), + current_saddle_point_state_.get_dual_gradient().data(), + problem_ptr->constraint_lower_bounds.data(), + problem_ptr->constraint_upper_bounds.data()), + reflected_dual_.data(), + dual_size_h_, + dual_reflected_projection(dual_step_size.data()), + stream_view_.value()); +} + template struct dual_reflected_projection_batch { HDI f_t @@ -1113,22 +1229,25 @@ void pdhg_solver_t::compute_next_primal_dual_solution_reflected( using f_t2 = typename type_2::type; + if (is_distributed_master()) { mgpu_engine_->sync_await_shards(stream_view_); } + // Compute next primal solution reflected. if (should_major) { graph_all.run(should_major, [&]() { + // Adds all the shards streams into the graph capture + if (is_distributed_master()) { mgpu_engine_->graph_capture_fork_to_shards(stream_view_); } + compute_At_y(); - if (!batch_mode_) { - cub::DeviceTransform::Transform( - cuda::std::make_tuple(current_saddle_point_state_.get_primal_solution().data(), - problem_ptr->objective_coefficients.data(), - current_saddle_point_state_.get_current_AtY().data(), - problem_ptr->variable_bounds.data()), - thrust::make_zip_iterator( - potential_next_primal_solution_.data(), dual_slack_.data(), reflected_primal_.data()), - primal_size_h_, - primal_reflected_major_projection(primal_step_size.data()), - stream_view_.value()); + + if (is_distributed_master()) { + mgpu_engine_->for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + sub_pdlp.pdhg_solver_.primal_reflected_major_projection_transform( + sub_pdlp.get_primal_step_size()); + }); + } else if (!batch_mode_) { + primal_reflected_major_projection_transform(primal_step_size); } else { cub::DeviceFor::Bulk( potential_next_primal_solution_.size(), @@ -1187,16 +1306,14 @@ void pdhg_solver_t::compute_next_primal_dual_solution_reflected( // Compute next dual compute_A_x(); - if (!batch_mode_) { - cub::DeviceTransform::Transform( - cuda::std::make_tuple(current_saddle_point_state_.get_dual_solution().data(), - current_saddle_point_state_.get_dual_gradient().data(), - problem_ptr->constraint_lower_bounds.data(), - problem_ptr->constraint_upper_bounds.data()), - thrust::make_zip_iterator(potential_next_dual_solution_.data(), reflected_dual_.data()), - dual_size_h_, - dual_reflected_major_projection(dual_step_size.data()), - stream_view_.value()); + if (is_distributed_master()) { + mgpu_engine_->for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + sub_pdlp.pdhg_solver_.dual_reflected_major_projection_transform( + sub_pdlp.get_dual_step_size()); + }); + } else if (!batch_mode_) { + dual_reflected_major_projection_transform(dual_step_size); } else { cub::DeviceFor::Bulk( potential_next_dual_solution_.size(), @@ -1217,10 +1334,18 @@ void pdhg_solver_t::compute_next_primal_dual_solution_reflected( print("potential_next_dual_solution_", potential_next_dual_solution_); print("reflected_dual_", reflected_dual_); #endif + + // Multi-GPU: close the fork by joining every shard stream back into + // the master stream so cudaStreamEndCapture sees a single graph + // spanning all streams. + if (is_distributed_master()) { mgpu_engine_->graph_capture_join_from_shards(stream_view_); } }); } else { graph_all.run(should_major, [&]() { + // Same reason as above, adds all the shards streams into the graph capture + if (is_distributed_master()) { mgpu_engine_->graph_capture_fork_to_shards(stream_view_); } + // Compute next primal compute_At_y(); @@ -1232,16 +1357,14 @@ void pdhg_solver_t::compute_next_primal_dual_solution_reflected( current_saddle_point_state_.get_current_AtY()); #endif - if (!batch_mode_) { - cub::DeviceTransform::Transform( - cuda::std::make_tuple(current_saddle_point_state_.get_primal_solution().data(), - problem_ptr->objective_coefficients.data(), - current_saddle_point_state_.get_current_AtY().data(), - problem_ptr->variable_bounds.data()), - reflected_primal_.data(), - primal_size_h_, - primal_reflected_projection(primal_step_size.data()), - stream_view_.value()); + if (is_distributed_master()) { + mgpu_engine_->for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + sub_pdlp.pdhg_solver_.primal_reflected_projection_transform( + sub_pdlp.get_primal_step_size()); + }); + } else if (!batch_mode_) { + primal_reflected_projection_transform(primal_step_size); } else { cub::DeviceFor::Bulk( reflected_primal_.size(), @@ -1301,16 +1424,13 @@ void pdhg_solver_t::compute_next_primal_dual_solution_reflected( // Compute next dual compute_A_x(); - if (!batch_mode_) { - cub::DeviceTransform::Transform( - cuda::std::make_tuple(current_saddle_point_state_.get_dual_solution().data(), - current_saddle_point_state_.get_dual_gradient().data(), - problem_ptr->constraint_lower_bounds.data(), - problem_ptr->constraint_upper_bounds.data()), - reflected_dual_.data(), - dual_size_h_, - dual_reflected_projection(dual_step_size.data()), - stream_view_.value()); + if (is_distributed_master()) { + mgpu_engine_->for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + sub_pdlp.pdhg_solver_.dual_reflected_projection_transform(sub_pdlp.get_dual_step_size()); + }); + } else if (!batch_mode_) { + dual_reflected_projection_transform(dual_step_size); } else { cub::DeviceFor::Bulk( reflected_dual_.size(), @@ -1328,8 +1448,13 @@ void pdhg_solver_t::compute_next_primal_dual_solution_reflected( #ifdef CUPDLP_DEBUG_MODE print("reflected_dual_", reflected_dual_); #endif + + if (is_distributed_master()) { mgpu_engine_->graph_capture_join_from_shards(stream_view_); } }); } + + // sync to master stream after the graph is captured + if (is_distributed_master()) { mgpu_engine_->sync_await_master(stream_view_); } } template diff --git a/cpp/src/pdlp/pdhg.hpp b/cpp/src/pdlp/pdhg.hpp index 1860f0e8f2..83f3b5bde4 100644 --- a/cpp/src/pdlp/pdhg.hpp +++ b/cpp/src/pdlp/pdhg.hpp @@ -24,6 +24,12 @@ #include namespace cuopt::mathematical_optimization::pdlp { + +// Forward-declared to avoid include cycle: multi_gpu_engine.hpp itself includes pdhg.hpp +// (engine calls per-shard pdhg compute_*). pdhg.cu does the full include. +template +struct multi_gpu_engine_t; + template class pdhg_solver_t { public: @@ -76,6 +82,41 @@ class pdhg_solver_t { void update_solution(cusparse_view_t& current_op_problem_evaluation_cusparse_view_); void refine_initial_primal_projection(const rmm::device_uvector& bound_rescaling); + // SpMV primitives. Public so the multi-GPU engine can drive them per-shard + // after halo-exchanging the relevant vector. + // + // If set_multi_gpu_engine() has been called, these dispatch to the engine + // (halo exchange + per-shard SpMV). Otherwise they run the single-GPU + // cusparse path on the local matrix. + void compute_At_y(); + void compute_A_x(); + void spmvop_At_y(); + void spmvop_A_x(); + + // Parameterized SpMVs used by the multi-GPU engine. Thin wrappers around + // cusparsespmv on this shard's local A / A^T + void spmv_At_into(cusparseDnVecDescr_t in_desc, cusparseDnVecDescr_t out_desc); + void spmv_A_into(cusparseDnVecDescr_t in_desc, cusparseDnVecDescr_t out_desc); + + // Pure cub-transform extractions. Allows for clearer containment of the calls and ensures + // the single-GPU vs distributed-GPU uses the same calls + void primal_reflected_major_projection_transform(rmm::device_uvector& primal_step_size); + void dual_reflected_major_projection_transform(rmm::device_uvector& dual_step_size); + void primal_reflected_projection_transform(rmm::device_uvector& primal_step_size); + void dual_reflected_projection_transform(rmm::device_uvector& dual_step_size); + + // Master PDLP wires the engine pointer here after the engine is built. Only + // the master's pdhg_solver_ holds a non-null engine; shards leave it null and + // run single-GPU SpMV on their local matrix. + void set_multi_gpu_engine(multi_gpu_engine_t* engine) { mgpu_engine_ = engine; } + multi_gpu_engine_t* get_mgpu_engine() const { return mgpu_engine_; } + + // True only on the master pdhg of a distributed run (the one wired to the + // engine, which orchestrates the shards). + // Shards report false. + // Single-GPU PDHG reports false. + bool is_distributed_master() const { return mgpu_engine_ != nullptr; } + i_t total_pdhg_iterations_; private: @@ -93,10 +134,6 @@ class pdhg_solver_t { void compute_primal_projection_with_gradient(rmm::device_uvector& primal_step_size); void compute_primal_projection(rmm::device_uvector& primal_step_size); - void compute_At_y(); - void compute_A_x(); - void spmvop_At_y(); - void spmvop_A_x(); bool batch_mode_{false}; raft::handle_t const* handle_ptr_{nullptr}; @@ -150,6 +187,10 @@ class pdhg_solver_t { rmm::device_uvector new_bounds_lower_; rmm::device_uvector new_bounds_upper_; cuda::fast_mod_div batch_size_divisor_; + + // Non-owning. Set on the master pdhg_solver_ in distributed mode; null + // (default) means single-GPU path. + multi_gpu_engine_t* mgpu_engine_{nullptr}; }; } // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/pdlp.cu b/cpp/src/pdlp/pdlp.cu index a936763fd4..488a2bcc6b 100644 --- a/cpp/src/pdlp/pdlp.cu +++ b/cpp/src/pdlp/pdlp.cu @@ -11,12 +11,16 @@ #include #include +#include +#include #include #include #include +#include #include #include "cuopt/mathematical_optimization/pdlp/solver_solution.hpp" +#include "distributed_pdlp/multi_gpu_engine.hpp" #include #include @@ -40,9 +44,11 @@ #include #include +#include #include #include #include +#include #include namespace cuopt::mathematical_optimization::pdlp { @@ -155,7 +161,8 @@ static size_t batch_size_handler(const pdlp_solver_settings_t& setting template pdlp_solver_t::pdlp_solver_t(mip::problem_t& op_problem, pdlp_solver_settings_t const& settings, - bool is_legacy_batch_mode) + bool is_legacy_batch_mode, + bool is_distributed_sub_pdlp) : original_batch_size_(batch_size_handler(settings)), climber_strategies_(original_batch_size_), batch_mode_(climber_strategies_.size() > 1), @@ -198,7 +205,9 @@ pdlp_solver_t::pdlp_solver_t(mip::problem_t& op_problem, op_problem_scaled_.reverse_constraints, &pdhg_solver_, settings_.hyper_params, - static_cast(original_batch_size_)}, + static_cast(original_batch_size_), + /*running_mip=*/false, + /*skip_initial_scaling=*/is_distributed_sub_pdlp}, average_op_problem_evaluation_cusparse_view_{handle_ptr_, op_problem, unscaled_primal_avg_solution_, @@ -260,7 +269,8 @@ pdlp_solver_t::pdlp_solver_t(mip::problem_t& op_problem, reusable_device_scalar_value_0_{f_t(0.0), stream_view_}, batch_solution_to_return_{pdlp_termination_status_t::TimeLimit, stream_view_}, best_primal_solution_so_far{pdlp_termination_status_t::TimeLimit, stream_view_}, - inside_mip_{false} + inside_mip_{false}, + is_distributed_sub_pdlp_{is_distributed_sub_pdlp} { cuopt_expects(!(settings_.first_primal_feasible && settings_.all_primal_feasible), error_type_t::ValidationError, @@ -349,7 +359,11 @@ pdlp_solver_t::pdlp_solver_t(mip::problem_t& op_problem, best_primal_quality_so_far_.primal_objective = (op_problem_scaled_.maximize) ? -std::numeric_limits::infinity() : std::numeric_limits::infinity(); - op_problem.check_problem_representation(true, false); + // On a distributed sub-solver, op_problem.coefficients (owned rows) and + // op_problem.reverse_coefficients (owned cols) are two independent slices + // of the global matrix, not transposes of each other; skip that check. + op_problem.check_problem_representation(/*check_transposed=*/!is_distributed_sub_pdlp_, + /*empty=*/false); if (batch_mode_) { batch_solution_to_return_.get_additional_termination_informations().resize( @@ -371,6 +385,223 @@ pdlp_solver_t::pdlp_solver_t(mip::problem_t& op_problem, } } +// ============================================================================ +// Distributed multi-GPU ctor. +// needs placeholder_problem to be a shape-0 problem +// reads the problem from mps_data_model directly +// builds internal attributes from the placeholder_problem +// builds the engine from the mps_data_model +template +pdlp_solver_t::pdlp_solver_t( + mip::problem_t& placeholder_problem, + cuopt::mathematical_optimization::io::mps_data_model_t const& mps, + pdlp_solver_settings_t const& settings) + // Makes all inner feilds of master 0 size + : pdlp_solver_t(placeholder_problem, settings, /*is_legacy_batch_mode=*/false) +{ + cuopt_expects(placeholder_problem.n_variables == 0 && placeholder_problem.n_constraints == 0 && + placeholder_problem.nnz == 0, + error_type_t::ValidationError, + "Distributed mGPU pdlp_solver_t ctor requires a shape-0 " + "placeholder problem (n_variables == n_constraints == nnz == 0)"); + cuopt_expects(settings.hyper_params.never_restart_to_average, + error_type_t::ValidationError, + "Distributed PDLP requires never_restart_to_average = true"); + const int distributed_pdlp_num_gpus = settings.distributed_pdlp_num_gpus; + CUOPT_LOG_INFO("Solving with distributed PDLP on %d GPU.", distributed_pdlp_num_gpus); + + if constexpr (!std::is_same_v) { + cuopt_expects( + false, error_type_t::ValidationError, "Distributed PDLP currently requires double precision"); + return; + } + // ----- 1. Read problem shape and bulk data directly from mps (host) ----- + const i_t n_vars = static_cast(mps.get_objective_coefficients().size()); + const i_t n_cstr = static_cast(mps.get_constraint_lower_bounds().size()); + const i_t nnz = static_cast(mps.get_constraint_matrix_values().size()); + cuopt_expects(n_vars > 0, + error_type_t::ValidationError, + "Distributed PDLP from mps requires a non-empty objective"); + cuopt_expects(n_cstr > 0, + error_type_t::ValidationError, + "Distributed PDLP from mps requires at least one constraint"); + cuopt_expects(static_cast(mps.get_constraint_matrix_offsets().size()) == n_cstr + 1, + error_type_t::ValidationError, + "mps constraint_matrix_offsets size must equal n_constraints + 1"); + cuopt_expects( + static_cast(mps.get_constraint_matrix_indices().size()) == nnz, + error_type_t::ValidationError, + "mps constraint_matrix_indices size must equal nnz (constraint_matrix_values size)"); + cuopt_expects(static_cast(mps.get_constraint_upper_bounds().size()) == n_cstr, + error_type_t::ValidationError, + "mps constraint_upper_bounds size must equal n_constraints"); + cuopt_expects(static_cast(mps.get_variable_lower_bounds().size()) == n_vars, + error_type_t::ValidationError, + "mps variable_lower_bounds size must equal n_variables"); + cuopt_expects(static_cast(mps.get_variable_upper_bounds().size()) == n_vars, + error_type_t::ValidationError, + "mps variable_upper_bounds size must equal n_variables"); + + // A (CSR) — mutable copies for the engine + partitioner consumers below. + std::vector h_A_row_offsets = mps.get_constraint_matrix_offsets(); + std::vector h_A_col_indices = mps.get_constraint_matrix_indices(); + std::vector h_A_values = mps.get_constraint_matrix_values(); + + // ----- 2. Transpose A -> A^T on the host (one-shot CSR transpose) ----- + // CSC(A) and CSR(A^T) share the same memory layout, so the CSC produced + // by csr_matrix_t::to_compressed_col IS the CSR of A^T. + // O(nnz + n_vars) counting sort, same as problem_t::compute_transpose. + namespace ds = cuopt::mathematical_optimization; + ds::csr_matrix_t A_csr(n_cstr, n_vars, nnz); + A_csr.row_start = h_A_row_offsets; + A_csr.j = h_A_col_indices; + A_csr.x = h_A_values; + ds::csc_matrix_t AT_as_csc(n_vars, n_cstr, nnz); + A_csr.to_compressed_col(AT_as_csc); + std::vector h_A_t_row_offsets = std::move(AT_as_csc.col_start); + std::vector h_A_t_col_indices = std::move(AT_as_csc.i); + std::vector h_A_t_values = std::move(AT_as_csc.x); + + // ----- 3. Partition ----- + std::vector parts; + { + partitioner_input_t partition_input; + partition_input.nb_cstr = n_cstr; + partition_input.nb_vars = n_vars; + partition_input.nb_parts = distributed_pdlp_num_gpus; + + // Resolve which partitioner to use. The public enum is validated at the + // parameter-set layer (int_parameters min/max), so we only need to map + // the three known values to the backend selector. + partitioner_kind_t kind; + switch (settings.distributed_pdlp_partitioner) { + case distributed_pdlp_partitioner_t::Auto: + kind = (distributed_pdlp_num_gpus == 1) ? partitioner_kind_t::Dummy + : partitioner_kind_t::KaMinPar; + break; + case distributed_pdlp_partitioner_t::KaMinPar: kind = partitioner_kind_t::KaMinPar; break; + case distributed_pdlp_partitioner_t::Dummy: kind = partitioner_kind_t::Dummy; break; + default: + cuopt_expects(false, + error_type_t::ValidationError, + "Unknown distributed_pdlp_partitioner value %d", + static_cast(settings.distributed_pdlp_partitioner)); + kind = partitioner_kind_t::Dummy; // unreachable; silences -Wmaybe-uninitialized + } + const bool needs_graph = (kind == partitioner_kind_t::KaMinPar); + if (needs_graph) { + // csr_host_view_t members are std::span — an owning + // std::vector converts implicitly. + partition_input.A.row_offsets = h_A_row_offsets; + partition_input.A.col_indices = h_A_col_indices; + partition_input.A.num_rows = n_cstr; + partition_input.A.num_cols = n_vars; + partition_input.A_t.row_offsets = h_A_t_row_offsets; + partition_input.A_t.col_indices = h_A_t_col_indices; + partition_input.A_t.num_rows = n_vars; + partition_input.A_t.num_cols = n_cstr; + // 0 => KaMinPar auto-detects and uses all hardware threads. + partition_input.nb_threads = 0; + } + const char* kind_name = (kind == partitioner_kind_t::Dummy) ? "dummy" + : (kind == partitioner_kind_t::KaMinPar) ? "kaminpar" + : "unknown"; + CUOPT_LOG_INFO( + "Partitioning %d constraints + %d variables into %d part(s) using the %s " + "partitioner", + n_cstr, + n_vars, + distributed_pdlp_num_gpus, + kind_name); + auto partitioner = make_partitioner(kind); + parts = partitioner->partition(partition_input); + } + + // ----- 4. Build per-rank data ----- + std::vector> sub_pdlp_rank_data = + create_rank_data_from_parts(parts, + h_A_row_offsets, + h_A_col_indices, + h_A_values, + h_A_t_row_offsets, + h_A_t_col_indices, + h_A_t_values, + settings.distributed_pdlp_num_gpus, + n_cstr, + n_vars, + nnz); + + // ----- 5. Per-shard settings ----- + pdlp_solver_settings_t sub_pdlp_settings = settings; + sub_pdlp_settings.num_gpus = 1; + sub_pdlp_settings.distributed_pdlp_num_gpus = 1; + // Disable automatic ruiz and pock-chambolle in the initial_scaling ctor: the + // distributed pipeline computes them via distributed_scaling using the + // GLOBAL problem. + sub_pdlp_settings.hyper_params.do_ruiz_scaling = false; + sub_pdlp_settings.hyper_params.do_pock_chambolle_scaling = false; + + // ----- 6. Construct the engine: NCCL comms + per-shard pdlp_solver_t ----- + multi_gpu_engine.emplace(std::move(sub_pdlp_rank_data), mps, sub_pdlp_settings); + + // Non-owning back-pointer so engine helpers (e.g. allreduce_sum_inplace_to_master) + // can apply a single pdlp_solver_t-shaped accessor lambda to master too. + multi_gpu_engine->set_master(this); + + // Wire the engine into master's pdhg_solver_; shards keep mgpu_engine_ == nullptr. + pdhg_solver_.set_multi_gpu_engine(&*multi_gpu_engine); + + // ----- 9. Resize master gather destinations to the full problem size ----- + pdhg_solver_.get_potential_next_primal_solution().resize(n_vars, stream_view_); + pdhg_solver_.get_potential_next_dual_solution().resize(n_cstr, stream_view_); + current_termination_strategy_.get_convergence_information().get_reduced_cost().resize( + n_vars, stream_view_); + primal_size_h_ = n_vars; + dual_size_h_ = n_cstr; + + // Distributed conergence_information::init_l2_norms + multi_gpu_engine->for_each_shard([](auto& shard) { + shard.sub_pdlp->get_current_termination_strategy() + .get_convergence_information() + .compute_owned_reference_norm_partials(shard.rank_data.owned_var_size, + shard.rank_data.owned_cstr_size); + }); + multi_gpu_engine->allreduce_sum_inplace([](pdlp_solver_t& sp) -> f_t* { + return sp.get_current_termination_strategy() + .get_convergence_information() + .l2_norm_primal_right_hand_side_data(); + }); + multi_gpu_engine->allreduce_sum_inplace([](pdlp_solver_t& sp) -> f_t* { + return sp.get_current_termination_strategy() + .get_convergence_information() + .l2_norm_primal_linear_objective_data(); + }); + multi_gpu_engine->for_each_shard([](auto& shard) { + shard.sub_pdlp->get_current_termination_strategy() + .get_convergence_information() + .sqrt_reference_norms_inplace(); + shard.stream.synchronize(); + }); + // Broadcast the values to the master + { + auto& s0 = *multi_gpu_engine->shards[0]; + auto& s0_conv = s0.sub_pdlp->get_current_termination_strategy().get_convergence_information(); + raft::device_setter guard(s0.device_id); + for (auto* ts : {¤t_termination_strategy_, &average_termination_strategy_}) { + auto& ci = ts->get_convergence_information(); + raft::copy(ci.l2_norm_primal_right_hand_side_data(), + s0_conv.l2_norm_primal_right_hand_side_data(), + 1, + stream_view_); + raft::copy(ci.l2_norm_primal_linear_objective_data(), + s0_conv.l2_norm_primal_linear_objective_data(), + 1, + stream_view_); + } + } + handle_ptr_->sync_stream(stream_view_); +} + template void pdlp_solver_t::set_initial_primal_weight(f_t initial_primal_weight) { @@ -1980,6 +2211,26 @@ void pdlp_solver_t::resize_and_swap_all_context_loop( RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); } +// delta = reflected - current, for both primal and dual, written into the +// saddle-point delta buffers. Shared by the single-GPU and per-shard +// (distributed) paths so the two only differ by which pdhg/stream they pass. +template +static void compute_primal_dual_deltas(pdhg_solver_t& pdhg, rmm::cuda_stream_view stream) +{ + cub::DeviceTransform::Transform( + cuda::std::make_tuple(pdhg.get_reflected_primal().data(), pdhg.get_primal_solution().data()), + pdhg.get_saddle_point_state().get_delta_primal().data(), + pdhg.get_primal_solution().size(), + cuda::std::minus{}, + stream); + cub::DeviceTransform::Transform( + cuda::std::make_tuple(pdhg.get_reflected_dual().data(), pdhg.get_dual_solution().data()), + pdhg.get_saddle_point_state().get_delta_dual().data(), + pdhg.get_dual_solution().size(), + cuda::std::minus{}, + stream); +} + template void pdlp_solver_t::compute_fixed_error(std::vector& has_restarted) { @@ -1988,54 +2239,105 @@ void pdlp_solver_t::compute_fixed_error(std::vector& has_restarte #ifdef CUPDLP_DEBUG_MODE printf("Computing compute_fixed_point_error \n"); #endif - cuopt_assert( - pdhg_solver_.get_reflected_primal().size() == primal_size_h_ * climber_strategies_.size(), - "reflected_primal_ size mismatch"); - cuopt_assert( - pdhg_solver_.get_reflected_dual().size() == dual_size_h_ * climber_strategies_.size(), - "reflected_dual_ size mismatch"); - cuopt_assert( - pdhg_solver_.get_primal_solution().size() == primal_size_h_ * climber_strategies_.size(), - "primal_solution_ size mismatch"); - cuopt_assert(pdhg_solver_.get_dual_solution().size() == dual_size_h_ * climber_strategies_.size(), - "dual_solution_ size mismatch"); - cuopt_assert(pdhg_solver_.get_saddle_point_state().get_delta_primal().size() == - primal_size_h_ * climber_strategies_.size(), - "delta_primal_ size mismatch"); - cuopt_assert(pdhg_solver_.get_saddle_point_state().get_delta_dual().size() == - dual_size_h_ * climber_strategies_.size(), - "delta_dual_ size mismatch"); - - // Computing the deltas + if (is_distributed_master()) { + // The master's own pdhg_solver_ buffers are built from a shape-0 placeholder and must be empty + cuopt_assert(pdhg_solver_.get_reflected_primal().size() == 0, + "master must not own a reflected_primal_ buffer (shards compute it)"); + cuopt_assert(pdhg_solver_.get_reflected_dual().size() == 0, + "master must not own a reflected_dual_ buffer (shards compute it)"); + cuopt_assert(pdhg_solver_.get_primal_solution().size() == 0, + "master must not own a primal_solution_ buffer (shards compute it)"); + cuopt_assert(pdhg_solver_.get_dual_solution().size() == 0, + "master must not own a dual_solution_ buffer (shards compute it)"); + cuopt_assert(pdhg_solver_.get_saddle_point_state().get_delta_primal().size() == 0, + "master must not own a delta_primal_ buffer (shards compute it)"); + cuopt_assert(pdhg_solver_.get_saddle_point_state().get_delta_dual().size() == 0, + "master must not own a delta_dual_ buffer (shards compute it)"); + } else { + cuopt_assert( + pdhg_solver_.get_reflected_primal().size() == primal_size_h_ * climber_strategies_.size(), + "reflected_primal_ size mismatch"); + cuopt_assert( + pdhg_solver_.get_reflected_dual().size() == dual_size_h_ * climber_strategies_.size(), + "reflected_dual_ size mismatch"); + cuopt_assert( + pdhg_solver_.get_primal_solution().size() == primal_size_h_ * climber_strategies_.size(), + "primal_solution_ size mismatch"); + cuopt_assert( + pdhg_solver_.get_dual_solution().size() == dual_size_h_ * climber_strategies_.size(), + "dual_solution_ size mismatch"); + cuopt_assert(pdhg_solver_.get_saddle_point_state().get_delta_primal().size() == + primal_size_h_ * climber_strategies_.size(), + "delta_primal_ size mismatch"); + cuopt_assert(pdhg_solver_.get_saddle_point_state().get_delta_dual().size() == + dual_size_h_ * climber_strategies_.size(), + "delta_dual_ size mismatch"); + } + + // Computing the deltas (delta = reflected - current) // TODO batch mdoe: this only works if everyone restarts - cub::DeviceTransform::Transform(cuda::std::make_tuple(pdhg_solver_.get_reflected_primal().data(), - pdhg_solver_.get_primal_solution().data()), - pdhg_solver_.get_saddle_point_state().get_delta_primal().data(), - pdhg_solver_.get_primal_solution().size(), - cuda::std::minus{}, - stream_view_.value()); - cub::DeviceTransform::Transform(cuda::std::make_tuple(pdhg_solver_.get_reflected_dual().data(), - pdhg_solver_.get_dual_solution().data()), - pdhg_solver_.get_saddle_point_state().get_delta_dual().data(), - pdhg_solver_.get_dual_solution().size(), - cuda::std::minus{}, - stream_view_.value()); + if (is_distributed_master()) { + multi_gpu_engine->for_each_shard([](auto& shard) { + compute_primal_dual_deltas(shard.sub_pdlp->pdhg_solver_, shard.stream.view()); + }); + } else { + compute_primal_dual_deltas(pdhg_solver_, stream_view_); + } auto& cusparse_view = pdhg_solver_.get_cusparse_view(); - // Sync to make sure all previous cuSparse operations are finished before setting the - // potential_next_dual_solution - RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); - // Make potential_next_dual_solution point towards reflected dual solution to reuse the code - RAFT_CUSPARSE_TRY(cusparseDnVecSetValues(cusparse_view.potential_next_dual_solution, - (void*)pdhg_solver_.get_reflected_dual().data())); + // Distributed compute_fixed_error second part + if (is_distributed_master()) { + // SpMV is the first operation in compute_interaction_and_movement so we can do halo before and + // call it naturally we then reduce the local dot products + multi_gpu_engine->halo_exchange_cstr( + [](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_reflected_dual(); + }); - if (batch_mode_) - RAFT_CUSPARSE_TRY(cusparseDnMatSetValues(cusparse_view.batch_potential_next_dual_solution, + multi_gpu_engine->for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + auto& sub_cv = sub_pdlp.pdhg_solver_.get_cusparse_view(); + + RAFT_CUSPARSE_TRY( + cusparseDnVecSetValues(sub_cv.potential_next_dual_solution, + (void*)sub_pdlp.pdhg_solver_.get_reflected_dual().data())); + + sub_pdlp.step_size_strategy_.compute_interaction_and_movement( + sub_pdlp.pdhg_solver_.get_primal_tmp_resource(), + sub_cv, + sub_pdlp.pdhg_solver_.get_saddle_point_state(), + shard.rank_data.owned_var_size, + shard.rank_data.owned_cstr_size); + + RAFT_CUSPARSE_TRY(cusparseDnVecSetValues( + sub_cv.potential_next_dual_solution, + (void*)sub_pdlp.pdhg_solver_.get_potential_next_dual_solution().data())); + }); + + multi_gpu_engine->allreduce_sum_inplace_to_master( + [](auto& sp) -> f_t* { return sp.step_size_strategy_.get_interaction().data(); }); + multi_gpu_engine->allreduce_sum_inplace_to_master([](auto& sp) -> f_t* { + return sp.step_size_strategy_.get_norm_squared_delta_primal().data(); + }); + multi_gpu_engine->allreduce_sum_inplace_to_master( + [](auto& sp) -> f_t* { return sp.step_size_strategy_.get_norm_squared_delta_dual().data(); }); + } else { + // Sync to make sure all previous cuSparse operations are finished before setting the + // potential_next_dual_solution + RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); + + // Make potential_next_dual_solution point towards reflected dual solution to reuse the code + RAFT_CUSPARSE_TRY(cusparseDnVecSetValues(cusparse_view.potential_next_dual_solution, (void*)pdhg_solver_.get_reflected_dual().data())); - step_size_strategy_.compute_interaction_and_movement( - pdhg_solver_.get_primal_tmp_resource(), cusparse_view, pdhg_solver_.get_saddle_point_state()); + if (batch_mode_) + RAFT_CUSPARSE_TRY(cusparseDnMatSetValues(cusparse_view.batch_potential_next_dual_solution, + (void*)pdhg_solver_.get_reflected_dual().data())); + + step_size_strategy_.compute_interaction_and_movement( + pdhg_solver_.get_primal_tmp_resource(), cusparse_view, pdhg_solver_.get_saddle_point_state()); + } if (batch_mode_) { const auto [grid_size, block_size] = kernel_config_from_batch_size(climber_strategies_.size()); @@ -2066,11 +2368,12 @@ void pdlp_solver_t::compute_fixed_error(std::vector& has_restarte // potential_next_dual_solution RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); - // Put back - RAFT_CUSPARSE_TRY( - cusparseDnVecSetValues(cusparse_view.potential_next_dual_solution, - (void*)pdhg_solver_.get_potential_next_dual_solution().data())); - + // Put back, already done in multi-gpu side + if (!is_distributed_master()) { + RAFT_CUSPARSE_TRY( + cusparseDnVecSetValues(cusparse_view.potential_next_dual_solution, + (void*)pdhg_solver_.get_potential_next_dual_solution().data())); + } if (batch_mode_) { RAFT_CUSPARSE_TRY( cusparseDnMatSetValues(cusparse_view.batch_potential_next_dual_solution, @@ -2309,7 +2612,11 @@ optimization_problem_solution_t pdlp_solver_t::run_solver(co std::cout << "Starting PDLP loop:" << std::endl; #endif - // TODO handle that properly + bool warm_start_was_given = settings_.get_pdlp_warm_start_data().is_populated(); + + // The four setup calls (compute_initial_step_size, compute_initial_primal_weight, + // scale_problem, create_spmv_op_plans) run unconditionally here. Each of them + // branches on is_distributed_master() at entry. if (settings_.hyper_params.compute_initial_step_size_before_scaling && !settings_.get_initial_step_size().has_value()) compute_initial_step_size(); @@ -2317,24 +2624,24 @@ optimization_problem_solution_t pdlp_solver_t::run_solver(co !settings_.get_initial_primal_weight().has_value()) compute_initial_primal_weight(); - initial_scaling_strategy_.scale_problem(); - if constexpr (std::is_same_v) { - if (!batch_mode_ && !pdhg_solver_.get_cusparse_view().mixed_precision_enabled_) { - pdhg_solver_.get_cusparse_view().create_spmv_op_plans( - settings_.hyper_params.use_reflected_primal_dual); - } - } + scale_problem(); + create_spmv_op_plans(); - // Update FP32 matrix copies for mixed precision SpMV after scaling - pdhg_solver_.get_cusparse_view().update_mixed_precision_matrices(); + // mixed precision and cusparse structure redirection are not supported in distributed + // as memory footprint is not currently a bottleneck in distributed + if (!settings_.use_distributed_pdlp) { + // Update FP32 matrix copies for mixed precision SpMV after scaling + pdhg_solver_.get_cusparse_view().update_mixed_precision_matrices(); - // Redirect cuSPARSE descriptors to use the original problem's structural data (offsets, indices), - // then free the duplicated structural vectors from the scaled copy to save device memory. - pdhg_solver_.get_cusparse_view().redirect_cusparse_csr_structure_pointers(*problem_ptr); - op_problem_scaled_.variables.resize(0, stream_view_); - op_problem_scaled_.offsets.resize(0, stream_view_); - op_problem_scaled_.reverse_constraints.resize(0, stream_view_); - op_problem_scaled_.reverse_offsets.resize(0, stream_view_); + // Redirect cuSPARSE descriptors to use the original problem's structural data (offsets, + // indices), then free the duplicated structural vectors from the scaled copy to save device + // memory. + pdhg_solver_.get_cusparse_view().redirect_cusparse_csr_structure_pointers(*problem_ptr); + op_problem_scaled_.variables.resize(0, stream_view_); + op_problem_scaled_.offsets.resize(0, stream_view_); + op_problem_scaled_.reverse_constraints.resize(0, stream_view_); + op_problem_scaled_.reverse_offsets.resize(0, stream_view_); + } if (!settings_.hyper_params.compute_initial_step_size_before_scaling && !settings_.get_initial_step_size().has_value()) @@ -2343,211 +2650,230 @@ optimization_problem_solution_t pdlp_solver_t::run_solver(co !settings_.get_initial_primal_weight().has_value()) compute_initial_primal_weight(); + // Distributed counterpart of the single-GPU + // step_size_strategy_.get_primal_and_dual_stepsizes() + if (settings_.use_distributed_pdlp) { + step_size_strategy_.get_primal_and_dual_stepsizes(primal_step_size_, dual_step_size_); + multi_gpu_engine->for_each_shard([&](auto& shard) { + auto& sub = *shard.sub_pdlp; + sub.step_size_strategy_.get_primal_and_dual_stepsizes(sub.primal_step_size_, + sub.dual_step_size_); + }); + multi_gpu_engine->sync_await_shards(stream_view_); + handle_ptr_->sync_stream(stream_view_); + } + + // Everything below (seed-from-settings, initial_k, get_primal_and_dual_stepsizes, + // initial primal/dual, projection, transpose, verbose prints, log header) + // still runs single-GPU only. Distributed rejects + // has_initial_{primal,dual}_solution() and warm-start data up front, and + // its per-shard primal/dual step sizes were derived above + if (!settings_.use_distributed_pdlp) { #ifdef PDLP_DEBUG_MODE - std::cout << "Initial Scaling done" << std::endl; + std::cout << "Initial Scaling done" << std::endl; #endif - - // Needs to be performed here before the below line to make sure the initial primal_weight / step - // size are used as previous point when potentially updating them in this next call - if (settings_.get_initial_step_size().has_value() || initial_step_size_.has_value()) { - if (initial_step_size_.has_value()) - thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), - step_size_.begin(), - step_size_.end(), - initial_step_size_.value()); - else - thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), - step_size_.begin(), - step_size_.end(), - settings_.get_initial_step_size().value()); - } - if (settings_.get_initial_primal_weight().has_value() || initial_primal_weight_.has_value()) { - if (initial_primal_weight_.has_value()) { - thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), - primal_weight_.begin(), - primal_weight_.end(), - initial_primal_weight_.value()); - if (is_cupdlpx_restart(settings_.hyper_params)) + // Needs to be performed here before the below line to make sure the initial primal_weight / + // step size are used as previous point when potentially updating them in this next call + if (settings_.get_initial_step_size().has_value() || initial_step_size_.has_value()) { + if (initial_step_size_.has_value()) + thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), + step_size_.begin(), + step_size_.end(), + initial_step_size_.value()); + else + thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), + step_size_.begin(), + step_size_.end(), + settings_.get_initial_step_size().value()); + } + if (settings_.get_initial_primal_weight().has_value() || initial_primal_weight_.has_value()) { + if (initial_primal_weight_.has_value()) { thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), - best_primal_weight_.begin(), - best_primal_weight_.end(), + primal_weight_.begin(), + primal_weight_.end(), initial_primal_weight_.value()); - } else { - thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), - primal_weight_.begin(), - primal_weight_.end(), - settings_.get_initial_primal_weight().value()); - if (is_cupdlpx_restart(settings_.hyper_params)) + if (is_cupdlpx_restart(settings_.hyper_params)) + thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), + best_primal_weight_.begin(), + best_primal_weight_.end(), + initial_primal_weight_.value()); + } else { thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), - best_primal_weight_.begin(), - best_primal_weight_.end(), + primal_weight_.begin(), + primal_weight_.end(), settings_.get_initial_primal_weight().value()); + if (is_cupdlpx_restart(settings_.hyper_params)) + thrust::uninitialized_fill(handle_ptr_->get_thrust_policy(), + best_primal_weight_.begin(), + best_primal_weight_.end(), + settings_.get_initial_primal_weight().value()); + } + } + if (initial_k_.has_value()) { + pdhg_solver_.total_pdhg_iterations_ = initial_k_.value(); + pdhg_solver_.get_d_total_pdhg_iterations().set_value_async(initial_k_.value(), stream_view_); + } + if (settings_.get_initial_pdlp_iteration().has_value()) { + total_pdlp_iterations_ = settings_.get_initial_pdlp_iteration().value(); + // This is meaningless in batch mode since pdhg step is never used, set it just to avoid + // assertions + pdhg_solver_.get_d_total_pdhg_iterations().set_value_async(total_pdlp_iterations_, + stream_view_); + pdhg_solver_.total_pdhg_iterations_ = total_pdlp_iterations_; + // Reset the fixed point error since at this pdlp iteration it is expected to already be + // initialized to some value + std::fill(restart_strategy_.initial_fixed_point_error_.begin(), + restart_strategy_.initial_fixed_point_error_.end(), + f_t(0.0)); + std::fill(restart_strategy_.fixed_point_error_.begin(), + restart_strategy_.fixed_point_error_.end(), + f_t(0.0)); } - } - if (initial_k_.has_value()) { - pdhg_solver_.total_pdhg_iterations_ = initial_k_.value(); - pdhg_solver_.get_d_total_pdhg_iterations().set_value_async(initial_k_.value(), stream_view_); - } - if (settings_.get_initial_pdlp_iteration().has_value()) { - total_pdlp_iterations_ = settings_.get_initial_pdlp_iteration().value(); - // This is meaningless in batch mode since pdhg step is never used, set it just to avoid - // assertions - pdhg_solver_.get_d_total_pdhg_iterations().set_value_async(total_pdlp_iterations_, - stream_view_); - pdhg_solver_.total_pdhg_iterations_ = total_pdlp_iterations_; - // Reset the fixed point error since at this pdlp iteration it is expected to already be - // initialized to some value - std::fill(restart_strategy_.initial_fixed_point_error_.begin(), - restart_strategy_.initial_fixed_point_error_.end(), - f_t(0.0)); - std::fill(restart_strategy_.fixed_point_error_.begin(), - restart_strategy_.fixed_point_error_.end(), - f_t(0.0)); - } - - // Only the primal_weight_ and step_size_ variables are initialized during the initial phase - // The associated primal/dual step_size (computed using the two firstly mentionned) are not - // initialized. This calls ensures the latter - // In the event of a given primal and dual solutions and if the option is toggled, calling the - // update primal_weight and step_size will also update the associated primal_step_size_, - // dual_step_size_. - // In summary: the below call is only mandatory at the beginning when - // computing/setting the initial primal weight and step size and if they are not recomputed later. - step_size_strategy_.get_primal_and_dual_stepsizes(primal_step_size_, dual_step_size_); + + // Only the primal_weight_ and step_size_ variables are initialized during the initial phase + // The associated primal/dual step_size (computed using the two firstly mentionned) are not + // initialized. This calls ensures the latter + // In the event of a given primal and dual solutions and if the option is toggled, calling the + // update primal_weight and step_size will also update the associated primal_step_size_, + // dual_step_size_. + // In summary: the below call is only mandatory at the beginning when + // computing/setting the initial primal weight and step size and if they are not recomputed + // later. + step_size_strategy_.get_primal_and_dual_stepsizes(primal_step_size_, dual_step_size_); #ifdef CUPDLP_DEBUG_MODE - if (initial_primal_.size() != 0 || initial_dual_.size() != 0) { - std::cout << "Initial primal and dual solution before scaling" << std::endl; - if (initial_primal_.size() != 0) { print("initial_primal_", initial_primal_); } - if (initial_dual_.size() != 0) { print("initial_dual_", initial_dual_); } - } + if (initial_primal_.size() != 0 || initial_dual_.size() != 0) { + std::cout << "Initial primal and dual solution before scaling" << std::endl; + if (initial_primal_.size() != 0) { print("initial_primal_", initial_primal_); } + if (initial_dual_.size() != 0) { print("initial_dual_", initial_dual_); } + } #endif - // If there is an initial primal or dual we should update the restart info as if there was a step - // that has happend - if (initial_primal_.size() != 0 || initial_dual_.size() != 0) { - update_primal_dual_solutions( - (initial_primal_.size() != 0) ? std::make_optional(&initial_primal_) : std::nullopt, - (initial_dual_.size() != 0) ? std::make_optional(&initial_dual_) : std::nullopt); - } + // If there is an initial primal or dual we should update the restart info as if there was a + // step that has happend + if (initial_primal_.size() != 0 || initial_dual_.size() != 0) { + update_primal_dual_solutions( + (initial_primal_.size() != 0) ? std::make_optional(&initial_primal_) : std::nullopt, + (initial_dual_.size() != 0) ? std::make_optional(&initial_dual_) : std::nullopt); + } #ifdef CUPDLP_DEBUG_MODE - std::cout << "Solution before projection" << std::endl; - print("pdhg_solver_.get_primal_solution()", pdhg_solver_.get_primal_solution()); - print("pdhg_solver_.get_dual_solution()", pdhg_solver_.get_dual_solution()); - print("pdhg_solver_.get_potential_next_primal_solution()", - pdhg_solver_.get_potential_next_primal_solution()); - print("pdhg_solver_.get_potential_next_dual_solution()", - pdhg_solver_.get_potential_next_dual_solution()); - print("restart_strategy_.last_restart_duality_gap_.primal_solution_", - restart_strategy_.last_restart_duality_gap_.primal_solution_); - print("restart_strategy_.last_restart_duality_gap_.dual_solution_", - restart_strategy_.last_restart_duality_gap_.dual_solution_); + std::cout << "Solution before projection" << std::endl; + print("pdhg_solver_.get_primal_solution()", pdhg_solver_.get_primal_solution()); + print("pdhg_solver_.get_dual_solution()", pdhg_solver_.get_dual_solution()); + print("pdhg_solver_.get_potential_next_primal_solution()", + pdhg_solver_.get_potential_next_primal_solution()); + print("pdhg_solver_.get_potential_next_dual_solution()", + pdhg_solver_.get_potential_next_dual_solution()); + print("restart_strategy_.last_restart_duality_gap_.primal_solution_", + restart_strategy_.last_restart_duality_gap_.primal_solution_); + print("restart_strategy_.last_restart_duality_gap_.dual_solution_", + restart_strategy_.last_restart_duality_gap_.dual_solution_); #endif - // Project initial primal solution - if (settings_.hyper_params.project_initial_primal) { - using f_t2 = typename type_2::type; - if (batch_mode_) { - // In batch mode variable_bounds are shared and only the bound rescaling is per climber. - // Apply it here too so the initial point is projected into the correct saacled space - cub::DeviceTransform::Transform( - cuda::std::make_tuple( + // Project initial primal solution + if (settings_.hyper_params.project_initial_primal) { + using f_t2 = typename type_2::type; + if (batch_mode_) { + // In batch mode variable_bounds are shared and only the bound rescaling is per climber. + // Apply it here too so the initial point is projected into the correct scaled space. + cub::DeviceTransform::Transform( + cuda::std::make_tuple( + pdhg_solver_.get_primal_solution().data(), + thrust::make_transform_iterator( + thrust::make_zip_iterator( + problem_wrap_container(op_problem_scaled_.variable_bounds), + batch_wrapped_container(initial_scaling_strategy_.get_bound_rescaling_vector(), + primal_size_h_)), + scale_bounds_by_scalar_op{})), pdhg_solver_.get_primal_solution().data(), - thrust::make_transform_iterator( - thrust::make_zip_iterator( - problem_wrap_container(op_problem_scaled_.variable_bounds), - batch_wrapped_container(initial_scaling_strategy_.get_bound_rescaling_vector(), - primal_size_h_)), - scale_bounds_by_scalar_op{})), - pdhg_solver_.get_primal_solution().data(), - pdhg_solver_.get_primal_solution().size(), - clamp(), - stream_view_.value()); - } else { - cub::DeviceTransform::Transform( - cuda::std::make_tuple(pdhg_solver_.get_primal_solution().data(), - problem_wrap_container(op_problem_scaled_.variable_bounds)), - pdhg_solver_.get_primal_solution().data(), - pdhg_solver_.get_primal_solution().size(), - clamp(), - stream_view_.value()); - } + pdhg_solver_.get_primal_solution().size(), + clamp(), + stream_view_.value()); + } else { + cub::DeviceTransform::Transform( + cuda::std::make_tuple(pdhg_solver_.get_primal_solution().data(), + problem_wrap_container(op_problem_scaled_.variable_bounds)), + pdhg_solver_.get_primal_solution().data(), + pdhg_solver_.get_primal_solution().size(), + clamp(), + stream_view_.value()); + } + + pdhg_solver_.refine_initial_primal_projection( + initial_scaling_strategy_.get_bound_rescaling_vector()); - pdhg_solver_.refine_initial_primal_projection( - initial_scaling_strategy_.get_bound_rescaling_vector()); - - if (!settings_.hyper_params.never_restart_to_average) { - cuopt_expects(!batch_mode_, - cuopt::error_type_t::ValidationError, - "Restart to average not supported in batch mode"); - cub::DeviceTransform::Transform( - cuda::std::make_tuple(unscaled_primal_avg_solution_.data(), - op_problem_scaled_.variable_bounds.data()), - unscaled_primal_avg_solution_.data(), - primal_size_h_, - clamp(), - stream_view_.value()); + if (!settings_.hyper_params.never_restart_to_average) { + cuopt_expects(!batch_mode_, + cuopt::error_type_t::ValidationError, + "Restart to average not supported in batch mode"); + cub::DeviceTransform::Transform( + cuda::std::make_tuple(unscaled_primal_avg_solution_.data(), + op_problem_scaled_.variable_bounds.data()), + unscaled_primal_avg_solution_.data(), + primal_size_h_, + clamp(), + stream_view_.value()); + } } - } #ifdef CUPDLP_DEBUG_MODE - std::cout << "Solution after projection" << std::endl; - print("pdhg_solver_.get_primal_solution()", pdhg_solver_.get_primal_solution()); - print("pdhg_solver_.get_dual_solution()", pdhg_solver_.get_dual_solution()); - print("pdhg_solver_.get_potential_next_primal_solution()", - pdhg_solver_.get_potential_next_primal_solution()); - print("pdhg_solver_.get_potential_next_dual_solution()", - pdhg_solver_.get_potential_next_dual_solution()); - print("restart_strategy_.last_restart_duality_gap_.primal_solution_", - restart_strategy_.last_restart_duality_gap_.primal_solution_); - print("restart_strategy_.last_restart_duality_gap_.dual_solution_", - restart_strategy_.last_restart_duality_gap_.dual_solution_); + std::cout << "Solution after projection" << std::endl; + print("pdhg_solver_.get_primal_solution()", pdhg_solver_.get_primal_solution()); + print("pdhg_solver_.get_dual_solution()", pdhg_solver_.get_dual_solution()); + print("pdhg_solver_.get_potential_next_primal_solution()", + pdhg_solver_.get_potential_next_primal_solution()); + print("pdhg_solver_.get_potential_next_dual_solution()", + pdhg_solver_.get_potential_next_dual_solution()); + print("restart_strategy_.last_restart_duality_gap_.primal_solution_", + restart_strategy_.last_restart_duality_gap_.primal_solution_); + print("restart_strategy_.last_restart_duality_gap_.dual_solution_", + restart_strategy_.last_restart_duality_gap_.dual_solution_); #endif - // Need to to tranpose primal solution to row format as there might be initial values or clamping - // Value may not be all 0 - if (batch_mode_) { - rmm::device_uvector dummy(0, stream_view_); - transpose_primal_dual_to_row( - pdhg_solver_.get_primal_solution(), pdhg_solver_.get_dual_solution(), dummy); - if (settings_.hyper_params.use_reflected_primal_dual) { - transpose_primal_dual_to_row(pdhg_solver_.get_potential_next_primal_solution(), - pdhg_solver_.get_potential_next_dual_solution(), - dummy); - transpose_primal_dual_to_row(restart_strategy_.last_restart_duality_gap_.primal_solution_, - restart_strategy_.last_restart_duality_gap_.dual_solution_, - dummy); + // Need to to tranpose primal solution to row format as there might be initial values or + // clamping Value may not be all 0 + if (batch_mode_) { + rmm::device_uvector dummy(0, stream_view_); + transpose_primal_dual_to_row( + pdhg_solver_.get_primal_solution(), pdhg_solver_.get_dual_solution(), dummy); + if (settings_.hyper_params.use_reflected_primal_dual) { + transpose_primal_dual_to_row(pdhg_solver_.get_potential_next_primal_solution(), + pdhg_solver_.get_potential_next_dual_solution(), + dummy); + transpose_primal_dual_to_row(restart_strategy_.last_restart_duality_gap_.primal_solution_, + restart_strategy_.last_restart_duality_gap_.dual_solution_, + dummy); + } + transpose_problem_fields(/*to_row=*/true); + } + + if (verbose) { + std::cout << "primal_size_h_ " << primal_size_h_ << " dual_size_h_ " << dual_size_h_ + << " nnz " << problem_ptr->nnz << std::endl; + std::cout << "Problem before scaling" << std::endl; + print_problem_info(problem_ptr->coefficients, + problem_ptr->objective_coefficients, + problem_ptr->combined_bounds); + std::cout << "Problem after scaling" << std::endl; + print_problem_info(op_problem_scaled_.coefficients, + op_problem_scaled_.objective_coefficients, + op_problem_scaled_.combined_bounds); + raft::print_device_vector("Initial step_size", step_size_.data(), 1, std::cout); + raft::print_device_vector("Initial primal_weight", primal_weight_.data(), 1, std::cout); + raft::print_device_vector("Initial primal_step_size", primal_step_size_.data(), 1, std::cout); + raft::print_device_vector("Initial dual_step_size", dual_step_size_.data(), 1, std::cout); } - transpose_problem_fields(/*to_row=*/true); - } - - if (verbose) { - std::cout << "primal_size_h_ " << primal_size_h_ << " dual_size_h_ " << dual_size_h_ << " nnz " - << problem_ptr->nnz << std::endl; - std::cout << "Problem before scaling" << std::endl; - print_problem_info( - problem_ptr->coefficients, problem_ptr->objective_coefficients, problem_ptr->combined_bounds); - std::cout << "Problem after scaling" << std::endl; - print_problem_info(op_problem_scaled_.coefficients, - op_problem_scaled_.objective_coefficients, - op_problem_scaled_.combined_bounds); - raft::print_device_vector("Initial step_size", step_size_.data(), 1, std::cout); - raft::print_device_vector("Initial primal_weight", primal_weight_.data(), 1, std::cout); - raft::print_device_vector("Initial primal_step_size", primal_step_size_.data(), 1, std::cout); - raft::print_device_vector("Initial dual_step_size", dual_step_size_.data(), 1, std::cout); - } #ifdef CUPDLP_DEBUG_MODE - raft::print_device_vector("Initial step_size", step_size_.data(), step_size_.size(), std::cout); - raft::print_device_vector( - "Initial primal_weight", primal_weight_.data(), primal_weight_.size(), std::cout); + raft::print_device_vector("Initial step_size", step_size_.data(), step_size_.size(), std::cout); + raft::print_device_vector( + "Initial primal_weight", primal_weight_.data(), primal_weight_.size(), std::cout); #endif - bool warm_start_was_given = settings_.get_pdlp_warm_start_data().is_populated(); - - if (!inside_mip_) { - CUOPT_LOG_INFO( - " Iter Primal Obj. Dual Obj. Gap Primal Res. Dual Res. Time"); + if (!inside_mip_) { + CUOPT_LOG_INFO( + " Iter Primal Obj. Dual Obj. Gap Primal Res. Dual Res. Time"); + } } while (true) { #ifdef CUPDLP_DEBUG_MODE @@ -2646,10 +2972,21 @@ optimization_problem_solution_t pdlp_solver_t::run_solver(co initial_scaling_strategy_.unscale_solutions(pdhg_solver_.get_primal_solution(), pdhg_solver_.get_dual_solution()); } else { - initial_scaling_strategy_.unscale_solutions( - pdhg_solver_.get_potential_next_primal_solution(), - pdhg_solver_.get_potential_next_dual_solution(), - pdhg_solver_.get_dual_slack()); + if (is_distributed_master()) { + // The only branch in cuPDLPx (Stable3) + multi_gpu_engine->for_each_shard([&](auto& shard) { + auto& sub = *shard.sub_pdlp; + sub.get_initial_scaling_strategy().unscale_solutions( + sub.pdhg_solver_.get_potential_next_primal_solution(), + sub.pdhg_solver_.get_potential_next_dual_solution(), + sub.pdhg_solver_.get_dual_slack()); + }); + } else { + initial_scaling_strategy_.unscale_solutions( + pdhg_solver_.get_potential_next_primal_solution(), + pdhg_solver_.get_potential_next_dual_solution(), + pdhg_solver_.get_dual_slack()); + } } #ifdef CUPDLP_DEBUG_MODE @@ -2681,10 +3018,21 @@ optimization_problem_solution_t pdlp_solver_t::run_solver(co initial_scaling_strategy_.scale_solutions(pdhg_solver_.get_primal_solution(), pdhg_solver_.get_dual_solution()); } else { - initial_scaling_strategy_.scale_solutions( - pdhg_solver_.get_potential_next_primal_solution(), - pdhg_solver_.get_potential_next_dual_solution(), - pdhg_solver_.get_dual_slack()); + if (is_distributed_master()) { + // The only branch in cuPDLPx (Stable3) + multi_gpu_engine->for_each_shard([&](auto& shard) { + auto& sub = *shard.sub_pdlp; + sub.get_initial_scaling_strategy().scale_solutions( + sub.pdhg_solver_.get_potential_next_primal_solution(), + sub.pdhg_solver_.get_potential_next_dual_solution(), + sub.pdhg_solver_.get_dual_slack()); + }); + } else { + initial_scaling_strategy_.scale_solutions( + pdhg_solver_.get_potential_next_primal_solution(), + pdhg_solver_.get_potential_next_dual_solution(), + pdhg_solver_.get_dual_slack()); + } } } @@ -2782,13 +3130,23 @@ optimization_problem_solution_t pdlp_solver_t::run_solver(co transpose_problem_fields(/*to_row=*/true); } } - halpern_update(); + if (is_distributed_master()) { + multi_gpu_engine->for_each_shard([&](auto& shard) { shard.sub_pdlp->halpern_update(); }); + } else { + halpern_update(); + } } ++total_pdlp_iterations_; ++internal_solver_iterations_; - if (settings_.hyper_params.never_restart_to_average) + if (settings_.hyper_params.never_restart_to_average) { restart_strategy_.increment_iteration_since_last_restart(); + if (is_distributed_master()) { + multi_gpu_engine->for_each_shard([&](auto& shard) { + shard.sub_pdlp->restart_strategy_.increment_iteration_since_last_restart(); + }); + } + } } return optimization_problem_solution_t{pdlp_termination_status_t::NumericalError, stream_view_}; @@ -2928,11 +3286,72 @@ void pdlp_solver_t::take_step([[maybe_unused]] i_t total_pdlp_iteratio // print("potential next dual", pdhg_solver_.get_potential_next_dual_solution()); } +template +void pdlp_solver_t::scale_problem() +{ + // Scale problem then free scratch buffers + raft::common::nvtx::range fun_scope("pdlp_solver_t::scale_problem"); + if (is_distributed_master()) { + multi_gpu_engine->distributed_scaling(settings_.hyper_params, primal_size_h_, inside_mip_); + + // Free per-shard scratch: no further scaling passes happen after this point. + multi_gpu_engine->for_each_shard([](auto& shard) { + auto& scaling = shard.sub_pdlp->get_initial_scaling_strategy(); + scaling.get_iteration_variable_scaling().resize(0, shard.stream.view()); + scaling.get_iteration_constraint_matrix_scaling().resize(0, shard.stream.view()); + }); + } else { + initial_scaling_strategy_.scale_problem(); + + // Free scratch: no further scaling passes happen after this point. + initial_scaling_strategy_.get_iteration_variable_scaling().resize(0, stream_view_); + initial_scaling_strategy_.get_iteration_constraint_matrix_scaling().resize(0, stream_view_); + } +} + +template +void pdlp_solver_t::create_spmv_op_plans() +{ + raft::common::nvtx::range fun_scope("pdlp_solver_t::create_spmv_op_plans"); + if (is_distributed_master()) { + // Distributed path: fan out the same per-shard cusparse_view call the + // single-GPU path would make. + multi_gpu_engine->for_each_shard([&](auto& shard) { + shard.sub_pdlp->pdhg_solver_.get_cusparse_view().create_spmv_op_plans( + settings_.hyper_params.use_reflected_primal_dual); + }); + multi_gpu_engine->sync_await_shards(stream_view_); + handle_ptr_->sync_stream(stream_view_); + return; + } + if constexpr (std::is_same_v) { + if (!batch_mode_ && !pdhg_solver_.get_cusparse_view().mixed_precision_enabled_) { + pdhg_solver_.get_cusparse_view().create_spmv_op_plans( + settings_.hyper_params.use_reflected_primal_dual); + } + } +} + template void pdlp_solver_t::compute_initial_step_size() { raft::common::nvtx::range fun_scope("compute_initial_step_size"); + // Shared knobs for the power-iteration path (both single-GPU and + // distributed) + constexpr f_t scaling_factor = f_t{0.998}; + constexpr int max_iterations = 5000; + constexpr f_t tolerance = f_t{1e-4}; + + if (is_distributed_master()) { + // Distributed dispatch: everything (sigma_max, deriving primal/dual + // step sizes from master's current primal_weight_, seeding master + + // all shards, syncs) lives inside distributed_compute_initial_step_size. + multi_gpu_engine->distributed_compute_initial_step_size( + settings_.hyper_params, dual_size_h_, scaling_factor, max_iterations, tolerance); + return; + } + if (!settings_.hyper_params.initial_step_size_max_singular_value) { // set stepsize relative to maximum absolute value of A rmm::device_scalar abs_max_element{0.0, stream_view_}; @@ -2965,23 +3384,14 @@ void pdlp_solver_t::compute_initial_step_size() // Sync since we are using local variable RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); } else { - constexpr i_t max_iterations = 5000; - constexpr f_t tolerance = 1e-4; - i_t m = op_problem_scaled_.n_constraints; i_t n = op_problem_scaled_.n_variables; - std::vector z(m); + std::vector z = make_singular_value_probe(static_cast(m)); rmm::device_uvector d_z(m, stream_view_); rmm::device_uvector d_q(m, stream_view_); rmm::device_uvector d_atq(n, stream_view_); - std::mt19937 gen(1); - std::normal_distribution dist(f_t(0.0), f_t(1.0)); - - for (int i = 0; i < m; ++i) - z[i] = dist(gen); - device_copy(d_z, z, stream_view_); rmm::device_scalar norm_q(stream_view_); @@ -3011,12 +3421,11 @@ void pdlp_solver_t::compute_initial_step_size() cuopt_assert(norm_q.value(stream_view_) != f_t(0), "norm q can't be 0"); // d_q *= 1 / norm_q - cub::DeviceTransform::Transform( - d_q.data(), - d_q.data(), - d_q.size(), - [norm_q = norm_q.data()] __device__(f_t d_q) { return d_q / *norm_q; }, - stream_view_.value()); + cub::DeviceTransform::Transform(d_q.data(), + d_q.data(), + d_q.size(), + divide_by_device_scalar_t{norm_q.data()}, + stream_view_.value()); // A_t_q = A_t @ d_q RAFT_CUSPARSE_TRY( @@ -3053,14 +3462,12 @@ void pdlp_solver_t::compute_initial_step_size() sigma_max_sq.data(), stream_view_.value())); - cub::DeviceTransform::Transform( - cuda::std::make_tuple(d_q.data(), d_z.data()), - d_q.data(), - d_q.size(), - [sigma_max_sq = sigma_max_sq.data()] __device__(f_t d_q, f_t d_z) { - return d_q * -(*sigma_max_sq) + d_z; - }, - stream_view_.value()); + // d_q := -sigma_max_sq * d_q + d_z + cub::DeviceTransform::Transform(cuda::std::make_tuple(d_q.data(), d_z.data()), + d_q.data(), + d_q.size(), + residual_fma_neg_scalar_t{sigma_max_sq.data()}, + stream_view_.value()); my_l2_norm(d_q, residual_norm, handle_ptr_); @@ -3070,8 +3477,7 @@ void pdlp_solver_t::compute_initial_step_size() printf("iter_count %d\n", sing_iters); #endif - constexpr f_t scaling_factor = 0.998; - const f_t step_size = scaling_factor / std::sqrt(sigma_max_sq.value(stream_view_)); + const f_t step_size = scaling_factor / std::sqrt(sigma_max_sq.value(stream_view_)); thrust::uninitialized_fill( handle_ptr_->get_thrust_policy(), step_size_.begin(), step_size_.end(), step_size); @@ -3121,6 +3527,15 @@ void pdlp_solver_t::compute_initial_primal_weight() { raft::common::nvtx::range fun_scope("compute_initial_primal_weight"); + if (is_distributed_master()) { + // Distributed dispatch: + // - short-circuit -> 1 + // - primal/dual step sizes from master's current step_size_, seeding + // master + all shards, syncs) + multi_gpu_engine->distributed_compute_initial_primal_weight(settings_.hyper_params); + return; + } + // Here we use the combined bounds of the op_problem_scaled which may or may not be scaled yet // based on pdlp config pdlp::combine_constraint_bounds(op_problem_scaled_, op_problem_scaled_.combined_bounds); diff --git a/cpp/src/pdlp/pdlp.cuh b/cpp/src/pdlp/pdlp.cuh index 3735db0cd4..492f03c610 100644 --- a/cpp/src/pdlp/pdlp.cuh +++ b/cpp/src/pdlp/pdlp.cuh @@ -8,10 +8,12 @@ #pragma once #include +#include #include #include #include +#include #include #include #include @@ -57,10 +59,18 @@ class pdlp_solver_t { * * @param[in] op_problem An mip::problem_t object with a * representation of a linear program + * @param[in] is_distributed_sub_pdlp true when constructed as a shard of a + * distributed solve. */ pdlp_solver_t(mip::problem_t& op_problem, pdlp_solver_settings_t const& settings, - bool is_batch_mode = false); + bool is_batch_mode = false, + bool is_distributed_sub_pdlp = false); + + // Distributed Solver Constructor + pdlp_solver_t(mip::problem_t& placeholder_problem, + cuopt::mathematical_optimization::io::mps_data_model_t const& mps, + pdlp_solver_settings_t const& settings); optimization_problem_solution_t run_solver(const timer_t& timer); @@ -97,6 +107,31 @@ class pdlp_solver_t { void compute_initial_step_size(); void compute_initial_primal_weight(); + // Thin dispatch wrappers used by run_solver so single-GPU and distributed + // callers hit the same call sites. + void scale_problem(); + void create_spmv_op_plans(); + + // Needed by multi-GPU to mutate them + mip::problem_t& get_op_problem_scaled() { return op_problem_scaled_; } + pdlp_initial_scaling_strategy_t& get_initial_scaling_strategy() + { + return initial_scaling_strategy_; + } + pdlp_restart_strategy_t& get_restart_strategy() { return restart_strategy_; } + + // Per-shard primal/dual step sizes are private state on pdlp_solver_t but + // are needed inside the multi-GPU dispatch paths that fan out a master cub + // call across all shards' pdhg_solver_t::*_transform methods. + rmm::device_uvector& get_primal_step_size() { return primal_step_size_; } + rmm::device_uvector& get_dual_step_size() { return dual_step_size_; } + // Multi-GPU restart broadcast needs to mirror master's primal_weight / + // best_primal_weight onto every shard after each cuPDLPx restart so that + // downstream shard-side restart machinery stays in sync with master. + rmm::device_uvector& get_primal_weight() { return primal_weight_; } + rmm::device_uvector& get_best_primal_weight() { return best_primal_weight_; } + rmm::device_uvector& get_step_size() { return step_size_; } + raft::handle_t const* get_handle_ptr() const { return handle_ptr_; } private: void print_termination_criteria(const timer_t& timer, bool is_average = false); @@ -184,10 +219,24 @@ class pdlp_solver_t { pdlp::adaptive_step_size_strategy_t step_size_strategy_; public: + // std::optional because multi_gpu_engine_t is non-default-constructible + // (collectively bootstraps NCCL, owns RMM resources). Stays nullopt in + // single-GPU mode; emplaced by the multi-GPU ctor. + std::optional> multi_gpu_engine; + // Inner solver pdlp::pdhg_solver_t pdhg_solver_; void halpern_update(); + // This solver is the distributed-PDLP master orchestrator iff it owns the + // multi-GPU engine. Shards (sub-solvers) leave the optional empty -> false. + // Single-GPU PDLP reports false. + bool is_distributed_master() const { return multi_gpu_engine.has_value(); } + + // True when this pdlp_solver_t was constructed as a per-shard sub-solver + // of a distributed solve (set once at construction via the ctor arg). + bool is_distributed_sub_pdlp() const { return is_distributed_sub_pdlp_; } + private: void compute_fixed_error(std::vector& has_restarted); @@ -246,6 +295,8 @@ class pdlp_solver_t { primal_quality_adapter_t best_primal_quality_so_far_; // Flag to indicate if solver is being called from MIP. No logging is done in this case. bool inside_mip_{false}; + // See is_distributed_sub_pdlp() above. Initialized from ctor arg. + bool is_distributed_sub_pdlp_{false}; }; } // namespace cuopt::mathematical_optimization::pdlp diff --git a/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cu b/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cu index 85bc703a05..5d3258502a 100644 --- a/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cu +++ b/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cu @@ -8,6 +8,8 @@ #include #include +#include +#include #include #include #include @@ -891,21 +893,24 @@ void pdlp_restart_strategy_t::cupdlpx_restart( should_restart.begin(), should_restart.end(), [](int restarted) { return restarted == 1; }), "If any, all should be true"); - // Computing the deltas - distance_squared_moved_from_last_restart_period( - pdhg_solver.get_potential_next_primal_solution(), - last_restart_duality_gap_.primal_solution_, - pdhg_solver.get_primal_tmp_resource(), - primal_size_h_, - 1, - last_restart_duality_gap_.primal_distance_traveled_); - distance_squared_moved_from_last_restart_period( - pdhg_solver.get_potential_next_dual_solution(), - last_restart_duality_gap_.dual_solution_, - pdhg_solver.get_dual_tmp_resource(), - dual_size_h_, - 1, - last_restart_duality_gap_.dual_distance_traveled_); + // Computing the distributed deltas + if (auto* engine = pdhg_solver.get_mgpu_engine()) { + engine->for_each_shard([&](auto& shard) { + auto& sub = *shard.sub_pdlp; + sub.get_restart_strategy().primal_dual_distance_squared_moved_from_last_restart_period( + sub.pdhg_solver_, shard.rank_data.owned_var_size, shard.rank_data.owned_cstr_size); + }); + + engine->allreduce_sum_inplace_to_master([](pdlp_solver_t& sp) -> f_t* { + return sp.get_restart_strategy().last_restart_duality_gap_.primal_distance_traveled_.data(); + }); + engine->allreduce_sum_inplace_to_master([](pdlp_solver_t& sp) -> f_t* { + return sp.get_restart_strategy().last_restart_duality_gap_.dual_distance_traveled_.data(); + }); + } else { + primal_dual_distance_squared_moved_from_last_restart_period( + pdhg_solver, primal_size_h_, dual_size_h_); + } auto view = make_cupdlpx_restart_view(last_restart_duality_gap_.primal_distance_traveled_, last_restart_duality_gap_.dual_distance_traveled_, @@ -952,30 +957,56 @@ void pdlp_restart_strategy_t::cupdlpx_restart( hyper_params_.restart_k_i, hyper_params_.restart_k_d, hyper_params_.restart_i_smooth); - primal_weight.set_element_async(0, primal_weight_value, stream_view_); - primal_step_size.set_element_async(0, primal_step_size_value, stream_view_); - dual_step_size.set_element_async(0, dual_step_size_value, stream_view_); - best_primal_weight.set_element_async(0, best_primal_weight_value, stream_view_); + if (auto* engine = pdhg_solver.get_mgpu_engine()) { + engine->set_scalar_on_master_and_shards( + primal_weight_value, [](auto& sp) { return sp.get_primal_weight().data(); }); + engine->set_scalar_on_master_and_shards( + primal_step_size_value, [](auto& sp) { return sp.get_primal_step_size().data(); }); + engine->set_scalar_on_master_and_shards( + dual_step_size_value, [](auto& sp) { return sp.get_dual_step_size().data(); }); + engine->set_scalar_on_master_and_shards( + best_primal_weight_value, [](auto& sp) { return sp.get_best_primal_weight().data(); }); + } else { + primal_weight.set_element_async(0, primal_weight_value, stream_view_); + primal_step_size.set_element_async(0, primal_step_size_value, stream_view_); + dual_step_size.set_element_async(0, dual_step_size_value, stream_view_); + best_primal_weight.set_element_async(0, best_primal_weight_value, stream_view_); + } } // TODO later batch mode: remove if you have per climber restart - raft::copy(last_restart_duality_gap_.primal_solution_.data(), - pdhg_solver.get_potential_next_primal_solution().data(), - last_restart_duality_gap_.primal_solution_.size(), - stream_view_); - raft::copy(pdhg_solver.get_primal_solution().data(), - pdhg_solver.get_potential_next_primal_solution().data(), - last_restart_duality_gap_.primal_solution_.size(), - stream_view_); - raft::copy(last_restart_duality_gap_.dual_solution_.data(), - pdhg_solver.get_potential_next_dual_solution().data(), - last_restart_duality_gap_.dual_solution_.size(), - stream_view_); - raft::copy(pdhg_solver.get_dual_solution().data(), - pdhg_solver.get_potential_next_dual_solution().data(), - last_restart_duality_gap_.dual_solution_.size(), - stream_view_); + // Small copy helper to use in both single-GPU and distributed paths. + auto commit_potential_next_as_last_restart = [](pdlp_restart_strategy_t& rest, + pdhg_solver_t& solver, + rmm::cuda_stream_view stream) { + raft::copy(rest.last_restart_duality_gap_.primal_solution_.data(), + solver.get_potential_next_primal_solution().data(), + rest.last_restart_duality_gap_.primal_solution_.size(), + stream); + raft::copy(solver.get_primal_solution().data(), + solver.get_potential_next_primal_solution().data(), + solver.get_primal_solution().size(), + stream); + raft::copy(rest.last_restart_duality_gap_.dual_solution_.data(), + solver.get_potential_next_dual_solution().data(), + rest.last_restart_duality_gap_.dual_solution_.size(), + stream); + raft::copy(solver.get_dual_solution().data(), + solver.get_potential_next_dual_solution().data(), + solver.get_dual_solution().size(), + stream); + }; + + if (auto* engine = pdhg_solver.get_mgpu_engine()) { + engine->for_each_shard([&](auto& shard) { + auto& sub = *shard.sub_pdlp; + commit_potential_next_as_last_restart( + sub.get_restart_strategy(), sub.pdhg_solver_, shard.stream.view()); + }); + } else { + commit_potential_next_as_last_restart(*this, pdhg_solver, stream_view_); + } #ifdef CUPDLP_DEBUG_MODE print("New last_restart_duality_gap_.primal_solution_", @@ -990,6 +1021,13 @@ void pdlp_restart_strategy_t::cupdlpx_restart( weighted_average_solution_.iterations_since_last_restart_ = 0; last_trial_fixed_point_error_[i] = std::numeric_limits::infinity(); } + + if (auto* engine = pdhg_solver.get_mgpu_engine()) { + engine->for_each_shard([&](auto& shard) { + shard.sub_pdlp->get_restart_strategy() + .weighted_average_solution_.iterations_since_last_restart_ = 0; + }); + } } template @@ -1265,6 +1303,26 @@ void pdlp_restart_strategy_t::distance_squared_moved_from_last_restart } } +template +void pdlp_restart_strategy_t::primal_dual_distance_squared_moved_from_last_restart_period( + pdhg_solver_t& pdhg_solver, i_t primal_size, i_t dual_size) +{ + distance_squared_moved_from_last_restart_period( + pdhg_solver.get_potential_next_primal_solution(), + last_restart_duality_gap_.primal_solution_, + pdhg_solver.get_primal_tmp_resource(), + primal_size, + 1, + last_restart_duality_gap_.primal_distance_traveled_); + distance_squared_moved_from_last_restart_period( + pdhg_solver.get_potential_next_dual_solution(), + last_restart_duality_gap_.dual_solution_, + pdhg_solver.get_dual_tmp_resource(), + dual_size, + 1, + last_restart_duality_gap_.dual_distance_traveled_); +} + template __global__ void compute_distance_traveled_last_restart_kernel( const typename localized_duality_gap_container_t::view_t duality_gap_view, diff --git a/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cuh b/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cuh index 03f9350314..68ef7503a0 100644 --- a/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cuh +++ b/cpp/src/pdlp/restart_strategy/pdlp_restart_strategy.cuh @@ -234,6 +234,14 @@ class pdlp_restart_strategy_t { i_t stride, rmm::device_uvector& distance_moved); + // Convenience wrapper: applies distance_squared_moved_from_last_restart_period + // to both primal and dual potential_next solutions of pdhg_solver, writing + // into last_restart_duality_gap_.primal_distance_traveled_ and + // last_restart_duality_gap_.dual_distance_traveled_. + // Used for distributed PDLP mirroring clarity + void primal_dual_distance_squared_moved_from_last_restart_period( + pdhg_solver_t& pdhg_solver, i_t primal_size, i_t dual_size); + void compute_primal_gradient(localized_duality_gap_container_t& duality_gap, cusparse_view_t& cusparse_view); diff --git a/cpp/src/pdlp/saddle_point.cu b/cpp/src/pdlp/saddle_point.cu index 9267a7ced6..b92fdc2fb3 100644 --- a/cpp/src/pdlp/saddle_point.cu +++ b/cpp/src/pdlp/saddle_point.cu @@ -37,8 +37,8 @@ saddle_point_state_t::saddle_point_state_t(raft::handle_t const* handl current_AtY_{batch_size * primal_size, handle_ptr->get_stream()}, next_AtY_{batch_size * primal_size, handle_ptr->get_stream()} { - EXE_CUOPT_EXPECTS(primal_size > 0, "Size of the primal problem must be larger than 0"); - EXE_CUOPT_EXPECTS(dual_size > 0, "Size of the dual problem must be larger than 0"); + EXE_CUOPT_EXPECTS(primal_size >= 0, "Size of the primal problem must be non-negative"); + EXE_CUOPT_EXPECTS(dual_size >= 0, "Size of the dual problem must be non-negative"); // Starting from all 0 thrust::fill( diff --git a/cpp/src/pdlp/solve.cu b/cpp/src/pdlp/solve.cu index 25b427fa9a..03fdaadfb5 100644 --- a/cpp/src/pdlp/solve.cu +++ b/cpp/src/pdlp/solve.cu @@ -756,6 +756,12 @@ static optimization_problem_solution_t run_pdlp_solver( const timer_t& timer, bool is_batch_mode) { + cuopt_expects(!settings.use_distributed_pdlp, + error_type_t::ValidationError, + "Distributed PDLP must be entered via solve_lp(mps_data_model, ...) " + "so the master GPU never materializes the full problem. Call sites " + "with a problem_t cannot dispatch to distributed mode."); + if (problem.n_constraints == 0) { CUOPT_LOG_CONDITIONAL_INFO( !settings.inside_mip, @@ -1944,7 +1950,7 @@ optimization_problem_solution_t solve_lp( // Declare result at outer scope so that result.reduced_problem (which may be // referenced by problem.original_problem_ptr) remains alive through the solve. - std::optional> result; + std::optional> result; if (run_presolve) { sort_csr(op_problem); @@ -1954,13 +1960,14 @@ optimization_problem_solution_t solve_lp( const double presolve_time_limit = std::max(1.0, std::min(0.1 * lp_timer.remaining_time(), 60.0)); presolver = std::make_unique>(); - result = presolver->apply(op_problem, - cuopt::mathematical_optimization::problem_category_t::LP, - settings.presolver, - settings.dual_postsolve, - settings.tolerances.absolute_primal_tolerance, - settings.tolerances.relative_primal_tolerance, - presolve_time_limit); + result = presolver->apply_presolve_from_op_problem( + op_problem, + cuopt::mathematical_optimization::problem_category_t::LP, + settings.presolver, + settings.dual_postsolve, + settings.tolerances.absolute_primal_tolerance, + settings.tolerances.relative_primal_tolerance, + presolve_time_limit); if (result->status == mip::third_party_presolve_status_t::INFEASIBLE) { return optimization_problem_solution_t( pdlp_termination_status_t::PrimalInfeasible, op_problem.get_handle_ptr()->get_stream()); @@ -1991,13 +1998,13 @@ optimization_problem_solution_t solve_lp( rmm::device_uvector empty_reduced_costs(0, op_problem.get_handle_ptr()->get_stream()); // Run postsolve to get the full solution - presolver->undo(empty_primal, - empty_dual, - empty_reduced_costs, - cuopt::mathematical_optimization::problem_category_t::LP, - false, // status_to_skip - settings.dual_postsolve, - op_problem.get_handle_ptr()->get_stream()); + presolver->undo_from_device(empty_primal, + empty_dual, + empty_reduced_costs, + cuopt::mathematical_optimization::problem_category_t::LP, + false, // status_to_skip + settings.dual_postsolve, + op_problem.get_handle_ptr()->get_stream()); // Create termination info with the objective from presolve typename optimization_problem_solution_t::additional_termination_information_t @@ -2062,13 +2069,13 @@ optimization_problem_solution_t solve_lp( cuopt::device_copy(solution.get_reduced_cost(), op_problem.get_handle_ptr()->get_stream()); bool status_to_skip = false; - presolver->undo(primal_solution, - dual_solution, - reduced_costs, - cuopt::mathematical_optimization::problem_category_t::LP, - status_to_skip, - settings.dual_postsolve, - op_problem.get_handle_ptr()->get_stream()); + presolver->undo_from_device(primal_solution, + dual_solution, + reduced_costs, + cuopt::mathematical_optimization::problem_category_t::LP, + status_to_skip, + settings.dual_postsolve, + op_problem.get_handle_ptr()->get_stream()); std::vector< typename optimization_problem_solution_t::additional_termination_information_t> @@ -2203,6 +2210,110 @@ mps_data_model_to_optimization_problem( return op_problem; } +template +cuopt::mathematical_optimization::io::mps_data_model_t op_problem_to_mps_data_model( + const optimization_problem_t& op_problem) +{ + raft::common::nvtx::range fun_scope("op_problem -> mps_data_model (D->H)"); + cuopt::mathematical_optimization::io::mps_data_model_t mps; + + mps.set_maximize(op_problem.get_sense()); + mps.set_objective_scaling_factor(op_problem.get_objective_scaling_factor()); + mps.set_objective_offset(op_problem.get_objective_offset()); + + if (!op_problem.get_problem_name().empty()) { + mps.set_problem_name(op_problem.get_problem_name()); + } + if (!op_problem.get_objective_name().empty()) { + mps.set_objective_name(op_problem.get_objective_name()); + } + if (!op_problem.get_variable_names().empty()) { + mps.set_variable_names(op_problem.get_variable_names()); + } + if (!op_problem.get_row_names().empty()) { mps.set_row_names(op_problem.get_row_names()); } + + const auto& d_coefficients = op_problem.get_constraint_matrix_values(); + const auto& d_offsets = op_problem.get_constraint_matrix_offsets(); + const auto& d_indices = op_problem.get_constraint_matrix_indices(); + const auto& d_obj_coeffs = op_problem.get_objective_coefficients(); + const auto& d_var_lb = op_problem.get_variable_lower_bounds(); + const auto& d_var_ub = op_problem.get_variable_upper_bounds(); + const auto& d_bounds = op_problem.get_constraint_bounds(); + const auto& d_row_types = op_problem.get_row_types(); + const auto& d_constr_lb = op_problem.get_constraint_lower_bounds(); + const auto& d_constr_ub = op_problem.get_constraint_upper_bounds(); + const auto& d_var_types = op_problem.get_variable_types(); + + std::vector h_coefficients(d_coefficients.size()); + std::vector h_offsets(d_offsets.size()); + std::vector h_indices(d_indices.size()); + std::vector h_obj_coeffs(d_obj_coeffs.size()); + std::vector h_var_lb(d_var_lb.size()); + std::vector h_var_ub(d_var_ub.size()); + std::vector h_bounds(d_bounds.size()); + std::vector h_row_types(d_row_types.size()); + std::vector h_constr_lb(d_constr_lb.size()); + std::vector h_constr_ub(d_constr_ub.size()); + std::vector h_var_types_enum(d_var_types.size()); + + auto stream = op_problem.get_handle_ptr()->get_stream(); + raft::copy(h_coefficients.data(), d_coefficients.data(), d_coefficients.size(), stream); + raft::copy(h_offsets.data(), d_offsets.data(), d_offsets.size(), stream); + raft::copy(h_indices.data(), d_indices.data(), d_indices.size(), stream); + raft::copy(h_obj_coeffs.data(), d_obj_coeffs.data(), d_obj_coeffs.size(), stream); + raft::copy(h_var_lb.data(), d_var_lb.data(), d_var_lb.size(), stream); + raft::copy(h_var_ub.data(), d_var_ub.data(), d_var_ub.size(), stream); + raft::copy(h_bounds.data(), d_bounds.data(), d_bounds.size(), stream); + raft::copy(h_row_types.data(), d_row_types.data(), d_row_types.size(), stream); + raft::copy(h_constr_lb.data(), d_constr_lb.data(), d_constr_lb.size(), stream); + raft::copy(h_constr_ub.data(), d_constr_ub.data(), d_constr_ub.size(), stream); + raft::copy(h_var_types_enum.data(), d_var_types.data(), d_var_types.size(), stream); + stream.synchronize(); + + if (!h_offsets.empty()) { + mps.set_csr_constraint_matrix( + std::span(h_coefficients.data(), h_coefficients.size()), + std::span(h_indices.data(), h_indices.size()), + std::span(h_offsets.data(), h_offsets.size())); + } else { + // set_csr_constraint_matrix rejects empty offsets — synthesize the [0] + // sentinel that downstream consumers expect for a zero-row problem. + std::vector empty_offsets{0}; + mps.set_csr_constraint_matrix( + {}, {}, std::span(empty_offsets.data(), empty_offsets.size())); + } + + if (!h_obj_coeffs.empty()) { + mps.set_objective_coefficients(std::span(h_obj_coeffs.data(), h_obj_coeffs.size())); + } + if (!h_var_lb.empty()) { + mps.set_variable_lower_bounds(std::span(h_var_lb.data(), h_var_lb.size())); + } + if (!h_var_ub.empty()) { + mps.set_variable_upper_bounds(std::span(h_var_ub.data(), h_var_ub.size())); + } + if (!h_bounds.empty()) { + mps.set_constraint_bounds(std::span(h_bounds.data(), h_bounds.size())); + } + if (!h_row_types.empty()) { + mps.set_row_types(std::span(h_row_types.data(), h_row_types.size())); + } + if (!h_constr_lb.empty()) { + mps.set_constraint_lower_bounds(std::span(h_constr_lb.data(), h_constr_lb.size())); + } + if (!h_constr_ub.empty()) { + mps.set_constraint_upper_bounds(std::span(h_constr_ub.data(), h_constr_ub.size())); + } + if (!h_var_types_enum.empty()) { + std::vector h_var_types_char(h_var_types_enum.size()); + std::transform( + h_var_types_enum.begin(), h_var_types_enum.end(), h_var_types_char.begin(), var_type_to_char); + mps.set_variable_types(h_var_types_char); + } + + return mps; +} + template optimization_problem_solution_t solve_lp( raft::handle_t const* handle_ptr, @@ -2211,8 +2322,275 @@ optimization_problem_solution_t solve_lp( bool problem_checking, bool use_pdlp_solver_mode) { + if (settings.use_distributed_pdlp) { + return solve_lp_distributed_from_mps( + handle_ptr, mps_data_model, settings, problem_checking, use_pdlp_solver_mode); + } auto op_problem = mps_data_model_to_optimization_problem(handle_ptr, mps_data_model); - return solve_lp(op_problem, settings, problem_checking, use_pdlp_solver_mode); + return solve_lp(op_problem, settings, problem_checking, use_pdlp_solver_mode, false); +} + +template +optimization_problem_solution_t solve_lp_distributed_from_mps( + raft::handle_t const* handle_ptr, + const cuopt::mathematical_optimization::io::mps_data_model_t& mps_data_model, + pdlp_solver_settings_t const& settings, + bool problem_checking, + bool use_pdlp_solver_mode) +{ + cuopt_expects(handle_ptr != nullptr, + error_type_t::ValidationError, + "solve_lp_distributed_from_mps: handle_ptr must not be null"); + cuopt_expects(settings.use_distributed_pdlp, + error_type_t::ValidationError, + "solve_lp_distributed_from_mps: settings.use_distributed_pdlp must be true"); + + pdlp_solver_settings_t settings_resolved = settings; + cuopt_expects(settings_resolved.method == method_t::PDLP, + error_type_t::ValidationError, + "Distributed MPS solve currently supports only method_t::PDLP"); + if (use_pdlp_solver_mode) { set_pdlp_solver_mode(settings_resolved); } + + if (settings_resolved.distributed_pdlp_num_gpus == -1) { + settings_resolved.distributed_pdlp_num_gpus = raft::device_setter::get_device_count(); + CUOPT_LOG_INFO( + "solve_lp_distributed_from_mps: distributed_pdlp_num_gpus == -1, auto-detected " + "%d visible CUDA device(s)", + settings_resolved.distributed_pdlp_num_gpus); + } + // PDLP precision validations (mirror the checks in run_pdlp; distributed + // path only supports the default-precision, non-batch double config). + cuopt_expects(settings_resolved.pdlp_precision == pdlp_precision_t::DefaultPrecision, + error_type_t::ValidationError, + "Distributed PDLP only supports DefaultPrecision (double)."); + cuopt_expects(!settings_resolved.inside_mip, + error_type_t::ValidationError, + "Distributed PDLP is not yet supported from inside MIP."); + // Reject initial solution and warm starts as they are not supported yes for distributed PDLP + cuopt_expects(!settings_resolved.has_initial_primal_solution() && + !settings_resolved.has_initial_dual_solution() && + !settings_resolved.get_pdlp_warm_start_data().is_populated(), + error_type_t::ValidationError, + "Distributed PDLP does not support initial primal/dual solutions or warm-start " + "data."); + // save_best_primal_so_far snapshots the current-or-average iterate mid-solve via + // record_best_primal_so_far. That path reads shard-local buffers + // (pdhg_solver_.get_primal_solution() / unscaled_primal_avg_solution_) that are never + // coherently assembled on the master in distributed mode, so the resulting snapshot would + // be silently wrong. Fail loudly instead. + cuopt_expects(!settings_resolved.save_best_primal_so_far, + error_type_t::ValidationError, + "Distributed PDLP does not support save_best_primal_so_far."); + // Distributed PDLP today only supports the Stable3-shaped hyper-param profile: + // - initial_step_size_max_singular_value = true (matches the sigma_max seeding + // driven by distributed_max_singular_value_squared in the setup), + // - initial_primal_weight_combined_bounds = false and bound_objective_rescaling = true + // (this is the profile where single-GPU compute_initial_primal_weight + // short-circuits to primal_weight = 1, which distributed_compute_initial_primal_weight + // mirrors verbatim). + // Any other profile would leave distributed setup silently divergent from + // single-GPU (either the step size or the primal weight would be seeded from + // a different formula than what single-GPU would compute). + // Fail early + cuopt_expects( + settings_resolved.hyper_params.initial_step_size_max_singular_value && + !settings_resolved.hyper_params.initial_primal_weight_combined_bounds && + settings_resolved.hyper_params.bound_objective_rescaling, + error_type_t::ValidationError, + "Distributed PDLP currently only supports the Stable3-shaped hyper-param profile " + "(initial_step_size_max_singular_value=true, initial_primal_weight_combined_bounds=false, " + "bound_objective_rescaling=true). Set pdlp_solver_mode = Stable3 (the default) or adjust " + "the hyper-params to match."); + + init_logger_t log(settings_resolved.log_file, settings_resolved.log_to_console); + print_version_info(); + init_handler(handle_ptr); + + const i_t n_vars = static_cast(mps_data_model.get_objective_coefficients().size()); + const i_t n_cstr = static_cast(mps_data_model.get_constraint_lower_bounds().size()); + const i_t nnz = static_cast(mps_data_model.get_constraint_matrix_values().size()); + CUOPT_LOG_INFO( + "Solving a problem with %d constraints, %d variables (%d integers), and %d " + "nonzeros (distributed mps-direct path)", + n_cstr, + n_vars, + 0, + nnz); + + auto lp_timer = cuopt::timer_t(settings_resolved.time_limit); + + if (settings_resolved.presolver == presolver_t::Default) { + settings_resolved.presolver = presolver_t::PSLP; + CUOPT_LOG_INFO("Using PSLP presolver"); + } + const bool run_presolve = settings_resolved.presolver != presolver_t::None; + + std::unique_ptr> presolver_ptr; + std::optional> host_res; + [[maybe_unused]] double presolve_time = 0.0; + + if (run_presolve) { + // mirroring single-GPU solve.cu + const double presolve_time_limit = + std::max(1.0, std::min(0.1 * lp_timer.remaining_time(), 60.0)); + + presolver_ptr = std::make_unique>(); + host_res = presolver_ptr->apply_presolve_from_mps_data( + mps_data_model, + cuopt::mathematical_optimization::problem_category_t::LP, + settings_resolved.presolver, + settings_resolved.dual_postsolve, + settings_resolved.tolerances.absolute_primal_tolerance, + settings_resolved.tolerances.relative_primal_tolerance, + presolve_time_limit); + + if (host_res->status == mip::third_party_presolve_status_t::INFEASIBLE) { + return optimization_problem_solution_t(pdlp_termination_status_t::PrimalInfeasible, + handle_ptr->get_stream()); + } + if (host_res->status == mip::third_party_presolve_status_t::UNBNDORINFEAS) { + return optimization_problem_solution_t( + pdlp_termination_status_t::UnboundedOrInfeasible, handle_ptr->get_stream()); + } + if (host_res->status == mip::third_party_presolve_status_t::UNBOUNDED) { + return optimization_problem_solution_t(pdlp_termination_status_t::DualInfeasible, + handle_ptr->get_stream()); + } + + // Presolve completely solved the problem. Mirroring single-GPU solve + if (host_res->reduced_problem.get_n_variables() == 0 && + host_res->reduced_problem.get_n_constraints() == 0) { + CUOPT_LOG_INFO("Presolve completely solved the problem"); + presolve_time = lp_timer.elapsed_time(); + CUOPT_LOG_INFO("%s presolve time: %.2fs", + settings_resolved.presolver == presolver_t::PSLP ? "PSLP" : "Papilo", + presolve_time); + + std::vector h_primal, h_dual, h_rc; + presolver_ptr->undo(h_primal, + h_dual, + h_rc, + cuopt::mathematical_optimization::problem_category_t::LP, + /*status_to_skip=*/false, + settings_resolved.dual_postsolve); + + auto primal_uv = cuopt::device_copy(h_primal, handle_ptr->get_stream()); + auto dual_uv = cuopt::device_copy(h_dual, handle_ptr->get_stream()); + auto rc_uv = cuopt::device_copy(h_rc, handle_ptr->get_stream()); + handle_ptr->sync_stream(); + + typename optimization_problem_solution_t::additional_termination_information_t + term_info; + term_info.primal_objective = host_res->reduced_problem.get_objective_offset(); + term_info.dual_objective = host_res->reduced_problem.get_objective_offset(); + term_info.number_of_steps_taken = 0; + term_info.solve_time = presolve_time; + term_info.l2_primal_residual = 0.0; + term_info.l2_dual_residual = 0.0; + term_info.gap = 0.0; + + std::vector< + typename optimization_problem_solution_t::additional_termination_information_t> + term_vec{term_info}; + std::vector status_vec{pdlp_termination_status_t::Optimal}; + + CUOPT_LOG_INFO("Status: Optimal Objective: %f", term_info.primal_objective); + return optimization_problem_solution_t(primal_uv, + dual_uv, + rc_uv, + mps_data_model.get_objective_name(), + mps_data_model.get_variable_names(), + mps_data_model.get_row_names(), + std::move(term_vec), + std::move(status_vec)); + } + + presolve_time = lp_timer.elapsed_time(); + CUOPT_LOG_INFO("%s presolve time: %.2fs", + settings_resolved.presolver == presolver_t::PSLP ? "PSLP" : "Papilo", + presolve_time); + } + + // mps_for_solver is what the distributed solver actually sees. + // the reduced + // problem when we ran presolve, the original otherwise. No data transits through device + const auto& mps_for_solver = run_presolve ? host_res->reduced_problem : mps_data_model; + + // -------------------------- DISTRIBUTED SOLVE -------------------------- + // Shape-0 placeholder: needed to build an empty pdlp_solver + cuopt::mathematical_optimization::optimization_problem_t placeholder_op(handle_ptr); + { + std::vector empty_offsets = {0}; + placeholder_op.set_csr_constraint_matrix( + nullptr, 0, nullptr, 0, empty_offsets.data(), static_cast(empty_offsets.size())); + } + placeholder_op.set_objective_offset(mps_for_solver.get_objective_offset()); + placeholder_op.set_objective_scaling_factor(mps_for_solver.get_objective_scaling_factor()); + placeholder_op.set_maximize(mps_for_solver.get_sense()); + mip::problem_t placeholder_problem(placeholder_op); + + pdlp::pdlp_solver_t solver(placeholder_problem, mps_for_solver, settings_resolved); + + auto sol = solver.run_solver(lp_timer); + + // Maximization post-processing (matches run_pdlp): + // PDLP internally solves the negated objective, so flip dual / reduced + // cost signs on the gathered solution before returning. + if (mps_for_solver.get_sense()) { + adjust_dual_solution_and_reduced_cost( + sol.get_dual_solution(), sol.get_reduced_cost(), handle_ptr->get_stream()); + handle_ptr->sync_stream(); + } + + // postsolve + if (run_presolve) { + auto h_primal = cuopt::host_copy(sol.get_primal_solution(), handle_ptr->get_stream()); + auto h_dual = cuopt::host_copy(sol.get_dual_solution(), handle_ptr->get_stream()); + auto h_rc = cuopt::host_copy(sol.get_reduced_cost(), handle_ptr->get_stream()); + handle_ptr->sync_stream(); + + presolver_ptr->undo(h_primal, + h_dual, + h_rc, + cuopt::mathematical_optimization::problem_category_t::LP, + /*status_to_skip=*/false, + settings_resolved.dual_postsolve); + + auto primal_uv = cuopt::device_copy(h_primal, handle_ptr->get_stream()); + auto dual_uv = cuopt::device_copy(h_dual, handle_ptr->get_stream()); + auto rc_uv = cuopt::device_copy(h_rc, handle_ptr->get_stream()); + handle_ptr->sync_stream(); + + auto term_vec = sol.get_additional_termination_informations(); + auto status_vec = sol.get_terminations_status(); + + sol = optimization_problem_solution_t(primal_uv, + dual_uv, + rc_uv, + std::move(sol.get_pdlp_warm_start_data()), + mps_data_model.get_objective_name(), + mps_data_model.get_variable_names(), + mps_data_model.get_row_names(), + std::move(term_vec), + std::move(status_vec)); + } + + sol.set_solve_time(lp_timer.elapsed_time()); + CUOPT_LOG_INFO("PDLP finished"); + if (sol.get_termination_status() != pdlp_termination_status_t::ConcurrentLimit) { + CUOPT_LOG_INFO("Status: %s Objective: %.8e Iterations: %d Time: %.3fs", + sol.get_termination_status_string().c_str(), + sol.get_objective_value(), + sol.get_additional_termination_information().number_of_steps_taken, + sol.get_solve_time()); + } + + if (settings_resolved.sol_file != "") { + CUOPT_LOG_INFO("Writing solution to file %s", settings_resolved.sol_file.c_str()); + sol.write_to_sol_file(settings_resolved.sol_file, handle_ptr->get_stream()); + } + + return sol; } // ============================================================================ @@ -2359,6 +2737,17 @@ std::unique_ptr> solve_lp( template optimization_problem_t mps_data_model_to_optimization_problem( \ raft::handle_t const* handle_ptr, \ const cuopt::mathematical_optimization::io::mps_data_model_t& data_model); \ + \ + template cuopt::mathematical_optimization::io::mps_data_model_t \ + op_problem_to_mps_data_model(const optimization_problem_t& op_problem); \ + \ + template optimization_problem_solution_t solve_lp_distributed_from_mps( \ + raft::handle_t const* handle_ptr, \ + const cuopt::mathematical_optimization::io::mps_data_model_t& mps_data_model, \ + pdlp_solver_settings_t const& settings, \ + bool problem_checking, \ + bool use_pdlp_solver_mode); \ + \ template void set_pdlp_solver_mode(pdlp_solver_settings_t& settings); #if MIP_INSTANTIATE_FLOAT @@ -2369,4 +2758,18 @@ INSTANTIATE(float) INSTANTIATE(double) #endif +// third_party_presolve_t (in mip_heuristics/presolve/) is built +// whenever PDLP_INSTANTIATE_FLOAT is on and depends on the float overloads of +// mps_data_model_to_optimization_problem and op_problem_to_mps_data_model. +// Make sure both symbols exist in PDLP-only float builds where +// MIP_INSTANTIATE_FLOAT is off. +#if PDLP_INSTANTIATE_FLOAT && !MIP_INSTANTIATE_FLOAT +template optimization_problem_t mps_data_model_to_optimization_problem( + raft::handle_t const* handle_ptr, + const cuopt::mathematical_optimization::io::mps_data_model_t& data_model); + +template cuopt::mathematical_optimization::io::mps_data_model_t +op_problem_to_mps_data_model(const optimization_problem_t& op_problem); +#endif + } // namespace cuopt::mathematical_optimization diff --git a/cpp/src/pdlp/solve.cuh b/cpp/src/pdlp/solve.cuh index dd249a825d..aed7a70ac2 100644 --- a/cpp/src/pdlp/solve.cuh +++ b/cpp/src/pdlp/solve.cuh @@ -33,6 +33,50 @@ cuopt::mathematical_optimization::optimization_problem_solution_t solv const timer_t& timer, bool is_batch_mode = false); +/** + * @brief Distributed-PDLP entry point that consumes the host MPS data model + * directly, partitioning it across GPUs without ever materializing the + * full problem on a single (master) GPU. + * + * Intended for problems whose `nnz` exceeds the memory of a single device. The + * master `pdlp_solver_t` is constructed from a shape-0 placeholder problem; the + * real work happens inside it, built straight from the host `mps_data_model`: + * 1. host-side graph partitioning off the MPS CSR, + * 2. per-shard host CSR slicing, + * 3. one shard pdlp_solver_t per GPU, while master holds only scalar metadata + * + gather buffers (no full A / A^T / scaled copies). + * It then runs the solver, gathers the solution to master, applies the + * maximization sign-flip on the dual / reduced cost when the sense is maximize, + * and returns the gathered solution. + * + * Resolves the `distributed_pdlp_num_gpus == -1` sentinel against the + * visible-device count. Several configurations are rejected up front + * (see @pre). + * + * @param handle_ptr Master raft handle (its stream owns the gather buffers and + * any master-side aggregator allocations). Must be non-null. + * @param mps_data_model Host-resident MPS data (CPU vectors only). + * @param settings User-supplied PDLP solver settings; the + * `distributed_pdlp_num_gpus == -1` sentinel is resolved + * here against the visible-device count. + * @param problem_checking Currently unused on this path (accepted for + * signature parity with the device-side entry points). + * @param use_pdlp_solver_mode When true, applies set_pdlp_solver_mode() to the + * resolved settings before solving. + * + * @pre `settings.use_distributed_pdlp == true`, `method == PDLP`, + * `pdlp_precision == DefaultPrecision`, not inside MIP, and no initial + * primal/dual or warm-start data. + */ +template +cuopt::mathematical_optimization::optimization_problem_solution_t +solve_lp_distributed_from_mps( + raft::handle_t const* handle_ptr, + const cuopt::mathematical_optimization::io::mps_data_model_t& mps_data_model, + pdlp_solver_settings_t const& settings, + bool problem_checking, + bool use_pdlp_solver_mode); + /** * @brief Entry point for batch PDLP. Solves multiple LPs sharing the same constraint * matrix structure in a single batched GPU run. diff --git a/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.cu b/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.cu index 5bc66dd7d0..f3e0d622b9 100644 --- a/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.cu +++ b/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.cu @@ -309,6 +309,24 @@ adaptive_step_size_strategy_t::get_norm_squared_delta_dual() const return norm_squared_delta_dual_; } +template +rmm::device_uvector& adaptive_step_size_strategy_t::get_interaction() +{ + return interaction_; +} + +template +rmm::device_uvector& adaptive_step_size_strategy_t::get_norm_squared_delta_primal() +{ + return norm_squared_delta_primal_; +} + +template +rmm::device_uvector& adaptive_step_size_strategy_t::get_norm_squared_delta_dual() +{ + return norm_squared_delta_dual_; +} + template void adaptive_step_size_strategy_t::set_valid_step_size(i_t valid) { @@ -346,8 +364,16 @@ template void adaptive_step_size_strategy_t::compute_interaction_and_movement( rmm::device_uvector& tmp_primal, cusparse_view_t& cusparse_view, - saddle_point_state_t& current_saddle_point_state) + saddle_point_state_t& current_saddle_point_state, + i_t owned_primal_size, + i_t owned_cstr_size) { + // mGPU needs to know owned size to restrict the reductions to the owned prefix + const i_t reduce_primal_size = + (owned_primal_size >= 0) ? owned_primal_size : current_saddle_point_state.get_primal_size(); + const i_t reduce_dual_size = + (owned_cstr_size >= 0) ? owned_cstr_size : current_saddle_point_state.get_dual_size(); + // QP would need this: // if iszero(problem.objective_matrix) // primal_objective_interaction = 0.0 @@ -426,7 +452,7 @@ void adaptive_step_size_strategy_t::compute_interaction_and_movement( // compute interaction (x'-x) . (A(y'-y)) RAFT_CUBLAS_TRY( raft::linalg::detail::cublasdot(handle_ptr_->get_cublas_handle(), - current_saddle_point_state.get_primal_size(), + reduce_primal_size, tmp_primal.data(), primal_stride, current_saddle_point_state.get_delta_primal().data(), @@ -444,7 +470,7 @@ void adaptive_step_size_strategy_t::compute_interaction_and_movement( // norm(delta_dual) ^ 2; RAFT_CUBLAS_TRY( raft::linalg::detail::cublasdot(handle_ptr_->get_cublas_handle(), - current_saddle_point_state.get_primal_size(), + reduce_primal_size, current_saddle_point_state.get_delta_primal().data(), primal_stride, current_saddle_point_state.get_delta_primal().data(), @@ -454,7 +480,7 @@ void adaptive_step_size_strategy_t::compute_interaction_and_movement( RAFT_CUBLAS_TRY( raft::linalg::detail::cublasdot(handle_ptr_->get_cublas_handle(), - current_saddle_point_state.get_dual_size(), + reduce_dual_size, current_saddle_point_state.get_delta_dual().data(), dual_stride, current_saddle_point_state.get_delta_dual().data(), diff --git a/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.hpp b/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.hpp index 1a31cd5c28..7e484abfe2 100644 --- a/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.hpp +++ b/cpp/src/pdlp/step_size_strategy/adaptive_step_size_strategy.hpp @@ -81,9 +81,20 @@ class adaptive_step_size_strategy_t { const rmm::device_uvector& get_norm_squared_delta_primal() const; const rmm::device_uvector& get_norm_squared_delta_dual() const; + // Mutable overloads — used by the multi-GPU path to NCCL-allreduce the + // per-shard scalar contributions in place and to mirror them back to the + // master step_size_strategy_. + rmm::device_uvector& get_interaction(); + rmm::device_uvector& get_norm_squared_delta_primal(); + rmm::device_uvector& get_norm_squared_delta_dual(); + + // owned_primal_size / owned_cstr_size are mGPU overrides. + // mGPU needs to know owned size to restrict the reductions to the owned prefix void compute_interaction_and_movement(rmm::device_uvector& tmp_primal, cusparse_view_t& cusparse_view, - saddle_point_state_t& current_saddle_point_state); + saddle_point_state_t& current_saddle_point_state, + i_t owned_primal_size = -1, + i_t owned_cstr_size = -1); void swap_context(const thrust::universal_host_pinned_vector>& swap_pairs); void resize_context(i_t new_size); diff --git a/cpp/src/pdlp/termination_strategy/convergence_information.cu b/cpp/src/pdlp/termination_strategy/convergence_information.cu index 966bf59724..27500d4c89 100644 --- a/cpp/src/pdlp/termination_strategy/convergence_information.cu +++ b/cpp/src/pdlp/termination_strategy/convergence_information.cu @@ -5,12 +5,16 @@ */ /* clang-format on */ +#include +#include #include #include #include #include #include +#include + #include #include @@ -208,6 +212,77 @@ void convergence_information_t::init_l2_norms() } } +template +void convergence_information_t::compute_owned_reference_norm_partials(i_t owned_var_size, + i_t owned_cstr_size) +{ + cuopt_assert(!batch_mode_, "owned reference-norm partials only used in non-batch mGPU mode"); + cuopt_assert(owned_var_size <= primal_size_h_, "owned_var_size must be <= primal_size_h_"); + cuopt_assert(owned_cstr_size <= dual_size_h_, "owned_cstr_size must be <= dual_size_h_"); + + // Σ objective[0:owned_var]² + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(handle_ptr_->get_cublas_handle(), + static_cast(owned_var_size), + problem_ptr->objective_coefficients.data(), + 1, + problem_ptr->objective_coefficients.data(), + 1, + l2_norm_primal_linear_objective_.data(), + stream_view_)); + + // rhs_sum_of_squares(lower[0:owned_cstr], upper[0:owned_cstr]) (no sqrt) + { + rmm::device_buffer d_temp_storage; + size_t bytes = 0; + auto zip_begin = thrust::make_zip_iterator(problem_ptr->constraint_lower_bounds.data(), + problem_ptr->constraint_upper_bounds.data()); + cub::DeviceReduce::TransformReduce(nullptr, + bytes, + zip_begin, + l2_norm_primal_right_hand_side_.data(), + static_cast(owned_cstr_size), + cuda::std::plus<>{}, + rhs_sum_of_squares_t{}, + f_t(0), + stream_view_); + d_temp_storage.resize(bytes, stream_view_); + cub::DeviceReduce::TransformReduce(d_temp_storage.data(), + bytes, + zip_begin, + l2_norm_primal_right_hand_side_.data(), + static_cast(owned_cstr_size), + cuda::std::plus<>{}, + rhs_sum_of_squares_t{}, + f_t(0), + stream_view_); + } + RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); +} + +template +void convergence_information_t::sqrt_reference_norms_inplace() +{ + cub::DeviceTransform::Transform(l2_norm_primal_linear_objective_.data(), + l2_norm_primal_linear_objective_.data(), + 1, + sqrt_func_t{}, + stream_view_); + cub::DeviceTransform::Transform(l2_norm_primal_right_hand_side_.data(), + l2_norm_primal_right_hand_side_.data(), + 1, + sqrt_func_t{}, + stream_view_); + // Broadcast slot [0] to all climbers (no-op outside batch mode). + thrust::fill(handle_ptr_->get_thrust_policy(), + l2_norm_primal_linear_objective_.begin(), + l2_norm_primal_linear_objective_.end(), + l2_norm_primal_linear_objective_.element(0, stream_view_)); + thrust::fill(handle_ptr_->get_thrust_policy(), + l2_norm_primal_right_hand_side_.begin(), + l2_norm_primal_right_hand_side_.end(), + l2_norm_primal_right_hand_side_.element(0, stream_view_)); +} + // --------------------------------------------------------------------------- // init_reduction_storage: allocate and size the temporary buffers used by // cub::DeviceReduce and cub::DeviceSegmentedReduce throughout solving. @@ -382,6 +457,141 @@ __global__ void compute_remaining_stats_kernel( raft::abs(convergence_information_view.dual_objective[idx]); } +template +template +void convergence_information_t::copy_scalar_from_shard0( + multi_gpu_engine_t& engine, Pick&& pick) +{ + // Sync shards with master stream to avoid race conditions before reading. + engine.sync_await_shards(stream_view_); + auto& s0 = *engine.shards[0]; + raft::device_setter guard(s0.device_id); + auto& s0_conv = s0.sub_pdlp->get_current_termination_strategy().get_convergence_information(); + raft::copy(pick(*this), pick(s0_conv), 1, stream_view_); +} + +template +void convergence_information_t::distributed_compute_primal_residual_and_objective( + multi_gpu_engine_t& engine, const pdlp_solver_settings_t& settings) +{ + cuopt_expects(!settings.per_constraint_residual, + error_type_t::ValidationError, + "per_constraint_residual is not yet supported in multi-GPU mode"); + cuopt_assert(!batch_mode_, "multi-GPU PDLP is not supported in batch mode"); + + // Prepare halo values in potential_next_primal_solution. + engine.halo_exchange_var([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_potential_next_primal_solution(); + }); + + // Per-shard primal residual + partial (owned) primal objective. + engine.for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + auto& sub_conv = sub_pdlp.get_current_termination_strategy().get_convergence_information(); + sub_conv.compute_primal_residual(sub_conv.op_problem_cusparse_view_, + sub_pdlp.pdhg_solver_.get_dual_tmp_resource(), + sub_pdlp.pdhg_solver_.get_potential_next_dual_solution()); + sub_conv.compute_primal_objective_owned_partial( + sub_pdlp.pdhg_solver_.get_potential_next_primal_solution(), shard.rank_data.owned_var_size); + }); + + // Reduce partial primal objectives across shards, mirror the result from + // shard 0 to master, then apply scaling/offset once on the reduced value + // (applying it per-shard would over-count the offset Nshards times). + engine.allreduce_sum_inplace([](pdlp_solver_t& sp) -> f_t* { + return sp.get_current_termination_strategy() + .get_convergence_information() + .get_primal_objective() + .data(); + }); + copy_scalar_from_shard0(engine, [](convergence_information_t& c) -> f_t* { + return c.primal_objective_.data(); + }); + apply_primal_objective_scaling_and_offset(); +} + +template +void convergence_information_t::distributed_compute_primal_residual_l2_norm( + multi_gpu_engine_t& engine) +{ + engine.distributed_l2_norm( + [](pdlp_solver_t& sp) -> rmm::device_uvector& { + return sp.get_current_termination_strategy().get_convergence_information().primal_residual_; + }, + [](pdlp_solver_t& sp) -> f_t* { + return sp.get_current_termination_strategy() + .get_convergence_information() + .l2_primal_residual_.data(); + }, + [](pdlp_shard_t& shard) -> i_t { return shard.rank_data.owned_cstr_size; }); + copy_scalar_from_shard0(engine, [](convergence_information_t& c) -> f_t* { + return c.l2_primal_residual_.data(); + }); +} + +template +void convergence_information_t::distributed_compute_dual_residual_and_objective( + multi_gpu_engine_t& engine) +{ + // 1) Halo-exchange potential_next_dual_solution on every shard so the + // A_T_shard @ y SpMV inside compute_dual_residual reads correct values + // in the cstr halo region + engine.halo_exchange_cstr([](pdlp_solver_t& p) -> rmm::device_uvector& { + return p.pdhg_solver_.get_potential_next_dual_solution(); + }); + + // Same primal_iterate fix as the primal block above: use the shard's + // (fresh, unscaled) potential_next_primal_solution, matching single-GPU + // cuPDLPx. + engine.for_each_shard([](auto& shard) { + auto& sub_pdlp = *shard.sub_pdlp; + auto& sub_conv = sub_pdlp.get_current_termination_strategy().get_convergence_information(); + sub_conv.compute_dual_residual(sub_conv.op_problem_cusparse_view_, + sub_pdlp.pdhg_solver_.get_primal_tmp_resource(), + sub_pdlp.pdhg_solver_.get_potential_next_primal_solution(), + sub_pdlp.pdhg_solver_.get_dual_slack()); + sub_conv.compute_dual_objective_owned_partial( + sub_pdlp.pdhg_solver_.get_potential_next_primal_solution(), + sub_pdlp.pdhg_solver_.get_dual_slack(), + shard.rank_data.owned_var_size, + shard.rank_data.owned_cstr_size); + }); + + // 4) Allreduce dual_objective_ across shards (sum, in place), mirror to + // master, then apply offset/scaling once + engine.allreduce_sum_inplace([](pdlp_solver_t& sp) -> f_t* { + return sp.get_current_termination_strategy() + .get_convergence_information() + .get_dual_objective() + .data(); + }); + copy_scalar_from_shard0(engine, [](convergence_information_t& c) -> f_t* { + return c.dual_objective_.data(); + }); + apply_dual_objective_scaling_and_offset(); +} + +template +void convergence_information_t::distributed_compute_dual_residual_l2_norm( + multi_gpu_engine_t& engine) +{ + // Same pattern as the primal L2 above, but the dual residual is var-shaped + // so we clip to owned_var_size. + engine.distributed_l2_norm( + [](pdlp_solver_t& sp) -> rmm::device_uvector& { + return sp.get_current_termination_strategy().get_convergence_information().dual_residual_; + }, + [](pdlp_solver_t& sp) -> f_t* { + return sp.get_current_termination_strategy() + .get_convergence_information() + .l2_dual_residual_.data(); + }, + [](pdlp_shard_t& shard) -> i_t { return shard.rank_data.owned_var_size; }); + copy_scalar_from_shard0(engine, [](convergence_information_t& c) -> f_t* { + return c.l2_dual_residual_.data(); + }); +} + template void convergence_information_t::compute_convergence_information( pdhg_solver_t& current_pdhg_solver, @@ -392,22 +602,29 @@ void convergence_information_t::compute_convergence_information( const rmm::device_uvector& objective_coefficients, const pdlp_solver_settings_t& settings) { - cuopt_assert( - primal_residual_.size() == dual_size_h_ * climber_strategies_.size(), - "primal_residual_ size must be equal to primal_size_h_ * climber_strategies_.size()"); - cuopt_assert(primal_iterate.size() == primal_size_h_ * climber_strategies_.size(), - "primal_iterate size must be equal to primal_size_h_ * climber_strategies_.size()"); - cuopt_assert(dual_residual_.size() == primal_size_h_ * climber_strategies_.size(), - "dual_residual_ size must be equal to primal_size_h_ * climber_strategies_.size()"); - cuopt_assert(dual_iterate.size() == dual_size_h_ * climber_strategies_.size(), - "dual_iterate size must be equal to dual_size_h_ * climber_strategies_.size()"); - cuopt_assert(l2_primal_residual_.size() == climber_strategies_.size(), - "l2_primal_residual_ size must be equal to climber_strategies_.size()"); + // Per-climber L2-norm scalars are uniform across single-GPU, master, and shards. cuopt_assert(l2_primal_residual_.size() == climber_strategies_.size(), "l2_primal_residual_ size must be equal to climber_strategies_.size()"); cuopt_assert(l2_dual_residual_.size() == climber_strategies_.size(), "l2_dual_residual_ size must be equal to climber_strategies_.size()"); - + if (current_pdhg_solver.is_distributed_master()) { + cuopt_assert(primal_residual_.size() == 0, + "master must not own a primal_residual_ buffer (shards compute it)"); + cuopt_assert(dual_residual_.size() == 0, + "master must not own a dual_residual_ buffer (shards compute it)"); + } else { + cuopt_assert( + primal_iterate.size() == primal_size_h_ * climber_strategies_.size(), + "primal_iterate size must be equal to primal_size_h_ * climber_strategies_.size()"); + cuopt_assert(dual_iterate.size() == dual_size_h_ * climber_strategies_.size(), + "dual_iterate size must be equal to dual_size_h_ * climber_strategies_.size()"); + cuopt_assert( + primal_residual_.size() == dual_size_h_ * climber_strategies_.size(), + "primal_residual_ size must be equal to dual_size_h_ * climber_strategies_.size()"); + cuopt_assert( + dual_residual_.size() == primal_size_h_ * climber_strategies_.size(), + "dual_residual_ size must be equal to primal_size_h_ * climber_strategies_.size()"); + } raft::common::nvtx::range fun_scope("compute_convergence_information"); #ifdef CUPDLP_DEBUG_MODE @@ -416,17 +633,25 @@ void convergence_information_t::compute_convergence_information( print("dual_slack", dual_slack); #endif - compute_primal_residual( - op_problem_cusparse_view_, current_pdhg_solver.get_dual_tmp_resource(), dual_iterate); - compute_primal_objective(primal_iterate); + if (current_pdhg_solver.is_distributed_master()) { + distributed_compute_primal_residual_and_objective(*current_pdhg_solver.get_mgpu_engine(), + settings); + } else { + compute_primal_residual( + op_problem_cusparse_view_, current_pdhg_solver.get_dual_tmp_resource(), dual_iterate); + compute_primal_objective(primal_iterate); + } #ifdef CUPDLP_DEBUG_MODE print("Primal Residual", primal_residual_); #endif - if (!batch_mode_) + // L2 Norm + if (current_pdhg_solver.is_distributed_master()) { + distributed_compute_primal_residual_l2_norm(*current_pdhg_solver.get_mgpu_engine()); + } else if (!batch_mode_) { my_l2_norm(primal_residual_, l2_primal_residual_, handle_ptr_); - else { + } else { segmented_sum_handler_.segmented_sum_helper( thrust::make_transform_iterator(primal_residual_.data(), power_two_func_t{}), l2_primal_residual_.data(), @@ -444,6 +669,7 @@ void convergence_information_t::compute_convergence_information( print("Absolute Primal Residual", l2_primal_residual_); #endif // If per_constraint_residual is false we still need to perform the l2 since it's used in kkt + // Not suported in mGPU if (settings.per_constraint_residual) { // Compute the linf of (residual_i - rel * b_i) if (settings.save_best_primal_so_far) { @@ -466,19 +692,25 @@ void convergence_information_t::compute_convergence_information( std::numeric_limits::lowest()); } - compute_dual_residual(op_problem_cusparse_view_, - current_pdhg_solver.get_primal_tmp_resource(), - primal_iterate, - dual_slack); - compute_dual_objective(dual_iterate, primal_iterate, dual_slack); + if (current_pdhg_solver.is_distributed_master()) { + distributed_compute_dual_residual_and_objective(*current_pdhg_solver.get_mgpu_engine()); + } else { + compute_dual_residual(op_problem_cusparse_view_, + current_pdhg_solver.get_primal_tmp_resource(), + primal_iterate, + dual_slack); + compute_dual_objective(dual_iterate, primal_iterate, dual_slack); + } #ifdef CUPDLP_DEBUG_MODE print("Dual Residual", dual_residual_); #endif - if (!batch_mode_) + if (current_pdhg_solver.is_distributed_master()) { + distributed_compute_dual_residual_l2_norm(*current_pdhg_solver.get_mgpu_engine()); + } else if (!batch_mode_) { my_l2_norm(dual_residual_, l2_dual_residual_, handle_ptr_); - else { + } else { segmented_sum_handler_.segmented_sum_helper( thrust::make_transform_iterator(dual_residual_.data(), power_two_func_t{}), l2_dual_residual_.data(), @@ -509,6 +741,8 @@ void convergence_information_t::compute_convergence_information( std::numeric_limits::lowest()); } + // In mGPU, full primal_objective and dual_objective already mirrored to master so no special + // behaviour const auto [grid_size, block_size] = kernel_config_from_batch_size(climber_strategies_.size()); compute_remaining_stats_kernel <<>>(this->view(), climber_strategies_.size()); @@ -615,6 +849,24 @@ __global__ void apply_objective_scaling_and_offset(raft::device_span object objective[idx] = objective_scaling_factor * (objective[idx] + objective_offsets[idx]); } +template +void convergence_information_t::compute_primal_objective_owned_partial( + rmm::device_uvector& primal_solution, i_t n_owned) +{ + raft::common::nvtx::range fun_scope("compute_primal_objective_owned_partial"); + cuopt_assert(!batch_mode_, "owned-partial primal objective is only used in non-batch mGPU mode"); + cuopt_assert(n_owned <= primal_size_h_, + "n_owned must be <= primal_size_h_ (owned slice is a prefix)"); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(handle_ptr_->get_cublas_handle(), + static_cast(n_owned), + primal_solution.data(), + primal_stride, + problem_ptr->objective_coefficients.data(), + primal_stride, + primal_objective_.data(), + stream_view_)); +} + template void convergence_information_t::compute_primal_objective( rmm::device_uvector& primal_solution) @@ -643,21 +895,25 @@ void convergence_information_t::compute_primal_objective( // Apply per-climber objective scaling and offset. objective_offsets_ is always populated // (defaults to the scalar problem offset replicated, or user-specified per-climber offsets). - { - const auto [grid_size, block_size] = kernel_config_from_batch_size(climber_strategies_.size()); - apply_objective_scaling_and_offset<<>>( - make_span(primal_objective_), - problem_ptr->presolve_data.objective_scaling_factor, - make_span(objective_offsets_), - climber_strategies_.size()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } + apply_primal_objective_scaling_and_offset(); #ifdef CUPDLP_DEBUG_MODE print("Primal objective", primal_objective_); #endif } +template +void convergence_information_t::apply_primal_objective_scaling_and_offset() +{ + const auto [grid_size, block_size] = kernel_config_from_batch_size(climber_strategies_.size()); + apply_objective_scaling_and_offset<<>>( + make_span(primal_objective_), + problem_ptr->presolve_data.objective_scaling_factor, + make_span(objective_offsets_), + climber_strategies_.size()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + template void convergence_information_t::compute_dual_residual( cusparse_view_t& cusparse_view, @@ -740,6 +996,50 @@ void convergence_information_t::compute_dual_residual( } } +template +void convergence_information_t::compute_dual_objective_owned_partial( + rmm::device_uvector& primal_solution, + rmm::device_uvector& dual_slack, + i_t n_owned_var, + i_t n_owned_cstr) +{ + raft::common::nvtx::range fun_scope("compute_dual_objective_owned_partial"); + cuopt_assert(!batch_mode_, "owned-partial dual objective is only used in non-batch mGPU mode"); + cuopt_assert(hyper_params_.use_reflected_primal_dual, + "owned-partial dual objective requires use_reflected_primal_dual"); + cuopt_assert(n_owned_var <= primal_size_h_, + "n_owned_var must be <= primal_size_h_ (owned slice is a prefix)"); + cuopt_assert(n_owned_cstr <= dual_size_h_, + "n_owned_cstr must be <= dual_size_h_ (owned slice is a prefix)"); + + // dual_dot_ = dot(dual_slack[0:n_owned_var], primal_solution[0:n_owned_var]) + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(handle_ptr_->get_cublas_handle(), + static_cast(n_owned_var), + dual_slack.data(), + primal_stride, + primal_solution.data(), + primal_stride, + dual_dot_.data(), + stream_view_)); + + // sum_primal_slack_ = Σ primal_slack_[0:n_owned_cstr] + // primal_slack_ is assumed populated for owned cstrs by a prior + // compute_primal_residual call on this same shard. + cub::DeviceReduce::Sum(rmm_tmp_buffer_.data(), + size_of_buffer_, + primal_slack_.data(), + sum_primal_slack_.data(), + static_cast(n_owned_cstr), + stream_view_); + + // dual_objective_ = dual_dot_ + sum_primal_slack_ (still a partial sum). + cub::DeviceTransform::Transform(cuda::std::make_tuple(dual_dot_.data(), sum_primal_slack_.data()), + dual_objective_.data(), + 1, + cuda::std::plus<>{}, + stream_view_); +} + template void convergence_information_t::compute_dual_objective( rmm::device_uvector& dual_solution, @@ -821,21 +1121,25 @@ void convergence_information_t::compute_dual_objective( } // Apply per-climber objective scaling and offset. - { - const auto [grid_size, block_size] = kernel_config_from_batch_size(climber_strategies_.size()); - apply_objective_scaling_and_offset<<>>( - make_span(dual_objective_), - problem_ptr->presolve_data.objective_scaling_factor, - make_span(objective_offsets_), - climber_strategies_.size()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } + apply_dual_objective_scaling_and_offset(); #ifdef CUPDLP_DEBUG_MODE print("Dual objective", dual_objective_); #endif } +template +void convergence_information_t::apply_dual_objective_scaling_and_offset() +{ + const auto [grid_size, block_size] = kernel_config_from_batch_size(climber_strategies_.size()); + apply_objective_scaling_and_offset<<>>( + make_span(dual_objective_), + problem_ptr->presolve_data.objective_scaling_factor, + make_span(objective_offsets_), + climber_strategies_.size()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + template void convergence_information_t::compute_reduced_cost_from_primal_gradient( const rmm::device_uvector& primal_gradient, const rmm::device_uvector& primal_solution) @@ -916,12 +1220,24 @@ const rmm::device_uvector& convergence_information_t::get_primal_ return primal_objective_; } +template +rmm::device_uvector& convergence_information_t::get_primal_objective() +{ + return primal_objective_; +} + template const rmm::device_uvector& convergence_information_t::get_dual_objective() const { return dual_objective_; } +template +rmm::device_uvector& convergence_information_t::get_dual_objective() +{ + return dual_objective_; +} + template const rmm::device_uvector& convergence_information_t::get_l2_dual_residual() const { diff --git a/cpp/src/pdlp/termination_strategy/convergence_information.hpp b/cpp/src/pdlp/termination_strategy/convergence_information.hpp index 41e5f06f26..266665abcb 100644 --- a/cpp/src/pdlp/termination_strategy/convergence_information.hpp +++ b/cpp/src/pdlp/termination_strategy/convergence_information.hpp @@ -52,7 +52,10 @@ class convergence_information_t { // Needed for kkt restart & debug prints const rmm::device_uvector& get_primal_objective() const; + // Non-const overload used by the multi-GPU branch to mirror / allreduce. + rmm::device_uvector& get_primal_objective(); const rmm::device_uvector& get_dual_objective() const; + rmm::device_uvector& get_dual_objective(); const rmm::device_uvector& get_l2_primal_residual() const; const rmm::device_uvector& get_l2_dual_residual() const; const rmm::device_uvector& get_relative_linf_primal_residual() const; @@ -66,6 +69,11 @@ class convergence_information_t { const rmm::device_uvector& get_l2_norm_primal_linear_objective() const; const rmm::device_uvector& get_l2_norm_primal_right_hand_side() const; + void compute_owned_reference_norm_partials(i_t owned_var_size, i_t owned_cstr_size); + void sqrt_reference_norms_inplace(); + f_t* l2_norm_primal_right_hand_side_data() { return l2_norm_primal_right_hand_side_.data(); } + f_t* l2_norm_primal_linear_objective_data() { return l2_norm_primal_linear_objective_.data(); } + struct view_t { i_t primal_size; i_t dual_size; @@ -123,12 +131,39 @@ class convergence_information_t { rmm::device_uvector& tmp_dual, [[maybe_unused]] const rmm::device_uvector& dual_iterate); + // Multi-GPU shard helper: writes a partial dot(c[0:n_owned], x[0:n_owned]) + // into primal_objective_ (no scaling, no offset). Master is responsible for + // allreduce SUM across shards and then applying scaling + offset once on the + // reduced value + void compute_primal_objective_owned_partial(rmm::device_uvector& primal_solution, + i_t n_owned); + + // Multi-GPU shard helper: writes a partial dual objective into + // dual_objective_ (no scaling, no offset). Computes + // dual_dot_ = dot(dual_slack[0:n_owned_var], primal_solution[0:n_owned_var]) + // sum_primal_slack_ = Σ primal_slack_[0:n_owned_cstr] + // dual_objective_ = dual_dot_ + sum_primal_slack_ + // primal_slack_ is assumed already populated by a prior per-shard + // compute_primal_residual call. Use only in the use_reflected_primal_dual + // path (the multi-GPU mode). + void compute_dual_objective_owned_partial(rmm::device_uvector& primal_solution, + rmm::device_uvector& dual_slack, + i_t n_owned_var, + i_t n_owned_cstr); + void swap_context(const thrust::universal_host_pinned_vector>& swap_pairs); void resize_context(i_t new_size); private: void compute_primal_objective(rmm::device_uvector& primal_solution); + // Applies per-climber objective scaling + offset to primal_objective_. + // Single-GPU path: called from compute_primal_objective right after the dot. + // Multi-GPU path: called on master once after allreduce of partial sums. + void apply_primal_objective_scaling_and_offset(); + // Same as above but for dual_objective_. + void apply_dual_objective_scaling_and_offset(); + void compute_dual_residual(cusparse_view_t& cusparse_view, rmm::device_uvector& tmp_primal, rmm::device_uvector& primal_solution, @@ -143,6 +178,20 @@ class convergence_information_t { void compute_reduced_costs_dual_objective_contribution(); + // ----- Distributed-PDLP sub-steps of compute_convergence_information ----- + // Each is the multi-GPU equivalent of the single-GPU block it replaces + void distributed_compute_primal_residual_and_objective( + multi_gpu_engine_t& engine, const pdlp_solver_settings_t& settings); + void distributed_compute_primal_residual_l2_norm(multi_gpu_engine_t& engine); + void distributed_compute_dual_residual_and_objective(multi_gpu_engine_t& engine); + void distributed_compute_dual_residual_l2_norm(multi_gpu_engine_t& engine); + + // Mirror one per-shard convergence scalar from shard 0 to master (after a + // cross-shard reduction has made every shard agree). `pick(conv)` selects the + // f_t* field; it is applied to both master (*this) and shard 0's instance. + template + void copy_scalar_from_shard0(multi_gpu_engine_t& engine, Pick&& pick); + // Ctor helpers — each handles both batch and non-batch internally. void init_objective_offsets(); void init_l2_norms(); diff --git a/cpp/src/pdlp/termination_strategy/termination_strategy.cu b/cpp/src/pdlp/termination_strategy/termination_strategy.cu index 1218c60fa4..13acee138c 100644 --- a/cpp/src/pdlp/termination_strategy/termination_strategy.cu +++ b/cpp/src/pdlp/termination_strategy/termination_strategy.cu @@ -5,6 +5,9 @@ */ /* clang-format on */ +#include + +#include #include #include #include @@ -195,6 +198,13 @@ pdlp_termination_strategy_t::get_convergence_information() const return convergence_information_; } +template +convergence_information_t& +pdlp_termination_strategy_t::get_convergence_information() +{ + return convergence_information_; +} + template const infeasibility_information_t& pdlp_termination_strategy_t::get_infeasibility_information() const @@ -545,12 +555,24 @@ pdlp_termination_strategy_t::fill_return_problem_solution( std::vector&& termination_status, bool deep_copy) { - cuopt_assert( - primal_iterate.size() == current_pdhg_solver.get_primal_size() * termination_status.size(), - "Primal iterate size mismatch"); - cuopt_assert( - dual_iterate.size() == current_pdhg_solver.get_dual_size() * termination_status.size(), - "Dual iterate size mismatch"); + // Skip the size checks on the distributed master: its pdhg_solver_ is built from a + // shape-0 placeholder while termination_strategy is built from the full problem size + if (!current_pdhg_solver.is_distributed_master()) { + cuopt_assert( + primal_iterate.size() == current_pdhg_solver.get_primal_size() * termination_status.size(), + "Primal iterate size mismatch"); + cuopt_assert( + dual_iterate.size() == current_pdhg_solver.get_dual_size() * termination_status.size(), + "Dual iterate size mismatch"); + } + + // In distributed PDLP, gather solutions from the shards to the master. + if (auto* engine = current_pdhg_solver.get_mgpu_engine()) { + const bool is_current_live_iterate = + (&primal_iterate == ¤t_pdhg_solver.get_potential_next_primal_solution()) || + (&primal_iterate == ¤t_pdhg_solver.get_primal_solution()); + if (is_current_live_iterate) { engine->gather_potential_next_solutions_to_master(); } + } typename convergence_information_t::view_t convergence_information_view = convergence_information_.view(); diff --git a/cpp/src/pdlp/termination_strategy/termination_strategy.hpp b/cpp/src/pdlp/termination_strategy/termination_strategy.hpp index 42d7cc038a..948f7674d0 100644 --- a/cpp/src/pdlp/termination_strategy/termination_strategy.hpp +++ b/cpp/src/pdlp/termination_strategy/termination_strategy.hpp @@ -187,6 +187,7 @@ class pdlp_termination_strategy_t { i_t get_optimal_solution_id() const; const convergence_information_t& get_convergence_information() const; + convergence_information_t& get_convergence_information(); const infeasibility_information_t& get_infeasibility_information() const; // Deep copy is used when save best primal so far is toggled diff --git a/cpp/src/pdlp/utils.cuh b/cpp/src/pdlp/utils.cuh index adc1432b39..91c59be15b 100644 --- a/cpp/src/pdlp/utils.cuh +++ b/cpp/src/pdlp/utils.cuh @@ -12,6 +12,9 @@ #include #include +#include +#include + #include #include #include @@ -32,6 +35,39 @@ namespace cuopt::mathematical_optimization::pdlp { +// Host-side probe vector z ~ Normal(0, 1) that seeds the power iteration +// for sigma_max(A) in compute_initial_step_size (single-GPU) and +// distributed_max_singular_value_squared. Fixed seed (1) so single-GPU and +// distributed runs on the same problem produce bit-identical initial iterates. +template +inline std::vector make_singular_value_probe(std::size_t size) +{ + std::vector out(size); + std::mt19937 gen(1); + std::normal_distribution dist(f_t(0.0), f_t(1.0)); + for (std::size_t i = 0; i < size; ++i) + out[i] = dist(gen); + return out; +} + +// Elementwise: v := v / *scalar. Named struct (rather than an extended HD +// lambda) so it can be used as a cub::DeviceTransform op inside templates whose +// template arguments are themselves local lambda types (which nvcc rejects). +template +struct divide_by_device_scalar_t { + f_t const* scalar; + __host__ __device__ f_t operator()(f_t v) const { return v / *scalar; } +}; + +// Elementwise: q := -*scalar * q + z. Used in the power-iteration residual +// update (single-GPU compute_initial_step_size and distributed +// distributed_max_singular_value_squared). Same named-struct rationale as above. +template +struct residual_fma_neg_scalar_t { + f_t const* scalar; + __host__ __device__ f_t operator()(f_t q, f_t z) const { return -(*scalar) * q + z; } +}; + template DI f_t deterministic_block_reduce(raft::device_span shared, f_t val) { @@ -264,6 +300,21 @@ struct rhs_sum_of_squares_t { } }; +// Per-element weighted square (used to compute the L2 norm of the weighted +// objective coefficients). +template +struct weighted_square_op { + f_t weight; + HDI f_t operator()(f_t v) const { return v * v * weight; } +}; + +// Convert a squared L2 norm to its bound/objective rescaling scalar, +// 1 / (sqrt(sum) + 1). Uses raft::sqrt which is __host__ __device__. +template +struct rescaling_from_squared_norm_op { + HDI f_t operator()(f_t sum) const { return f_t(1.0) / (raft::sqrt(sum) + f_t(1.0)); } +}; + template void inline combine_constraint_bounds(const mip::problem_t& op_problem, rmm::device_uvector& combined_bounds) diff --git a/cpp/tests/linear_programming/CMakeLists.txt b/cpp/tests/linear_programming/CMakeLists.txt index bc057db1e2..ea3292394d 100644 --- a/cpp/tests/linear_programming/CMakeLists.txt +++ b/cpp/tests/linear_programming/CMakeLists.txt @@ -15,6 +15,14 @@ ConfigureTest(PDLP_TEST ${CMAKE_CURRENT_SOURCE_DIR}/pdlp_test.cu LABELS numopt) +# Multi-GPU distributed PDLP tests. Binary name ends in _MG_TEST so that +# ci/test_cpp_multi_gpu.sh picks it up via its *_MG_TEST glob. Each test +# calls GTEST_SKIP when fewer than 2 GPUs are visible so the binary is +# harmless on single-GPU runners. +ConfigureTest(PDLP_MG_TEST + ${CMAKE_CURRENT_SOURCE_DIR}/pdlp_distributed_test.cu + LABELS numopt) + # ################################################################################################## # - MPS / LP parser tests -------------------------------------------------------------------------- ConfigureTest(MPS_PARSER_TEST diff --git a/cpp/tests/linear_programming/pdlp_distributed_test.cu b/cpp/tests/linear_programming/pdlp_distributed_test.cu new file mode 100644 index 0000000000..74c35ddb34 --- /dev/null +++ b/cpp/tests/linear_programming/pdlp_distributed_test.cu @@ -0,0 +1,143 @@ +/* clang-format off */ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +/* clang-format on */ + +// Multi-GPU distributed PDLP parity tests. +// Binary name PDLP_MG_TEST matches the *_MG_TEST glob in ci/test_cpp_multi_gpu.sh. + +#include "utilities/pdlp_test_utilities.cuh" + +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include + +#include +#include + +namespace cuopt::mathematical_optimization::test { + +// Solve `mps_rel_path` with the single-GPU PDLP ("base") and with distributed PDLP +// (num_gpus = -1 => auto-detect), then assert the distributed run matches the base run on +// everything meaningful: termination status, step count (within 15%), primal/dual objective, +// and the full primal/dual solution vectors. All value comparisons use a loose relative tolerance. +// Pass disable_presolve=true for tiny instances that PSLP would solve entirely, which would +// otherwise bypass the distributed solver path we want to exercise. +static void expect_distributed_matches_base(raft::handle_t const& handle, + std::string const& mps_rel_path, + bool fixed_mps_format = false, + bool disable_presolve = false) +{ + constexpr double loose_rel = 1e-3; + auto near_rel = [](double a, double b, double rel) { + return std::fabs(a - b) <= rel * (1.0 + std::fabs(a)); + }; + + auto path = make_path_absolute(mps_rel_path); + io::mps_data_model_t problem = io::read_mps(path, fixed_mps_format); + + pdlp_solver_settings_t base_settings{}; + base_settings.method = method_t::PDLP; + if (disable_presolve) { base_settings.presolver = presolver_t::None; } + + auto base_op = mps_data_model_to_optimization_problem(&handle, problem); + auto base = solve_lp(base_op, base_settings); + + pdlp_solver_settings_t dist_settings = base_settings; + dist_settings.use_distributed_pdlp = true; + dist_settings.distributed_pdlp_num_gpus = -1; + auto dist = solve_lp(&handle, problem, dist_settings); + + ASSERT_EQ(static_cast(base.get_termination_status()), CUOPT_TERMINATION_STATUS_OPTIMAL) + << mps_rel_path << ": base did not reach optimal"; + EXPECT_EQ(static_cast(dist.get_termination_status()), + static_cast(base.get_termination_status())) + << mps_rel_path << ": distributed termination status differs from base"; + + const auto& base_info = base.get_additional_termination_information(); + const auto& dist_info = dist.get_additional_termination_information(); + + EXPECT_TRUE(near_rel(base_info.primal_objective, dist_info.primal_objective, loose_rel)) + << mps_rel_path << ": primal objective base=" << base_info.primal_objective + << " distributed=" << dist_info.primal_objective; + EXPECT_TRUE(near_rel(base_info.dual_objective, dist_info.dual_objective, loose_rel)) + << mps_rel_path << ": dual objective base=" << base_info.dual_objective + << " distributed=" << dist_info.dual_objective; + + const int base_steps = base_info.number_of_steps_taken; + const int dist_steps = dist_info.number_of_steps_taken; + const int max_steps = std::max(base_steps, dist_steps); + const int step_diff = std::max(base_steps, dist_steps) - std::min(base_steps, dist_steps); + EXPECT_LE(static_cast(step_diff), 0.15 * max_steps) + << mps_rel_path << ": step counts differ by >15% (base=" << base_steps + << ", distributed=" << dist_steps << ")"; + + auto base_primal = cuopt::host_copy(base.get_primal_solution(), handle.get_stream()); + auto dist_primal = cuopt::host_copy(dist.get_primal_solution(), handle.get_stream()); + ASSERT_EQ(base_primal.size(), dist_primal.size()) << mps_rel_path << ": primal size mismatch"; + for (std::size_t i = 0; i < base_primal.size(); ++i) { + EXPECT_TRUE(near_rel(base_primal[i], dist_primal[i], loose_rel)) + << mps_rel_path << ": primal[" << i << "] base=" << base_primal[i] + << " distributed=" << dist_primal[i]; + } + + auto base_dual = cuopt::host_copy(base.get_dual_solution(), handle.get_stream()); + auto dist_dual = cuopt::host_copy(dist.get_dual_solution(), handle.get_stream()); + ASSERT_EQ(base_dual.size(), dist_dual.size()) << mps_rel_path << ": dual size mismatch"; + for (std::size_t i = 0; i < base_dual.size(); ++i) { + EXPECT_TRUE(near_rel(base_dual[i], dist_dual[i], loose_rel)) + << mps_rel_path << ": dual[" << i << "] base=" << base_dual[i] + << " distributed=" << dist_dual[i]; + } +} + +TEST(pdlp_class, distributed_parity_afiro) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires >=2 GPUs, found " << raft::device_setter::get_device_count(); + } + const raft::handle_t handle{}; + expect_distributed_matches_base(handle, "linear_programming/afiro_original.mps", true); +} + +// Maximization LP: exercises the sign-flipping paths through normalization / +// presolve / solution translation on the distributed side. +TEST(pdlp_class, distributed_parity_cod105_max) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires >=2 GPUs, found " << raft::device_setter::get_device_count(); + } + const raft::handle_t handle{}; + expect_distributed_matches_base(handle, "mip/cod105_max.mps"); +} + +TEST(pdlp_class, distributed_parity_graph40_40) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires >=2 GPUs, found " << raft::device_setter::get_device_count(); + } + const raft::handle_t handle{}; + expect_distributed_matches_base(handle, "linear_programming/graph40-40/graph40-40.mps"); +} + +TEST(pdlp_class, distributed_parity_ex10) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires >=2 GPUs, found " << raft::device_setter::get_device_count(); + } + const raft::handle_t handle{}; + expect_distributed_matches_base(handle, "linear_programming/ex10/ex10.mps"); +} + +} // namespace cuopt::mathematical_optimization::test diff --git a/cpp/tests/linear_programming/unit_tests/presolve_test.cu b/cpp/tests/linear_programming/unit_tests/presolve_test.cu index 0c76de192a..3ca9b5833a 100644 --- a/cpp/tests/linear_programming/unit_tests/presolve_test.cu +++ b/cpp/tests/linear_programming/unit_tests/presolve_test.cu @@ -515,13 +515,13 @@ TEST_P(dual_crush_round_trip, kkt_check) // Step 1: Presolve with a single presolver instance (same one used for crush later) sort_csr(op_problem); mip::third_party_presolve_t presolver; - auto result = presolver.apply(op_problem, - problem_category_t::LP, - presolver_t::Papilo, - /*dual_postsolve=*/true, - /*abs_tol=*/1e-6, - /*rel_tol=*/1e-9, - /*time_limit=*/60.0); + auto result = presolver.apply_presolve_from_op_problem(op_problem, + problem_category_t::LP, + presolver_t::Papilo, + /*dual_postsolve=*/true, + /*abs_tol=*/1e-6, + /*rel_tol=*/1e-9, + /*time_limit=*/60.0); ASSERT_TRUE(result.status == mip::third_party_presolve_status_t::REDUCED || result.status == mip::third_party_presolve_status_t::UNCHANGED); @@ -569,7 +569,8 @@ TEST_P(dual_crush_round_trip, kkt_check) rc_sol = cuopt::device_copy(reduced_solution.get_reduced_cost(), stream); } - presolver.undo(primal_sol, dual_sol, rc_sol, problem_category_t::LP, false, true, stream); + presolver.undo_from_device( + primal_sol, dual_sol, rc_sol, problem_category_t::LP, false, true, stream); auto x_orig = host_copy(primal_sol, stream); auto y_orig = host_copy(dual_sol, stream); @@ -772,13 +773,13 @@ TEST_P(crush_warmstart, round_trip) // Step 1: Presolve sort_csr(op_problem); mip::third_party_presolve_t presolver; - auto result = presolver.apply(op_problem, - problem_category_t::LP, - presolver_t::Papilo, - /*dual_postsolve=*/true, - /*abs_tol=*/1e-6, - /*rel_tol=*/1e-9, - /*time_limit=*/60.0); + auto result = presolver.apply_presolve_from_op_problem(op_problem, + problem_category_t::LP, + presolver_t::Papilo, + /*dual_postsolve=*/true, + /*abs_tol=*/1e-6, + /*rel_tol=*/1e-9, + /*time_limit=*/60.0); ASSERT_TRUE(result.status == mip::third_party_presolve_status_t::REDUCED || result.status == mip::third_party_presolve_status_t::UNCHANGED); @@ -817,7 +818,8 @@ TEST_P(crush_warmstart, round_trip) } auto rc_sol = cuopt::device_copy(z_red, stream); - presolver.undo(primal_sol, dual_sol, rc_sol, problem_category_t::LP, false, true, stream); + presolver.undo_from_device( + primal_sol, dual_sol, rc_sol, problem_category_t::LP, false, true, stream); auto x_orig = host_copy(primal_sol, stream); auto y_orig = host_copy(dual_sol, stream); diff --git a/cpp/tests/mip/presolve_test.cu b/cpp/tests/mip/presolve_test.cu index 9f9b65d6e0..7067479093 100644 --- a/cpp/tests/mip/presolve_test.cu +++ b/cpp/tests/mip/presolve_test.cu @@ -39,14 +39,15 @@ TEST(problem, find_implied_integers) auto mps_data_model = cuopt::mathematical_optimization::io::read_mps(path, false); auto op_problem = mps_data_model_to_optimization_problem(&handle_, mps_data_model); auto presolver = std::make_unique>(); - auto result = presolver->apply(op_problem, - cuopt::mathematical_optimization::problem_category_t::MIP, - cuopt::mathematical_optimization::presolver_t::Papilo, - false, - 1e-6, - 1e-12, - 20, - 1); + auto result = presolver->apply_presolve_from_op_problem( + op_problem, + cuopt::mathematical_optimization::problem_category_t::MIP, + cuopt::mathematical_optimization::presolver_t::Papilo, + false, + 1e-6, + 1e-12, + 20, + 1); ASSERT_NE(result.status, mip::third_party_presolve_status_t::INFEASIBLE); ASSERT_NE(result.status, mip::third_party_presolve_status_t::UNBNDORINFEAS); @@ -144,7 +145,7 @@ TEST(gf2_presolve, uses_compact_constraint_indices) problem.set_constraint_upper_bounds(constraint_ub.data(), constraint_ub.size()); auto presolver = std::make_unique>(); - auto result = presolver->apply( + auto result = presolver->apply_presolve_from_op_problem( problem, problem_category_t::MIP, presolver_t::Papilo, false, 1e-6, 1e-12, 20, 1); EXPECT_EQ(result.status, mip::third_party_presolve_status_t::REDUCED); diff --git a/dependencies.yaml b/dependencies.yaml index 22a5752514..82d6a908db 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -266,6 +266,7 @@ dependencies: - libboost-devel - cpp-argparse - &pyyaml pyyaml>=6.0.0 + - nccl >=2.19 - tbb-devel - zlib - bzip2