improve inplace performance

Summary:
!ci_branch_mk2

When inplace open, poprithm memory graph will repeatedly constructed to query alias (a view) of a tensor.

This is time consuming. The naive idea is to reuse the the alias model and poprithm graph grower to save
compilation time.

Reviewers: #popart, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, jamesn, shirazb

Reviewed By: #popart, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, jamesn, shirazb

Subscribers: markk, shirazb

Maniphest Tasks: T74421

Differential Revision: https://phabricator.sourcevertex.net/D85549

Author: yiakwy-xpu-ml-framework-team

willow/include/popart/ir.hpp

@@ -12,6 +12,8 @@
 #include <string>
 #include <typeindex>
 #include <vector>
+#include <popart/alias/aliasmodel.hpp>
+#include <popart/alias/aliasmodelgrower.hpp>
 #include <popart/bimap.hpp>
 #include <popart/dataflow.hpp>
 #include <popart/inputshapeinfo.hpp>

willow/include/popart/op.hpp

@@ -16,6 +16,7 @@
 #include <utility>
 #include <vector>
 #include <poprithms/memory/inplace/proposal.hpp>
+#include <popart/alias/aliasmodelgrower.hpp>
 #include <popart/attributes.hpp>
 #include <popart/basicoptionals.hpp>
 #include <popart/bwdgraphinfo.hpp>
