Skip to content

Commit

Permalink
use block structure when evaluating residuals
Browse files Browse the repository at this point in the history
  • Loading branch information
@RSchwan
RSchwan committed Feb 3, 2025
1 parent d880dc0 commit 1511a89
Show file tree
Hide file tree
Showing 6 changed files with 145 additions and 31 deletions.
21 changes: 21 additions & 0 deletions include/piqp/dense/kkt.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,27 @@ class KKT : public KKTSystem<T> {
#endif
}

// z = alpha * P * x
void eval_P_x(const T& alpha, const CVecRef<T>& x, VecRef<T> z)
{
z.noalias() = alpha * data.P_utri * x;
z.noalias() += data.P_utri.transpose().template triangularView<Eigen::StrictlyLower>() * (alpha * x);
}

// zn = beta_n * yn + alpha_n * A * xn, zt = beta_t * yt + alpha_t * A^T * xt
void eval_A_xn_and_AT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt)
{
zn.noalias() = beta_n * yn + alpha_n * data.AT.transpose() * xn;
zt.noalias() = beta_t * yt + alpha_t * data.AT * xt;
}

// zn = beta_n * yn + alpha_n * G * xn, zt = beta_t * yt + alpha_t * G^T * xt
void eval_G_xn_and_GT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt)
{
zn.noalias() = beta_n * yn + alpha_n * data.GT.transpose() * xn;
zt.noalias() = beta_t * yt + alpha_t * data.GT * xt;
}

Mat<T>& internal_kkt_mat()
{
return kkt_mat;
Expand Down
7 changes: 7 additions & 0 deletions include/piqp/kkt_system.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,13 @@ class KKTSystem
VecRef<T> delta_z, VecRef<T> delta_z_lb, VecRef<T> delta_z_ub,
VecRef<T> delta_s, VecRef<T> delta_s_lb, VecRef<T> delta_s_ub) = 0;

// z = alpha * P * x
virtual void eval_P_x(const T& alpha, const CVecRef<T>& x, VecRef<T> z) = 0;
// zn = beta_n * yn + alpha_n * A * xn, zt = beta_t * yt + alpha_t * A^T * xt
virtual void eval_A_xn_and_AT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt) = 0;
// zn = beta_n * yn + alpha_n * G * xn, zt = beta_t * yt + alpha_t * G^T * xt
virtual void eval_G_xn_and_GT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt) = 0;

virtual void print_info() {};
};

Expand Down
18 changes: 12 additions & 6 deletions include/piqp/solver.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -835,8 +835,7 @@ class SolverBase
using std::abs;

// first part of dual residual and infeasibility calculation (used in cost calculation)
rx_nr.noalias() = -m_data.P_utri * m_result.x;
rx_nr.noalias() -= m_data.P_utri.transpose().template triangularView<Eigen::StrictlyLower>() * m_result.x;
m_kkt->eval_P_x(T(-1), m_result.x, rx_nr);
m_result.info.dual_rel_inf = m_preconditioner.unscale_dual_res(rx_nr).template lpNorm<Eigen::Infinity>();

// calculate primal cost, dual cost, and duality gap
Expand Down Expand Up @@ -866,11 +865,18 @@ class SolverBase
m_result.info.dual_obj = m_preconditioner.unscale_cost(m_result.info.dual_obj);
m_result.info.duality_gap = m_preconditioner.unscale_cost(m_result.info.duality_gap);

// ry_nr = -A * x
// dx = A^T * y, use dx as a temporary
m_kkt->eval_A_xn_and_AT_xt(T(-1), T(1), m_result.x, m_result.y, T(0), T(0), ry_nr, dx, ry_nr, dx);
// rz_nr = -G * x
// dx += G^T * z
m_kkt->eval_G_xn_and_GT_xt(T(-1), T(1), m_result.x, m_result.z, T(0), T(1), rz_nr, dx, rz_nr, dx);

