JIT: Make strength reduced IV updates amenable to post-indexed addressing

On arm64 have strength reduction try to insert IV updates after the last
use if that last use is a legal insertion point. This often allows the
backend to use post-indexed addressing to combine the use with the IV
update.

Author: jakobbotsch

src/coreclr/jit/inductionvariableopts.cpp

@@ -1341,7 +1341,11 @@ class StrengthReductionContext
     bool        CheckAdvancedCursors(ArrayStack<CursorInfo>* cursors, ScevAddRec** nextIV);
     bool        StaysWithinManagedObject(ArrayStack<CursorInfo>* cursors, ScevAddRec* addRec);
     bool        TryReplaceUsesWithNewPrimaryIV(ArrayStack<CursorInfo>* cursors, ScevAddRec* iv);
-    BasicBlock* FindUpdateInsertionPoint(ArrayStack<CursorInfo>* cursors);
+    BasicBlock* FindUpdateInsertionPoint(ArrayStack<CursorInfo>* cursors, Statement** afterStmt);
+    BasicBlock* FindPostUseUpdateInsertionPoint(ArrayStack<CursorInfo>* cursors,
+                                                BasicBlock*             backEdgeDominator,
+                                                Statement**             afterStmt);
+    bool        InsertionPointPostDominatesUses(BasicBlock* insertionPoint, ArrayStack<CursorInfo>* cursors);
 
     bool StressProfitability()
     {
@@ -2000,7 +2004,8 @@ bool StrengthReductionContext::TryReplaceUsesWithNewPrimaryIV(ArrayStack<CursorI
         return false;
     }
 
-    BasicBlock* insertionPoint = FindUpdateInsertionPoint(cursors);
+    Statement*  afterStmt;
+    BasicBlock* insertionPoint = FindUpdateInsertionPoint(cursors, &afterStmt);
     if (insertionPoint == nullptr)
     {
         JITDUMP("    Skipping: could not find a legal insertion point for the new IV update\n");
@@ -2032,7 +2037,14 @@ bool StrengthReductionContext::TryReplaceUsesWithNewPrimaryIV(ArrayStack<CursorI
         m_comp->gtNewOperNode(GT_ADD, iv->Type, m_comp->gtNewLclVarNode(newPrimaryIV, iv->Type), stepValue);
     GenTree*   stepStore = m_comp->gtNewTempStore(newPrimaryIV, nextValue);
     Statement* stepStmt  = m_comp->fgNewStmtFromTree(stepStore);
-    m_comp->fgInsertStmtNearEnd(insertionPoint, stepStmt);
+    if (afterStmt != nullptr)
+    {
+        m_comp->fgInsertStmtAfter(insertionPoint, afterStmt, stepStmt);
+    }
+    else
+    {
+        m_comp->fgInsertStmtNearEnd(insertionPoint, stepStmt);
+    }
 
     JITDUMP("    Inserting step statement in " FMT_BB "\n", insertionPoint->bbNum);
     DISPSTMT(stepStmt);
@@ -2084,22 +2096,27 @@ bool StrengthReductionContext::TryReplaceUsesWithNewPrimaryIV(ArrayStack<CursorI
 // of a new primary IV introduced by strength reduction.
 //
 // Parameters:
-//   cursors - The list of cursors pointing to uses that are being replaced by
-//             the new IV
+//   cursors   - The list of cursors pointing to uses that are being replaced by
+//               the new IV
+//   afterStmt - [out] Statement to insert the update after. Set to nullptr if
+//               update should be inserted near the end of the block.
 //
 // Returns:
 //   Basic block; the insertion point is the end (before a potential
 //   terminator) of this basic block. May return null if no insertion point
 //   could be found.
 //
-BasicBlock* StrengthReductionContext::FindUpdateInsertionPoint(ArrayStack<CursorInfo>* cursors)
+BasicBlock* StrengthReductionContext::FindUpdateInsertionPoint(ArrayStack<CursorInfo>* cursors, Statement** afterStmt)
 {
+    *afterStmt = nullptr;
+
     // Find insertion point. It needs to post-dominate all uses we are going to
     // replace and it needs to dominate all backedges.
     // TODO-CQ: Canonicalizing backedges would make this simpler and work in
     // more cases.
 
     BasicBlock* insertionPoint = nullptr;
+
     for (FlowEdge* backEdge : m_loop->BackEdges())
     {
         if (insertionPoint == nullptr)
@@ -2112,6 +2129,18 @@ BasicBlock* StrengthReductionContext::FindUpdateInsertionPoint(ArrayStack<Cursor
         }
     }
 
+#ifdef TARGET_ARM64
+    // For arm64 we try to place the IV update after a use if possible. This
+    // sets the backend up for post-indexed addressing mode.
+    BasicBlock* postUseInsertionPoint = FindPostUseUpdateInsertionPoint(cursors, insertionPoint, afterStmt);
+    if (postUseInsertionPoint != nullptr)
+    {
+        JITDUMP("    Found a legal insertion point after a last use of the IV in " FMT_BB " after " FMT_STMT "\n",
+                postUseInsertionPoint->bbNum, (*afterStmt)->GetID());
+        return postUseInsertionPoint;
+    }
+#endif
+
     while ((insertionPoint != nullptr) && m_loop->ContainsBlock(insertionPoint) &&
            m_loop->MayExecuteBlockMultipleTimesPerIteration(insertionPoint))
     {
@@ -2123,6 +2152,124 @@ BasicBlock* StrengthReductionContext::FindUpdateInsertionPoint(ArrayStack<Cursor
         return nullptr;
     }
 
+    if (!InsertionPointPostDominatesUses(insertionPoint, cursors))
+    {
+        return nullptr;
+    }
+
+    JITDUMP("    Found a legal insertion point in " FMT_BB "\n", insertionPoint->bbNum);
+    return insertionPoint;
+}
+
+//------------------------------------------------------------------------
+// FindPostUseUpdateInsertionPoint: Try finding an insertion point for the IV
+// update that is right after one of the uses of it.
+//
+// Parameters:
+//   cursors           - The list of cursors pointing to uses that are being replaced by
+//                       the new IV
+//   backEdgeDominator - A basic block that dominates all backedges
+//   afterStmt         - [out] Statement to insert the update after, if the
+//                       return value is non-null.
+//
+// Returns:
+//   nullptr if no such insertion point could be found. Otherwise returns the
+//   basic block and statement after which the update can be inserted.
+//
+BasicBlock* StrengthReductionContext::FindPostUseUpdateInsertionPoint(ArrayStack<CursorInfo>* cursors,
+                                                                      BasicBlock*             backEdgeDominator,
+                                                                      Statement**             afterStmt)
+{
+    BitVecTraits poTraits = m_loop->GetDfsTree()->PostOrderTraits();
+
+#ifdef DEBUG
+    // We will be relying on the fact that the cursors are ordered in a useful
+    // way here: loop locals are visited in post order within each basic block,
+    // meaning that "cursors" has the last uses first for each basic block.
+    // Assert that here.
+
+    BitVec seenBlocks(BitVecOps::MakeEmpty(&poTraits));
+    for (int i = 1; i < cursors->Height(); i++)
+    {
+        CursorInfo& prevCursor = cursors->BottomRef(i - 1);
+        CursorInfo& cursor     = cursors->BottomRef(i);
+
+        if (cursor.Block != prevCursor.Block)
+        {
+            assert(BitVecOps::TryAddElemD(&poTraits, seenBlocks, prevCursor.Block->bbPostorderNum));
+            continue;
+        }
+
+        Statement* curStmt = cursor.Stmt;
+        while ((curStmt != nullptr) && (curStmt != prevCursor.Stmt))
+        {
+            curStmt = curStmt->GetNextStmt();
+        }
+
+        assert(curStmt == prevCursor.Stmt);
+    }
+#endif
+
+    BitVec blocksWithUses(BitVecOps::MakeEmpty(&poTraits));
+    for (int i = 0; i < cursors->Height(); i++)
+    {
+        CursorInfo& cursor = cursors->BottomRef(i);
+        BitVecOps::AddElemD(&poTraits, blocksWithUses, cursor.Block->bbPostorderNum);
+    }
+
+    while ((backEdgeDominator != nullptr) && m_loop->ContainsBlock(backEdgeDominator))
+    {
+        if (!BitVecOps::IsMember(&poTraits, blocksWithUses, backEdgeDominator->bbPostorderNum))
+        {
+            backEdgeDominator = backEdgeDominator->bbIDom;
+            continue;
+        }
+
+        if (m_loop->MayExecuteBlockMultipleTimesPerIteration(backEdgeDominator))
+        {
+            return nullptr;
+        }
+
+        for (int i = 0; i < cursors->Height(); i++)
+        {
+            CursorInfo& cursor = cursors->BottomRef(i);
+            if (cursor.Block != backEdgeDominator)
+            {
+                continue;
+            }
+
+            if (!InsertionPointPostDominatesUses(cursor.Block, cursors))
+            {
+                return nullptr;
+            }
+
+            *afterStmt = cursor.Stmt;
+            return cursor.Block;
+        }
+    }
+
+    return nullptr;
+}
+
+//------------------------------------------------------------------------
+// InsertionPointPostDominatesUses: Check if a basic block post-dominates all
+// locations specified by the cursors.
+//
+// Parameters:
+//   insertionPoint - The insertion point
+//   cursors        - Cursors specifying locations
+//
+// Returns:
+//   True if so.
+//
+// Remarks:
+//   For cursors inside "insertionPoint", the function expects that the
+//   insertion point is _after_ the use, except if the use is in a terminator
+//   statement.
+//
+bool StrengthReductionContext::InsertionPointPostDominatesUses(BasicBlock*             insertionPoint,
+                                                               ArrayStack<CursorInfo>* cursors)
+{
     for (int i = 0; i < cursors->Height(); i++)
     {
         CursorInfo& cursor = cursors->BottomRef(i);
@@ -2131,19 +2278,19 @@ BasicBlock* StrengthReductionContext::FindUpdateInsertionPoint(ArrayStack<Cursor
         {
             if (insertionPoint->HasTerminator() && (cursor.Stmt == insertionPoint->lastStmt()))
             {
-                return nullptr;
+                return false;
             }
         }
         else
         {
             if (!m_loop->IsPostDominatedOnLoopIteration(cursor.Block, insertionPoint))
             {
-                return nullptr;
+                return false;
             }
         }
     }
 
-    return insertionPoint;
+    return true;
 }
 
 //------------------------------------------------------------------------