ITF export: Implement index batching support

SeQuant/core/eval/eval_node_compare.hpp

@@ -4,6 +4,7 @@
 #include <SeQuant/core/hash.hpp>
 #include <SeQuant/core/utility/macros.hpp>
 #include <SeQuant/core/utility/tensor.hpp>
+#include <SeQuant/external/bliss/graph.hh>
 
 #include <cstddef>
 #include <unordered_map>
@@ -34,6 +35,10 @@ struct TreeNodeEqualityComparator {
   /// Trait used by the C++ STL allowing heterogenous lookups
   using is_transparent = void;
 
+  TreeNodeEqualityComparator() = default;
+  TreeNodeEqualityComparator(std::vector<Index> indices)
+      : block_comparator_(std::move(indices)) {}
+
   bool operator()(const TreeNode *lhs, const TreeNode *rhs) const {
     return (*this)(*lhs, *rhs);
   }
@@ -84,8 +89,7 @@ struct TreeNodeEqualityComparator {
       const Tensor &lhs_tensor = lhs->as_tensor();
       const Tensor &rhs_tensor = rhs->as_tensor();
 
-      TensorBlockEqualComparator cmp;
-      if (!cmp(lhs_tensor, rhs_tensor)) {
+      if (!block_comparator_(lhs_tensor, rhs_tensor)) {
         return false;
       }
     }
@@ -119,6 +123,9 @@ struct TreeNodeEqualityComparator {
 
     return true;
   }
+
+ private:
+  IndexSpecificTensorBlockEqualComparator block_comparator_;
 };
 
 /// A map between (sub)tree hashes and how often they have been found

SeQuant/core/export/context.cpp

@@ -74,16 +74,16 @@ bool operator!=(Usage usage, UsageSet set) { return set != usage; }
 
 ExportContext::ExportContext(TensorStrategyMap tensorMap,
                              VariableStrategyMap variableMap)
-    : m_tensorStrategies({{GLOBAL, std::move(tensorMap)}}),
-      m_variableStrategies({{GLOBAL, std::move(variableMap)}}) {}
+    : m_tensorStrategies({{ID_GLOBAL, std::move(tensorMap)}}),
+      m_variableStrategies({{ID_GLOBAL, std::move(variableMap)}}) {}
 
 ExportContext::ExportContext(VariableStrategyMap map)
-    : m_variableStrategies({{GLOBAL, std::move(map)}}) {}
+    : m_variableStrategies({{ID_GLOBAL, std::move(map)}}) {}
 
 ExportContext::~ExportContext() = default;
 
 LoadStrategy ExportContext::loadStrategy(const Tensor &tensor) const {
-  if (auto map_iter = m_tensorStrategies.find(GLOBAL);
+  if (auto map_iter = m_tensorStrategies.find(ID_GLOBAL);
       map_iter != m_tensorStrategies.end()) {
     auto iter = map_iter->second.find(tensor);
     if (iter != map_iter->second.end()) {
@@ -105,7 +105,7 @@ LoadStrategy ExportContext::loadStrategy(const Tensor &tensor) const {
 }
 
 LoadStrategy ExportContext::loadStrategy(const Variable &variable) const {
-  if (auto map_iter = m_variableStrategies.find(GLOBAL);
+  if (auto map_iter = m_variableStrategies.find(ID_GLOBAL);
       map_iter != m_variableStrategies.end()) {
     auto iter = map_iter->second.find(variable);
     if (iter != map_iter->second.end()) {
@@ -130,7 +130,7 @@ void ExportContext::setLoadStrategy(
     const Tensor &tensor, LoadStrategy strategy,
     const std::optional<std::size_t> &expression_id) {
   std::size_t id = expression_id.value_or(
-      has_current_expression_id() ? current_expression_id() : GLOBAL);
+      has_current_expression_id() ? current_expression_id() : ID_GLOBAL);
 
   auto iter = m_tensorStrategies[id].find(tensor);
 
@@ -145,7 +145,7 @@ void ExportContext::setLoadStrategy(
     const Variable &variable, LoadStrategy strategy,
     const std::optional<std::size_t> &expression_id) {
   std::size_t id = expression_id.value_or(
-      has_current_expression_id() ? current_expression_id() : GLOBAL);
+      has_current_expression_id() ? current_expression_id() : ID_GLOBAL);
 
   auto iter = m_variableStrategies[id].find(variable);
 
@@ -157,7 +157,7 @@ void ExportContext::setLoadStrategy(
 }
 
 ZeroStrategy ExportContext::zeroStrategy(const Tensor &tensor) const {
-  if (auto map_iter = m_tensorStrategies.find(GLOBAL);
+  if (auto map_iter = m_tensorStrategies.find(ID_GLOBAL);
       map_iter != m_tensorStrategies.end()) {
     auto iter = map_iter->second.find(tensor);
     if (iter != map_iter->second.end()) {
@@ -179,7 +179,7 @@ ZeroStrategy ExportContext::zeroStrategy(const Tensor &tensor) const {
 }
 
 ZeroStrategy ExportContext::zeroStrategy(const Variable &variable) const {
-  if (auto map_iter = m_variableStrategies.find(GLOBAL);
+  if (auto map_iter = m_variableStrategies.find(ID_GLOBAL);
       map_iter != m_variableStrategies.end()) {
     auto iter = map_iter->second.find(variable);
     if (iter != map_iter->second.end()) {
@@ -204,7 +204,7 @@ void ExportContext::setZeroStrategy(
     const Tensor &tensor, ZeroStrategy strategy,
     const std::optional<std::size_t> &expression_id) {
   std::size_t id = expression_id.value_or(
-      has_current_expression_id() ? current_expression_id() : GLOBAL);
+      has_current_expression_id() ? current_expression_id() : ID_GLOBAL);
 
   auto iter = m_tensorStrategies[id].find(tensor);
 
@@ -219,7 +219,7 @@ void ExportContext::setZeroStrategy(
     const Variable &variable, ZeroStrategy strategy,
     const std::optional<std::size_t> &expression_id) {
   std::size_t id = expression_id.value_or(
-      has_current_expression_id() ? current_expression_id() : GLOBAL);
+      has_current_expression_id() ? current_expression_id() : ID_GLOBAL);
 
   auto iter = m_variableStrategies[id].find(variable);
 
@@ -264,4 +264,9 @@ void ExportContext::clear_current_expression_id() {
   m_currentExpressionID.reset();
 }
 
+std::vector<Index> ExportContext::batch_indices(
+    std::optional<std::size_t>) const {
+  return {};
+}
+
 }  // namespace sequant

SeQuant/core/export/context.hpp

@@ -2,13 +2,15 @@
 #define SEQUANT_CORE_EXPORT_CONTEXT_HPP
 
 #include <SeQuant/core/expr.hpp>
+#include <SeQuant/core/index.hpp>
 #include <SeQuant/core/utility/tensor.hpp>
 
 #include <limits>
 #include <map>
 #include <optional>
 #include <string>
 #include <type_traits>
+#include <vector>
 
 namespace sequant {
 
@@ -199,9 +201,17 @@ class ExportContext {
   /// Resets the ID of the current expression
   virtual void clear_current_expression_id();
 
- private:
-  static constexpr std::size_t GLOBAL = std::numeric_limits<std::size_t>::max();
+  /// @param id The ID of the relevant expression
+  /// @returns The list of indices that the given expression should be batched
+  /// over
+  virtual std::vector<Index> batch_indices(
+      std::optional<std::size_t> id = {}) const;
+
+ protected:
+  static constexpr std::size_t ID_GLOBAL =
+      std::numeric_limits<std::size_t>::max();
 
+ private:
   std::map<std::size_t, TensorStrategyMap> m_tensorStrategies;
   std::map<std::size_t, VariableStrategyMap> m_variableStrategies;
   std::optional<std::string> m_currentSection;

SeQuant/core/export/export.hpp

@@ -30,6 +30,17 @@ namespace sequant {
 
 namespace detail {
 
+/// A collection of various meta-data that the preprocessing stage will collect
+struct PreprocessResult {
+  std::unordered_map<std::size_t, ExprPtr> scalarFactors;
+  std::set<Index> indices;
+  std::map<Tensor, UsageSet, TensorBlockLessThanComparator> tensors;
+  std::map<Variable, UsageSet> variables;
+
+  std::map<Tensor, std::size_t, TensorBlockLessThanComparator> tensorReferences;
+  std::map<Variable, std::size_t> variableReferences;
+};
+
 /// Visitor objects that will steer code generation while visiting a given
 /// expression/evaluation tree by triggering the corresponding callbacks in the
 /// provided Generator objects.
@@ -44,7 +55,19 @@ class GenerationVisitor {
   /// need to be multiplied with the result before storing the node
   GenerationVisitor(Generator<Context> &generator, Context &ctx,
                     const std::unordered_map<NodeID, ExprPtr> &scalarFactors)
-      : m_generator(generator), m_ctx(ctx), m_scalarFactors(scalarFactors) {}
+      : m_generator(generator), m_ctx(ctx), m_scalarFactors(scalarFactors) {
+    if (m_generator.supports_index_batching()) {
+      // Make tensor comparator aware of the list of batched indices
+      std::vector<Index> batchIndices =
+          m_ctx.batch_indices(m_ctx.current_expression_id());
+      if (!batchIndices.empty()) {
+        IndexSpecificTensorBlockLessThanComparator cmp =
+            m_tensorUses.key_comp();
+        cmp.set_indices(std::move(batchIndices));
+        m_tensorUses = decltype(m_tensorUses)(std::move(cmp));
+      }
+    }
+  }
 
   void operator()(const ExportNode<NodeData> &node, TreeTraversal context) {
     // Note the context for leaf nodes is always TreeTraversal::Any
@@ -304,23 +327,13 @@ class GenerationVisitor {
   Generator<Context> &m_generator;
   Context &m_ctx;
   const std::unordered_map<NodeID, ExprPtr> &m_scalarFactors;
-  std::map<Tensor, std::size_t, TensorBlockLessThanComparator> m_tensorUses;
+  std::map<Tensor, std::size_t, IndexSpecificTensorBlockLessThanComparator>
+      m_tensorUses;
   std::map<Variable, std::size_t> m_variableUses;
 
   std::optional<NodeID> m_rootID;
 };
 
-/// A collection of various meta-data that the preprocessing stage will collect
-struct PreprocessResult {
-  std::unordered_map<std::size_t, ExprPtr> scalarFactors;
-  std::set<Index> indices;
-  std::map<Tensor, UsageSet, TensorBlockLessThanComparator> tensors;
-  std::map<Variable, UsageSet> variables;
-
-  std::map<Tensor, std::size_t, TensorBlockLessThanComparator> tensorReferences;
-  std::map<Variable, std::size_t> variableReferences;
-};
-
 /// Removes explicitly represented scalar factors from the provided tree and
 /// instead stores them separately. This yields a much more compact tree and
 /// makes subsequent visiting easier as scalar factors should simply be
@@ -409,8 +422,9 @@ bool prune_scalar_factor(ExportNode<T> &node, PreprocessResult &result,
 /// Renames the given Tensor to a name that doesn't collide with any currently
 /// loaded object. This may reuse previously used/declares tensors.
 bool rename(Tensor &tensor, PreprocessResult &result);
-/// Renames the given Variable to a name that doesn't collide with any currently
-/// loaded object. This may reuse previously used/declares variables.
+/// Renames the given Variable to a name that doesn't collide with any
+/// currently loaded object. This may reuse previously used/declares
+/// variables.
 bool rename(Variable &variable, PreprocessResult &result);
 
 /// Preprocesses the given expression
@@ -423,8 +437,9 @@ void preprocess(ExprType expr, ExportContext &ctx, Node &node,
 
   bool storeExpr = false;
 
-  // TODO: find a way to pass usage information to this call so that indices of
-  // tensors that are only used as an intermediate can be more easily reordered
+  // TODO: find a way to pass usage information to this call so that indices
+  // of tensors that are only used as an intermediate can be more easily
+  // reordered
   storeExpr |= ctx.rewrite(expr);
 
   if (node.leaf()) {
@@ -532,8 +547,8 @@ void track_usage(const EvalNode<T> &node, PreprocessResult &result) {
   }
 }
 
-/// @returns Whether the given node may be pruned from its parent in order to be
-/// represented implicitly rather than by explicit occurrence in the tree
+/// @returns Whether the given node may be pruned from its parent in order to
+/// be represented implicitly rather than by explicit occurrence in the tree
 template <typename T>
 bool may_prune(const EvalNode<T> &tree) {
   // Tree must represent a product and must itself not be a leaf (pruning that
@@ -565,17 +580,18 @@ bool may_prune(const EvalNode<T> &tree) {
 ///
 /// Preprocesses the provided binary tree by
 /// - removing explicit appearances of scalar leafs. We don't want them to
-///   be represented in the tree. Instead, we keep track of them in a different
-///   way in order to be able to give scalar factors alongside the actual
-///   tensor contraction they are supposed to scale (this is necessary
+///   be represented in the tree. Instead, we keep track of them in a
+///   different way in order to be able to give scalar factors alongside the
+///   actual tensor contraction they are supposed to scale (this is necessary
 ///   as there are backends which only support scaling in this context)
 /// - rebalance the tree such that for any given non-leaf node, its left
 ///   subtree is always larger (or equally large) than its right one.
 ///   This ensures that we have to have the least amount of tensors loaded
 ///   at the same time, when generating code for a backend which only supports
 ///   stack-like memory allocations (e.g. when A is allocated before B, B
 ///   must be deleted before A can be deleted).
-/// - Rename intermediate tensors that have the same name and describe the same
+/// - Rename intermediate tensors that have the same name and describe the
+/// same
 ///   tensor block, which are required as two separate entities at the same
 ///   time when evaluating the tree (thus a single tensor object is
 ///   insufficient).
@@ -603,8 +619,9 @@ class PreprocessVisitor {
           prune_redundant_intermediate(tree);
         }
 
-        // It is important to track_usage AFTER prune_redundant_intermediate as
-        // the latter might change the result expression on the current node
+        // It is important to track_usage AFTER prune_redundant_intermediate
+        // as the latter might change the result expression on the current
+        // node
         track_usage(tree, m_result);
         break;
       case TreeTraversal::PostOrder:
@@ -729,25 +746,40 @@ class PreprocessVisitor {
   }
 
   void release_used_terms(ExportNode<T> &node) {
-    // Mark tensors/variables as no longer in use
-    if (node.left()->is_tensor()) {
-      const Tensor &tensor = node.left()->as_tensor();
+    auto handle_tensor = [&](const Tensor &tensor) {
       SEQUANT_ASSERT(m_result.tensorReferences[tensor] > 0);
       m_result.tensorReferences[tensor]--;
-    } else if (node.left()->is_variable()) {
-      const Variable &variable = node.left()->as_variable();
+    };
+    auto handle_variable = [&](const Variable &variable) {
       SEQUANT_ASSERT(m_result.variableReferences[variable] > 0);
       m_result.variableReferences[variable]--;
+    };
+
+    // Mark tensors/variables as no longer in use
+    if (node.left()->is_tensor()) {
+      handle_tensor(node.left()->as_tensor());
+    } else if (node.left()->is_variable()) {
+      handle_variable(node.left()->as_variable());
+    } else if (node.left()->is_power()) {
+      const Power &power = node.left()->as_power();
+      if (power.base().is<Tensor>()) {
+        handle_tensor(power.base().as<Tensor>());
+      } else if (power.base().is<Variable>()) {
+        handle_variable(power.base().as<Variable>());
+      }
     }
 
     if (node.right()->is_tensor()) {
-      const Tensor &tensor = node.right()->as_tensor();
-      SEQUANT_ASSERT(m_result.tensorReferences[tensor] > 0);
-      m_result.tensorReferences[tensor]--;
+      handle_tensor(node.right()->as_tensor());
     } else if (node.right()->is_variable()) {
-      const Variable &variable = node.right()->as_variable();
-      SEQUANT_ASSERT(m_result.variableReferences[variable] > 0);
-      m_result.variableReferences[variable]--;
+      handle_variable(node.right()->as_variable());
+    } else if (node.right()->is_power()) {
+      const Power &power = node.right()->as_power();
+      if (power.base().is<Tensor>()) {
+        handle_tensor(power.base().as<Tensor>());
+      } else if (power.base().is<Variable>()) {
+        handle_variable(power.base().as<Variable>());
+      }
     }
   }
 
@@ -797,6 +829,7 @@ void export_expression(ExportNode<T> &expression, Generator<Context> &generator,
 
   detail::GenerationVisitor<T, Context> visitor(generator, ctx,
                                                 pp_result.scalarFactors);
+
   expression.visit(
       [&visitor](const FullBinaryNode<T> &node, TreeTraversal context) {
         visitor(node, context);
@@ -842,8 +875,8 @@ void declare_all(const Range &range, Generator<Context> &generator,
 }
 
 /// Combines the known T from the range of
-/// PreprocessResults and clears the respective fields of the individual result
-/// objects.
+/// PreprocessResults and clears the respective fields of the individual
+/// result objects.
 /// @returns The combined set of known objects
 template <typename T, typename Compare = std::less<T>, typename Range>
   requires std::ranges::range<Range> &&