feat: moved operator generation to parse

RivinHD · RivinHD · commit 20401dc9da20 · 2025-07-29T17:42:34.000Z
diff --git a/docs_sphinx/chapters/einsum_trees.rst b/docs_sphinx/chapters/einsum_trees.rst
@@ -42,7 +42,9 @@ This structure holds one node of the tree, its possible children, dimension size
     struct EinsumNode
     {
       NodeType type;
-      float *tensor;
+      int32_t input_tensor_index = -1;
+      float *tensor = nullptr;
+      mini_jit::TensorOperation tensor_op;
 
       // Always filled — dims of the output tensor
       std::vector<int64_t> output_dim_ids;
@@ -75,12 +77,12 @@ This structure holds one node of the tree, its possible children, dimension size
       std::string _to_string(uint depth, std::string connection, std::string depthString) const;
     };
 
-Then, we implemented the logic to parse the string into a set of nodes in the ``parse_tree_no_optimization()`` method. This method also indicates whether
+Then, we implemented the logic to parse the string into a set of nodes in the ``parse_tree_no_optimization(bool)`` method. This method also indicates whether
 the parsing was successful, ``ErrorParse``.
 
 .. code-block:: cpp
 
-    ErrorParse parse_tree_no_optimization();
+    ErrorParse parse_tree_no_optimization(bool build_operators);
 
     // AND
 
diff --git a/docs_sphinx/submissions/report_25_06_12.rst b/docs_sphinx/submissions/report_25_06_12.rst
@@ -87,12 +87,12 @@ This structure holds one node of the tree, its possible children, dimension size
       std::string _to_string(uint depth, std::string connection, std::string depthString) const;
     };
 
-Then, we implemented the logic to parse the string into a set of nodes in the ``parse_tree_no_optimization()`` method. This method also indicates whether
+Then, we implemented the logic to parse the string into a set of nodes in the ``parse_tree_no_optimization(bool)`` method. This method also indicates whether
 the parsing was successful, ``ErrorParse``.
 
 .. code-block:: cpp
 
-    ErrorParse parse_tree_no_optimization();
+    ErrorParse parse_tree_no_optimization(bool build_operators);
 
     // AND
 
