From ef91eacd51aeeeea8965cccdb7feae22088edfe8 Mon Sep 17 00:00:00 2001 From: Bo Tang Date: Mon, 26 Oct 2026 22:31:24 -0400 Subject: [PATCH 1/2] Fix: plan the reduction workspace from the reductions it serves --- CMakeLists.txt | 3 ++ internal/feasibility_polish.h | 4 ++ internal/infeasibility.h | 8 ++- internal/reduce.h | 32 ++++++++---- internal/termination.h | 8 +++ src/feasibility_polish.cu | 15 ++++-- src/infeasibility.cu | 74 ++++++++++++++++++--------- src/reduce.cu | 14 ++---- src/solver.cu | 47 ++++++++++++----- src/termination.cu | 62 +++++++++++++++-------- test/test_reduce_workspace.cu | 95 +++++++++++++++++++++++++++++++++++ 11 files changed, 280 insertions(+), 82 deletions(-) create mode 100644 test/test_reduce_workspace.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index 0c27288..fcfa052 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/internal/feasibility_polish.h b/internal/feasibility_polish.h index 2a3d1ee..3f559d4 100644 --- a/internal/feasibility_polish.h +++ b/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 diff --git a/internal/infeasibility.h b/internal/infeasibility.h index 5a3a4fc..7890cce 100644 --- a/internal/infeasibility.h +++ b/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 } diff --git a/internal/reduce.h b/internal/reduce.h index 1dcf130..bfed45c 100644 --- a/internal/reduce.h +++ b/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 void reduce_sum(pdhg_solver_state_t *state, int n, F f, double *out) +template 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(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 void reduce_max(pdhg_solver_state_t *state, int n, F f, double *out) +template 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(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 diff --git a/internal/termination.h b/internal/termination.h index a1dde82..6d3597d 100644 --- a/internal/termination.h +++ b/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); diff --git a/src/feasibility_polish.cu b/src/feasibility_polish.cu index c59af21..2b52142 100644 --- a/src/feasibility_polish.cu +++ b/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) diff --git a/src/infeasibility.cu b/src/infeasibility.cu index 68fc80a..e4d0748 100644 --- a/src/infeasibility.cu +++ b/src/infeasibility.cu @@ -150,7 +150,7 @@ __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; @@ -179,11 +179,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 +199,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 +385,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 +433,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->prod_m + (size_t)c * m; + 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->prod_n + (size_t)c * n; + 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 +478,13 @@ void launch_abs_products(pdhg_solver_state_t *state, bool primal_side, bool dual { abs_kernel<<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<<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); } } diff --git a/src/reduce.cu b/src/reduce.cu index e865a06..cbf810e 100644 --- a/src/reduce.cu +++ b/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)); } diff --git a/src/solver.cu b/src/solver.cu index 2bdcb59..d4c5269 100644 --- a/src/solver.cu +++ b/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, bool detect_infeasibility); 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); @@ -161,6 +163,8 @@ cupdlpx_result_t *optimize(const pdhg_parameters_t *params, const lp_problem_t * state->infeas_ctx = infeasibility_context_create(state); } + create_reduction_storage(state, detect_infeasibility); + // residuals of the starting point; also runs every check kernel once before the capture display_iteration_header(params); compute_fixed_point_error_device(state); @@ -656,7 +660,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, @@ -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,19 @@ static void initialize_step_size_and_primal_weight(pdhg_solver_state_t *state, c state->best_primal_weight = state->primal_weight; } +static void create_reduction_storage(pdhg_solver_state_t *state, bool detect_infeasibility) +{ + 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); + if (detect_infeasibility) + plan_infeasibility_reductions(state, &bytes); + reduce_temp_storage_create(state, bytes); +} + void pdhg_solver_state_free(pdhg_solver_state_t *state) { if (state == NULL) diff --git a/src/termination.cu b/src/termination.cu index 7125fb1..6e55ffc 100644 --- a/src/termination.cu +++ b/src/termination.cu @@ -105,6 +105,26 @@ void fetch_device_scalars(pdhg_solver_state_t *state) CUDA_CHECK(cudaStreamSynchronize(state->stream)); } +void launch_residual_reductions(const pdhg_solver_state_t *state, + size_t *required_bytes, + const double *objective_vector, + const double *primal_solution) +{ + device_scalars_t *out = state->scalars_d; + const int m = state->num_constraints, n = state->num_variables; + const double *primal_residual = state->primal_residual, *primal_slack = state->primal_slack; + const double *dual_residual = state->dual_residual, *dual_slack = state->dual_slack; + const double *objective = objective_vector, *x = primal_solution; + reduce_sum( + state, required_bytes, m, DEVICE_FN(i) { return primal_residual[i] * primal_residual[i]; }, &out->primal_res_sq); + reduce_max(state, required_bytes, m, DEVICE_FN(i) { return fabs(primal_residual[i]); }, &out->primal_res_max); + reduce_sum(state, required_bytes, m, DEVICE_FN(i) { return primal_slack[i]; }, &out->primal_slack_sum); + reduce_sum(state, required_bytes, n, DEVICE_FN(j) { return dual_residual[j] * dual_residual[j]; }, &out->dual_res_sq); + reduce_max(state, required_bytes, n, DEVICE_FN(j) { return fabs(dual_residual[j]); }, &out->dual_res_max); + reduce_sum(state, required_bytes, n, DEVICE_FN(j) { return objective[j] * x[j]; }, &out->objective_dot); + reduce_sum(state, required_bytes, n, DEVICE_FN(j) { return dual_slack[j] * x[j]; }, &out->dual_slack_dot); +} + void launch_residual_check(pdhg_solver_state_t *state, const double *objective_vector, const double *primal_solution, @@ -129,18 +149,7 @@ void launch_residual_check(pdhg_solver_state_t *state, state->num_constraints, state->num_variables); - device_scalars_t *out = state->scalars_d; - const int m = state->num_constraints, n = state->num_variables; - const double *primal_residual = state->primal_residual, *primal_slack = state->primal_slack; - const double *dual_residual = state->dual_residual, *dual_slack = state->dual_slack; - const double *objective = objective_vector, *x = primal_solution; - reduce_sum(state, m, DEVICE_FN(i) { return primal_residual[i] * primal_residual[i]; }, &out->primal_res_sq); - reduce_max(state, m, DEVICE_FN(i) { return fabs(primal_residual[i]); }, &out->primal_res_max); - reduce_sum(state, m, DEVICE_FN(i) { return primal_slack[i]; }, &out->primal_slack_sum); - reduce_sum(state, n, DEVICE_FN(j) { return dual_residual[j] * dual_residual[j]; }, &out->dual_res_sq); - reduce_max(state, n, DEVICE_FN(j) { return fabs(dual_residual[j]); }, &out->dual_res_max); - reduce_sum(state, n, DEVICE_FN(j) { return objective[j] * x[j]; }, &out->objective_dot); - reduce_sum(state, n, DEVICE_FN(j) { return dual_slack[j] * x[j]; }, &out->dual_slack_dot); + launch_residual_reductions(state, NULL, objective_vector, primal_solution); } void compute_residual_device(pdhg_solver_state_t *state) @@ -182,6 +191,25 @@ void compute_residual(pdhg_solver_state_t *state, norm_type_t optimality_norm) compute_residual_host(state, optimality_norm); } +void launch_fixed_point_reductions(const pdhg_solver_state_t *state, size_t *required_bytes) +{ + // the squared norms of the two steps and the interaction term of the + // primal step with A^T times the dual step + device_scalars_t *out = state->scalars_d; + const int m = state->num_constraints, n = state->num_variables; + const double *delta_primal = state->delta_primal_solution, *delta_dual = state->delta_dual_solution; + const double *delta_dual_product = state->delta_dual_product; + reduce_sum(state, required_bytes, m, DEVICE_FN(i) { return delta_dual[i] * delta_dual[i]; }, &out->fixed_point_dual_sq); + reduce_sum( + state, required_bytes, n, DEVICE_FN(j) { return delta_primal[j] * delta_primal[j]; }, &out->fixed_point_primal_sq); + reduce_sum( + state, + required_bytes, + n, + DEVICE_FN(j) { return delta_dual_product[j] * delta_primal[j]; }, + &out->fixed_point_cross); +} + // Queues the fixed point error of the current iterate. void compute_fixed_point_error_device(pdhg_solver_state_t *state) { @@ -193,15 +221,7 @@ void compute_fixed_point_error_device(pdhg_solver_state_t *state) cupdlpx_spmv_ATx(state->sparse_handle, state->spmv_ctx, state->delta_dual_solution, state->delta_dual_product); - // the squared norms of the two steps and the interaction term of the - // primal step with A^T times the dual step - device_scalars_t *out = state->scalars_d; - const int m = state->num_constraints, n = state->num_variables; - const double *delta_primal = state->delta_primal_solution, *delta_dual = state->delta_dual_solution; - const double *delta_dual_product = state->delta_dual_product; - reduce_sum(state, m, DEVICE_FN(i) { return delta_dual[i] * delta_dual[i]; }, &out->fixed_point_dual_sq); - reduce_sum(state, n, DEVICE_FN(j) { return delta_primal[j] * delta_primal[j]; }, &out->fixed_point_primal_sq); - reduce_sum(state, n, DEVICE_FN(j) { return delta_dual_product[j] * delta_primal[j]; }, &out->fixed_point_cross); + launch_fixed_point_reductions(state, NULL); } void compute_fixed_point_error_host(pdhg_solver_state_t *state) diff --git a/test/test_reduce_workspace.cu b/test/test_reduce_workspace.cu new file mode 100644 index 0000000..33203ef --- /dev/null +++ b/test/test_reduce_workspace.cu @@ -0,0 +1,95 @@ +/* +Copyright 2025 Haihao Lu + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +#include "reduce.h" + +#include +#include + +void launch_test_reductions( + const pdhg_solver_state_t *state, size_t *required_bytes, int n, const double *input, double *out); +static int expect_results(const double *out, double sum, double max_abs, double sum_sq); + +int main() +{ + const int n = 12000; + pdhg_solver_state_t state = {}; + CUDA_CHECK(cudaStreamCreate(&state.stream)); + double *input_h = (double *)safe_malloc(n * sizeof(double)); + for (int i = 0; i < n; ++i) + input_h[i] = (double)(i % 3) - 1.0; + double *input, *out; + CUDA_CHECK(cudaMalloc(&input, n * sizeof(double))); + CUDA_CHECK(cudaMalloc(&out, 3 * sizeof(double))); + CUDA_CHECK(cudaMemcpy(input, input_h, n * sizeof(double), cudaMemcpyHostToDevice)); + free(input_h); + + size_t bytes = 0; + launch_test_reductions(&state, &bytes, n, input, out); + reduce_temp_storage_create(&state, bytes); + + cudaGraph_t graph; + cudaGraphExec_t graphExec; + CUDA_CHECK(cudaStreamBeginCapture(state.stream, cudaStreamCaptureModeGlobal)); + launch_test_reductions(&state, NULL, n, input, out); + CUDA_CHECK(cudaStreamEndCapture(state.stream, &graph)); + CUDA_CHECK(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0)); + int failed = 0; + for (int replay = 0; replay < 2; ++replay) + { + CUDA_CHECK(cudaGraphLaunch(graphExec, state.stream)); + CUDA_CHECK(cudaStreamSynchronize(state.stream)); + failed |= expect_results(out, 0.0, 1.0, 8000.0); + } + launch_test_reductions(&state, NULL, 0, input, out); + CUDA_CHECK(cudaStreamSynchronize(state.stream)); + failed |= expect_results(out, 0.0, 0.0, 0.0); + + CUDA_CHECK(cudaGraphExecDestroy(graphExec)); + CUDA_CHECK(cudaGraphDestroy(graph)); + CUDA_CHECK(cudaFree(state.reduce_temp_d)); + CUDA_CHECK(cudaFree(input)); + CUDA_CHECK(cudaFree(out)); + CUDA_CHECK(cudaStreamDestroy(state.stream)); + if (failed) + return 1; + printf("Planned reduction workspace checks passed.\n"); + return 0; +} + +void launch_test_reductions( + const pdhg_solver_state_t *state, size_t *required_bytes, int n, const double *input, double *out) +{ + reduce_sum(state, required_bytes, n, DEVICE_FN(i) { return input[i]; }, out); + reduce_max(state, required_bytes, n, DEVICE_FN(i) { return fabs(input[i]); }, out + 1); + reduce_sum(state, required_bytes, n, DEVICE_FN(i) { return input[i] * input[i]; }, out + 2); +} + +static int expect_results(const double *out, double sum, double max_abs, double sum_sq) +{ + double expected[3] = {sum, max_abs, sum_sq}; + double actual[3]; + CUDA_CHECK(cudaMemcpy(actual, out, sizeof(actual), cudaMemcpyDeviceToHost)); + for (int k = 0; k < 3; ++k) + { + if (actual[k] != expected[k]) + { + fprintf(stderr, "Reduction %d: expected %.17g, got %.17g\n", k, expected[k], actual[k]); + return 1; + } + } + return 0; +} From c3a10c2f1ab4784a31dd45e1c66fd92f239fd7c9 Mon Sep 17 00:00:00 2001 From: ZedongPeng Date: Fri, 30 Oct 2026 15:16:48 -0400 Subject: [PATCH 2/2] Fix: create the reduction workspace with the solver state --- src/infeasibility.cu | 13 +++++++------ src/solver.cu | 15 +++++++++------ 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/infeasibility.cu b/src/infeasibility.cu index e4d0748..c10322b 100644 --- a/src/infeasibility.cu +++ b/src/infeasibility.cu @@ -155,10 +155,11 @@ void launch_candidate_reductions(const pdhg_solver_state_t *state, size_t *requi 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; @@ -438,7 +439,7 @@ void launch_primal_abs_reductions(const pdhg_solver_state_t *state, size_t *requ const int m = state->num_constraints; for (int c = 0; c < INFEAS_NUM_CANDIDATES; c++) { - const double *p = state->infeas_ctx->prod_m + (size_t)c * m; + const double *p = state->infeas_ctx ? state->infeas_ctx->prod_m + (size_t)c * m : NULL; reduce_max( state, required_bytes, @@ -453,7 +454,7 @@ void launch_dual_abs_reductions(const pdhg_solver_state_t *state, size_t *requir const int n = state->num_variables; for (int c = 0; c < INFEAS_NUM_CANDIDATES; c++) { - const double *p = state->infeas_ctx->prod_n + (size_t)c * n; + 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); } diff --git a/src/solver.cu b/src/solver.cu index d4c5269..46a7c87 100644 --- a/src/solver.cu +++ b/src/solver.cu @@ -56,7 +56,7 @@ static cupdlpx_result_t *create_result_from_state(pdhg_solver_state_t *state, co 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, bool detect_infeasibility); +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); @@ -163,8 +163,6 @@ cupdlpx_result_t *optimize(const pdhg_parameters_t *params, const lp_problem_t * state->infeas_ctx = infeasibility_context_create(state); } - create_reduction_storage(state, detect_infeasibility); - // residuals of the starting point; also runs every check kernel once before the capture display_iteration_header(params); compute_fixed_point_error_device(state); @@ -669,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; } @@ -1093,7 +1093,11 @@ static void initialize_step_size_and_primal_weight(pdhg_solver_state_t *state, c state->best_primal_weight = state->primal_weight; } -static void create_reduction_storage(pdhg_solver_state_t *state, bool detect_infeasibility) +// 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); @@ -1101,8 +1105,7 @@ static void create_reduction_storage(pdhg_solver_state_t *state, bool detect_inf launch_restart_reductions(state, &bytes); launch_polish_fixed_point_reductions(state, &bytes, true); launch_polish_fixed_point_reductions(state, &bytes, false); - if (detect_infeasibility) - plan_infeasibility_reductions(state, &bytes); + plan_infeasibility_reductions(state, &bytes); reduce_temp_storage_create(state, bytes); }