Refactor: move finite-bound computation

Author: LucasBoTang

src/preconditioner.cu

@@ -27,16 +27,12 @@ limitations under the License.
 __global__ void scale_variables_kernel(double *__restrict__ objective_vector,
                                        double *__restrict__ variable_lower_bound,
                                        double *__restrict__ variable_upper_bound,
-                                       double *__restrict__ variable_lower_bound_finite_val,
-                                       double *__restrict__ variable_upper_bound_finite_val,
                                        double *__restrict__ initial_primal_solution,
                                        const double *__restrict__ variable_rescaling,
                                        const double *__restrict__ inverse_variable_rescaling,
                                        int num_variables);
 __global__ void scale_constraints_kernel(double *__restrict__ constraint_lower_bound,
                                          double *__restrict__ constraint_upper_bound,
-                                         double *__restrict__ constraint_lower_bound_finite_val,
-                                         double *__restrict__ constraint_upper_bound_finite_val,
                                          double *__restrict__ initial_dual_solution,
                                          const double *__restrict__ constraint_rescaling,
                                          const double *__restrict__ inverse_constraint_rescaling,
@@ -70,8 +66,6 @@ __global__ void reduce_bound_norm_sq_kernel(
 __global__ void scale_bounds_kernel(
     double *__restrict__ constraint_lower_bound,
     double *__restrict__ constraint_upper_bound,
-    double *__restrict__ constraint_lower_bound_finite_val,
-    double *__restrict__ constraint_upper_bound_finite_val,
     double *__restrict__ initial_dual_solution,
     int num_constraints,
     double constraint_scale,
@@ -80,8 +74,6 @@ __global__ void scale_objective_kernel(
     double *__restrict__ objective_vector,
     double *__restrict__ variable_lower_bound,
     double *__restrict__ variable_upper_bound,
-    double *__restrict__ variable_lower_bound_finite_val,
-    double *__restrict__ variable_upper_bound_finite_val,
     double *__restrict__ initial_primal_solution,
     int num_variables,
     double constraint_scale,
@@ -108,8 +100,6 @@ static void scale_problem(
         state->objective_vector,
         state->variable_lower_bound,
         state->variable_upper_bound,
-        state->variable_lower_bound_finite_val,
-        state->variable_upper_bound_finite_val,
         state->initial_primal_solution,
         variable_rescaling,
         inverse_variable_rescaling,
@@ -118,8 +108,6 @@ static void scale_problem(
     scale_constraints_kernel<<<state->num_blocks_dual, THREADS_PER_BLOCK>>>(
         state->constraint_lower_bound,
         state->constraint_upper_bound,
-        state->constraint_lower_bound_finite_val,
-        state->constraint_upper_bound_finite_val,
         state->initial_dual_solution,
         constraint_rescaling,
         inverse_constraint_rescaling,
@@ -248,8 +236,6 @@ static void bound_objective_rescaling(
     scale_bounds_kernel<<<state->num_blocks_dual, THREADS_PER_BLOCK>>>(
         state->constraint_lower_bound,
         state->constraint_upper_bound,
-        state->constraint_lower_bound_finite_val,
-        state->constraint_upper_bound_finite_val,
         state->initial_dual_solution,
         num_constraints,
         constraint_scale,
@@ -259,8 +245,6 @@ static void bound_objective_rescaling(
         state->objective_vector,
         state->variable_lower_bound,
         state->variable_upper_bound,
-        state->variable_lower_bound_finite_val,
-        state->variable_upper_bound_finite_val,
         state->initial_primal_solution,
         num_variables,
         constraint_scale,
@@ -326,8 +310,6 @@ rescale_info_t *rescale_problem(
 __global__ void scale_variables_kernel(double *__restrict__ objective_vector,
                                        double *__restrict__ variable_lower_bound,
                                        double *__restrict__ variable_upper_bound,
-                                       double *__restrict__ variable_lower_bound_finite_val,
-                                       double *__restrict__ variable_upper_bound_finite_val,
                                        double *__restrict__ initial_primal_solution,
                                        const double *__restrict__ variable_rescaling,
                                        const double *__restrict__ inverse_variable_rescaling,
@@ -341,15 +323,11 @@ __global__ void scale_variables_kernel(double *__restrict__ objective_vector,
     objective_vector[j] *= inv_dj;
     variable_lower_bound[j] *= dj;
     variable_upper_bound[j] *= dj;
-    variable_lower_bound_finite_val[j] *= dj;
-    variable_upper_bound_finite_val[j] *= dj;
     initial_primal_solution[j] *= dj;
 }
 
 __global__ void scale_constraints_kernel(double *__restrict__ constraint_lower_bound,
                                          double *__restrict__ constraint_upper_bound,
-                                         double *__restrict__ constraint_lower_bound_finite_val,
-                                         double *__restrict__ constraint_upper_bound_finite_val,
                                          double *__restrict__ initial_dual_solution,
                                          const double *__restrict__ constraint_rescaling,
                                          const double *__restrict__ inverse_constraint_rescaling,
@@ -362,8 +340,6 @@ __global__ void scale_constraints_kernel(double *__restrict__ constraint_lower_b
     double ei = constraint_rescaling[i];
     constraint_lower_bound[i] *= inv_ei;
     constraint_upper_bound[i] *= inv_ei;
-    constraint_lower_bound_finite_val[i] *= inv_ei;
-    constraint_upper_bound_finite_val[i] *= inv_ei;
     initial_dual_solution[i] *= ei;
 }
 
@@ -504,8 +480,6 @@ __global__ void reduce_bound_norm_sq_kernel(
 __global__ void scale_bounds_kernel(
     double *__restrict__ constraint_lower_bound,
     double *__restrict__ constraint_upper_bound,
-    double *__restrict__ constraint_lower_bound_finite_val,
-    double *__restrict__ constraint_upper_bound_finite_val,
     double *__restrict__ initial_dual_solution,
     int num_constraints,
     double constraint_scale,
@@ -516,17 +490,13 @@ __global__ void scale_bounds_kernel(
         return;
     constraint_lower_bound[i] *= constraint_scale;
     constraint_upper_bound[i] *= constraint_scale;
-    constraint_lower_bound_finite_val[i] *= constraint_scale;
-    constraint_upper_bound_finite_val[i] *= constraint_scale;
     initial_dual_solution[i] *= objective_scale;
 }
 
 __global__ void scale_objective_kernel(
     double *__restrict__ objective_vector,
     double *__restrict__ variable_lower_bound,
     double *__restrict__ variable_upper_bound,
-    double *__restrict__ variable_lower_bound_finite_val,
-    double *__restrict__ variable_upper_bound_finite_val,
     double *__restrict__ initial_primal_solution,
     int num_variables,
     double constraint_scale,
@@ -537,8 +507,6 @@ __global__ void scale_objective_kernel(
         return;
     variable_lower_bound[j] *= constraint_scale;
     variable_upper_bound[j] *= constraint_scale;
-    variable_lower_bound_finite_val[j] *= constraint_scale;
-    variable_upper_bound_finite_val[j] *= constraint_scale;
     objective_vector[j] *= objective_scale;
     initial_primal_solution[j] *= constraint_scale;
 }

src/solver.cu

@@ -37,6 +37,11 @@ __global__ void build_transpose_map(const int *__restrict__ A_row_ind,
                                     const int *__restrict__ At_col_ind,
                                     int nnz,
                                     int *__restrict__ A_to_At);
+__global__ void fill_finite_bounds_kernel(const double *__restrict__ lower_bound,
+                                          const double *__restrict__ upper_bound,
+                                          double *__restrict__ lower_bound_finite_val,
+                                          double *__restrict__ upper_bound_finite_val,
+                                          int num_elements);
 __global__ void compute_next_pdhg_primal_solution_kernel(const double *current_primal,
                                                          double *reflected_primal,
                                                          const double *dual_product,
@@ -402,6 +407,25 @@ static pdhg_solver_state_t *initialize_solver_state(
     rescale_info->var_rescale = NULL;
     rescale_info_free(rescale_info);
 
+    CUDA_CHECK(cudaMalloc(&state->constraint_lower_bound_finite_val, con_bytes));
+    CUDA_CHECK(cudaMalloc(&state->constraint_upper_bound_finite_val, con_bytes));
+    CUDA_CHECK(cudaMalloc(&state->variable_lower_bound_finite_val, var_bytes));
+    CUDA_CHECK(cudaMalloc(&state->variable_upper_bound_finite_val, var_bytes));
+
+    fill_finite_bounds_kernel<<<state->num_blocks_dual, THREADS_PER_BLOCK>>>(
+        state->constraint_lower_bound,
+        state->constraint_upper_bound,
+        state->constraint_lower_bound_finite_val,
+        state->constraint_upper_bound_finite_val,
+        n_cons);
+
+    fill_finite_bounds_kernel<<<state->num_blocks_primal, THREADS_PER_BLOCK>>>(
+        state->variable_lower_bound,
+        state->variable_upper_bound,
+        state->variable_lower_bound_finite_val,
+        state->variable_upper_bound_finite_val,
+        n_vars);
+
     CUDA_CHECK(cudaFree(state->constraint_matrix->row_ind));
     state->constraint_matrix->row_ind = NULL;
     CUDA_CHECK(cudaFree(state->constraint_matrix_t->row_ind));
@@ -571,6 +595,24 @@ __global__ void build_transpose_map(
     A_to_At[k] = pos;
 }
 
+__global__ void fill_finite_bounds_kernel(
+    const double *__restrict__ lb,
+    const double *__restrict__ ub,
+    double *__restrict__ lb_finite,
+    double *__restrict__ ub_finite,
+    int num_elements)
+{
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= num_elements)
+        return;
+
+    double Li = lb[i];
+    double Ui = ub[i];
+
+    lb_finite[i] = isfinite(Li) ? Li : 0.0;
+    ub_finite[i] = isfinite(Ui) ? Ui : 0.0;
+}
+
 __global__ void compute_next_pdhg_primal_solution_kernel(
     const double *current_primal, double *reflected_primal,
     const double *dual_product, const double *objective, const double *var_lb,