Add support for sparsity exploitation for SOCs of large dimensionality

cpp/include/cuopt/linear_programming/constants.h

@@ -48,6 +48,7 @@
 #define CUOPT_ORDERING                             "ordering"
 #define CUOPT_BARRIER_DUAL_INITIAL_POINT           "barrier_dual_initial_point"
 #define CUOPT_BARRIER_ITERATIVE_REFINEMENT         "barrier_iterative_refinement"
+#define CUOPT_BARRIER_SOC_THRESHOLD                "barrier_soc_threshold"
 #define CUOPT_BARRIER_STEP_SCALE                   "barrier_step_scale"
 #define CUOPT_ELIMINATE_DENSE_COLUMNS              "eliminate_dense_columns"
 #define CUOPT_CUDSS_DETERMINISTIC                  "cudss_deterministic"

cpp/include/cuopt/linear_programming/pdlp/solver_settings.hpp

@@ -283,6 +283,7 @@ class pdlp_solver_settings_t {
   bool eliminate_dense_columns{true};
   pdlp_precision_t pdlp_precision{pdlp_precision_t::DefaultPrecision};
   bool barrier_iterative_refinement{true};
+  i_t barrier_soc_threshold{5};
   f_t barrier_step_scale{0.9};
   bool save_best_primal_so_far{false};
   /**

cpp/src/barrier/iterative_refinement.hpp

@@ -24,7 +24,6 @@
 
 #include <algorithm>
 #include <cmath>
-#include <cstdio>
 #include <limits>
 #include <vector>
 

cpp/src/barrier/second_order_cone_reduction.cuh

@@ -10,8 +10,8 @@
 #include <utilities/copy_helpers.hpp>
 #include <utilities/macros.cuh>
 
+#include <cub/block/block_reduce.cuh>
 #include <cub/device/device_reduce.cuh>
-#include <cub/device/device_segmented_reduce.cuh>
 #include <cub/warp/warp_reduce.cuh>
 
 #include <rmm/device_buffer.hpp>
@@ -24,7 +24,6 @@
 
 #include <thrust/iterator/constant_iterator.h>
 #include <thrust/iterator/discard_iterator.h>
-#include <thrust/iterator/permutation_iterator.h>
 
 #include <algorithm>
 #include <concepts>
@@ -46,14 +45,22 @@ __global__ void __launch_bounds__(warps_per_cta* raft::WarpSize)
                               OutputIt output,
                               value_t init);
 
+template <std::integral i_t, typename value_t, int block_dim, typename InputIt, typename OutputIt>
+__global__ void __launch_bounds__(block_dim)
+  block_per_cone_reduce_kernel(InputIt input,
+                               raft::device_span<const i_t> medium_cone_ids,
+                               raft::device_span<const std::size_t> cone_offsets,
+                               OutputIt output,
+                               value_t init);
+
 /**
  * Segmented-sum dispatcher for packed second-order cone vectors.
  *
  * Cone dimensions are fixed for a solve, so the constructor partitions cone
  * ids once by reduction strategy. Each call then reuses those partitions:
- * small cones use one warp per cone, medium cones use CUB DeviceSegmentedReduce,
+ * small cones use one warp per cone, medium cones use one block per cone,
  * and large cones use CUB DeviceReduce one cone at a time. The object owns the
- * CUB workspace for those medium/large paths. Call `prepare_workspace` once
+ * CUB workspace for the large-cone path. Call `prepare_workspace` once
  * before using a CUB-backed path.
  */
 template <std::integral i_t, i_t warp_cone_dim = 64, i_t large_cone_cutoff = 32768>
@@ -69,7 +76,7 @@ struct segmented_sum_t {
   std::vector<i_t> large_cone_ids;
   std::vector<i_t> large_cone_dimensions;
 
-  // Maximum CUB temporary storage needed by prepared medium/large reductions.
+  // Maximum CUB temporary storage needed by prepared large reductions.
   std::size_t cub_workspace_bytes = 0;
   rmm::device_buffer cub_workspace;
 
@@ -80,24 +87,6 @@ struct segmented_sum_t {
     auto input  = thrust::make_constant_iterator(value_t{});
     auto output = thrust::make_discard_iterator();
 
-    if (!medium_cone_ids.is_empty()) {
-      const auto medium_begin_offsets =
-        thrust::make_permutation_iterator(cone_offsets.data(), medium_cone_ids.begin());
-      const auto medium_end_offsets =
-        thrust::make_permutation_iterator(cone_offsets.data() + 1, medium_cone_ids.begin());
-
-      std::size_t temp_storage_bytes = 0;
-      RAFT_CUDA_TRY(cub::DeviceSegmentedReduce::Sum(nullptr,
-                                                    temp_storage_bytes,
-                                                    input,
-                                                    output,
-                                                    medium_cone_ids.size(),
-                                                    medium_begin_offsets,
-                                                    medium_end_offsets,
-                                                    stream.value()));
-      cub_workspace_bytes = std::max(cub_workspace_bytes, temp_storage_bytes);
-    }
-
     for (std::size_t i = 0; i < large_cone_ids.size(); ++i) {
       std::size_t temp_storage_bytes = 0;
       RAFT_CUDA_TRY(cub::DeviceReduce::Sum(nullptr,
@@ -109,7 +98,7 @@ struct segmented_sum_t {
       cub_workspace_bytes = std::max(cub_workspace_bytes, temp_storage_bytes);
     }
 
-    if (cub_workspace.size() < cub_workspace_bytes) {
+    if (cub_workspace_bytes > 0 && cub_workspace.size() < cub_workspace_bytes) {
       cub_workspace.resize(cub_workspace_bytes, stream);
     }
   }
@@ -138,31 +127,17 @@ struct segmented_sum_t {
     }
 
     if (!medium_cone_ids.is_empty()) {
-      cuopt_assert(cub_workspace_bytes > 0 && cub_workspace.size() >= cub_workspace_bytes,
-                   "segmented_sum_t::prepare_workspace must be called before reducing medium or "
-                   "large cones");
-
-      const auto medium_output = thrust::make_permutation_iterator(output, medium_cone_ids.begin());
-      const auto medium_begin_offsets =
-        thrust::make_permutation_iterator(cone_offsets.data(), medium_cone_ids.begin());
-      const auto medium_end_offsets =
-        thrust::make_permutation_iterator(cone_offsets.data() + 1, medium_cone_ids.begin());
-
-      std::size_t temp_storage_bytes = cub_workspace_bytes;
-      RAFT_CUDA_TRY(cub::DeviceSegmentedReduce::Sum(cub_workspace.data(),
-                                                    temp_storage_bytes,
-                                                    input,
-                                                    medium_output,
-                                                    medium_cone_ids.size(),
-                                                    medium_begin_offsets,
-                                                    medium_end_offsets,
-                                                    stream.value()));
+      constexpr int medium_block_dim = 256;
+      const auto n_medium            = medium_cone_ids.size();
+      block_per_cone_reduce_kernel<i_t, value_t, medium_block_dim>
+        <<<n_medium, medium_block_dim, 0, stream.value()>>>(
+          input, cuopt::make_span(medium_cone_ids), cone_offsets, output, init);
+      RAFT_CUDA_TRY(cudaPeekAtLastError());
     }
 
     if (!large_cone_ids.empty()) {
       cuopt_assert(cub_workspace_bytes > 0 && cub_workspace.size() >= cub_workspace_bytes,
-                   "segmented_sum_t::prepare_workspace must be called before reducing medium or "
-                   "large cones");
+                   "segmented_sum_t::prepare_workspace must be called before reducing large cones");
 
       for (std::size_t i = 0; i < large_cone_ids.size(); ++i) {
         std::size_t temp_storage_bytes = cub_workspace_bytes;
@@ -258,4 +233,55 @@ __global__ void __launch_bounds__(warps_per_cta* raft::WarpSize)
   if (lane_id == 0) { output[cone] = sum; }
 }
 
+template <std::integral i_t, typename value_t, int block_dim, typename InputIt, typename OutputIt>
+__global__ void __launch_bounds__(block_dim)
+  block_per_cone_reduce_kernel(InputIt input,
+                               raft::device_span<const i_t> medium_cone_ids,
+                               raft::device_span<const std::size_t> cone_offsets,
+                               OutputIt output,
+                               value_t init)
+{
+  static_assert(block_dim > 0);
+  static_assert(block_dim <= 1024);
+
+  constexpr int items_per_thread = 4;
+
+  using block_reduce_t = cub::BlockReduce<value_t, block_dim>;
+  __shared__ typename block_reduce_t::TempStorage temp_storage;
+
+  const auto slot = blockIdx.x;
+  if (slot >= medium_cone_ids.size()) { return; }
+
+  const auto cone = medium_cone_ids[slot];
+  const auto off  = cone_offsets[cone];
+  const auto dim  = cone_offsets[cone + 1] - off;
+
+  value_t acc[items_per_thread];
+#pragma unroll
+  for (int k = 0; k < items_per_thread; ++k) {
+    acc[k] = value_t{0};
+  }
+
+  // Thread t owns elements t, t+block_dim, ..., t+(items_per_thread-1)*block_dim
+  // within each tile of `block_dim * items_per_thread` consecutive entries, so
+  // every warp load stays coalesced.
+  const std::size_t tile = static_cast<std::size_t>(block_dim) * items_per_thread;
+  for (std::size_t tile_start = 0; tile_start < dim; tile_start += tile) {
+#pragma unroll
+    for (int k = 0; k < items_per_thread; ++k) {
+      const std::size_t idx = tile_start + threadIdx.x + static_cast<std::size_t>(k) * block_dim;
+      if (idx < dim) { acc[k] = acc[k] + input[off + idx]; }
+    }
+  }
+
+  value_t sum = init;
+#pragma unroll
+  for (int k = 0; k < items_per_thread; ++k) {
+    sum = sum + acc[k];
+  }
+
+  sum = block_reduce_t(temp_storage).Sum(sum);
+  if (threadIdx.x == 0) { output[cone] = sum; }
+}
+
 }  // namespace cuopt::linear_programming::dual_simplex

cpp/src/dual_simplex/simplex_solver_settings.hpp

@@ -69,6 +69,7 @@ struct simplex_solver_settings_t {
       eliminate_dense_columns(true),
       barrier_iterative_refinement(true),
       barrier_step_scale(0.9),
+      barrier_soc_threshold(5),
       num_gpus(1),
       folding(-1),
       augmented(0),
@@ -157,6 +158,7 @@ struct simplex_solver_settings_t {
   bool eliminate_dense_columns;       // true to eliminate dense columns from A*D*A^T
   bool barrier_iterative_refinement;  // true to use iterative refinement for barrier method
   f_t barrier_step_scale;             // step scale for barrier method
+  i_t barrier_soc_threshold;          // SOC dimension above which rank-2 sparse scaling is used
   int num_gpus;   // Number of GPUs to use (maximum of 2 gpus are supported at the moment)
   i_t folding;    // -1 automatic, 0 don't fold, 1 fold
   i_t augmented;  // -1 automatic, 0 to solve with ADAT, 1 to solve with augmented system

cpp/src/math_optimization/solver_settings.cu

@@ -132,6 +132,7 @@ solver_settings_t<i_t, f_t>::solver_settings_t() : pdlp_settings(), mip_settings
     {CUOPT_DUALIZE, &pdlp_settings.dualize, -1, 1, -1},
     {CUOPT_ORDERING, &pdlp_settings.ordering, -1, 1, -1},
     {CUOPT_BARRIER_DUAL_INITIAL_POINT, &pdlp_settings.barrier_dual_initial_point, -1, 1, -1},
+    {CUOPT_BARRIER_SOC_THRESHOLD, &pdlp_settings.barrier_soc_threshold, 0, 1000000, 5, "SOC dimension above which rank-2 sparse KKT expansion is used"},
     {CUOPT_MIP_CUT_PASSES, &mip_settings.max_cut_passes, -1, std::numeric_limits<i_t>::max(), 10},
     {CUOPT_MIP_MIXED_INTEGER_ROUNDING_CUTS, &mip_settings.mir_cuts, -1, 1, -1},
     {CUOPT_MIP_MIXED_INTEGER_GOMORY_CUTS, &mip_settings.mixed_integer_gomory_cuts, -1, 1, -1},

cpp/src/pdlp/solve.cu

@@ -511,6 +511,7 @@ run_barrier(dual_simplex::user_problem_t<i_t, f_t>& user_problem,
   barrier_settings.crossover                       = settings.crossover;
   barrier_settings.eliminate_dense_columns         = settings.eliminate_dense_columns;
   barrier_settings.barrier_iterative_refinement    = settings.barrier_iterative_refinement;
+  barrier_settings.barrier_soc_threshold           = settings.barrier_soc_threshold;
   barrier_settings.barrier_step_scale              = settings.barrier_step_scale;
   barrier_settings.cudss_deterministic             = settings.cudss_deterministic;
   barrier_settings.barrier_relaxed_feasibility_tol = settings.tolerances.relative_primal_tolerance;

cpp/tests/socp/CMakeLists.txt

@@ -5,6 +5,7 @@
 
 ConfigureTest(SOCP_TEST
     ${CMAKE_CURRENT_SOURCE_DIR}/second_order_cone_kernels.cu
+    ${CMAKE_CURRENT_SOURCE_DIR}/sparse_augmented_kkt_test.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/solve_barrier_socp.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/general_quadratic_test.cu
     LABELS numopt)