Lots of performance upgrades

examples/MachineLearning/GravityNN/network/IndexContainer.hpp

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <map>
+#include <unordered_map>
 #include <vector>
 #include <gravity/types.h>
 #include <gravity/var.h>
@@ -34,13 +34,10 @@ class IndexSet {
     }
 
     indices& operator[](const std::string& name) {
-        if (this->_indices.count(name) == 0) {
-            throw std::runtime_error("Unknown index set: " + name);
-        }
         return this->_indices.at(name);
     }
 
-    std::map<std::string, indices> _indices;
+    std::unordered_map<std::string, indices> _indices;
     size_t row_id = 0;
     size_t row_id2 = 0;
 };
@@ -59,7 +56,7 @@ class IndexContainer {
     }
 
     void add(OType op, std::vector<std::vector<std::string>> names) {
-        if (this->_indices.count(op) > 0) {
+        if (this->_indices.find(op) != this->_indices.end()) {
             return;
         }
 
@@ -82,5 +79,5 @@ class IndexContainer {
     indices hidden_states, y_ids;
 
     // Map from operator type to index set
-    std::map<OType, IndexSet> _indices;
+    std::unordered_map<OType, IndexSet> _indices;
 };

examples/MachineLearning/GravityNN/network/Layers/LayerBase.hpp

@@ -15,6 +15,8 @@ class Layer {
 public:
     Layer(const onnx::NodeProto& node, Tensors& tensors) {
         this->name = node.name();
+        this->opname = node.op_type();
+
         for (const auto& input : node.input()) {
             this->inputs.push_back(&tensors.at(input));
         }
@@ -112,13 +114,15 @@ class Layer {
 
     double range_lower = HMIN;
     double range_upper = HMAX;
+
+    std::string opname;
 };
 
 class UnsupportedLayer: public Layer {
 public:
     UnsupportedLayer(const onnx::NodeProto& node, Tensors& tensors) : Layer(node, tensors) {
         this->operator_type = _unsupported;
-        std::cout << "WARNING: Unsupported layer: " << node.op_type() << std::endl;
+        std::cout << "WARNING: Unsupported layer: " << node.op_type() << " | " << node.name() << std::endl;
     }
 
     std::vector<std::vector<std::string>> get_indices() const override {

examples/MachineLearning/GravityNN/network/Layers/Linear.hpp

@@ -354,28 +354,33 @@ class Conv : public Layer {
             inds["Constr"].add(this->name + "," + to_string(j));
         }
 
+        auto& oInd = inds["Out"];
+        auto& wInd = inds["W"];
+        auto& iInd = inds["In"];
+        auto& bInd = inds["B"];
+
         for (int ob = 0; ob < this->Y->shape[0]; ob++) {
             for (int oh = 0; oh < this->out_h; oh++) {
                 for (int ow = 0; ow < this->out_w; ow++) {
                     for (int oc = 0; oc < this->out_c; oc++) {
-                        inds["Out"].add_ref(this->Y->strkey(ob, oc, oh, ow));
+                        oInd.add_ref(this->Y->strkey(ob, oc, oh, ow));
                         for (int kh = 0; kh < this->kern_h; kh++) {
                             for (int kw = 0; kw < this->kern_w; kw++) {
                                 for (int kc = 0; kc < this->kern_c; kc++) {
                                     int h_ind = (this->strides[0]*oh + this->dilations[0]*kh - this->pads[0]);
                                     int w_ind = (this->strides[1]*ow + this->dilations[1]*kw - this->pads[3]);
                                     if ((h_ind < this->inp_h) && (h_ind >= 0) && (w_ind < this->inp_w) && (w_ind >= 0)) {
-                                        inds["W"].add_in_row(inds.row_id, this->W->strkey(oc, kc, kh, kw));
-                                        inds["In"].add_in_row(inds.row_id, this->X->strkey(ob, kc, h_ind, w_ind));
+                                        wInd.add_in_row(inds.row_id, this->W->strkey(oc, kc, kh, kw));
+                                        iInd.add_in_row(inds.row_id, this->X->strkey(ob, kc, h_ind, w_ind));
                                     }
                                 }
                             }
                         }
                         // Add bias
                         if (this->B) {
-                            inds["B"].add_in_row(inds.row_id, this->B->strkey(oc));
+                            bInd.add_in_row(inds.row_id, this->B->strkey(oc));
                         } else {
-                            inds["B"].add_empty_row();
+                            bInd.add_empty_row();
                         }
                         inds.row_id++;
                     }

examples/MachineLearning/GravityNN/network/Layers/Unary.hpp

@@ -156,6 +156,7 @@ class Input : public Layer {
 public:
     Input(const onnx::ValueInfoProto& node, Tensors& tensors): Layer(node, tensors) {
         operator_type = _input;
+        this->opname = "Input";
 
         // TODO: Remove!
         this->range_lower = -1000.0;

examples/MachineLearning/GravityNN/network/NeuralNet.hpp

@@ -129,7 +129,8 @@ class NeuralNet {
         this->x.in(this->indices.hidden_states);
         this->y.in(this->indices.y_ids);
 
-        this->initialize_state(x, y);
+        // This is taking far too long on ConvNets
+        // this->initialize_state(x, y);
 
         this->NN.add(this->x);
         this->NN.add(this->y);
@@ -281,7 +282,7 @@ class NeuralNet {
                 continue;
             }
 
-            std::cout << " - " << l->name << std::endl;
+            std::cout << " - " << l->opname << std::endl;
             l->add_constraints(this->NN, this->indices(l->operator_type), this->indices.w, this->x, this->y);
         }
     }
@@ -300,4 +301,4 @@ class NeuralNet {
     param<> x_lb, x_ub;
     var<> x;
     var<int> y;
-};
+};

examples/MachineLearning/GravityNN/network/Tensor.hpp

@@ -105,13 +105,13 @@ class Tensor {
     std::string strkey(std::vector<size_t> indices) const {
         size_t flat = this->flatten_index(indices);
         this->_boundcheck(flat);
-        return this->name + "," + std::to_string(flat);
+        return Tensor::_strkey(this->name, flat);
     }
 
     std::string strkey(size_t idx) const {
         size_t flat = idx;
         this->_boundcheck(flat);
-        return this->name + "," + std::to_string(flat);
+        return Tensor::_strkey(this->name, flat);
     }
 
     std::string strkey(size_t i, size_t j) const {
@@ -120,7 +120,7 @@ class Tensor {
         }
         size_t flat = i * this->shape[1] + j;
         this->_boundcheck(flat);
-        return this->name + "," + std::to_string(flat);
+        return Tensor::_strkey(this->name, flat);
     }
 
     std::string strkey(size_t i, size_t j, size_t k) const {
@@ -129,7 +129,7 @@ class Tensor {
         }
         size_t flat = i * this->shape[1] * this->shape[2] + j * this->shape[2] + k;
         this->_boundcheck(flat);
-        return this->name + "," + std::to_string(flat);
+        return Tensor::_strkey(this->name, flat);
     }
 
     std::string strkey(size_t i, size_t j, size_t k, size_t l) const {
@@ -138,13 +138,17 @@ class Tensor {
         }
         size_t flat = i * this->shape[1] * this->shape[2] * this->shape[3] + j * this->shape[2] * this->shape[3] + k * this->shape[3] + l;
         this->_boundcheck(flat);
-        return this->name + "," + std::to_string(flat);
+        return Tensor::_strkey(this->name, flat);
     }
 
     size_t flatten_index(const std::vector<size_t>& indices) const {
         return Tensor::flatten_index(indices, this->shape);
     }
 
+    static std::string _strkey(std::string name, size_t idx) {
+        return name.append(",").append(std::to_string(idx));
+    }
+
     static size_t flatten_index(const std::vector<size_t>& indices, const std::vector<size_t>& shape) {
         if (indices.size() != shape.size()) {
             throw std::runtime_error("Cannot flatten index with " + std::to_string(indices.size()) + " indices for a tensor with " + std::to_string(shape.size()) + " dimensions");

include/gravity/param.h

@@ -334,7 +334,7 @@ namespace gravity {
             return _intype;
         }
 
-        shared_ptr<map<string,size_t>> get_keys_map() const {
+        shared_ptr<std::unordered_map<std::string, size_t>> get_keys_map() const {
             if(_indices){
                 return _indices->_keys_map;
             }
@@ -1679,7 +1679,7 @@ namespace gravity {
             res._type = matrix_;
             res._ids = make_shared<vector<vector<size_t>>>();
             string key = "", invariant_key="";
-            map<string,size_t> invariant_map;
+            std::unordered_map<std::string, size_t> invariant_map;
             int row_id = 0;
             string prev_key = "";
             if(is_indexed()){/* If ids has key references, use those */

include/gravity/types.h

@@ -18,6 +18,7 @@
 #include <iostream>
 #include <algorithm>
 #include <complex>      // std::complex
+#include <unordered_map>
 
 inline int nthOccurrence(const std::string& str, const std::string& findMe, int nth)
 {
@@ -251,7 +252,7 @@ class indices{
     shared_ptr<vector<size_t>>              _dim = nullptr;/*<< A vector storing the dimension of sub-indices */
     shared_ptr<vector<string>>              _keys = nullptr; /*<< A vector storing all the keys */
 
-    shared_ptr<map<string,size_t>>          _keys_map = nullptr; /*<< A map storing all the indices, the size_t number indicates the right position in the _keys vector */
+    shared_ptr<std::unordered_map<std::string, size_t>>          _keys_map = nullptr; /*<< A map storing all the indices, the size_t number indicates the right position in the _keys vector */
 
     set<size_t>                             _excluded_keys; /*<< A set storing all indices that should be excluded */
     shared_ptr<vector<vector<size_t>>>      _ids = nullptr;
@@ -269,7 +270,7 @@ class indices{
 
     indices(string name){
         _name = name;
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _dim->resize(2,0);
@@ -282,22 +283,22 @@ class indices{
 
     indices(int name){
         _name = to_string(name);
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _dim->resize(2,0);
     }
 
     indices(){
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _dim->resize(2,0);
     }
 
     indices(const ordered_pairs& pairs){
         auto n = pairs._keys.size();
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _keys->resize(n);
@@ -1051,7 +1052,7 @@ class indices{
         indices = {forward<string>(args)...};
         indices.push_front(idx1);
         auto n = indices.size();
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _keys->resize(n);
         _dim = make_shared<vector<size_t>>();
@@ -1151,7 +1152,7 @@ class indices{
             cpy._keys = make_shared<vector<string>>(*_keys);
         }
         if(_keys_map){
-            cpy._keys_map = make_shared<map<string,size_t>>(*_keys_map);
+            cpy._keys_map = make_shared<std::unordered_map<std::string, size_t>>(*_keys_map);
         }
         cpy._time_extended = _time_extended;
         cpy._time_pos = _time_pos;
@@ -1182,7 +1183,7 @@ class indices{
 
     //        template<typename Tobj>
     //        indices(const vector<Tobj*>& vec){
-    //            _keys_map = make_shared<map<string,size_t>>();
+    //            _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
     //            _keys = make_shared<vector<string>>();
     //            size_t i = 0;
     //            for (auto idx:vec) {
@@ -1196,7 +1197,7 @@ class indices{
     //
     //        template<typename Tobj>
     //        indices(const vector<Tobj>& vec){
-    //            _keys_map = make_shared<map<string,size_t>>();
+    //            _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
     //            _keys = make_shared<vector<string>>();
     //            size_t i = 0;
     //            for (auto idx:vec) {
@@ -1210,7 +1211,7 @@ class indices{
 
     template<typename Tobj>
     indices(const vector<Tobj*>& vec, bool include_inactive = false){
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _dim->resize(1);
@@ -1230,7 +1231,7 @@ class indices{
 
     template<typename Tobj>
     indices(const vector<Tobj>& vec, bool include_inactive = false){
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _dim->resize(1);
@@ -1283,7 +1284,7 @@ class indices{
         //            if (vecs.size()==2) {
         //                _type = matrix_;
         //            }
-        _keys_map = make_shared<map<string,size_t>>();
+        _keys_map = make_shared<std::unordered_map<std::string, size_t>>();
         _keys = make_shared<vector<string>>();
         _dim = make_shared<vector<size_t>>();
         _dim->resize(vecs.size());
@@ -1540,7 +1541,7 @@ class indices{
 
     void remove_excluded(){
         _ids = nullptr;
-        map<string,size_t> new_keys_map;
+        std::unordered_map<std::string, size_t> new_keys_map;
         for(auto &key_id: _excluded_keys){
             auto key = _keys->at(key_id);
             _keys_map->erase(key);
@@ -1569,7 +1570,7 @@ class indices{
 
     indices exclude(string key){
         auto res =  *this;
-        res._keys_map = make_shared<map<string,size_t>>(*_keys_map);
+        res._keys_map = make_shared<std::unordered_map<std::string, size_t>>(*_keys_map);
         res._keys = make_shared<vector<string>>(*_keys);
         res._excluded_keys.insert(res._keys_map->at(key));
         //            if(!is_indexed()){
@@ -1661,7 +1662,7 @@ indices union_ids(const indices& ids1, Args&&... args) {
     indices res("Union(");
     auto nb_entries = ids1.get_nb_entries();
     if(!ids1.is_indexed()){ //if the index set is not indexed, we assume the rest are not indexed as well
-        res._keys_map = make_shared<map<string,size_t>>(*ids1._keys_map);
+        res._keys_map = make_shared<std::unordered_map<std::string, size_t>>(*ids1._keys_map);
         res._keys = make_shared<vector<string>>(*ids1._keys);
         for (size_t i= 1; i < all_ids.size(); i++) {
             auto ids = all_ids[i];

src/utils.cpp

@@ -42,7 +42,7 @@ gravity::indices gravity::range(size_t p1 ,size_t p2){
     indices res("range("+to_string(p1)+"<"+to_string(p2)+")");
     auto n = p2 - p1 + 1;
     assert(n >= 0);
-    res._keys_map = make_shared<map<string,size_t>>();
+    res._keys_map = make_shared<std::unordered_map<std::string, size_t>>();
     res._keys = make_shared<vector<string>>();
     res._keys->resize(n);
     res._dim = make_shared<vector<size_t>>();