JIT: track memory loop dependence of trees during value numbering

src/coreclr/jit/compiler.cpp

@@ -2015,10 +2015,11 @@ void Compiler::compInit(ArenaAllocator*       pAlloc,
     }
 #endif // DEBUG
 
-    vnStore               = nullptr;
-    m_opAsgnVarDefSsaNums = nullptr;
-    fgSsaPassesCompleted  = 0;
-    fgVNPassesCompleted   = 0;
+    vnStore                    = nullptr;
+    m_opAsgnVarDefSsaNums      = nullptr;
+    m_nodeToLoopMemoryBlockMap = nullptr;
+    fgSsaPassesCompleted       = 0;
+    fgVNPassesCompleted        = 0;
 
     // check that HelperCallProperties are initialized
 

src/coreclr/jit/compiler.h

@@ -5053,6 +5053,33 @@ class Compiler
         return m_opAsgnVarDefSsaNums;
     }
 
+    // This map tracks nodes whose value numbers explicitly or implicitly depend on memory states.
+    // The map provides the entry block of the most closely enclosing loop that
+    // defines the memory region accessed when defining the nodes's VN.
+    //
+    // This information should be consulted when considering hoisting node out of a loop, as the VN
+    // for the node will only be valid within the indicated loop.
+    //
+    // It is not fine-grained enough to track memory dependence within loops, so cannot be used
+    // for more general code motion.
+    //
+    // If a node does not have an entry in the map we currently assume the VN is not memory dependent
+    // and so memory does not constrain hoisting.
+    //
+    typedef JitHashTable<GenTree*, JitPtrKeyFuncs<GenTree>, BasicBlock*> NodeToLoopMemoryBlockMap;
+    NodeToLoopMemoryBlockMap* m_nodeToLoopMemoryBlockMap;
+    NodeToLoopMemoryBlockMap* GetNodeToLoopMemoryBlockMap()
+    {
+        if (m_nodeToLoopMemoryBlockMap == nullptr)
+        {
+            m_nodeToLoopMemoryBlockMap = new (getAllocator()) NodeToLoopMemoryBlockMap(getAllocator());
+        }
+        return m_nodeToLoopMemoryBlockMap;
+    }
+
+    void optRecordLoopMemoryDependence(GenTree* tree, BasicBlock* block, ValueNum memoryVN);
+    void optCopyLoopMemoryDependence(GenTree* fromTree, GenTree* toTree);
+
     // Requires value numbering phase to have completed. Returns the value number ("gtVN") of the
     // "tree," EXCEPT in the case of GTF_VAR_USEASG, because the tree node's gtVN member is the
     // "use" VN. Performs a lookup into the map of (use asg tree -> def VN.) to return the "def's"
@@ -9965,6 +9992,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 
     BasicBlock* compCurBB;   // the current basic block in process
     Statement*  compCurStmt; // the current statement in process
+    GenTree*    compCurTree; // the current tree in process
 
     //  The following is used to create the 'method JIT info' block.
     size_t compInfoBlkSize;

src/coreclr/jit/optimizer.cpp

@@ -5486,6 +5486,9 @@ void Compiler::optPerformHoistExpr(GenTree* origExpr, unsigned lnum)
     // so clear the RegNum if it was set in the original expression
     hoistExpr->ClearRegNum();
 
+    // Copy any loop memory dependence.
+    optCopyLoopMemoryDependence(origExpr, hoistExpr);
+
     // At this point we should have a cloned expression, marked with the GTF_MAKE_CSE flag
     assert(hoistExpr != origExpr);
     assert(hoistExpr->gtFlags & GTF_MAKE_CSE);
@@ -6036,6 +6039,106 @@ bool Compiler::optIsProfitableToHoistableTree(GenTree* tree, unsigned lnum)
     return true;
 }
 