diff --git a/src/main/EinsumTree.cpp b/src/main/EinsumTree.cpp
@@ -14,7 +14,7 @@ mini_jit::EinsumTree::EinsumTree(const std::string &tree_str, const std::vector<
 {
 }
 
-mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::parse_tree_no_optimization()
+mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::parse_tree_no_optimization(bool build_operators)
 {
   if (root != nullptr)
   {
@@ -27,6 +27,16 @@ mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::parse_tree_no_optimizatio
   tensorIndex = 0;
   assign_tensor_indices(root);
 
+  if (error_parse != ErrorParse::None)
+  {
+    return error_parse;
+  }
+
+  if (build_operators)
+  {
+    error_parse = generate_operators();
+  }
+
   return error_parse;
 }
 
@@ -496,20 +506,7 @@ mini_jit::EinsumTree::ErrorExecute mini_jit::EinsumTree::execute_node(const std:
       node->tensor = new float[node->get_size(dim_sizes)]();
     }
 
-    mini_jit::TensorOperation tensor_op;
-    TensorConfig config = lower_node(node);
-    TensorOperation::error_t error_setup = tensor_op.setup(config);
-    error = parse_setup_error(error_setup);
-
-    if (error != ErrorExecute::None)
-    {
-      return error;
-    }
-
-#ifdef SAVE_JITS_TO_FILE
-    tensor_op.write_kernel_to_file(node->name());
-#endif  // SAVE_JITS_TO_FILE
-    tensor_op.execute(node->left->tensor, nullptr, node->tensor);
+    node->tensor_op.execute(node->left->tensor, nullptr, node->tensor);
   }
   else if (node->type == NodeType::Contraction)
   {
@@ -539,20 +536,7 @@ mini_jit::EinsumTree::ErrorExecute mini_jit::EinsumTree::execute_node(const std:
       node->tensor = new float[node->get_size(dim_sizes)]();
     }
 
-    mini_jit::TensorOperation tensor_op;
-    TensorConfig config = lower_node(node);
-    TensorOperation::error_t error_setup = tensor_op.setup(config);
-    error = parse_setup_error(error_setup);
-
-    if (error != ErrorExecute::None)
-    {
-      return error;
-    }
-
-#ifdef SAVE_JITS_TO_FILE
-    tensor_op.write_kernel_to_file(node->name());
-#endif  // SAVE_JITS_TO_FILE
-    tensor_op.execute(node->left->tensor, node->right->tensor, node->tensor);
+    node->tensor_op.execute(node->left->tensor, node->right->tensor, node->tensor);
   }
   else
   {
@@ -598,19 +582,19 @@ std::vector<int64_t> mini_jit::EinsumTree::get_output_dims(const std::vector<int
   return dims;
 }
 
-mini_jit::EinsumTree::ErrorExecute mini_jit::EinsumTree::parse_setup_error(TensorOperation::error_t error)
+mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::parse_setup_error(TensorOperation::error_t error)
 {
   if (error == TensorOperation::error_t::success)
   {
-    return ErrorExecute::None;
+    return ErrorParse::None;
   }
 
   uint32_t error_num = static_cast<uint32_t>(error) + 100;
 
   release_assert(error_num >= 101, "Expected error_num to be larger equal than 101.");
   release_assert(error_num <= 115, "Expected error_num to be less equal than 115.");
 
-  return static_cast<ErrorExecute>(error_num);
+  return static_cast<ErrorParse>(error_num);
 }
 
 bool mini_jit::EinsumTree::is_unit_stride_n(EinsumNode *node)
@@ -673,7 +657,7 @@ void mini_jit::EinsumTree::conditional_swap(mini_jit::EinsumTree::EinsumNode *no
 
 mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::parse_tree()
 {
-  ErrorParse error = parse_tree_no_optimization();
+  ErrorParse error = parse_tree_no_optimization(false);
 
   if (error != ErrorParse::None)
   {
@@ -682,6 +666,77 @@ mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::parse_tree()
 
   optimize(root);
 
+  error = generate_operators();
+
+  return error;
+}
+
+mini_jit::EinsumTree::ErrorParse mini_jit::EinsumTree::generate_operators()
+{
+  if (root == nullptr)
+  {
+    std::cerr << "EinsumTree: Cannot execute, root is null." << std::endl;
+    return ErrorParse::InvalidRoot;
+  }
+
+  ErrorParse error = ErrorParse::None;
+  std::vector<EinsumNode *> stack = {root};
+
+  while (stack.size() > 0)
+  {
+    EinsumNode *node = stack.back();
+    stack.pop_back();
+
+    if (node->type == NodeType::Leaf)
+    {
+      continue;
+    }
+    else if (node->type == NodeType::Transposition)
+    {
+      release_assert(node->left != nullptr, "Expected the left child of contraction to be a valid pointer.");
+      release_assert(node->right == nullptr, "Expected the right child of contraction to be a nullptr.");
+
+      stack.push_back(node->left);
+
+      TensorConfig config = lower_node(node);
+      TensorOperation::error_t error_setup = node->tensor_op.setup(config);
+      error = parse_setup_error(error_setup);
+
+      if (error != ErrorParse::None)
+      {
+        return error;
+      }
+
+#ifdef SAVE_JITS_TO_FILE
+      node->tensor_op.write_kernel_to_file(node->name());
+#endif  // SAVE_JITS_TO_FILE
+    }
+    else if (node->type == NodeType::Contraction)
+    {
+      release_assert(node->left != nullptr, "Expected the left child of contraction to be a valid pointer.");
+      release_assert(node->right != nullptr, "Expected the right child of contraction to be a valid pointer.");
+
+      stack.push_back(node->left);
+      stack.push_back(node->right);
+
+      TensorConfig config = lower_node(node);
+      TensorOperation::error_t error_setup = node->tensor_op.setup(config);
+      error = parse_setup_error(error_setup);
+
+      if (error != ErrorParse::None)
+      {
+        return error;
+      }
+
+#ifdef SAVE_JITS_TO_FILE
+      node->tensor_op.write_kernel_to_file(node->name());
+#endif  // SAVE_JITS_TO_FILE
+    }
+    else
+    {
+      release_assert(false, "Found unhandled einsum tree node type.");
+    }
+  }
   return ErrorParse::None;
 }
 
diff --git a/src/main/EinsumTree.h b/src/main/EinsumTree.h
@@ -24,15 +24,7 @@ namespace mini_jit
       ExpectedDimensionList = 5,
       NotAllowedToParseAgain = 6,
       UndefinedNode = 7,
-    };
-
-    enum class ErrorExecute
-    {
-      None = 0,
-      InvalidRoot = 1,
-      NotEnoughInputTensors = 2,
-      TooManyInputTensors = 3,
-      NullPtrAsInputTensor = 5,
+      InvalidRoot = 8,
 
       err_wrong_dtype = 101,
       err_wrong_dimension = 102,
@@ -51,6 +43,15 @@ namespace mini_jit
       err_shared_required_for_parallel_execution = 115,
     };
 
+    enum class ErrorExecute
+    {
+      None = 0,
+      InvalidRoot = 1,
+      NotEnoughInputTensors = 2,
+      TooManyInputTensors = 3,
+      NullPtrAsInputTensor = 5,
+    };
+
     enum class NodeType
     {
       Leaf,
@@ -63,6 +64,7 @@ namespace mini_jit
       NodeType type;
       int32_t input_tensor_index = -1;
       float *tensor = nullptr;
+      mini_jit::TensorOperation tensor_op;
 
       // Always filled — dims of the output tensor
       std::vector<int64_t> output_dim_ids;
@@ -198,7 +200,7 @@ namespace mini_jit
      * @param error The error code from TensorOperation.
      * @return An ErrorExecute enum representing the parsed error.
      */
-    ErrorExecute parse_setup_error(TensorOperation::error_t error);
+    ErrorParse parse_setup_error(TensorOperation::error_t error);
 
     // Cleanup
     /**
@@ -231,6 +233,13 @@ namespace mini_jit
      */
     int32_t findMDim(EinsumNode *Node);
 
+    /**
+     * @brief Generates the operator to the parsed einsum tree.
+     *
+     * @return ErrorParse indicating the result of the parsing operation.
+     */
+    ErrorParse generate_operators();
+
   public:
     EinsumTree(const std::string &tree_str);
     EinsumTree(const std::string &tree_str, const std::vector<int64_t> &sorted_dim_sizes);
@@ -248,9 +257,10 @@ namespace mini_jit
     /**
      * Parses the einsum tree string and builds the tree structure.
      *
+     * @param build_operators indicates if the operators should be generate with the parse.
      * @return ErrorParse indicating the result of the parsing operation.
      */
-    ErrorParse parse_tree_no_optimization();
+    ErrorParse parse_tree_no_optimization(bool build_operators = true);
 
     /**
      * Parses the einsum tree string, builds the tree structure and optimizes the tree.
diff --git a/src/test/EinsumTree.test.cpp b/src/test/EinsumTree.test.cpp
@@ -24,7 +24,7 @@ TEST_CASE("Test einsum tree parser simple example", "[einsumtree][parse][correct
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization(false);
   REQUIRE(err == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -47,7 +47,7 @@ TEST_CASE("Test einsum tree parser first example", "[einsumtree][parse][correctn
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization(false);
   REQUIRE(err == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -70,7 +70,7 @@ TEST_CASE("Test einsum tree parser second example", "[einsumtree][parse][correct
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err_parse = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err_parse = tree.parse_tree_no_optimization(false);
   REQUIRE(err_parse == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -236,7 +236,7 @@ TEST_CASE("Test einsum tree optimize swap", "[einsumtree][optimize][correctness]
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization(false);
   REQUIRE(err == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -264,7 +264,7 @@ TEST_CASE("Test einsum tree optimize reorder left", "[einsumtree][optimize][corr
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization(false);
   REQUIRE(err == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -300,7 +300,7 @@ TEST_CASE("Test einsum tree optimize reorder right", "[einsumtree][optimize][cor
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization(false);
   REQUIRE(err == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -338,7 +338,7 @@ TEST_CASE("Test einsum tree optimize", "[einsumtree][optimize][correctness]")
 
   EinsumTree tree(tree_str, sorted_dim_sizes);
 
-  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization();
+  mini_jit::EinsumTree::ErrorParse err = tree.parse_tree_no_optimization(false);
   REQUIRE(err == mini_jit::EinsumTree::ErrorParse::None);
   REQUIRE(tree.get_root() != nullptr);
   INFO(tree.get_root()->to_string());
@@ -420,7 +420,7 @@ TEST_CASE("Test einsum tree optimize and execute first example", "[einsumtree][e
                                       "   └─ 0,5\n";
   REQUIRE_THAT(expected_optimization, Catch::Matchers::Equals(tree.get_root()->to_string(), Catch::CaseSensitive::Yes));
 
-  tree_no_optimization.parse_tree_no_optimization();
+  tree_no_optimization.parse_tree_no_optimization(false);
   INFO("No Optimization");
   INFO(tree_no_optimization.get_root()->to_string());
 
@@ -483,7 +483,7 @@ TEST_CASE("Test einsum tree optimize and execute second example", "[einsumtree][
                                       "   └─ 1,4,7,8\n";
   REQUIRE_THAT(expected_optimization, Catch::Matchers::Equals(tree.get_root()->to_string(), Catch::CaseSensitive::Yes));
 
-  tree_no_optimization.parse_tree_no_optimization();
+  tree_no_optimization.parse_tree_no_optimization(false);
   INFO("No Optimization");
   INFO(tree_no_optimization.get_root()->to_string());
 
@@ -550,7 +550,7 @@ TEST_CASE("Test einsum tree optimize and execute third example", "[einsumtree][e
 
   REQUIRE_THAT(expected_optimization, Catch::Matchers::Equals(tree.get_root()->to_string(), Catch::CaseSensitive::Yes));
 
-  tree_no_optimization.parse_tree_no_optimization();
+  tree_no_optimization.parse_tree_no_optimization(false);
   INFO("No Optimization");
   INFO(tree_no_optimization.get_root()->to_string());
 
