Geometry Nodes: move up destruct instructions in procedure

This implements an optimization pass for multi-function procedures.
It optimizes memory reuse by moving destruct instructions up.
For more details see the in-code comment.

In very large fields with many short lived intermediate values, this change
can improve performance 3-4x. Furthermore, in such cases, peak memory
consumption is reduced significantly (e.g. 100x lower peak memory usage).

Differential Revision: https://developer.blender.org/D13548

5 files changed, 160 insertions, 4 deletions

diff --git a/source/blender/functions/CMakeLists.txt b/source/blender/functions/CMakeLists.txt
index 63c11164275..9cfaf3eabea 100644
--- a/source/blender/functions/CMakeLists.txt
+++ b/source/blender/functions/CMakeLists.txt
@@ -38,6 +38,7 @@ set(SRC
   intern/multi_function_procedure.cc
   intern/multi_function_procedure_builder.cc
   intern/multi_function_procedure_executor.cc
+  intern/multi_function_procedure_optimization.cc
 
   FN_cpp_type.hh
   FN_cpp_type_make.hh
@@ -59,6 +60,7 @@ set(SRC
   FN_multi_function_procedure.hh
   FN_multi_function_procedure_builder.hh
   FN_multi_function_procedure_executor.hh
+  FN_multi_function_procedure_optimization.hh
   FN_multi_function_signature.hh
 )
 
diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index 57c2753b951..cf96eff62bd 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -53,9 +53,6 @@
 
 #include "FN_generic_virtual_array.hh"
 #include "FN_multi_function_builder.hh"
-#include "FN_multi_function_procedure.hh"
-#include "FN_multi_function_procedure_builder.hh"
-#include "FN_multi_function_procedure_executor.hh"
 
 namespace blender::fn {
 
diff --git a/source/blender/functions/FN_multi_function_procedure_optimization.hh b/source/blender/functions/FN_multi_function_procedure_optimization.hh
new file mode 100644
index 00000000000..e5ffc12b241
--- /dev/null
+++ b/source/blender/functions/FN_multi_function_procedure_optimization.hh
@@ -0,0 +1,61 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#pragma once
+
+/** \file
+ * \ingroup fn
+ *
+ * A #MFProcedure optimization pass takes an existing procedure and changes it in a way that
+ * improves its performance when executed.
+ *
+ * Oftentimes it would also be possible to implement a specific optimization directly during
+ * construction of the initial #MFProcedure. There is a trade-off between doing that or just
+ * building a "simple" procedure and then optimizing it uses separate optimization passes.
+ * - Doing optimizations directly during construction is typically faster than doing it as a
+ *   separate pass. However, it would be much harder to turn the optimization off when it is not
+ *   necessary, making the construction potentially slower in those cases.
+ * - Doing optimizations directly would also make code more complex, because it mixes the logic
+ *   that generates the procedure from some other data with optimization decisions.
+ * - Having a separate pass allows us to use it in different places when necessary.
+ * - Having a separate pass allows us to enable and disable it easily to better understand its
+ *   impact on performance.
+ */
+
+#include "FN_multi_function_procedure.hh"
+
+namespace blender::fn::procedure_optimization {
+
+/**
+ * When generating a procedure, destruct instructions (#MFDestructInstruction) have to be inserted
+ * for all variables that are not outputs. Often the simplest approach is to add these instructions
+ * at the very end. However, when the procedure is executed this is not optimal, because many more
+ * variables are initialized at the same time than necessary. This inhibits the reuse of memory
+ * buffers which decreases performance and increases memory use.
+ *
+ * This optimization pass moves destruct instructions up in the procedure. The goal is to destruct
+ * each variable right after its last use.
+ *
+ * For simplicity, and because this is the most common use case, this optimization currently only
+ * works on a single chain of instructions. Destruct instructions are not moved across branches.
+ *
+ * \param procedure The procedure that should be optimized.
+ * \param block_end_instr The instruction that points to the last instruction within a linear chain
+ *   of instructions. The algorithm moves instructions backward starting at this instruction.
+ */
+void move_destructs_up(MFProcedure &procedure, MFInstruction &block_end_instr);
+
+}  // namespace blender::fn::procedure_optimization
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index 3274af4a7be..a014fd113e4 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -21,6 +21,10 @@
 #include "BLI_vector_set.hh"
 
 #include "FN_field.hh"
+#include "FN_multi_function_procedure.hh"
+#include "FN_multi_function_procedure_builder.hh"
+#include "FN_multi_function_procedure_executor.hh"
+#include "FN_multi_function_procedure_optimization.hh"
 
 namespace blender::fn {
 
@@ -251,7 +255,9 @@ static void build_multi_function_procedure_for_fields(MFProcedure &procedure,
     builder.add_destruct(*variable);
   }
 
-  builder.add_return();
+  MFReturnInstruction &return_instr = builder.add_return();
+
+  procedure_optimization::move_destructs_up(procedure, return_instr);
 
   // std::cout << procedure.to_dot() << "\n";
   BLI_assert(procedure.validate());
diff --git a/source/blender/functions/intern/multi_function_procedure_optimization.cc b/source/blender/functions/intern/multi_function_procedure_optimization.cc
new file mode 100644
index 00000000000..f220c85e535
--- /dev/null
+++ b/source/blender/functions/intern/multi_function_procedure_optimization.cc
@@ -0,0 +1,90 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#include "FN_multi_function_procedure_optimization.hh"
+
+namespace blender::fn::procedure_optimization {
+
+void move_destructs_up(MFProcedure &procedure, MFInstruction &block_end_instr)
+{
+  /* A mapping from a variable to its destruct instruction. */
+  Map<MFVariable *, MFDestructInstruction *> destruct_instructions;
+  MFInstruction *current_instr = &block_end_instr;
+  while (true) {
+    MFInstructionType instr_type = current_instr->type();
+    switch (instr_type) {
+      case MFInstructionType::Destruct: {
+        MFDestructInstruction &destruct_instr = static_cast<MFDestructInstruction &>(
+            *current_instr);
+        MFVariable *variable = destruct_instr.variable();
+        if (variable == nullptr) {
+          continue;
+        }
+        /* Remember this destruct instruction so that it can be moved up later on when the last use
+         * of the variable is found. */
+        destruct_instructions.add(variable, &destruct_instr);
+        break;
+      }
+      case MFInstructionType::Call: {
+        MFCallInstruction &call_instr = static_cast<MFCallInstruction &>(*current_instr);
+        /* For each variable, place the corresponding remembered destruct instruction right after
+         * this call instruction. */
+        for (MFVariable *variable : call_instr.params()) {
+          if (variable == nullptr) {
+            continue;
+          }
+          MFDestructInstruction *destruct_instr = destruct_instructions.pop_default(variable,
+                                                                                    nullptr);
+          if (destruct_instr == nullptr) {
+            continue;
+          }
+
+          /* Unlink destruct instruction from previous position. */
+          MFInstruction *after_destruct_instr = destruct_instr->next();
+          while (!destruct_instr->prev().is_empty()) {
+            /* Do a copy of the cursor here, because `destruct_instr->prev()` changes when
+             * #set_next is called below. */
+            const MFInstructionCursor cursor = destruct_instr->prev()[0];
+            cursor.set_next(procedure, after_destruct_instr);
+          }
+
+          /* Insert destruct instruction in new position. */
+          MFInstruction *next_instr = call_instr.next();
+          call_instr.set_next(destruct_instr);
+          destruct_instr->set_next(next_instr);
+        }
+        break;
+      }
+      default: {
+        break;
+      }
+    }
+
+    const Span<MFInstructionCursor> prev_cursors = current_instr->prev();
+    if (prev_cursors.size() != 1) {
+      /* Stop when there is some branching before this instruction. */
+      break;
+    }
+    const MFInstructionCursor &prev_cursor = prev_cursors[0];
+    current_instr = prev_cursor.instruction();
+    if (current_instr == nullptr) {
+      /* Stop when there is no previous instruction. E.g. when this is the first instruction.  */
+      break;
+    }
+  }
+}
+
+}  // namespace blender::fn::procedure_optimization