// dual residual and infeasibility calculation
rx_nr.noalias() -= m_data.c;
m_result.info.dual_rel_inf = (std::max)(m_result.info.dual_rel_inf, m_preconditioner.unscale_dual_res(m_data.c).template lpNorm<Eigen::Infinity>());
dx.noalias() = m_data.AT * m_result.y; // use dx as a temporary
dx.noalias() += m_data.GT * m_result.z;
// dx.noalias() = m_data.AT * m_result.y; // use dx as a temporary
// dx.noalias() += m_data.GT * m_result.z;
for (isize i = 0; i < m_data.n_lb; i++)
{
dx(m_data.x_lb_idx(i)) -= m_data.x_lb_scaling(i) * m_result.z_lb(i);
Expand All @@ -883,12 +889,12 @@ class SolverBase
rx_nr.noalias() -= dx;

// primal residual and infeasibility calculation
ry_nr.noalias() = -m_data.AT.transpose() * m_result.x;
// ry_nr.noalias() = -m_data.AT.transpose() * m_result.x;
m_result.info.primal_rel_inf = m_preconditioner.unscale_primal_res_eq(ry_nr).template lpNorm<Eigen::Infinity>();
ry_nr.noalias() += m_data.b;
m_result.info.primal_rel_inf = (std::max)(m_result.info.primal_rel_inf, m_preconditioner.unscale_primal_res_eq(m_data.b).template lpNorm<Eigen::Infinity>());

rz_nr.noalias() = -m_data.GT.transpose() * m_result.x;
// rz_nr.noalias() = -m_data.GT.transpose() * m_result.x;
m_result.info.primal_rel_inf = (std::max)(m_result.info.primal_rel_inf, m_preconditioner.unscale_primal_res_ineq(rz_nr).template lpNorm<Eigen::Infinity>());
rz_nr.noalias() += m_data.h - m_result.s;
m_result.info.primal_rel_inf = (std::max)(m_result.info.primal_rel_inf, m_preconditioner.unscale_primal_res_ineq(m_data.h).template lpNorm<Eigen::Infinity>());
Expand Down
20 changes: 10 additions & 10 deletions include/piqp/sparse/blocksparse/block_vec.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,13 +117,13 @@ struct BlockVec
int block_size = x[block_idx].rows();
for (int inner_idx = 0; inner_idx < block_size; inner_idx++)
{
other(i++) = BLASFEO_DVECEL(this->x[block_idx].ref(), inner_idx);
other(i++) += BLASFEO_DVECEL(this->x[block_idx].ref(), inner_idx);
}
}
while (i < other.rows())
{
other(i++) = 0;
}
// while (i < other.rows())
// {
// other(i++) = 0;
// }
}

template <typename Derived1, typename Derived2>
Expand All @@ -135,13 +135,13 @@ struct BlockVec
int block_size = x[block_idx].rows();
for (int inner_idx = 0; inner_idx < block_size; inner_idx++)
{
other(perm_inv(i++)) = BLASFEO_DVECEL(this->x[block_idx].ref(), inner_idx);
other(perm_inv(i++)) += BLASFEO_DVECEL(this->x[block_idx].ref(), inner_idx);
}
}
while (i < other.rows())
{
other(perm_inv(i++)) = 0;
}
// while (i < other.rows())
// {
// other(perm_inv(i++)) = 0;
// }
}
};

Expand Down
89 changes: 74 additions & 15 deletions include/piqp/sparse/blocksparse_stage_kkt.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
block_delta_z.assign(delta_z, GT.perm_inv);

// block_rhs_x = P * block_delta_x, P is symmetric and only the lower triangular part of P is accessed
block_symv_l(P, block_delta_x, block_rhs_x);
block_symv_l(1.0, P, block_delta_x, block_rhs_x);
// block_rhs_x += AT * block_delta_y
// block_rhs_y = A * block_delta_x
block_t_gemv_nt(1.0, 1.0, AT, block_delta_y, block_delta_x, 1.0, 0.0, block_rhs_x, block_rhs_y, block_rhs_x, block_rhs_y);
Expand All @@ -239,10 +239,10 @@ class BlocksparseStageKKT : public KKTSystem<T>
block_t_gemv_nt(1.0, 1.0, GT, block_delta_z, block_delta_x, 1.0, 0.0, block_rhs_x, block_rhs_z, block_rhs_x, block_rhs_z);

rhs_x.setZero();
rhs_y.setZero();
rhs_z.setZero();
block_rhs_x.load(rhs_x);
rhs_y.setZero();
block_rhs_y.load(rhs_y, AT.perm_inv);
rhs_z.setZero();
block_rhs_z.load(rhs_z, GT.perm_inv);

rhs_x.noalias() += m_rho * delta_x;
Expand Down Expand Up @@ -332,8 +332,11 @@ class BlocksparseStageKKT : public KKTSystem<T>
// block_delta_z = G * block_delta_x
block_t_gemv_t(1.0, GT, block_delta_x, 0.0, block_delta_z, block_delta_z);

