Allow more inserts before reIndexTopology

torch/csrc/jit/ir/ir.cpp

@@ -1283,13 +1283,20 @@ void Node::assignTopoPosition() {
 
     // insert between two existing nodes
   } else {
-    const auto posBetween = prevPos + (nextPos - prevPos) / 2;
-    if (posBetween == prevPos) {
+    int64_t remaining = nextPos - prevPos;
+    AT_ASSERT(remaining > 0);
+    if (remaining == 1) {
       // There was no room
       owningBlock()->reIndexTopology();
       return;
     }
-    topo_position_ = posBetween;
+    int64_t predicted_future_insertions = 0;
+    if (next() == graph_->insertPoint()) {
+      predicted_future_insertions = graph_->predicted_insert_count_++;
+    }
+    topo_position_ = prevPos +
+        std::max(int64_t(1), remaining / (2 + predicted_future_insertions));
+    AT_ASSERT(prevPos < topo_position_ && topo_position_ < nextPos);
   }
 }
 

torch/csrc/jit/ir/ir.h

@@ -1206,6 +1206,7 @@ struct Graph : std::enable_shared_from_this<Graph> {
   // when insertNode() is called, the node is inserted before this node
   // by default this is set to append to the top level block
   Node* insert_before_;
+  int64_t predicted_insert_count_ = 0;
 
   c10::optional<size_t> op_version_;
 
@@ -1403,14 +1404,15 @@ struct Graph : std::enable_shared_from_this<Graph> {
   // set where nodes are inserted to append to the end of this block
   void setInsertPoint(Block* b) {
     AT_ASSERT(b->owningGraph() == this);
-    insert_before_ = b->return_node();
+    setInsertPoint(b->return_node());
   }
   // set where nodes are inserted to insert _before_ this node
   // for implementation simplicity we only support inserting before a node for
   // now
   void setInsertPoint(Node* n) {
     AT_ASSERT(n->owningGraph() == this && n->inBlockList());
     insert_before_ = n;
+    predicted_insert_count_ = 0;
   }
   Node* insertPoint() {
     return insert_before_;