@@ -1678,6 +1679,16 @@ class Op : public Vertex {
    */
   bool inputUnmodifiable(InIndex in) const;
 
+  /**
+   * Check if the input index is unmodifiable or aliases an unmodifiable tensor
+   * with given poprithm graph.
+   *
+   * \param in The input index to check.
+   * \returns `true` if any connected variable tensor has a non-empty alias
+   *      chain and is unmodifiable, `false` otherwise.
+   */
+  bool inputUnmodifiableFor(InIndex in, const AliasModel *popMem) const;
+
   /**
    * Check if output is modified by any consumer.
    *
@@ -1687,6 +1698,15 @@ class Op : public Vertex {
    */
   bool hasAliasedModifiers(OutIndex out) const;
 
+  /**
+   * Check if output is modified by any consumer with the given poprithm graph.
+   *
+   * \param out The output index to check.
+   * \returns `true` if any consumer of any aliased tensor downstream modifies
+   *              a non-empty region, `false` otherwise.
+   */
+  bool hasAliasedModifiersFor(OutIndex out, const AliasModel *popMem) const;
+
   // Helper functions for probing graph structure.
   /**
    * Check if the graph is a parent of the op.

willow/include/popart/tensor.hpp

@@ -11,6 +11,7 @@
 #include <string>
 #include <utility>
 #include <vector>
+#include <popart/alias/aliasmodelgrower.hpp>
 #include <popart/dataflow.hpp>
 #include <popart/error.hpp>
 #include <popart/names.hpp>
@@ -221,6 +222,11 @@ class Tensor : public Vertex {
   // Returns true if the tensor or any of it's aliases fulfill the predicate
   bool anyAlias(std::function<bool(Tensor *)> predicate) const;
 
+  // Returns true if the tensor or any of it's aliases fulfill the predicate in
+  // the given poprithm memory graph
+  bool anyAliasFor(std::function<bool(Tensor *)> predicate,
+                   const AliasModel &popMem) const;
+
   void setTensorDataFromCopyOf(const void *src, std::size_t size);
   void setTensorDataFromViewOf(void *src, std::size_t size);
   void setTensorDataByEmplaceOf(std::vector<char> &&data);
@@ -291,6 +297,12 @@ class Tensor : public Vertex {
    */
   std::set<Op *, POpCmp> getInplaceModifiers() const;
 
+  /**
+   * Find operations that modify a tensor with the given poprithm graph
+   * \return All operations that (direct and indirectly) modify this tensor
+   */
+  std::set<Op *, POpCmp> getInplaceModifiersFor(const AliasModel *popMem) const;
+
   // Backtrack through input and parent graph tensors in order to get data from
   // initializer tensors (if they exist).
   // When ops are performed on initializers (e.g. slice), the
@@ -311,6 +323,10 @@ class Tensor : public Vertex {
 
   int getBatchAxisFromOp(Op *, bool, int) const;
 
+  bool anyAliasImpl(std::function<bool(Tensor *)> predicate,
+                    const AliasModel &popMem,
+                    const char *scopeDesc) const;
+
   const TensorDebugInfo di;
 
   /**

willow/src/ir.cpp

@@ -8,7 +8,6 @@
 #include <graphfromlosstolossupdater.hpp>
 #include <onnxutil.hpp>
 #include <poprithms/logging/timepartitionlogger.hpp>
-#include <popart/alias/aliasmodelgrower.hpp>
 #include <popart/ces/constexpr.hpp>
 #include <popart/devicemanager.hpp>
 #include <popart/error.hpp>
@@ -48,7 +47,6 @@
 #include <poprithms/memory/inplace/proposal.hpp>
 #include <poprithms/memory/inplace/result.hpp>
 #include <poprithms/util/typedinteger.hpp>
-#include <popart/alias/aliasmodel.hpp>
 #include <popart/dotvisualizer.hpp>
 #include <popart/op/copyvarupdate.hpp>
 #include <popart/op/ipucopy.hpp>
@@ -157,6 +155,55 @@ std::ostream &operator<<(std::ostream &ost, const OpsBeforeKey &o) {
   return ost;
 }
 
+struct PopMemGrower {
+  std::unique_ptr<AliasModel> popMem_;
+  std::unique_ptr<AliasModelGrower> aliasModelGrower_;
+
+  using Proposal      = poprithms::memory::inplace::Proposal;
+  using OpeningStatus = poprithms::memory::inplace::OpeningStatus;
+
+  PopMemGrower();
+
+  void init();
+
+  void growFullGraph(Graph &graph,
+                     DataDependenciesOnly dep = DataDependenciesOnly::No) {
+    aliasModelGrower_->growFullGraph(graph, dep);
+  }
+
+  AliasModel &getAliasModelRef() { return *(popMem_.get()); }
+
+  Proposal mapInplaceProposal(Op *op, OperatorIdentifier identifier) {
+    return op->mapInplaceProposal(*popMem_, identifier);
+  }
+
+  OpeningStatus tryOpening(Proposal &proposal) {
+    return popMem_->g.tryOpening(
+        proposal,
+        poprithms::memory::inplace::CheckParallelWriteable::No,
+        poprithms::memory::inplace::AllowMultiGateAlias::No);
+  }
+
+  void reset(Graph &graph) {
+    init();
+    growFullGraph(graph, DataDependenciesOnly::Yes);
+  }
+};
+
+PopMemGrower::PopMemGrower()
+    : popMem_(std::make_unique<AliasModel>()),
+      aliasModelGrower_(std::make_unique<AliasModelGrower>(*popMem_)) {}
+
+void PopMemGrower::init() {
+  if (popMem_ == nullptr) {
+    popMem_           = std::make_unique<AliasModel>();
+    aliasModelGrower_ = std::make_unique<AliasModelGrower>(*popMem_);
+  } else {
+    popMem_.reset(new AliasModel());
+    aliasModelGrower_.reset(new AliasModelGrower(*popMem_));
+  }
+}
+
 poprithms::logging::TimePartitionLogger &Ir::timePartitionLogger() const {
   return *timePartitionLogger_;
 }
@@ -3466,12 +3513,16 @@ void Ir::applyUpdateInplacePrioritiesForIpu() {
 }
 
 void Ir::applyInplacePattern(Graph &graph) {
-
   // The decision of where topological constraints need to be inserted is made
   // by a poprithms Graph whose Ops mirror those in \a graph.
-  AliasModel popMem;
-  AliasModelGrower aliasModelGrower{popMem};
-  aliasModelGrower.growFullGraph(graph, DataDependenciesOnly::No);
+  // Create poprithm memroy graph growers for this graph
+  PopMemGrower popMemGrowerOfSubgraph;
+  PopMemGrower popMemGrowerOfTensor;
+
+  popMemGrowerOfSubgraph.growFullGraph(graph, DataDependenciesOnly::No);
+  popMemGrowerOfTensor.growFullGraph(graph, DataDependenciesOnly::Yes);
+
+  AliasModel &popMem = *(popMemGrowerOfSubgraph.popMem_.get());
 
   Inplace inplace;
 
@@ -3591,6 +3642,8 @@ void Ir::applyInplacePattern(Graph &graph) {
         continue;
       }
 
+      auto tProposal = popMemGrowerOfTensor.mapInplaceProposal(op, identifier);
+
       // Convert poprithms topological constraints into popart constraints
       OpsBeforeKey newTopoCons;
       for (auto from_to : result.constraints()) {
@@ -3695,14 +3748,22 @@ void Ir::applyInplacePattern(Graph &graph) {
           };
 
           bool restoreInplaceIn =
-              op->input->tensor(in_index.first)->anyAlias(restoreInplaceTensor);
-          bool restoreInplaceOut = op->output->tensor(out_index.first)
-                                       ->anyAlias(restoreInplaceTensor);
+              op->input->tensor(in_index.first)
+                  ->anyAliasFor(restoreInplaceTensor,
+                                popMemGrowerOfTensor.getAliasModelRef());
+          bool restoreInplaceOut =
+              op->output->tensor(out_index.first)
+                  ->anyAliasFor(restoreInplaceTensor,
+                                popMemGrowerOfTensor.getAliasModelRef());
 
           bool conflictIn =
-              op->input->tensor(in_index.first)->anyAlias(isConflictTensor);
+              op->input->tensor(in_index.first)
+                  ->anyAliasFor(isConflictTensor,
+                                popMemGrowerOfTensor.getAliasModelRef());
           bool conflictOut =
-              op->output->tensor(out_index.first)->anyAlias(isConflictTensor);
+              op->output->tensor(out_index.first)
+                  ->anyAliasFor(isConflictTensor,
+                                popMemGrowerOfTensor.getAliasModelRef());
 
           // Check that no conflict tensors, through aliasing, can be consumed
           // by a RestoreInplaceOp
@@ -3726,10 +3787,13 @@ void Ir::applyInplacePattern(Graph &graph) {
 
           // Unmodifiable
           // 1. Is the input unmodifiable?
-          bool unmodifiable = op->inputUnmodifiable(in_index.first);
+          bool unmodifiable = op->inputUnmodifiableFor(
+              in_index.first, popMemGrowerOfTensor.popMem_.get());
           // 2. Does it indirectly modify this tensor and alias it?
           bool indirectModify =
-              (op->hasAliasedModifiers(out_index.first) && opAliases);
+              (opAliases &&
+               op->hasAliasedModifiersFor(out_index.first,
+                                          popMemGrowerOfTensor.popMem_.get()));
           // 3. Does it directly modify a weight?
           bool directModify = inplaceOp->modifiesIndex(in_index.first);
           // If ((1 and 2) or 3) : do not inplace.
@@ -3748,7 +3812,8 @@ void Ir::applyInplacePattern(Graph &graph) {
 
           if ((indirectModify || directModify) &&
               op->input->tensor(in_index.first)
-                  ->anyAlias(isImplicitRecomputeTensor)) {
+                  ->anyAliasFor(isImplicitRecomputeTensor,
+                                popMemGrowerOfTensor.getAliasModelRef())) {
             logging::pattern::trace("[Inplacing] Not inplacing {} with {} as "
                                     "it would be modified by a recomputation "
                                     "{} -> {} ",
@@ -3881,7 +3946,8 @@ void Ir::applyInplacePattern(Graph &graph) {
         };
 
         for (auto tensor : currentInsOuts) {
-          tensor->anyAlias(populateConsumersInIndices);
+          tensor->anyAliasFor(populateConsumersInIndices,
+                              popMemGrowerOfTensor.getAliasModelRef());
         }
 
         for (const auto &consumerInIndices : consumersInIndices) {
@@ -3944,6 +4010,15 @@ void Ir::applyInplacePattern(Graph &graph) {
         // The Op in graph has changed, mirror the change in the poprithms
         // Graph
         popMem.update(id, opOutput->getProducer()->id);
+
+        const auto status = popMemGrowerOfTensor.tryOpening(tProposal);
+        if (status != PopMemGrower::OpeningStatus::Valid) {
+          popMemGrowerOfTensor.reset(graph);
+        } else {
+          // The Op in graph has changed, mirror the change in the poprithms
+          // Graph
+          popMemGrowerOfTensor.popMem_->update(id, opOutput->getProducer()->id);
+        }
       }
     }
   }

willow/src/op.cpp

@@ -1352,12 +1352,26 @@ bool Op::producesGraphOutput() const {
 }
 
 bool Op::inputUnmodifiable(InIndex in) const {
+  return inputUnmodifiableFor(in, nullptr);
+}
+
+bool Op::inputUnmodifiableFor(InIndex in, const AliasModel *popMemPtr) const {
   auto t = input->tensor(in);
   // If the input tensor itself, or any of it's aliases, are unmodifiable
-  return t->anyAlias([](Tensor *tensor) { return tensor->isUnmodifiable(); });
+  if (popMemPtr == nullptr) {
+    return t->anyAlias([](Tensor *tensor) { return tensor->isUnmodifiable(); });
+  } else {
+    return t->anyAliasFor(
+        [](Tensor *tensor) { return tensor->isUnmodifiable(); }, *popMemPtr);
+  }
 }
 
 bool Op::hasAliasedModifiers(OutIndex out) const {
+  return hasAliasedModifiersFor(out, nullptr);
+}
+
+bool Op::hasAliasedModifiersFor(OutIndex out,
+                                const AliasModel *popMemPtr) const {
   auto t = output->tensor(out);
 
   auto checkConsumers = [](Tensor *t_in) {
@@ -1372,7 +1386,11 @@ bool Op::hasAliasedModifiers(OutIndex out) const {
     return false;
   };
 
-  return t->anyAlias(checkConsumers);
+  if (popMemPtr == nullptr) {
+    return t->anyAlias(checkConsumers);
+  } else {
+    return t->anyAliasFor(checkConsumers, *popMemPtr);
+  }
 }
 
 bool Op::isParentOf(const Op *op) const {

willow/src/tensor.cpp

@@ -307,8 +307,14 @@ view::Regions Tensor::modifiedRegionsByOps(std::vector<Op *> ops,
 }
 
 std::set<Op *, POpCmp> Tensor::getInplaceModifiers() const {
+  return getInplaceModifiersFor(nullptr);
+}
+
+std::set<Op *, POpCmp>
+Tensor::getInplaceModifiersFor(const AliasModel *popMem) const {
   std::set<Op *, POpCmp> ops;
-  anyAlias([&ops](Tensor *t) {
+
+  auto cb = [&ops](Tensor *t) {
     auto consumers = t->consumers.getOps();
     for (auto c : consumers) {
       if (c->modifies()) {
@@ -317,7 +323,13 @@ std::set<Op *, POpCmp> Tensor::getInplaceModifiers() const {
     }
     // Continue until all aliases have been visited
     return false;
-  });
+  };
+
+  if (popMem == nullptr) {
+    anyAlias(cb);
+  } else {
+    anyAliasFor(cb, *popMem);
+  }
   return ops;
 }
 
@@ -1122,9 +1134,6 @@ bool Tensor::isRootAnchor() const { return graph.getIr().isRootAnchor(id); }
 bool Tensor::anyAlias(std::function<bool(Tensor *)> predicate) const {
 
   constexpr const char *const ctxt{"Tensor::anyAlias"};
-  logging::ir::trace("{} for Tensor {},", ctxt, str());
-
-  auto scopedStopwatch = getIr().timePartitionLogger().scopedStopwatch(ctxt);
 
   // First check if this tensor itself satisfies the predicate. If so, we need
   // not bother constructing a poprithms graph to check for alias tensors.
@@ -1138,6 +1147,34 @@ bool Tensor::anyAlias(std::function<bool(Tensor *)> predicate) const {
   AliasModelGrower aliasModelGrower{popMem};
   aliasModelGrower.growPartialGraph(graph, id, DataDependenciesOnly::Yes);
 
+  return anyAliasImpl(predicate, popMem, ctxt);
+}
+
+bool Tensor::anyAliasFor(std::function<bool(Tensor *)> predicate,
+                         const AliasModel &popMem) const {
+
+  constexpr const char *const ctxt{
+      "Tensor::anyAlias with prebuild AliasModel and AliasModelGrower."};
+
+  return anyAliasImpl(predicate, popMem, ctxt);
+}
+
+bool Tensor::anyAliasImpl(std::function<bool(Tensor *)> predicate,
+                          const AliasModel &popMem,
+                          const char *scopeDesc) const {
+
+  logging::ir::trace("{} for Tensor {},", scopeDesc, str());
+
+  auto scopedStopwatch =
+      getIr().timePartitionLogger().scopedStopwatch(scopeDesc);
+
+  // First check if this tensor itself satisfies the predicate. If so, we need
+  // not bother constructing a poprithms graph to check for alias tensors.
+  Tensor *t = graph.getTensors().get(id);
+  if (predicate(t)) {
+    return true;
+  }
+
   // Get the identifier used to represent this tensor in poprithms.
   auto poprithmsTensorId = popMem.getPoprithmsTensorId(id);