Fix finding optimal contraction order (WIP)

include/search_tree.hpp

@@ -7,98 +7,82 @@
 #include <vector>
 #include <map>
 #include <string>
+#include <bitset>
+#include <memory>
+#include <unordered_map>
 #include "UniTensor.hpp"
 
-#ifdef BACKEND_TORCH
-#else
-
 namespace cytnx {
-  /// @cond
-  class PsudoUniTensor {
+
+  using IndexSet = std::bitset<128>;
+
+  class PseudoUniTensor {
    public:
-    // UniTensor utensor; //don't worry about copy, because everything are references in cytnx!
+    bool isLeaf;
+
+    // Leaf node data
     std::vector<std::string> labels;
     std::vector<cytnx_uint64> shape;
     bool is_assigned;
-    PsudoUniTensor *left;
-    PsudoUniTensor *right;
-    PsudoUniTensor *root;
+    cytnx_uint64 tensorIndex;
+
+    // Internal node data
+    std::unique_ptr<PseudoUniTensor> left;
+    std::unique_ptr<PseudoUniTensor> right;
+
     cytnx_float cost;
     cytnx_uint64 ID;
-
     std::string accu_str;
 
-    PsudoUniTensor()
-        : is_assigned(false), left(nullptr), right(nullptr), root(nullptr), cost(0), ID(0){};
-    PsudoUniTensor(const PsudoUniTensor &rhs) {
-      this->left = rhs.left;
-      this->right = rhs.right;
-      this->root = rhs.root;
-      this->labels = rhs.labels;
-      this->shape = rhs.shape;
-      this->is_assigned = rhs.is_assigned;
-      this->cost = rhs.cost;
-      this->accu_str = rhs.accu_str;
-      this->ID = rhs.ID;
-    }
-    PsudoUniTensor &operator==(const PsudoUniTensor &rhs) {
-      this->left = rhs.left;
-      this->right = rhs.right;
-      this->root = rhs.root;
-      this->labels = rhs.labels;
-      this->shape = rhs.shape;
-      this->is_assigned = rhs.is_assigned;
-      this->cost = rhs.cost;
-      this->accu_str = rhs.accu_str;
-      this->ID = rhs.ID;
-      return *this;
-    }
-    void from_utensor(const UniTensor &in_uten) {
-      this->labels = in_uten.labels();
-      this->shape = in_uten.shape();
-      this->is_assigned = true;
-    }
-    void clear_utensor() {
-      this->is_assigned = false;
-      this->labels.clear();
-      this->shape.clear();
-      this->ID = 0;
-      this->cost = 0;
-      this->accu_str = "";
-    }
-    void set_ID(const cytnx_int64 &ID) { this->ID = ID; }
+    // Constructors
+    explicit PseudoUniTensor(cytnx_uint64 index = 0)
+        : isLeaf(true), tensorIndex(index), is_assigned(false), cost(0), ID(1ULL << index), 
+          accu_str(std::to_string(index)) {}
+
+    PseudoUniTensor(std::unique_ptr<PseudoUniTensor> l, std::unique_ptr<PseudoUniTensor> r)
+        : isLeaf(false), left(std::move(l)), right(std::move(r)), cost(0), ID(0) {}
+
+    // Copy and move constructors and assignment operators
+    PseudoUniTensor(const PseudoUniTensor& rhs);
+    PseudoUniTensor(PseudoUniTensor&& rhs) noexcept;
+    PseudoUniTensor& operator=(const PseudoUniTensor& rhs);
+    PseudoUniTensor& operator=(PseudoUniTensor&& rhs) noexcept;
+    ~PseudoUniTensor() = default;
+
+    void from_utensor(const UniTensor& in_uten);
+    void clear_utensor();
   };
 
+  namespace OptimalTreeSolver {
+    std::unique_ptr<PseudoUniTensor> solve(const std::vector<PseudoUniTensor>& tensors,
+                                           bool verbose = false);
+  }
+
   class SearchTree {
    public:
-    std::vector<std::vector<PsudoUniTensor>> nodes_container;
-    // std::vector<PsudoUniTensor> nodes_container; // this contains intermediate layer.
-    std::vector<PsudoUniTensor> base_nodes;  // this is the button layer.
-
-    SearchTree(){};
-    SearchTree(const SearchTree &rhs) {
-      this->nodes_container = rhs.nodes_container;
-      this->base_nodes = rhs.base_nodes;
-    }
-    SearchTree &operator==(const SearchTree &rhs) {
-      this->nodes_container = rhs.nodes_container;
-      this->base_nodes = rhs.base_nodes;
-      return *this;
-    }
+
+    std::vector<PseudoUniTensor> base_nodes;
 
-    // clear all the elements in the whole tree.
+    SearchTree() = default;
     void clear() {
-      nodes_container.clear();
-      base_nodes.clear();
-      // nodes_container.reserve(1024);
+        root_ptr.reset();
+        base_nodes.clear();
     }
-    // clear all the intermediate layer, leave all the base_nodes intact.
-    // and reset the root pointer on the base ondes
-    void reset_search_order() { nodes_container.clear(); }
+    void reset_search_order() { root_ptr.reset(); }
     void search_order();
+
+    std::vector<std::vector<PseudoUniTensor*>> get_root() const {
+        return {{root_ptr.get()}};
+    }
+
+   private:
+    std::unique_ptr<PseudoUniTensor> root_ptr;
   };
-  /// @endcond
+
+  // Helper functions declarations
+  cytnx_float get_cost(const PseudoUniTensor& t1, const PseudoUniTensor& t2);
+  PseudoUniTensor pContract(PseudoUniTensor& t1, PseudoUniTensor& t2);
+
 }  // namespace cytnx
-#endif
 
 #endif  // CYTNX_SEARCH_TREE_H_

src/RegularNetwork.cpp

@@ -230,24 +230,24 @@ namespace cytnx {
   vector<pair<cytnx_int64, cytnx_int64>> CtTree_to_eisumpath(ContractionTree CtTree,
                                                              vector<string> tns) {
     vector<pair<cytnx_int64, cytnx_int64>> path;
-    stack<Node *> stk;
-    Node *root = &(CtTree.nodes_container.back());
+    stack<shared_ptr<Node>> stk;
+    shared_ptr<Node> root = make_shared<Node>(CtTree.nodes_container.back());
     int ly = 0;
     bool ict;
     do {
       while ((root != nullptr)) {
-        if (root->right != nullptr) stk.push(root->right);
+        if (root->right != nullptr) stk.push(make_shared<Node>(*(root->right)));
         stk.push(root);
-        root = root->left;
+        root = make_shared<Node>(*(root->left));
       }
       root = stk.top();
       stk.pop();
       ict = true;
       if ((root->right != nullptr) && !stk.empty()) {
-        if (stk.top() == root->right) {
+        if (stk.top()->name == root->right->name) {
           stk.pop();
           stk.push(root);
-          root = root->right;
+          root = make_shared<Node>(*(root->right));
           ict = false;
         }
       }
@@ -943,7 +943,7 @@ namespace cytnx {
       Stree.base_nodes[t].accu_str = this->names[t];
     }
     Stree.search_order();
-    return Stree.nodes_container.back()[0].accu_str;
+    return Stree.get_root().back()[0]->accu_str;
   }
 
   UniTensor RegularNetwork::Launch() {
@@ -969,28 +969,28 @@ namespace cytnx {
 
       // 2. contract using postorder traversal:
       // cout << this->CtTree.nodes_container.size() << endl;
-      stack<Node *> stk;
-      Node *root = &(this->CtTree.nodes_container.back());
+      stack<shared_ptr<Node>> stk;
+      shared_ptr<Node> root = make_shared<Node>(this->CtTree.nodes_container.back());
       int ly = 0;
       bool ict;
 
       do {
         // move the lmost
         while ((root != nullptr)) {
-          if (root->right != nullptr) stk.push(root->right);
+          if (root->right != nullptr) stk.push(make_shared<Node>(*(root->right)));
           stk.push(root);
-          root = root->left;
+          root = make_shared<Node>(*(root->left));
         }
 
         root = stk.top();
         stk.pop();
         // cytnx_error_msg(stk.size()==0,"[eRROR]","\n");
         ict = true;
         if ((root->right != nullptr) && !stk.empty()) {
-          if (stk.top() == root->right) {
+          if (stk.top()->name == root->right->name) {
             stk.pop();
             stk.push(root);
-            root = root->right;
+            root = make_shared<Node>(*(root->right));
             ict = false;
           }
         }
@@ -1100,28 +1100,28 @@ namespace cytnx {
       }
       // 2. contract using postorder traversal:
       // cout << this->CtTree.nodes_container.size() << endl;
-      stack<Node *> stk;
-      Node *root = &(this->CtTree.nodes_container.back());
+      stack<shared_ptr<Node>> stk;
+      shared_ptr<Node> root = make_shared<Node>(this->CtTree.nodes_container.back());
       int ly = 0;
       bool ict;
 
       do {
         // move the lmost
         while ((root != nullptr)) {
-          if (root->right != nullptr) stk.push(root->right);
+          if (root->right != nullptr) stk.push(make_shared<Node>(*(root->right)));
           stk.push(root);
-          root = root->left;
+          root = make_shared<Node>(*(root->left));
         }
 
         root = stk.top();
         stk.pop();
         // cytnx_error_msg(stk.size()==0,"[eRROR]","\n");
         ict = true;
         if ((root->right != nullptr) && !stk.empty()) {
-          if (stk.top() == root->right) {
+          if (stk.top()->name == root->right->name) {
             stk.pop();
             stk.push(root);
-            root = root->right;
+            root = make_shared<Node>(*(root->right));
             ict = false;
           }
         }