delta_x.setZero();
block_delta_x.load(delta_x);
delta_y.setZero();
block_delta_y.load(delta_y, AT.perm_inv);
delta_z.setZero();
block_delta_z.load(delta_z, GT.perm_inv);

delta_y.noalias() -= delta_inv * rhs_y;
Expand Down Expand Up @@ -363,14 +366,70 @@ class BlocksparseStageKKT : public KKTSystem<T>
m_s_ub.head(data.n_ub).array() * delta_z_ub.head(data.n_ub).array());
}

// z = alpha * P * x
void eval_P_x(const T& alpha, const CVecRef<T>& x, VecRef<T> z)
{
BlockVec& block_x = tmp1_x_block;
BlockVec& block_z = tmp2_x_block;

block_x.assign(x);
// block_z = alpha * P * block_x, P is symmetric and only the lower triangular part of P is accessed
block_symv_l(alpha, P, block_x, block_z);

z.setZero();
block_z.load(z);
}

// zn = beta_n * yn + alpha_n * A * xn, zt = beta_t * yt + alpha_t * A^T * xt
void eval_A_xn_and_AT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt)
{
BlockVec& block_xn = tmp1_x_block;
BlockVec& block_xt = tmp1_y_block;
BlockVec& block_zn = tmp2_y_block;
BlockVec& block_zt = tmp2_x_block;

block_xn.assign(xn);
block_xt.assign(xt, AT.perm_inv);

// block_zt = alpha_t * AT * block_xt
// block_zn = alpha_n * A * block_xn
block_t_gemv_nt(alpha_t, alpha_n, AT, block_xt, block_xn, 0.0, 0.0, block_zt, block_zn, block_zt, block_zn);

zn.noalias() = beta_n * yn;
block_zn.load(zn, AT.perm_inv);
zt.noalias() = beta_t * yt;
block_zt.load(zt);
}

// zn = beta_n * yn + alpha_n * G * xn, zt = beta_t * yt + alpha_t * G^T * xt
void eval_G_xn_and_GT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt)
{
BlockVec& block_xn = tmp1_x_block;
BlockVec& block_xt = tmp1_z_block;
BlockVec& block_zn = tmp2_z_block;
BlockVec& block_zt = tmp2_x_block;

block_xn.assign(xn);
block_xt.assign(xt, GT.perm_inv);

// block_zt = alpha_t * GT * block_xt
// block_zn = alpha_n * G * block_xn
block_t_gemv_nt(alpha_t, alpha_n, GT, block_xt, block_xn, 0.0, 0.0, block_zt, block_zn, block_zt, block_zn);

zn.noalias() = beta_n * yn;
block_zn.load(zn, GT.perm_inv);
zt.noalias() = beta_t * yt;
block_zt.load(zt);
}

void print_info()
{
std::size_t N = block_info.size();
piqp_print("block sizes:");
for (std::size_t i = 0; i < N - 1; i++) {
piqp_print(" %zd", block_info[i].width);
piqp_print(" %d", block_info[i].width);
}
piqp_print("\narrow width: %zd\n", block_info[N - 1].width);
piqp_print("\narrow width: %d\n", block_info[N - 1].width);
}

protected:
Expand Down Expand Up @@ -1278,8 +1337,8 @@ class BlocksparseStageKKT : public KKTSystem<T>
blasfeo_dpotrf_l(arrow_width, kkt_factor.D[N-1]->ref(), 0, 0, kkt_factor.D[N-1]->ref(), 0, 0);
}

