Port sparse inequality constraint Jacobian to GPU (#40)

* I am not going to lie, Cursor agent heavily helped me with this.

Replace PETSc-based inequality Jacobian with GPU RAJA kernels

Move the inequality constraint Jacobian computation for the HiOp sparse
GPU solver entirely to the device, eliminating the per-iteration host
back and forth (copy to host, PETSc compute, MatGetRow extraction, values
copy back to device). Elimiate PETSc use from this part of the code.

Three RAJA kernels now compute directly into device memory:
- Generator set-point constraints (AGC)
- Voltage-reactive-power bounds (FIXED_WITHIN_QBOUNDS)
- Line flow limits (Sf^2/St^2 derivatives + slack variables)

Supporting changes:
- Analytical NNZ counting replaces PETSc MatGetInfo at solver setup
- New device-side parameter fields (apf, vs, xpdevidx, xslackidx,
  bus-to-gen mapping) added to *ParamsRajaHiop structs
- Sparse position indices assigned at model setup for all three
  contribution types

Includes validation test (test_ineqjac_gpu) that solves with IPOPT,
then compares PETSc and GPU Jacobian values at the converged solution.
Optional -benchmark flag for performance comparison.

Made-with: Cursor


---------

Co-authored-by: kswirydo <kasia.swirydowicz@gmail.com>
Co-authored-by: pelesh <pelesh@users.noreply.github.com>
Co-authored-by: Nicholson Koukpaizan <koukpaizannk@ornl.gov>

Author: 3 people

src/opflow/CMakeLists.txt

@@ -20,6 +20,7 @@ if(EXAGO_ENABLE_RAJA)
     set(OPFLOW_FORM_SRC
         ${OPFLOW_FORM_SRC} model/power_bal_hiop/pbpolrajahiopsparse.cpp
         model/power_bal_hiop/pbpolrajahiopsparsekernels.cpp
+        model/power_bal_hiop/pbpolrajahiopsparse_gpu.cpp
     )
   endif()
 endif()
@@ -39,8 +40,8 @@ set_source_files_properties(${OPFLOW_SRC} PROPERTIES LANGUAGE CXX)
 if(EXAGO_ENABLE_RAJA AND EXAGO_ENABLE_CUDA)
   set_source_files_properties(
     model/power_bal_hiop/pbpolrajahiopkernels.cpp
-    model/power_bal_hiop/pbpolrajahiopsparsekernels.cpp PROPERTIES LANGUAGE
-                                                                   CUDA
+    model/power_bal_hiop/pbpolrajahiopsparsekernels.cpp
+    model/power_bal_hiop/pbpolrajahiopsparse_gpu.cpp PROPERTIES LANGUAGE CUDA
   )
 endif()
 

src/opflow/model/power_bal_hiop/paramsrajahiop.cpp

@@ -27,6 +27,11 @@ int BUSParamsRajaHiop::destroy(OPFLOW opflow) {
     h_allocator_.deallocate(jacsp_idx);
     h_allocator_.deallocate(jacsq_idx);
   }
+  h_allocator_.deallocate(ispv);
+  h_allocator_.deallocate(gineqidx);
+  h_allocator_.deallocate(ineqjacsp_idx);
+  h_allocator_.deallocate(genoffset);
+  h_allocator_.deallocate(ngenONbus);
 
 #ifdef EXAGO_ENABLE_GPU
   d_allocator_.deallocate(isref_dev_);
@@ -46,6 +51,11 @@ int BUSParamsRajaHiop::destroy(OPFLOW opflow) {
     d_allocator_.deallocate(jacsp_idx_dev_);
     d_allocator_.deallocate(jacsq_idx_dev_);
   }
+  d_allocator_.deallocate(ispv_dev_);
+  d_allocator_.deallocate(gineqidx_dev_);
+  d_allocator_.deallocate(ineqjacsp_idx_dev_);
+  d_allocator_.deallocate(genoffset_dev_);
+  d_allocator_.deallocate(ngenONbus_dev_);
 #endif
 
   return 0;
@@ -80,6 +90,11 @@ int BUSParamsRajaHiop::copy(OPFLOW opflow) {
     resmgr.copy(jacsq_idx_dev_, jacsq_idx);
     resmgr.copy(powerimbalance_penalty_dev_, powerimbalance_penalty);
   }
+  resmgr.copy(ispv_dev_, ispv);
+  resmgr.copy(gineqidx_dev_, gineqidx);
+  resmgr.copy(ineqjacsp_idx_dev_, ineqjacsp_idx);
+  resmgr.copy(genoffset_dev_, genoffset);
+  resmgr.copy(ngenONbus_dev_, ngenONbus);
 #else
   isref_dev_ = isref;
   isisolated_dev_ = isisolated;
@@ -96,6 +111,11 @@ int BUSParamsRajaHiop::copy(OPFLOW opflow) {
   jacsp_idx_dev_ = jacsp_idx;
   jacsq_idx_dev_ = jacsq_idx;
   powerimbalance_penalty_dev_ = powerimbalance_penalty;
+  ispv_dev_ = ispv;
+  gineqidx_dev_ = gineqidx;
+  ineqjacsp_idx_dev_ = ineqjacsp_idx;
+  genoffset_dev_ = genoffset;
+  ngenONbus_dev_ = ngenONbus;
 #endif
   return 0;
 }
@@ -132,18 +152,31 @@ int BUSParamsRajaHiop::allocate(OPFLOW opflow) {
     jacsp_idx = paramAlloc<int>(h_allocator_, nbus);
     jacsq_idx = paramAlloc<int>(h_allocator_, nbus);
   }
+  ispv = paramAlloc<int>(h_allocator_, nbus);
+  gineqidx = paramAlloc<int>(h_allocator_, nbus);
+  ineqjacsp_idx = paramAlloc<int>(h_allocator_, nbus);
+  genoffset = paramAlloc<int>(h_allocator_, nbus);
+  ngenONbus = paramAlloc<int>(h_allocator_, nbus);
 
   /* Memzero arrays */
   resmgr.memset(isref, 0, nbus * sizeof(int));
   resmgr.memset(ispvpq, 0, nbus * sizeof(int));
   resmgr.memset(isisolated, 0, nbus * sizeof(int));
+  resmgr.memset(ispv, 0, nbus * sizeof(int));
+  resmgr.memset(gineqidx, 0, nbus * sizeof(int));
+  resmgr.memset(ineqjacsp_idx, 0, nbus * sizeof(int));
+  resmgr.memset(genoffset, 0, nbus * sizeof(int));
+  resmgr.memset(ngenONbus, 0, nbus * sizeof(int));
 
+  int genoff = 0;
   for (int i = 0; i < nbus; i++) {
     bus = &ps->bus[i];
     loc = bus->startxVloc;
 
     xidx[i] = opflow->idxn2sd_map[loc];
     gidx[i] = bus->starteqloc;
+    genoffset[i] = genoff;
+    genoff += bus->ngenON;
 
     if (bus->ide == REF_BUS)
       isref[i] = 1;
@@ -152,6 +185,12 @@ int BUSParamsRajaHiop::allocate(OPFLOW opflow) {
     else
       ispvpq[i] = 1;
 
+    if (bus->ide == PV_BUS)
+      ispv[i] = 1;
+
+    ngenONbus[i] = bus->ngenON;
+    gineqidx[i] = bus->startineqloc;
+
     if (opflow->genbusvoltagetype == FIXED_AT_SETPOINT) {
       if (bus->ide == REF_BUS || bus->ide == PV_BUS) {
         /* Hold voltage at reference and PV buses */
@@ -200,6 +239,11 @@ int BUSParamsRajaHiop::allocate(OPFLOW opflow) {
     jacsp_idx_dev_ = paramAlloc<int>(d_allocator_, nbus);
     jacsq_idx_dev_ = paramAlloc<int>(d_allocator_, nbus);
   }
+  ispv_dev_ = paramAlloc<int>(d_allocator_, nbus);
+  gineqidx_dev_ = paramAlloc<int>(d_allocator_, nbus);
+  ineqjacsp_idx_dev_ = paramAlloc<int>(d_allocator_, nbus);
+  genoffset_dev_ = paramAlloc<int>(d_allocator_, nbus);
+  ngenONbus_dev_ = paramAlloc<int>(d_allocator_, nbus);
 #endif
   return 0;
 }
@@ -231,6 +275,8 @@ int LINEParamsRajaHiop::copy(OPFLOW opflow) {
     resmgr.copy(gineqidx_dev_, gineqidx);
     resmgr.copy(gbineqidx_dev_, gbineqidx);
     resmgr.copy(linelimidx_dev_, linelimidx);
+    resmgr.copy(ineqjacsp_idx_dev_, ineqjacsp_idx);
+    resmgr.copy(xslackidx_dev_, xslackidx);
   }
 #else
   Gff_dev_ = Gff;
@@ -250,6 +296,8 @@ int LINEParamsRajaHiop::copy(OPFLOW opflow) {
     gineqidx_dev_ = gineqidx;
     gbineqidx_dev_ = gbineqidx;
     linelimidx_dev_ = linelimidx;
+    ineqjacsp_idx_dev_ = ineqjacsp_idx;
+    xslackidx_dev_ = xslackidx;
   }
 #endif
   return 0;
@@ -277,6 +325,8 @@ int LINEParamsRajaHiop::destroy(OPFLOW opflow) {
     h_allocator_.deallocate(gineqidx);
     h_allocator_.deallocate(gbineqidx);
     h_allocator_.deallocate(linelimidx);
+    h_allocator_.deallocate(ineqjacsp_idx);
+    h_allocator_.deallocate(xslackidx);
   }
 
 #ifdef EXAGO_ENABLE_GPU
@@ -301,6 +351,8 @@ int LINEParamsRajaHiop::destroy(OPFLOW opflow) {
     d_allocator_.deallocate(gineqidx_dev_);
     d_allocator_.deallocate(gbineqidx_dev_);
     d_allocator_.deallocate(linelimidx_dev_);
+    d_allocator_.deallocate(ineqjacsp_idx_dev_);
+    d_allocator_.deallocate(xslackidx_dev_);
   }
 #endif
 
@@ -348,6 +400,8 @@ int LINEParamsRajaHiop::allocate(OPFLOW opflow) {
     linelimidx = paramAlloc<int>(h_allocator_, nlinelim);
     gineqidx = paramAlloc<int>(h_allocator_, nlinelim);
     gbineqidx = paramAlloc<int>(h_allocator_, nlinelim);
+    ineqjacsp_idx = paramAlloc<int>(h_allocator_, nlinelim);
+    xslackidx = paramAlloc<int>(h_allocator_, nlinelim);
   }
 
   PetscInt j = 0;
@@ -391,6 +445,9 @@ int LINEParamsRajaHiop::allocate(OPFLOW opflow) {
       gbineqidx[j] = opflow->nconeq + line->startineqloc;
       gineqidx[j] = line->startineqloc;
       linelimidx[j] = linei;
+      if (opflow->allow_lineflow_violation) {
+        xslackidx[j] = opflow->idxn2sd_map[line->startxslackloc];
+      }
       j++;
     }
 
@@ -420,6 +477,8 @@ int LINEParamsRajaHiop::allocate(OPFLOW opflow) {
     gineqidx_dev_ = paramAlloc<int>(d_allocator_, nlinelim);
     gbineqidx_dev_ = paramAlloc<int>(d_allocator_, nlinelim);
     linelimidx_dev_ = paramAlloc<int>(d_allocator_, nlinelim);
+    ineqjacsp_idx_dev_ = paramAlloc<int>(d_allocator_, nlinelim);
+    xslackidx_dev_ = paramAlloc<int>(d_allocator_, nlinelim);
   }
 #endif
   return 0;
@@ -558,7 +617,10 @@ int GENParamsRajaHiop::destroy(OPFLOW opflow) {
   h_allocator_.deallocate(qt);
   h_allocator_.deallocate(qb);
   h_allocator_.deallocate(isrenewable);
+  h_allocator_.deallocate(apf);
+  h_allocator_.deallocate(vs);
   h_allocator_.deallocate(xidx);
+  h_allocator_.deallocate(xpdevidx);
   h_allocator_.deallocate(gidxbus);
   h_allocator_.deallocate(eqjacspbus_idx);
   h_allocator_.deallocate(eqjacsqbus_idx);
@@ -581,7 +643,10 @@ int GENParamsRajaHiop::destroy(OPFLOW opflow) {
   d_allocator_.deallocate(qt_dev_);
   d_allocator_.deallocate(qb_dev_);
   d_allocator_.deallocate(isrenewable_dev_);
+  d_allocator_.deallocate(apf_dev_);
+  d_allocator_.deallocate(vs_dev_);
   d_allocator_.deallocate(xidx_dev_);
+  d_allocator_.deallocate(xpdevidx_dev_);
   d_allocator_.deallocate(gidxbus_dev_);
   d_allocator_.deallocate(eqjacspbus_idx_dev_);
   d_allocator_.deallocate(eqjacsqbus_idx_dev_);
@@ -615,8 +680,11 @@ int GENParamsRajaHiop::copy(OPFLOW opflow) {
   resmgr.copy(qt_dev_, qt);
   resmgr.copy(qb_dev_, qb);
   resmgr.copy(isrenewable_dev_, isrenewable);
+  resmgr.copy(apf_dev_, apf);
+  resmgr.copy(vs_dev_, vs);
 
   resmgr.copy(xidx_dev_, xidx);
+  resmgr.copy(xpdevidx_dev_, xpdevidx);
   resmgr.copy(gidxbus_dev_, gidxbus);
 
   resmgr.copy(eqjacspbus_idx_dev_, eqjacspbus_idx);
@@ -639,7 +707,10 @@ int GENParamsRajaHiop::copy(OPFLOW opflow) {
   qt_dev_ = qt;
   qb_dev_ = qb;
   isrenewable_dev_ = isrenewable;
+  apf_dev_ = apf;
+  vs_dev_ = vs;
   xidx_dev_ = xidx;
+  xpdevidx_dev_ = xpdevidx;
   gidxbus_dev_ = gidxbus;
   eqjacspbus_idx_dev_ = eqjacspbus_idx;
   eqjacsqbus_idx_dev_ = eqjacsqbus_idx;
@@ -682,8 +753,11 @@ int GENParamsRajaHiop::allocate(OPFLOW opflow) {
   qt = paramAlloc<double>(h_allocator_, ngenON);
   qb = paramAlloc<double>(h_allocator_, ngenON);
   isrenewable = paramAlloc<int>(h_allocator_, ngenON);
+  apf = paramAlloc<double>(h_allocator_, ngenON);
+  vs = paramAlloc<double>(h_allocator_, ngenON);
 
   xidx = paramAlloc<int>(h_allocator_, ngenON);
+  xpdevidx = paramAlloc<int>(h_allocator_, ngenON);
   gidxbus = paramAlloc<int>(h_allocator_, ngenON);
 
   eqjacspbus_idx = paramAlloc<int>(h_allocator_, ngenON);
@@ -720,11 +794,16 @@ int GENParamsRajaHiop::allocate(OPFLOW opflow) {
       qt[geni] = gen->qt;
       qb[geni] = gen->qb;
       isrenewable[geni] = (int)gen->isrenewable;
+      apf[geni] = gen->apf;
+      vs[geni] = gen->vs;
       if (opflow->has_gensetpoint) {
         pgs[geni] = gen->pgs;
       }
 
       xidx[geni] = opflow->idxn2sd_map[loc];
+      xpdevidx[geni] = (opflow->has_gensetpoint && !gen->isrenewable)
+                           ? opflow->idxn2sd_map[gen->startxpdevloc]
+                           : -1;
       gidxbus[geni] = gloc;
       if (opflow->has_gensetpoint) {
         geqidxgen[geni] = gen->starteqloc;
@@ -748,8 +827,11 @@ int GENParamsRajaHiop::allocate(OPFLOW opflow) {
   qt_dev_ = paramAlloc<double>(d_allocator_, ngenON);
   qb_dev_ = paramAlloc<double>(d_allocator_, ngenON);
   isrenewable_dev_ = paramAlloc<int>(d_allocator_, ngenON);
+  apf_dev_ = paramAlloc<double>(d_allocator_, ngenON);
+  vs_dev_ = paramAlloc<double>(d_allocator_, ngenON);
 
   xidx_dev_ = paramAlloc<int>(d_allocator_, ngenON);
+  xpdevidx_dev_ = paramAlloc<int>(d_allocator_, ngenON);
   gidxbus_dev_ = paramAlloc<int>(d_allocator_, ngenON);
 
   eqjacspbus_idx_dev_ = paramAlloc<int>(d_allocator_, ngenON);

src/opflow/model/power_bal_hiop/paramsrajahiop.h

@@ -28,7 +28,12 @@ struct BUSParamsRajaHiop {
                       vector */
   int *jacsp_idx;  /* Location number in the sparse Jacobian for Pimb */
   int *jacsq_idx;  /* Location number in the sparse Jacobian for Qimb */
-  int *hesssp_idx; /* Location number in the Hessian */
+  int *hesssp_idx; /* KS: Hessian indices */
+  int *ispv;       /* KS: ispv[i] = 1 if bus is PV bus */
+  int *gineqidx;   /* KS: starting position of bus ineq constraints */
+  int *ineqjacsp_idx; /* KS: index in flat sparse ineq Jacobian array */
+  int *genoffset;     /* KS: Offset into flattened gen array for this bus */
+  int *ngenONbus;     /* KS: Number of ON generators on this bus */
 
   // Device data
   int *isref_dev_;      /* isref[i] = 1 if bus is reference bus */
@@ -46,9 +51,14 @@ struct BUSParamsRajaHiop {
                          X vector */
   int *gidx_dev_; /* starting locations for bus balance equations in constraint
                      vector */
-  int *jacsp_idx_dev_;  /* Location number in the sparse Jacobian for Pimb */
-  int *jacsq_idx_dev_;  /* Location number in the sparse Jacobian for Qimb */
-  int *hesssp_idx_dev_; /* Location number in the Hessian */
+  int *jacsp_idx_dev_;     /* Location number in the sparse Jacobian for Pimb */
+  int *jacsq_idx_dev_;     /* Location number in the sparse Jacobian for Qimb */
+  int *hesssp_idx_dev_;    /* Location number in the Hessian */
+  int *ispv_dev_;          /* KS: dev counterpart of ispv */
+  int *gineqidx_dev_;      /* KS: dev counterpart of gineqidx */
+  int *ineqjacsp_idx_dev_; /* KS: device counterpart of ineqjacsp_idx_ */
+  int *genoffset_dev_;     /* KS: dev counterpart of genoffset */
+  int *ngenONbus_dev_;     /* KS: dev counterpart of ngenONbus */
 
   int allocate(OPFLOW);
   int destroy(OPFLOW);
@@ -72,9 +82,12 @@ struct GENParamsRajaHiop {
   double *qt;         /* min. reactive power gen. limits */
   double *qb;         /* max. reactive power gen. limits */
   double *pgs;        /* real power output setpoint */
+  double *apf;        /* generator AGC participation factor */
+  double *vs;         /* voltage setpoint */
   int *isrenewable;   /* Is renewable generator? */
 
-  int *xidx; /* starting locations in X vector */
+  int *xidx;     /* starting locations in X vector */
+  int *xpdevidx; /* KS: tarting locations of deviation variables in X vector */
   int *
       gidxbus; /* starting locations in constraint vector for bus constraints */
   int *geqidxgen;    /* starting locations in equality constraint vector for gen
@@ -104,9 +117,12 @@ struct GENParamsRajaHiop {
   double *qt_dev_;         /* min. reactive power gen. limits */
   double *qb_dev_;         /* max. reactive power gen. limits */
   double *pgs_dev_;        /* real power output setpoint */
+  double *apf_dev_;        /* KS: device counterpart of apf */
+  double *vs_dev_;         /* KS: device counterpart of vs */
   int *isrenewable_dev_;   /* Is renewable generator? */
 
   int *xidx_dev_;        /* starting locations in X vector */
+  int *xpdevidx_dev_;    /* KS: device coutnerpart of xpdevidx*/
   int *gidxbus_dev_;     /* starting locations in constraint vector for bus
                             constraints */
   int *geqidxgen_dev_;   /* starting locations in equality constraint vector for
@@ -143,7 +159,7 @@ struct LOADParamsRajaHiop {
   double *pl;               /* active power demand */
   double *ql;               /* reactive power demand */
   double *loadloss_penalty; /* Penalty for load loss */
-  int *xidx;                /* starting location in X vector */
+  int *xidx;                /* KS: starting location in X vector */
   int *gidx;                /* starting location in constraint vector */
 
   /* The following members are only used with HIOP */
@@ -192,9 +208,11 @@ struct LINEParamsRajaHiop {
                     contribution in constraints vector */
   int *gineqidx; /* Starting location to insert contribution to inequality
                     constraint */
-  int *gbineqidx;  /* Starting location to insert contribution to inequality
-                      constraint bound */
-  int *linelimidx; /* Indices for subset of lines that have finite limits */
+  int *gbineqidx;     /* Starting location to insert contribution to inequality
+                         constraint bound */
+  int *linelimidx;    /* Indices for subset of lines that have finite limits */
+  int *ineqjacsp_idx; /* KS: Position in flat sparse ineq Jacobian array */
+  int *xslackidx;     /* Starting location of slack variables in X vector */
 
   // Device data
   double *Gff_dev_;    /* From side self conductance */
@@ -218,6 +236,8 @@ struct LINEParamsRajaHiop {
                           constraint bound */
   int *
       linelimidx_dev_; /* Indices for subset of lines that have finite limits */
+  int *ineqjacsp_idx_dev_; /* KS: Position in flat sparse ineq Jacobian array */
+  int *xslackidx_dev_; /* Starting location of slack variables in X vector */
 
   int allocate(OPFLOW);
   int destroy(OPFLOW);
@@ -248,6 +268,9 @@ struct PbpolModelRajaHiop : public _p_FormPBPOLRAJAHIOP {
   LINEParamsRajaHiop lineparams;
   BUSParamsRajaHiop busparams;
 
+  int agc_xidx; /* KS: X-vector index for the AGC delta-P variable
+                   (ps->startxloc) */
+
   // Arrays to store Jacobian and Hessian indices and entries on CPU (used with
   // GPU sparse model)
   int *i_jaceq,