From 3f0daa0b01f746b6f60dcfaf55108b048962cbb6 Mon Sep 17 00:00:00 2001
From: nychiang <sorakid507@gmail.com>
Date: Sat, 11 Mar 2023 02:05:43 -0800
Subject: [PATCH] correct RAJA implementation

---
 src/LinAlg/hiopVectorRajaImpl.hpp       | 31 +++++++++++++------------
 src/Optimization/hiopNlpFormulation.cpp |  8 ++++++-
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/src/LinAlg/hiopVectorRajaImpl.hpp b/src/LinAlg/hiopVectorRajaImpl.hpp
index be5b796b0..b3165486a 100644
--- a/src/LinAlg/hiopVectorRajaImpl.hpp
+++ b/src/LinAlg/hiopVectorRajaImpl.hpp
@@ -2390,10 +2390,10 @@ void hiopVectorRaja<MEM, POL>::process_constraints_local(const hiopVector& gl_ve
   size_type n_in = dl_vec.get_local_size();
   size_type n_cons = n_eq + n_in;
   
-  hiopVectorInt* idx_eq_cumsum = LinearAlgebraFactory::create_vector_int(mem_space_, n_cons);
-  hiopVectorInt* idx_in_cumsum = LinearAlgebraFactory::create_vector_int(mem_space_, n_cons);
-  index_type* find_eq = idx_eq_cumsum->local_data();
-  index_type* find_in = idx_in_cumsum->local_data();
+  hiopVectorInt* find_eq = LinearAlgebraFactory::create_vector_int(mem_space_, n_cons);
+  hiopVectorInt* find_in = LinearAlgebraFactory::create_vector_int(mem_space_, n_cons);
+  index_type* idx_eq_cumsum = find_eq->local_data();
+  index_type* idx_in_cumsum = find_in->local_data();
 
   RAJA::ReduceSum< hiop_raja_reduce, int > sum_n_bnds_low(0);
   RAJA::ReduceSum< hiop_raja_reduce, int > sum_n_bnds_upp(0);
@@ -2405,11 +2405,11 @@ void hiopVectorRaja<MEM, POL>::process_constraints_local(const hiopVector& gl_ve
     RAJA_LAMBDA(RAJA::Index_type i)
     {
       if(gl[i] == gu[i]) {
-        find_eq[i] = 1;
-        find_in[i] = 0;
+        idx_eq_cumsum[i] = 1;
+        idx_in_cumsum[i] = 0;
       } else {
-        find_eq[i] = 0;
-        find_in[i] = 1;
+        idx_eq_cumsum[i] = 0;
+        idx_in_cumsum[i] = 1;
       }
     }
   );
@@ -2421,8 +2421,6 @@ void hiopVectorRaja<MEM, POL>::process_constraints_local(const hiopVector& gl_ve
   //     (0,1,1)  -- (1,1,2)  after scan
   // map [1]         [0,2]
 
-  index_type* nnz_cumsum = idx_cumsum_->local_data();
-  index_type v_n_local = v.n_local_;
   RAJA::forall<hiop_raja_exec>(
     RAJA::RangeSegment(0, n_cons),
     RAJA_LAMBDA(RAJA::Index_type i)
@@ -2464,9 +2462,10 @@ void hiopVectorRaja<MEM, POL>::process_constraints_local(const hiopVector& gl_ve
         } else {
           assert(idx_in_cumsum[i] == idx_in_cumsum[i-1] + 1);
           int in_idx = idx_in_cumsum[i] - 1;
-          incon_map[in_idx] = cons_type[i];
+          incon_type[in_idx] = cons_type[i];
           dl[in_idx] = gl[i];
           du[in_idx] = gu[i];
+          incon_map[in_idx] = i;
 
           if(gl[i]>-1e20) {
             idl[in_idx] = 1.;
@@ -2488,11 +2487,13 @@ void hiopVectorRaja<MEM, POL>::process_constraints_local(const hiopVector& gl_ve
     }
   );
 
-  n_bnds_low = sum_n_bnds_low.get();
-  n_bnds_upp = sum_n_bnds_upp.get();
-  n_bnds_lu = sum_n_bnds_lu.get();
+  n_ineq_low = sum_n_bnds_low.get();
+  n_ineq_upp = sum_n_bnds_upp.get();
+  n_ineq_lu = sum_n_bnds_lu.get();
+
+  delete find_eq;
+  delete find_in;  	
 
-  return true;
 }
 
 
diff --git a/src/Optimization/hiopNlpFormulation.cpp b/src/Optimization/hiopNlpFormulation.cpp
index 76fb5ba14..2e4ec4be1 100644
--- a/src/Optimization/hiopNlpFormulation.cpp
+++ b/src/Optimization/hiopNlpFormulation.cpp
@@ -524,7 +524,13 @@ bool hiopNlpFormulation::process_constraints()
                                     cons_ineq_type_,
                                     cons_type);
   
-  
+ 
+  hiopVectorIntSeq cons_eq_mapping_host(n_cons_eq_);
+  hiopVectorIntSeq cons_ineq_mapping_host(n_cons_ineq_);
+ 
+  cons_eq_mapping_->copy_to_vectorseq(cons_eq_mapping_host);
+  cons_ineq_mapping_->copy_to_vectorseq(cons_ineq_mapping_host);
+
   /* delete the temporary buffers */
   delete gl; 
   delete gu;