+//------------------------------------------------------------------------
+// optRecordLoopMemoryDependence: record that tree's value number
+//   is dependent on a particular memory VN
+//
+// Arguments:
+//   tree -- tree in question
+//   block -- block containing tree
+//   memoryVN -- VN for a "map" from a select operation encounterd
+//     while computing the tree's VN
+//
+// Notes:
+//   Only tracks trees in loops, and memory updates in the same loop nest.
+//   So this is a coarse-grained dependence that is only usable for
+//   hoisting tree out of its enclosing loops.
+//
+void Compiler::optRecordLoopMemoryDependence(GenTree* tree, BasicBlock* block, ValueNum memoryVN)
+{
+    // If tree is not in a loop, we don't need to track its loop dependence.
+    //
+    unsigned const loopNum = block->bbNatLoopNum;
+
+    if (loopNum == BasicBlock::NOT_IN_LOOP)
+    {
+        return;
+    }
+
+    // Find the loop associated with this memory VN.
+    //
+    unsigned const updateLoopNum = vnStore->LoopOfVN(memoryVN);
+
+    if (updateLoopNum == BasicBlock::NOT_IN_LOOP)
+    {
+        // memoryVN defined outside of any loop, we can ignore.
+        //
+        JITDUMP("      ==> Not updating loop memory dependence of [%06u], memory " FMT_VN " not defined in a loop\n",
+                dspTreeID(tree), memoryVN);
+        return;
+    }
+
+    // If the update block is not the the header of a loop containing
+    // block, we can also ignore the update.
+    //
+    if (!optLoopContains(updateLoopNum, loopNum))
+    {
+        JITDUMP("      ==> Not updating loop memory dependence of [%06u]/" FMT_LP ", memory " FMT_VN "/" FMT_LP
+                " is not defined in an enclosing loop\n",
+                dspTreeID(tree), loopNum, memoryVN, updateLoopNum);
+        return;
+    }
+
+    // If we already have a recorded a loop entry block for this
+    // tree, see if the new update is for a more closely nested
+    // loop.
+    //
+    NodeToLoopMemoryBlockMap* const map      = GetNodeToLoopMemoryBlockMap();
+    BasicBlock*                     mapBlock = nullptr;
+
+    if (map->Lookup(tree, &mapBlock))
+    {
+        unsigned const mapLoopNum = mapBlock->bbNatLoopNum;
+
+        // If the update loop contains the existing map loop,
+        // the existing map loop is more constraining. So no
+        // update needed.
+        //
+        if (optLoopContains(updateLoopNum, mapLoopNum))
+        {
+            JITDUMP("      ==> Not updating loop memory dependence of [%06u]; alrady constrained to " FMT_LP
+                    " nested in " FMT_LP "\n",
+                    dspTreeID(tree), mapLoopNum, updateLoopNum);
+            return;
+        }
+    }
+
+    // MemoryVN now describes the most constraining loop memory dependence
+    // we know of. Update the map.
+    //
+    JITDUMP("      ==> Updating loop memory dependence of [%06u] to " FMT_LP "\n", dspTreeID(tree), updateLoopNum);
+    map->Set(tree, optLoopTable[updateLoopNum].lpEntry);
+}
+
+//------------------------------------------------------------------------
+// optCopyLoopMemoryDependence: record that tree's loop memory dependence
+//   is the same as some other tree.
+//
+// Arguments:
+//   fromTree -- tree to copy dependence from
+//   toTree -- tree in question
+//
+void Compiler::optCopyLoopMemoryDependence(GenTree* fromTree, GenTree* toTree)
+{
+    NodeToLoopMemoryBlockMap* const map      = GetNodeToLoopMemoryBlockMap();
+    BasicBlock*                     mapBlock = nullptr;
+
+    if (map->Lookup(fromTree, &mapBlock))
+    {
+        map->Set(toTree, mapBlock);
+    }
+}
+
 //------------------------------------------------------------------------
 // optHoistLoopBlocks: Hoist invariant expression out of the loop.
 //
