Fix undersized reduction workspac

CMakeLists.txt

@@ -308,6 +308,9 @@ if(CUPDLPX_BUILD_TESTS)
     enable_testing()
     file(GLOB TEST_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/test/*.c" "${CMAKE_CURRENT_SOURCE_DIR}/test/*.cu")
     foreach(TEST_SRC ${TEST_SOURCES})
+        if(USE_HIP AND TEST_SRC MATCHES "\\.cu$")
+            set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP)
+        endif()
         get_filename_component(TEST_NAME ${TEST_SRC} NAME_WE)
         add_executable(${TEST_NAME} ${TEST_SRC})
         target_link_libraries(${TEST_NAME} PRIVATE cupdlpx_core)

internal/feasibility_polish.h

@@ -11,6 +11,10 @@ extern "C"
 
     void compute_polish_fixed_point_error_device(pdhg_solver_state_t *state, bool is_primal_polish);
 
+    void launch_polish_fixed_point_reductions(const pdhg_solver_state_t *state,
+                                              size_t *required_bytes,
+                                              bool is_primal_polish);
+
 #ifdef __cplusplus
 }
 #endif

internal/infeasibility.h

@@ -33,12 +33,18 @@ extern "C"
 
     void infeasibility_context_destroy(infeas_ctx_t *ctx);
 
+    // Raises required_bytes to the reductions of both rounds, even when a
+    // second-round side is not needed initially; see reduce.h.
+    void plan_infeasibility_reductions(const pdhg_solver_state_t *state, size_t *required_bytes);
+
     void compute_infeasibility_device(pdhg_solver_state_t *state);
 
     void compute_infeasibility_host(pdhg_solver_state_t *state, double objective_margin);
 
-    void launch_candidate_reductions(pdhg_solver_state_t *state, int c);
+    void launch_candidate_reductions(const pdhg_solver_state_t *state, size_t *required_bytes, int c);
     void launch_abs_products(pdhg_solver_state_t *state, bool primal_side, bool dual_side);
+    void launch_primal_abs_reductions(const pdhg_solver_state_t *state, size_t *required_bytes);
+    void launch_dual_abs_reductions(const pdhg_solver_state_t *state, size_t *required_bytes);
 
 #ifdef __cplusplus
 }

internal/reduce.h

@@ -29,35 +29,49 @@
 // The element function of a reduction: DEVICE_FN(i) { return ...; }
 #define DEVICE_FN(i) [=] __host__ __device__(int i)
 
-// Temporary storage of cub::DeviceReduce, sized once for the longest vector and
-// shared by every call. cub rejects a call that needs more.
-void reduce_temp_storage_create(pdhg_solver_state_t *state);
+// Temporary storage of cub::DeviceReduce, shared by every call. cub rejects a
+// call that needs more, and captured graphs keep its address, so it is sized
+// once by planning every reduction.
+void reduce_temp_storage_create(pdhg_solver_state_t *state, size_t bytes);
+
+// With required_bytes set, a reduction only raises it to the storage it needs.
+// Planning calls the same functions as execution, so cub sees the same types.
+// On Windows, nvcc rejects these lambdas in functions with internal linkage, so
+// functions that define reductions are not static.
 
 // out = sum over i < n of f(i).
-template <typename F> void reduce_sum(pdhg_solver_state_t *state, int n, F f, double *out)
+template <typename F> void reduce_sum(const pdhg_solver_state_t *state, size_t *required_bytes, int n, F f, double *out)
 {
     if (n <= 0)
     {
-        CUDA_CHECK(cudaMemsetAsync(out, 0, sizeof(double), state->stream));
+        if (!required_bytes)
+            CUDA_CHECK(cudaMemsetAsync(out, 0, sizeof(double), state->stream));
         return;
     }
     auto it = thrust::make_transform_iterator(thrust::counting_iterator<int>(0), f);
+    void *temp = required_bytes ? NULL : state->reduce_temp_d;
     size_t bytes = state->reduce_temp_bytes;
-    CUDA_CHECK(cub::DeviceReduce::Sum(state->reduce_temp_d, bytes, it, out, n, state->stream));
+    CUDA_CHECK(cub::DeviceReduce::Sum(temp, bytes, it, out, n, state->stream));
+    if (required_bytes && bytes > *required_bytes)
+        *required_bytes = bytes;
 }
 
 // out = max over i < n of f(i); 0 when n is 0, since every maximum taken here is
 // nonnegative.
-template <typename F> void reduce_max(pdhg_solver_state_t *state, int n, F f, double *out)
+template <typename F> void reduce_max(const pdhg_solver_state_t *state, size_t *required_bytes, int n, F f, double *out)
 {
     if (n <= 0)
     {
-        CUDA_CHECK(cudaMemsetAsync(out, 0, sizeof(double), state->stream));
+        if (!required_bytes)
+            CUDA_CHECK(cudaMemsetAsync(out, 0, sizeof(double), state->stream));
         return;
     }
     auto it = thrust::make_transform_iterator(thrust::counting_iterator<int>(0), f);
+    void *temp = required_bytes ? NULL : state->reduce_temp_d;
     size_t bytes = state->reduce_temp_bytes;
-    CUDA_CHECK(cub::DeviceReduce::Max(state->reduce_temp_d, bytes, it, out, n, state->stream));
+    CUDA_CHECK(cub::DeviceReduce::Max(temp, bytes, it, out, n, state->stream));
+    if (required_bytes && bytes > *required_bytes)
+        *required_bytes = bytes;
 }
 
 #endif // CUPDLPX_REDUCE_H

internal/termination.h

@@ -48,6 +48,14 @@ extern "C"
                                const double *constraint_lower_bound_finite,
                                const double *constraint_upper_bound_finite);
 
+    // The reductions of the residual and fixed point checks; see reduce.h for
+    // required_bytes.
+    void launch_residual_reductions(const pdhg_solver_state_t *state,
+                                    size_t *required_bytes,
+                                    const double *objective_vector,
+                                    const double *primal_solution);
+    void launch_fixed_point_reductions(const pdhg_solver_state_t *state, size_t *required_bytes);
+
     void compute_residual_device(pdhg_solver_state_t *state);
 
     void compute_residual_host(pdhg_solver_state_t *state, norm_type_t optimality_norm);

src/feasibility_polish.cu

@@ -205,6 +205,16 @@ static void compute_polish_residual_host(pdhg_solver_state_t *state, norm_type_t
     }
 }
 
+void launch_polish_fixed_point_reductions(const pdhg_solver_state_t *state,
+                                          size_t *required_bytes,
+                                          bool is_primal_polish)
+{
+    const double *delta = is_primal_polish ? state->delta_primal_solution : state->delta_dual_solution;
+    int len = is_primal_polish ? state->num_variables : state->num_constraints;
+    double *out = is_primal_polish ? &state->scalars_d->fixed_point_primal_sq : &state->scalars_d->fixed_point_dual_sq;
+    reduce_sum(state, required_bytes, len, DEVICE_FN(i) { return delta[i] * delta[i]; }, out);
+}
+
 // The fixed point error of a polishing iterate: only its own side of the
 // movement term, without the square root (the restart test compares ratios).
 void compute_polish_fixed_point_error_device(pdhg_solver_state_t *state, bool is_primal_polish)
@@ -214,10 +224,7 @@ void compute_polish_fixed_point_error_device(pdhg_solver_state_t *state, bool is
                                 state->reflected_primal_solution,
                                 state->pdhg_dual_solution,
                                 state->reflected_dual_solution);
-    const double *delta = is_primal_polish ? state->delta_primal_solution : state->delta_dual_solution;
-    int len = is_primal_polish ? state->num_variables : state->num_constraints;
-    double *out = is_primal_polish ? &state->scalars_d->fixed_point_primal_sq : &state->scalars_d->fixed_point_dual_sq;
-    reduce_sum(state, len, DEVICE_FN(i) { return delta[i] * delta[i]; }, out);
+    launch_polish_fixed_point_reductions(state, NULL, is_primal_polish);
 }
 
 static void compute_polish_fixed_point_error_host(pdhg_solver_state_t *state, bool is_primal_polish)

src/infeasibility.cu

@@ -150,15 +150,16 @@ __global__ void form_dual_candidates_kernel(candidate_directions_t directions,
 
 // First-round scalars of candidate c. Violations are reduced in the rescaled
 // space, which decides the certificate, and in original units, which are reported.
-void launch_candidate_reductions(pdhg_solver_state_t *state, int c)
+void launch_candidate_reductions(const pdhg_solver_state_t *state, size_t *required_bytes, int c)
 {
     infeas_ctx_t *ctx = state->infeas_ctx;
     const int n = state->num_variables, m = state->num_constraints;
     infeas_candidate_scalars_t *out = &state->scalars_d->infeas[c];
-    const double *d = ctx->D + (size_t)c * n; // the primal ray
-    const double *y = ctx->Y + (size_t)c * m; // the dual ray
-    const double *Ad = ctx->prod_m + (size_t)c * m;
-    const double *ATy = ctx->prod_n + (size_t)c * n;
+    // NULL while planning before the context exists; a size query never reads them
+    const double *d = ctx ? ctx->D + (size_t)c * n : NULL; // the primal ray
+    const double *y = ctx ? ctx->Y + (size_t)c * m : NULL; // the dual ray
+    const double *Ad = ctx ? ctx->prod_m + (size_t)c * m : NULL;
+    const double *ATy = ctx ? ctx->prod_n + (size_t)c * n : NULL;
     const double *objective = state->objective_vector;
     const double *variable_rescaling = state->variable_rescaling;
     const double *constraint_rescaling = state->constraint_rescaling;
@@ -179,11 +180,11 @@ void launch_candidate_reductions(pdhg_solver_state_t *state, int c)
     { return primal_ray_row_violation(Ad[i], isfinite(constraint_lower[i]), isfinite(constraint_upper[i])); };
     auto row_violation_orig = DEVICE_FN(i) { return row_violation(i) * constraint_rescaling[i]; };
 
-    reduce_max(state, n, d_abs_orig, &out->primal_ray.max_abs);
-    reduce_sum(state, n, objective_term, &out->primal_ray.obj);
-    reduce_sum(state, n, objective_term_abs, &out->primal_ray.obj_abs);
-    reduce_max(state, m, row_violation, &out->primal_ray.viol_max_scaled);
-    reduce_max(state, m, row_violation_orig, &out->primal_ray.viol_max_orig);
+    reduce_max(state, required_bytes, n, d_abs_orig, &out->primal_ray.max_abs);
+    reduce_sum(state, required_bytes, n, objective_term, &out->primal_ray.obj);
+    reduce_sum(state, required_bytes, n, objective_term_abs, &out->primal_ray.obj_abs);
+    reduce_max(state, required_bytes, m, row_violation, &out->primal_ray.viol_max_scaled);
+    reduce_max(state, required_bytes, m, row_violation_orig, &out->primal_ray.viol_max_orig);
 
     // dual ray y and its reduced cost r = -A^T y
     auto y_abs_orig = DEVICE_FN(i) { return fabs(y[i]) / (constraint_rescaling[i] * objective_rescaling); };
@@ -199,13 +200,13 @@ void launch_candidate_reductions(pdhg_solver_state_t *state, int c)
     { return dual_ray_reduced_cost_violation(-ATy[j], isfinite(variable_lower[j]), isfinite(variable_upper[j])); };
     auto reduced_cost_violation_orig = DEVICE_FN(j) { return reduced_cost_violation(j) * variable_rescaling[j]; };
 
-    reduce_max(state, m, y_abs_orig, &out->dual_ray.max_abs);
-    reduce_sum(state, m, con_bounds_term, &out->dual_ray.obj_con_bounds);
-    reduce_sum(state, n, var_bounds_term, &out->dual_ray.obj_var_bounds);
-    reduce_sum(state, m, con_bounds_term_abs, &out->dual_ray.obj_con_bounds_abs);
-    reduce_sum(state, n, var_bounds_term_abs, &out->dual_ray.obj_var_bounds_abs);
-    reduce_max(state, n, reduced_cost_violation, &out->dual_ray.viol_max_scaled);
-    reduce_max(state, n, reduced_cost_violation_orig, &out->dual_ray.viol_max_orig);
+    reduce_max(state, required_bytes, m, y_abs_orig, &out->dual_ray.max_abs);
+    reduce_sum(state, required_bytes, m, con_bounds_term, &out->dual_ray.obj_con_bounds);
+    reduce_sum(state, required_bytes, n, var_bounds_term, &out->dual_ray.obj_var_bounds);
+    reduce_sum(state, required_bytes, m, con_bounds_term_abs, &out->dual_ray.obj_con_bounds_abs);
+    reduce_sum(state, required_bytes, n, var_bounds_term_abs, &out->dual_ray.obj_var_bounds_abs);
+    reduce_max(state, required_bytes, n, reduced_cost_violation, &out->dual_ray.viol_max_scaled);
+    reduce_max(state, required_bytes, n, reduced_cost_violation_orig, &out->dual_ray.viol_max_orig);
 }
 
 static void spmm_buffer(
@@ -385,7 +386,7 @@ void compute_infeasibility_device(pdhg_solver_state_t *state)
     // everything that needs D, Y or the products, before the second round overwrites them
     for (int c = 0; c < K; c++)
     {
-        launch_candidate_reductions(state, c);
+        launch_candidate_reductions(state, NULL, c);
     }
 }
 
@@ -433,6 +434,40 @@ evaluate_objective_margins(const pdhg_solver_state_t *state, double objective_ma
     }
 }
 
+void launch_primal_abs_reductions(const pdhg_solver_state_t *state, size_t *required_bytes)
+{
+    const int m = state->num_constraints;
+    for (int c = 0; c < INFEAS_NUM_CANDIDATES; c++)
+    {
+        const double *p = state->infeas_ctx ? state->infeas_ctx->prod_m + (size_t)c * m : NULL;
+        reduce_max(
+            state,
+            required_bytes,
+            m,
+            DEVICE_FN(i) { return p[i]; },
+            &state->scalars_d->infeas[c].primal_ray.product_abs_max);
+    }
+}
+
+void launch_dual_abs_reductions(const pdhg_solver_state_t *state, size_t *required_bytes)
+{
+    const int n = state->num_variables;
+    for (int c = 0; c < INFEAS_NUM_CANDIDATES; c++)
+    {
+        const double *p = state->infeas_ctx ? state->infeas_ctx->prod_n + (size_t)c * n : NULL;
+        reduce_max(
+            state, required_bytes, n, DEVICE_FN(j) { return p[j]; }, &state->scalars_d->infeas[c].dual_ray.product_abs_max);
+    }
+}
+
+void plan_infeasibility_reductions(const pdhg_solver_state_t *state, size_t *required_bytes)
+{
+    for (int c = 0; c < INFEAS_NUM_CANDIDATES; c++)
+        launch_candidate_reductions(state, required_bytes, c);
+    launch_primal_abs_reductions(state, required_bytes);
+    launch_dual_abs_reductions(state, required_bytes);
+}
+
 // Second round for the sides that need it: D becomes |D| in place and prod_m
 // becomes |A||D|; likewise on the dual side.
 void launch_abs_products(pdhg_solver_state_t *state, bool primal_side, bool dual_side)
@@ -444,21 +479,13 @@ void launch_abs_products(pdhg_solver_state_t *state, bool primal_side, bool dual
     {
         abs_kernel<<<state->num_blocks_primal * K, THREADS_PER_BLOCK, 0, state->stream>>>(ctx->D, ctx->D, n * K);
         spmm(state->sparse_handle, ctx->matAbsA, ctx->dnD, ctx->dnProdM, ctx->buf_absA);
-        for (int c = 0; c < K; c++)
-        {
-            const double *p = ctx->prod_m + (size_t)c * m;
-            reduce_max(state, m, DEVICE_FN(i) { return p[i]; }, &state->scalars_d->infeas[c].primal_ray.product_abs_max);
-        }
+        launch_primal_abs_reductions(state, NULL);
     }
     if (dual_side)
     {
         abs_kernel<<<state->num_blocks_dual * K, THREADS_PER_BLOCK, 0, state->stream>>>(ctx->Y, ctx->Y, m * K);
         spmm(state->sparse_handle, ctx->matAbsAT, ctx->dnY, ctx->dnProdN, ctx->buf_absAT);
-        for (int c = 0; c < K; c++)
-        {
-            const double *p = ctx->prod_n + (size_t)c * n;
-            reduce_max(state, n, DEVICE_FN(j) { return p[j]; }, &state->scalars_d->infeas[c].dual_ray.product_abs_max);
-        }
+        launch_dual_abs_reductions(state, NULL);
     }
 }
 

src/reduce.cu

@@ -18,16 +18,8 @@ limitations under the License.
 
 #include "reduce.h"
 
-void reduce_temp_storage_create(pdhg_solver_state_t *state)
+void reduce_temp_storage_create(pdhg_solver_state_t *state, size_t bytes)
 {
-    int n_max = state->num_variables > state->num_constraints ? state->num_variables : state->num_constraints;
-    if (n_max < 1)
-        n_max = 1;
-    const double *in = NULL;
-    double *out = NULL;
-    size_t sum_bytes = 0, max_bytes = 0;
-    CUDA_CHECK(cub::DeviceReduce::Sum(NULL, sum_bytes, in, out, n_max, state->stream));
-    CUDA_CHECK(cub::DeviceReduce::Max(NULL, max_bytes, in, out, n_max, state->stream));
-    state->reduce_temp_bytes = sum_bytes > max_bytes ? sum_bytes : max_bytes;
-    CUDA_CHECK(cudaMalloc(&state->reduce_temp_d, state->reduce_temp_bytes));
+    state->reduce_temp_bytes = bytes;
+    CUDA_CHECK(cudaMalloc(&state->reduce_temp_d, bytes));
 }

src/solver.cu

@@ -54,7 +54,9 @@ __global__ void rescale_solution_kernel(double *__restrict__ primal_solution,
                                         int n_cons);
 static cupdlpx_result_t *create_result_from_state(pdhg_solver_state_t *state, const lp_problem_t *original_problem);
 void perform_restart(pdhg_solver_state_t *state, const pdhg_parameters_t *params);
+void launch_restart_reductions(const pdhg_solver_state_t *state, size_t *required_bytes);
 static void initialize_step_size_and_primal_weight(pdhg_solver_state_t *state, const pdhg_parameters_t *params);
+static void create_reduction_storage(pdhg_solver_state_t *state);
 static pdhg_solver_state_t *initialize_solver_state(const lp_problem_t *working_problem,
                                                     const pdhg_parameters_t *params,
                                                     objective_sense_t original_objective_sense);
@@ -656,7 +658,6 @@ static pdhg_solver_state_t *initialize_solver_state(const lp_problem_t *working_
     CUDA_CHECK(cudaStreamCreate(&state->stream));
     CUSPARSE_CHECK(cusparseSetStream(state->sparse_handle, state->stream));
     CUBLAS_CHECK(cublasSetStream(state->blas_handle, state->stream));
-    reduce_temp_storage_create(state);
 
     state->spmv_ctx = cupdlpx_spmv_ctx_create(state->sparse_handle,
                                               state->constraint_matrix,
@@ -666,6 +667,8 @@ static pdhg_solver_state_t *initialize_solver_state(const lp_problem_t *working_
                                               state->pdhg_dual_solution,
                                               state->dual_product);
 
+    create_reduction_storage(state);
+
     return state;
 }
 
@@ -968,17 +971,7 @@ void perform_restart(pdhg_solver_state_t *state, const pdhg_parameters_t *params
                                 state->initial_dual_solution,
                                 state->pdhg_dual_solution);
 
-    const double *delta_primal = state->delta_primal_solution, *delta_dual = state->delta_dual_solution;
-    reduce_sum(
-        state,
-        state->num_variables,
-        DEVICE_FN(j) { return delta_primal[j] * delta_primal[j]; },
-        &state->scalars_d->restart_primal_sq);
-    reduce_sum(
-        state,
-        state->num_constraints,
-        DEVICE_FN(i) { return delta_dual[i] * delta_dual[i]; },
-        &state->scalars_d->restart_dual_sq);
+    launch_restart_reductions(state, NULL);
     fetch_device_scalars(state);
     double primal_dist = sqrt(state->scalars_h->restart_primal_sq);
     double dual_dist = sqrt(state->scalars_h->restart_dual_sq);
@@ -1042,6 +1035,23 @@ void perform_restart(pdhg_solver_state_t *state, const pdhg_parameters_t *params
     sync_step_sizes_to_gpu(state);
 }
 
+void launch_restart_reductions(const pdhg_solver_state_t *state, size_t *required_bytes)
+{
+    const double *delta_primal = state->delta_primal_solution, *delta_dual = state->delta_dual_solution;
+    reduce_sum(
+        state,
+        required_bytes,
+        state->num_variables,
+        DEVICE_FN(j) { return delta_primal[j] * delta_primal[j]; },
+        &state->scalars_d->restart_primal_sq);
+    reduce_sum(
+        state,
+        required_bytes,
+        state->num_constraints,
+        DEVICE_FN(i) { return delta_dual[i] * delta_dual[i]; },
+        &state->scalars_d->restart_dual_sq);
+}
+
 static void initialize_step_size_and_primal_weight(pdhg_solver_state_t *state, const pdhg_parameters_t *params)
 {
     if (state->constraint_matrix->num_nonzeros == 0)
@@ -1083,6 +1093,22 @@ static void initialize_step_size_and_primal_weight(pdhg_solver_state_t *state, c
     state->best_primal_weight = state->primal_weight;
 }
 
+// Sizes the reduction workspace by planning every reduction a solve may run.
+// The infeasibility reductions are included whether or not detection is on,
+// so the size depends on no setting; they are planned before the context
+// exists, which a size query never reads.
+static void create_reduction_storage(pdhg_solver_state_t *state)
+{
+    size_t bytes = 0;
+    launch_residual_reductions(state, &bytes, state->objective_vector, state->pdhg_primal_solution);
+    launch_fixed_point_reductions(state, &bytes);
+    launch_restart_reductions(state, &bytes);
+    launch_polish_fixed_point_reductions(state, &bytes, true);
+    launch_polish_fixed_point_reductions(state, &bytes, false);
+    plan_infeasibility_reductions(state, &bytes);
+    reduce_temp_storage_create(state, bytes);
+}
+
 void pdhg_solver_state_free(pdhg_solver_state_t *state)
 {
     if (state == NULL)