// z = sA * x
void block_symv_l(BlockKKT& sA, BlockVec& x, BlockVec& z)
// z = alpha * sA * x
void block_symv_l(double alpha, BlockKKT& sA, BlockVec& x, BlockVec& z)
{
std::size_t N = block_info.size();
I arrow_width = block_info.back().width;
Expand All @@ -1289,8 +1348,8 @@ class BlocksparseStageKKT : public KKTSystem<T>
int m = sA.D[i]->rows();
assert(x.x[i].rows() == m && "size mismatch");
assert(z.x[i].rows() == m && "size mismatch");
// z_i = D_i * x_i, D_i is symmetric and only the lower triangular part of D_i is accessed
blasfeo_dsymv_l(m, 1.0, sA.D[i]->ref(), 0, 0, x.x[i].ref(), 0, 0.0, z.x[i].ref(), 0, z.x[i].ref(), 0);
// z_i = alpha * D_i * x_i, D_i is symmetric and only the lower triangular part of D_i is accessed
blasfeo_dsymv_l(m, alpha, sA.D[i]->ref(), 0, 0, x.x[i].ref(), 0, 0.0, z.x[i].ref(), 0, z.x[i].ref(), 0);
} else {
z.x[i].setZero();
}
Expand All @@ -1304,9 +1363,9 @@ class BlocksparseStageKKT : public KKTSystem<T>
assert(x.x[i+1].rows() == m && "size mismatch");
assert(z.x[i+1].rows() == m && "size mismatch");
assert(z.x[i].rows() == n && "size mismatch");
// z_{i+1} += B_i * x_i
// z_i += B_i^T * x_{i+1}
blasfeo_dgemv_nt(m, n, 1.0, 1.0, sA.B[i]->ref(), 0, 0, x.x[i].ref(), 0, x.x[i+1].ref(), 0, 1.0, 1.0, z.x[i+1].ref(), 0, z.x[i].ref(), 0, z.x[i+1].ref(), 0, z.x[i].ref(), 0);
// z_{i+1} += alpha * B_i * x_i
// z_i += alpha * B_i^T * x_{i+1}
blasfeo_dgemv_nt(m, n, alpha, alpha, sA.B[i]->ref(), 0, 0, x.x[i].ref(), 0, x.x[i+1].ref(), 0, 1.0, 1.0, z.x[i+1].ref(), 0, z.x[i].ref(), 0, z.x[i+1].ref(), 0, z.x[i].ref(), 0);
}
}
if (arrow_width > 0)
Expand All @@ -1320,9 +1379,9 @@ class BlocksparseStageKKT : public KKTSystem<T>
assert(z.x[N-1].rows() == m && "size mismatch");
assert(x.x[N-1].rows() == m && "size mismatch");
assert(z.x[i].rows() == n && "size mismatch");
// z_{N-1} += E_i * x_i
// z_i += E_i^T * x_{N-1}
blasfeo_dgemv_nt(m, n, 1.0, 1.0, sA.E[i]->ref(), 0, 0, x.x[i].ref(), 0, x.x[N-1].ref(), 0, 1.0, 1.0, z.x[N-1].ref(), 0, z.x[i].ref(), 0, z.x[N-1].ref(), 0, z.x[i].ref(), 0);
// z_{N-1} += alpha * E_i * x_i
// z_i += alpha * E_i^T * x_{N-1}
blasfeo_dgemv_nt(m, n, alpha, alpha, sA.E[i]->ref(), 0, 0, x.x[i].ref(), 0, x.x[N-1].ref(), 0, 1.0, 1.0, z.x[N-1].ref(), 0, z.x[i].ref(), 0, z.x[N-1].ref(), 0, z.x[i].ref(), 0);
}
}
}
Expand Down
21 changes: 21 additions & 0 deletions include/piqp/sparse/kkt.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -360,6 +360,27 @@ class KKT : public KKTSystem<T>, public KKTImpl<KKT<T, I, Mode, Ordering>, T, I,
#endif
}

// z = alpha * P * x
void eval_P_x(const T& alpha, const CVecRef<T>& x, VecRef<T> z)
{
z.noalias() = alpha * data.P_utri * x;
z.noalias() += alpha * data.P_utri.transpose().template triangularView<Eigen::StrictlyLower>() * x;
}

// zn = beta_n * yn + alpha_n * A * xn, zt = beta_t * yt + alpha_t * A^T * xt
void eval_A_xn_and_AT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt)
{
zn.noalias() = beta_n * yn + alpha_n * data.AT.transpose() * xn;
zt.noalias() = beta_t * yt + alpha_t * data.AT * xt;
}

// zn = beta_n * yn + alpha_n * G * xn, zt = beta_t * yt + alpha_t * G^T * xt
void eval_G_xn_and_GT_xt(const T& alpha_n, const T& alpha_t, const CVecRef<T>& xn, const CVecRef<T>& xt, const T& beta_n, const T& beta_t, const CVecRef<T>& yn, const CVecRef<T>& yt, VecRef<T> zn, VecRef<T> zt)
{
zn.noalias() = beta_n * yn + alpha_n * data.GT.transpose() * xn;
zt.noalias() = beta_t * yt + alpha_t * data.GT * xt;
}

SparseMat<T, I>& internal_kkt_mat()
{
return PKPt;
Expand Down

0 comments on commit 1511a89

Please sign in to comment.