@@ -6093,19 +6196,63 @@ void Compiler::optHoistLoopBlocks(unsigned loopNum, ArrayStack<BasicBlock*>* blo
         bool IsTreeVNInvariant(GenTree* tree)
         {
             ValueNum vn = tree->gtVNPair.GetLiberal();
-            if (m_compiler->vnStore->IsVNConstant(vn))
+            bool     vnIsInvariant =
+                m_compiler->optVNIsLoopInvariant(vn, m_loopNum, &m_hoistContext->m_curLoopVnInvariantCache);
+
+            // Even though VN is invariant in the loop (say a constant) its value may depend on position
+            // of tree, so for loop hoisting we must also check that any memory read by tree
+            // is also invariant in the loop.
+            //
+            if (vnIsInvariant)
+            {
+                vnIsInvariant = IsTreeLoopMemoryInvariant(tree);
+            }
+            return vnIsInvariant;
+        }
+
+        //------------------------------------------------------------------------
+        // IsTreeLoopMemoryInvariant: determine if the value number of tree
+        //   is dependent on the tree being executed within the current loop
+        //
+        // Arguments:
+        //   tree -- tree in question
+        //
+        // Returns:
+        //   true if tree could be evaluated just before loop and get the
+        //   same value.
+        //
+        // Note:
+        //   Calls are optimistically assumed to be invariant.
+        //   Caller must do their own analysis for these tree types.
+        //
+        bool IsTreeLoopMemoryInvariant(GenTree* tree)
+        {
+            if (tree->IsCall())
             {
-                // It is unsafe to allow a GT_CLS_VAR that has been assigned a constant.
-                // The logic in optVNIsLoopInvariant would consider it to be loop-invariant, even
-                // if the assignment of the constant to the GT_CLS_VAR was inside the loop.
+                // Calls are handled specially by hoisting, and loop memory dependence
+                // must be checked by other means.
                 //
-                if (tree->OperIs(GT_CLS_VAR))
+                return true;
+            }
+
+            NodeToLoopMemoryBlockMap* const map            = m_compiler->GetNodeToLoopMemoryBlockMap();
+            BasicBlock*                     loopEntryBlock = nullptr;
+            if (map->Lookup(tree, &loopEntryBlock))
+            {
+                for (MemoryKind memoryKind : allMemoryKinds())
                 {
-                    return false;
+                    ValueNum loopMemoryVN =
+                        m_compiler->GetMemoryPerSsaData(loopEntryBlock->bbMemorySsaNumIn[memoryKind])
+                            ->m_vnPair.GetLiberal();
+                    if (!m_compiler->optVNIsLoopInvariant(loopMemoryVN, m_loopNum,
+                                                          &m_hoistContext->m_curLoopVnInvariantCache))
+                    {
+                        return false;
+                    }
                 }
             }
 
-            return m_compiler->optVNIsLoopInvariant(vn, m_loopNum, &m_hoistContext->m_curLoopVnInvariantCache);
+            return true;
         }
 
     public:
@@ -6576,6 +6723,20 @@ bool Compiler::optVNIsLoopInvariant(ValueNum vn, unsigned lnum, VNToBoolMap* loo
         {
             for (unsigned i = 0; i < funcApp.m_arity; i++)
             {
+                // 4th arg of mapStore identifies the loop where the store happens.
+                //
+                if (funcApp.m_func == VNF_MapStore)
+                {
+                    assert(funcApp.m_arity == 4);
+
+                    if (i == 3)
+                    {
+                        const unsigned vnLoopNum = funcApp.m_args[3];
+                        res                      = !optLoopContains(lnum, vnLoopNum);
+                        break;
+                    }
+                }
+
                 // TODO-CQ: We need to either make sure that *all* VN functions
                 // always take VN args, or else have a list of arg positions to exempt, as implicitly
                 // constant.

src/coreclr/jit/valuenum.cpp

@@ -2136,7 +2136,7 @@ ValueNum ValueNumStore::VNForFunc(
     assert(arg0VN == VNNormalValue(arg0VN));
     assert(arg1VN == VNNormalValue(arg1VN));
     assert(arg2VN == VNNormalValue(arg2VN));
-    assert(arg3VN == VNNormalValue(arg3VN));
+    assert((func == VNF_MapStore) || (arg3VN == VNNormalValue(arg3VN)));
     assert(VNFuncArity(func) == 4);
 
     ValueNum resultVN;
@@ -2176,12 +2176,15 @@ ValueNum ValueNumStore::VNForFunc(
 
 ValueNum ValueNumStore::VNForMapStore(var_types typ, ValueNum arg0VN, ValueNum arg1VN, ValueNum arg2VN)
 {
-    ValueNum result = VNForFunc(typ, VNF_MapStore, arg0VN, arg1VN, arg2VN);
+    BasicBlock* const            bb      = m_pComp->compCurBB;
+    BasicBlock::loopNumber const loopNum = bb->bbNatLoopNum;
+    ValueNum const               result  = VNForFunc(typ, VNF_MapStore, arg0VN, arg1VN, arg2VN, loopNum);
+
 #ifdef DEBUG
     if (m_pComp->verbose)
     {
-        printf("    VNForMapStore(" FMT_VN ", " FMT_VN ", " FMT_VN "):%s returns ", arg0VN, arg1VN, arg2VN,
-               varTypeName(typ));
+        printf("    VNForMapStore(" FMT_VN ", " FMT_VN ", " FMT_VN "):%s in " FMT_BB " returns ", arg0VN, arg1VN,
+               arg2VN, varTypeName(typ), bb->bbNum);
         m_pComp->vnPrint(result, 1);
         printf("\n");
     }
@@ -2315,6 +2318,8 @@ ValueNum ValueNumStore::VNForMapSelectWork(
                             ") ==> " FMT_VN ".\n",
                             funcApp.m_args[0], arg0VN, funcApp.m_args[1], funcApp.m_args[2], arg1VN, funcApp.m_args[2]);
 #endif
+
+                    m_pComp->optRecordLoopMemoryDependence(m_pComp->compCurTree, m_pComp->compCurBB, funcApp.m_args[0]);
                     return funcApp.m_args[2];
                 }
                 // i # j ==> select(store(m, i, v), j) == select(m, j)
@@ -4346,7 +4351,7 @@ var_types ValueNumStore::TypeOfVN(ValueNum vn)
 }
 
 //------------------------------------------------------------------------
-// LoopOfVN: If the given value number is VNF_MemOpaque, return
+// LoopOfVN: If the given value number is VNF_MemOpaque or VNF_MapStore, return
 //    the loop number where the memory update occurs, otherwise returns MAX_LOOP_NUM.
 //
 // Arguments:
@@ -4359,9 +4364,16 @@ var_types ValueNumStore::TypeOfVN(ValueNum vn)
 BasicBlock::loopNumber ValueNumStore::LoopOfVN(ValueNum vn)
 {
     VNFuncApp funcApp;
-    if (GetVNFunc(vn, &funcApp) && (funcApp.m_func == VNF_MemOpaque))
+    if (GetVNFunc(vn, &funcApp))
     {
-        return (BasicBlock::loopNumber)funcApp.m_args[0];
+        if (funcApp.m_func == VNF_MemOpaque)
+        {
+            return (BasicBlock::loopNumber)funcApp.m_args[0];
+        }
+        else if (funcApp.m_func == VNF_MapStore)
+        {
+            return (BasicBlock::loopNumber)funcApp.m_args[3];
+        }
     }
 
     return BasicBlock::MAX_LOOP_NUM;
@@ -5707,13 +5719,18 @@ void ValueNumStore::vnDumpMapStore(Compiler* comp, VNFuncApp* mapStore)
     ValueNum mapVN    = mapStore->m_args[0]; // First arg is the map id
     ValueNum indexVN  = mapStore->m_args[1]; // Second arg is the index
     ValueNum newValVN = mapStore->m_args[2]; // Third arg is the new value
+    unsigned loopNum  = mapStore->m_args[3]; // Fourth arg is the loop num
 
     comp->vnPrint(mapVN, 0);
     printf("[");
     comp->vnPrint(indexVN, 0);
     printf(" := ");
     comp->vnPrint(newValVN, 0);
     printf("]");
+    if (loopNum != BasicBlock::NOT_IN_LOOP)
+    {
+        printf("@" FMT_LP, loopNum);
+    }
 }
 
 void ValueNumStore::vnDumpMemOpaque(Compiler* comp, VNFuncApp* memOpaque)
@@ -6593,7 +6610,10 @@ void Compiler::fgValueNumberBlock(BasicBlock* blk)
 
         for (GenTree* const tree : stmt->TreeList())
         {
+            // Set up ambient var referring to current tree.
+            compCurTree = tree;
             fgValueNumberTree(tree);
+            compCurTree = nullptr;
         }
 
 #ifdef DEBUG

src/coreclr/jit/valuenumfuncs.h

@@ -7,7 +7,7 @@
 
 // clang-format off
 ValueNumFuncDef(MemOpaque, 1, false, false, false)  // Args: 0: loop num
-ValueNumFuncDef(MapStore, 3, false, false, false)   // Args: 0: map, 1: index (e. g. field handle), 2: value being stored.
+ValueNumFuncDef(MapStore, 4, false, false, false)   // Args: 0: map, 1: index (e. g. field handle), 2: value being stored, 3: loop num.
 ValueNumFuncDef(MapSelect, 2, false, false, false)  // Args: 0: map, 1: key.
 
 ValueNumFuncDef(FieldSeq, 2, false, false, false)   // Sequence (VN of null == empty) of (VN's of) field handles.