Add Jacobi Precondition in Inner Solving

README.md

@@ -99,6 +99,7 @@ Solver Parameters:
 | --inner_init_tol | double | Initial tolerance for the inner solver. | 1e-3 |
 | --inner_min_tol | double | Minimum tolerance for the inner solver. | 1e-9 |
 | --presolve | int | Enable (1) or disable (0) presolve. | 1 |
+| --no_diag_precond | flag | Disable the Jacobi diagonal preconditioner used in the inner subproblem (enabled by default). | false |
 
 **Distributed Options** (only available when built with `-DPDHCG_COMPILE_DISTRIBUTED=ON`):
 | Option | Type | Description | Default |

distributed/distributed_solver.cu

@@ -433,8 +433,16 @@ pdhcg_result_t *distributed_optimize(const pdhg_parameters_t *params, const qp_p
             row_perm = (int *)malloc(working_problem->num_constraints * sizeof(int));
             col_perm = (int *)malloc(working_problem->num_variables * sizeof(int));
 
-            generate_random_permutation(working_problem->num_constraints, row_perm);
-            generate_random_permutation(working_problem->num_variables, col_perm);
+            if (params->permute_method == FULL_RANDOM_PERMUTATION)
+            {
+                generate_random_permutation(working_problem->num_variables, col_perm);
+                generate_random_permutation(working_problem->num_constraints, row_perm);
+            }
+            else if (params->permute_method == BLOCK_RANDOM_PERMUTATION)
+            {
+                generate_block_permutation(working_problem->num_variables, params->permute_block_size, col_perm);
+                generate_block_permutation(working_problem->num_constraints, params->permute_block_size, row_perm);
+            }
 
             permuted_problem = permute_problem_return_new(working_problem, row_perm, col_perm);
             working_problem = permuted_problem;

distributed/distributed_utils.cu

@@ -979,11 +979,12 @@ void print_distributed_params(const pdhg_parameters_t *params)
 
     if (params->grid_size.decided)
     {
-        printf("  Grid Size         : %d x %d (Rows x Cols)\n", params->grid_size.row_dims, params->grid_size.col_dims);
+        printf(
+            "  Grid Size          : %d x %d (Rows x Cols)\n", params->grid_size.row_dims, params->grid_size.col_dims);
     }
     else
     {
-        printf("  Grid Size         : Auto-detect (implementation dependent)\n");
+        printf("  Grid Size          : Auto-detect (implementation dependent)\n");
     }
 
     printf("  Partition Method   : ");

docs/c/types.md

@@ -235,6 +235,7 @@ typedef struct {
   norm_type_t optimality_norm;
   inner_solver_parameters_t inner_solver_parameters;
   bool presolve;
+  bool diag_jacobi_precond;
   partition_method_t partition_method;
   permute_method_t permute_method;
   grid_size_t grid_size;

docs/index.md

@@ -28,7 +28,7 @@ Where:
 - **GPU Acceleration**: Fully leverages NVIDIA CUDA for extreme-scale QP problems
 - **Flexible Problem Structure**: Supports sparse quadratic terms, low-rank quadratic terms, or both
 - **High Performance**: Competitive with commercial solvers on large-scale problems
-- **SpMVOp Auto-Detection**: Automatically uses cuSPARSE SpMVOp (fused SpMV+elementwise operations) on CUDA 13+ while falling back to standard SpMV on CUDA 12.x
+- **SpMVOp Auto-Detection**: Automatically uses cuSPARSE SpMVOp on CUDA 13+ while falling back to standard SpMV on CUDA 12.x
 - **Multi-GPU Distributed Solving**: Supports parallel solving across multiple GPUs via MPI and NCCL (optional, enabled at compile time)
 
 ## Quick Links

docs/installation.md

@@ -10,7 +10,7 @@
 !!! note "CUDA Version and SpMVOp"
     PDHCG automatically detects your CUDA version at compile time:
 
-    - **CUDA 13+**: Uses cuSPARSE **SpMVOp** (fused SpMV + elementwise operations) for improved performance.
+    - **CUDA 13+**: Uses cuSPARSE **SpMVOp** for improved performance.
     - **CUDA 12.x**: Falls back to the standard **SpMV** API. No manual intervention is required.
 
 ## C++ Executable
@@ -51,6 +51,9 @@ When enabled, the solver binary automatically detects whether it is launched wit
 
 ## Python Package
 
+!!! note "Multi-GPU support"
+    The Python interface currently supports single-GPU solving only. For multi-GPU distributed solving, build the C++ executable with `-DPDHCG_COMPILE_DISTRIBUTED=ON` and launch it via `mpirun`.
+
 ### From PyPI (Recommended)
 
 ```bash

docs/python/parameters.md

@@ -40,6 +40,7 @@ m.setParams(TimeLimit=3600, LogLevel=1)
 | `InnerIterLimit` | int | 1000 | Max iterations for inner CG solver |
 | `InnerInitTol` | float | 1e-3 | Initial tolerance for inner solver |
 | `InnerMinTol` | float | 1e-9 | Minimum tolerance for inner solver |
+| `DiagJacobiPrecond` | bool | True | Use the Jacobi diagonal preconditioner in the inner subproblem. Set to `False` to disable. |
 
 ### Singular Value Estimation
 

include/pdhcg_types.h

@@ -151,6 +151,7 @@ extern "C"
         norm_type_t optimality_norm;
         inner_solver_parameters_t inner_solver_parameters;
         bool presolve;
+        bool diag_jacobi_precond;
         partition_method_t partition_method;
         permute_method_t permute_method;
         grid_size_t grid_size;

internal/internal_types.h

@@ -70,6 +70,13 @@ typedef struct
     double *gradient;
     double *direction;
     double *scalar_buffer;
+    bool precond_enabled;
+    double *diag_h_static;
+    double *m_diag;
+    double *m_inv;
+    double *Ms_buffer;
+    double cached_inv_tau;
+    double tol_scale;
 } bb_step_size_t;
 
 typedef struct

internal/pdhcg_kernels.cuh

@@ -170,6 +170,18 @@ extern "C"
 
     __global__ void compute_bb_alpha_safeguard_kernel(const double *d_norm_gtg, const double *d_tmp, double *d_alpha);
 
+    __global__ void compute_bb_alpha_M_kernel(const double *d_stMs, const double *d_tmp, double *d_alpha);
+
+    __global__ void scalar_sqrt_copy_kernel(const double *src, double *dst);
+
+    __global__ void
+    compute_csr_diag_kernel(const int *row_ptr, const int *col_ind, const double *val, double *diag, int num_rows);
+
+    __global__ void compute_csr_row_sq_norm_kernel(const int *row_ptr, const double *val, double *out, int num_rows);
+
+    __global__ void refresh_inner_precond_kernel(
+        const double *diag_h_static, double inv_tau, double *m_diag, double *m_inv, int n_vars);
+
     __global__ void primal_gradient_descent_kernel_bb_init(const double *dual_product,
                                                            double *gradient,
                                                            double *direction,
@@ -200,6 +212,28 @@ extern "C"
                                                       const double *d_alpha,
                                                       const int n_vars);
 
+    __global__ void primal_gradient_descent_kernel_bb_init_precond(const double *dual_product,
+                                                                   double *gradient,
+                                                                   double *direction,
+                                                                   const double *current_primal_solution,
+                                                                   double *pdhg_primal_solution,
+                                                                   const double *objective_vector,
+                                                                   const double *objective_product,
+                                                                   const double *var_lb,
+                                                                   const double *var_ub,
+                                                                   const double *m_inv,
+                                                                   const double stepsize,
+                                                                   const int n_vars);
+
+    __global__ void primal_bb_update_direction_kernel_precond(double *pdhg_primal_solution,
+                                                              const double *gradient,
+                                                              double *direction,
+                                                              const double *var_lb,
+                                                              const double *var_ub,
+                                                              const double *m_inv,
+                                                              const double *d_alpha,
+                                                              const int n_vars);
+
     __global__ void primal_bb_final_kernel(const double *current_primal_solution,
                                            const double *pdhg_primal_solution,
                                            double *reflected_primal_solution,

python/README.md

@@ -20,9 +20,12 @@ It provides a high-level, Pythonic API for constructing, modifying, and solving
 !!! note "CUDA Version and SpMVOp"
     PDHCG automatically detects your CUDA version at compile time:
 
-    - **CUDA 13+**: Uses cuSPARSE **SpMVOp** (fused SpMV + elementwise operations) for improved performance.
+    - **CUDA 13+**: Uses cuSPARSE **SpMVOp** for improved performance.
     - **CUDA 12.x**: Falls back to the standard **SpMV** API. No manual intervention is required.
 
+!!! note "Multi-GPU support"
+    The Python interface currently supports single-GPU solving only. For multi-GPU distributed solving, build the C++ executable with `-DPDHCG_COMPILE_DISTRIBUTED=ON` and launch it via `mpirun` (see the [main README](../README.md)).
+
 ### Install
 Install from PyPI:
 
@@ -189,6 +192,7 @@ Below is a list of commonly used parameters, their internal keys, and descriptio
 | `InnerIterLimit` | `inner_iter_limit` | int | `1000` | Maximum number of iterations for the inner solver. |
 | `InnerInitTol` | `inner_init_tol` | float | `1e-3` | Initial tolerance for the inner solver. |
 | `InnerMinTol` | `inner_min_tol` | float | `1e-9` | Minimum tolerance for the inner solver. |
+| `DiagJacobiPrecond` | `diag_jacobi_precond` | bool | `True` | Whether to use the Jacobi diagonal preconditioner in the inner subproblem. Set to `False` to disable. |
 
 They can be set in multiple ways:
 

python/pdhcg/PDHCG.py

@@ -61,6 +61,7 @@
     "InnerIterLimit": "inner_iter_limit",
     "InnerInitTol": "inner_init_tol",
     "InnerMinTol": "inner_min_tol",
+    "DiagJacobiPrecond": "diag_jacobi_precond",
     # presolve
     "Presolve": "presolve",
 }

python_bindings/_core_bindings.cpp

@@ -296,6 +296,7 @@ static py::dict get_default_params_py()
     d["inner_iter_limit"] = p.inner_solver_parameters.iteration_limit;
     d["inner_init_tol"] = p.inner_solver_parameters.initial_tolerance;
     d["inner_min_tol"] = p.inner_solver_parameters.min_tolerance;
+    d["diag_jacobi_precond"] = p.diag_jacobi_precond;
 
     return d;
 }
@@ -381,6 +382,7 @@ static void parse_params_from_python(py::object params_obj, pdhg_parameters_t *p
     geti("inner_iter_limit", p->inner_solver_parameters.iteration_limit);
     getf("inner_init_tol", p->inner_solver_parameters.initial_tolerance);
     getf("inner_min_tol", p->inner_solver_parameters.min_tolerance);
+    getb("diag_jacobi_precond", p->diag_jacobi_precond);
 
     // presolve
     getb("presolve", p->presolve);

src/cli.c

@@ -169,6 +169,9 @@ void print_usage(const char *prog_name)
     fprintf(stderr, "      --inner_init_tol     Initial tolerance for the inner solver (default: 1e-3).\n");
     fprintf(stderr, "      --inner_min_tol      Minimum tolerance for the inner solver (default: 1e-9).\n");
     fprintf(stderr, "      --presolve <int>     Enable (1) or disable (0) presolve (default: 1).\n");
+    fprintf(
+        stderr,
+        "      --no_diag_precond    Disable Jacobi diagonal preconditioner for inner subproblem (default: enabled).\n");
 
 #ifdef PDHCG_COMPILE_DISTRIBUTED
     fprintf(stderr, "\nDistributed Options (MPI & NCCL):\n");
@@ -206,6 +209,7 @@ int run_pdhcg(int argc, char *argv[])
                                            {"inner_init_tol", required_argument, 0, 1016},
                                            {"inner_min_tol", required_argument, 0, 1017},
                                            {"presolve", required_argument, 0, 1018},
+                                           {"no_diag_precond", no_argument, 0, 1019},
                                            {0, 0, 0, 0}};
 
     int opt;
@@ -285,6 +289,9 @@ int run_pdhcg(int argc, char *argv[])
             case 1018:
                 params.presolve = (atoi(optarg) != 0);
                 break;
+            case 1019:
+                params.diag_jacobi_precond = false;
+                break;
             case '?':
                 return 1;
         }
@@ -376,6 +383,7 @@ int run_d_pdhcg(int argc, char *argv[])
                                            {"inner_init_tol", required_argument, 0, 1016},
                                            {"inner_min_tol", required_argument, 0, 1017},
                                            {"presolve", required_argument, 0, 1018},
+                                           {"no_diag_precond", no_argument, 0, 1019},
                                            {"grid_size", required_argument, 0, 2001},
                                            {"partition_method", required_argument, 0, 2002},
                                            {"permute_method", required_argument, 0, 2003},
@@ -462,6 +470,9 @@ int run_d_pdhcg(int argc, char *argv[])
             case 1018:
                 params.presolve = (atoi(optarg) != 0);
                 break;
+            case 1019:
+                params.diag_jacobi_precond = false;
+                break;
             case 2001: // --grid_size r,c
             {
                 int r, c;

src/pdhcg_kernels.cu

@@ -277,6 +277,65 @@ __global__ void compute_bb_alpha_safeguard_kernel(const double *d_norm_gtg, cons
     *d_alpha = (*d_norm_gtg * *d_norm_gtg) / *d_tmp;
 }
 
+__global__ void compute_bb_alpha_M_kernel(const double *d_stMs, const double *d_tmp, double *d_alpha)
+{
+    *d_alpha = *d_stMs / *d_tmp;
+}
+
+__global__ void scalar_sqrt_copy_kernel(const double *src, double *dst)
+{
+    *dst = sqrt(*src);
+}
+
+__global__ void
+compute_csr_diag_kernel(const int *row_ptr, const int *col_ind, const double *val, double *diag, int num_rows)
+{
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < num_rows)
+    {
+        double sum = 0.0;
+        int start = row_ptr[i];
+        int end = row_ptr[i + 1];
+        for (int k = start; k < end; ++k)
+        {
+            if (col_ind[k] == i)
+                sum += val[k];
+        }
+        diag[i] = sum;
+    }
+}
+
+__global__ void compute_csr_row_sq_norm_kernel(const int *row_ptr, const double *val, double *out, int num_rows)
+{
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < num_rows)
+    {
+        double sum = 0.0;
+        int start = row_ptr[i];
+        int end = row_ptr[i + 1];
+        for (int k = start; k < end; ++k)
+        {
+            double v = val[k];
+            sum += v * v;
+        }
+        out[i] = sum;
+    }
+}
+
+__global__ void
+refresh_inner_precond_kernel(const double *diag_h_static, double inv_tau, double *m_diag, double *m_inv, int n_vars)
+{
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n_vars)
+    {
+        double m = diag_h_static[i] + inv_tau;
+        if (m <= 0.0)
+            m = 1.0;
+        m_diag[i] = m;
+        m_inv[i] = 1.0 / m;
+    }
+}
+
 __global__ void primal_gradient_descent_kernel_bb_init(const double *dual_product,
                                                        double *gradient,
                                                        double *direction,
@@ -357,6 +416,53 @@ __global__ void primal_bb_final_kernel(const double *current_primal_solution,
     }
 }
 
+__global__ void primal_gradient_descent_kernel_bb_init_precond(const double *dual_product,
+                                                               double *gradient,
+                                                               double *direction,
+                                                               const double *current_primal_solution,
+                                                               double *pdhg_primal_solution,
+                                                               const double *objective_vector,
+                                                               const double *objective_product,
+                                                               const double *var_lb,
+                                                               const double *var_ub,
+                                                               const double *m_inv,
+                                                               const double stepsize,
+                                                               const int n_vars)
+{
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n_vars)
+    {
+        double current_grad = objective_product[i] + objective_vector[i] - dual_product[i];
+        double current_primal_sol = current_primal_solution[i];
+        double next_primal_sol = current_primal_sol - stepsize * m_inv[i] * current_grad;
+        next_primal_sol = fmax(var_lb[i], fmin(next_primal_sol, var_ub[i]));
+        pdhg_primal_solution[i] = next_primal_sol;
+        gradient[i] = current_grad;
+        direction[i] = next_primal_sol - current_primal_sol;
+    }
+}
+
+__global__ void primal_bb_update_direction_kernel_precond(double *pdhg_primal_solution,
+                                                          const double *gradient,
+                                                          double *direction,
+                                                          const double *var_lb,
+                                                          const double *var_ub,
+                                                          const double *m_inv,
+                                                          const double *d_alpha,
+                                                          const int n_vars)
+{
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n_vars)
+    {
+        double alpha = *d_alpha;
+        double cur_sol = pdhg_primal_solution[i];
+        double next_sol = cur_sol - alpha * m_inv[i] * gradient[i];
+        next_sol = fmax(var_lb[i], fmin(next_sol, var_ub[i]));
+        direction[i] = next_sol - cur_sol;
+        pdhg_primal_solution[i] = next_sol;
+    }
+}
+
 __global__ void compute_lp_residual_kernel(double *primal_residual,
                                            const double *primal_product,
                                            const double *constraint_lower_bound,