JIT: Introduce VNWalkPhis and use in VNNeverNegative as an example

src/coreclr/jit/valuenum.cpp

@@ -1635,7 +1635,20 @@ bool ValueNumStore::IsKnownNonNull(ValueNum vn)
     }
 
     VNFuncApp funcAttr;
-    return GetVNFunc(vn, &funcAttr) && (s_vnfOpAttribs[funcAttr.m_func] & VNFOA_KnownNonNull) != 0;
+    if (!GetVNFunc(vn, &funcAttr))
+    {
+        return false;
+    }
+
+    if ((s_vnfOpAttribs[funcAttr.m_func] & VNFOA_KnownNonNull) != 0)
+    {
+        return true;
+    }
+
+    // TODO: we can recognize more non-null idioms here, e.g.
+    // ADD(IsKnownNonNull(op1), smallCns), etc.
+
+    return false;
 }
 
 bool ValueNumStore::IsSharedStatic(ValueNum vn)
@@ -3110,6 +3123,22 @@ class SmallValueNumSet
         }
     }
 
+    bool Lookup(ValueNum vn)
+    {
+        if (m_numElements <= ArrLen(m_inlineElements))
+        {
+            for (unsigned i = 0; i < m_numElements; i++)
+            {
+                if (m_inlineElements[i] == vn)
+                {
+                    return true;
+                }
+            }
+            return false;
+        }
+        return m_set->Lookup(vn);
+    }
+
     void Add(Compiler* comp, ValueNum vn)
     {
         if (m_numElements <= ArrLen(m_inlineElements))
@@ -6483,6 +6512,24 @@ bool ValueNumStore::IsVNNeverNegative(ValueNum vn)
 #endif
 #endif // FEATURE_HW_INTRINSICS
 
+            case VNF_PhiDef:
+            case VNF_PhiMemoryDef:
+            {
+                // Inspect all PHI args - if all of them are non-negative,
+                // then the PHI itself is non-negative too (even if it has cycles).
+                auto phiVisitor = [this](ValueNum phiVN) -> PhiArgWalkResult {
+                    // Bail out if the type is not integral
+                    if (!varTypeIsIntegral(TypeOfVN(phiVN)))
+                    {
+                        return PhiArgWalkResult::Abort;
+                    }
+                    return IsVNNeverNegative(phiVN) ? PhiArgWalkResult::Continue : PhiArgWalkResult::Abort;
+                };
+
+                // maxDepth is set to 2 to improve JIT TP
+                return VNWalkPhis(vn, funcApp, phiVisitor, /* maxDepth */ 2) == PhiDefWalkResult::Completed;
+            }
+
             default:
                 break;
         }
@@ -15034,3 +15081,149 @@ void ValueNumStore::PeelOffsets(ValueNum* vn, target_ssize_t* offset)
         }
     }
 }
+
+//--------------------------------------------------------------------------------
+// VNWalkPhis: Walk the given PhiDef/PhiMemoryDef and call phiArgVisitor on each real VN.
+//
+// Arguments:
+//    phiDefVN      - The VN of the PhiDef
+//    phiDefFunc    - VNFuncApp for phiDefVN since caller typically has it anyway (to save TP).
+//    phiArgVisitor - The callback function to call on each Phi argument.
+//    maxDepth      - Maximum depth to handle recursive PhiDef chains. -1 means no limits.
+//    vnKind        - The kind of VN to use (Conservative or Liberal)
+//
+// Return Value:
+//    PhiDefWalkResult::InvalidPhiDef  if any PhiDef didn't allow us to make any assumptions
+//    PhiDefWalkResult::HitLimitations if we hit the maxDepth limit (for TP reasons)
+//    PhiDefWalkResult::Aborted        if any phiArgVisitor returned PhiArgWalkResult::Abort
+//    PhiDefWalkResult::Completed      if the walk was successful (all phiArgVisitors returned
+//                                     PhiArgWalkResult::Continue) and the phiArgVisitor was called at least once
+//
+template <typename TPhiArgVisitorFunc>
+ValueNumStore::PhiDefWalkResult ValueNumStore::VNWalkPhis(
+    ValueNum phiDefVN, VNFuncApp phiDefFunc, TPhiArgVisitorFunc phiArgVisitor, int maxDepth, ValueNumKind vnKind)
+{
+    int              realVNsVisited = 0;
+    SmallValueNumSet hashSet;
+    PhiDefWalkResult result =
+        VNWalkPhisInternal(phiDefVN, phiDefFunc, phiArgVisitor, maxDepth, vnKind, realVNsVisited, hashSet);
+    if ((result == PhiDefWalkResult::Completed) && (realVNsVisited == 0))
+    {
+        // We somehow didn't visit any real VN - bail out.
+        return PhiDefWalkResult::InvalidPhiDef;
+    }
+    return result;
+}
+
+//--------------------------------------------------------------------------------
+// VNWalkPhisInternal: internal worker for VNWalkPhis, see comments in VNWalkPhis
+//
+template <typename TPhiArgVisitorFunc>
+ValueNumStore::PhiDefWalkResult ValueNumStore::VNWalkPhisInternal(ValueNum           phiDefVN,
+                                                                  VNFuncApp          phiDefFunc,
+                                                                  TPhiArgVisitorFunc phiArgVisitor,
+                                                                  int                maxDepth,
+                                                                  ValueNumKind       vnKind,
+                                                                  int&               realVNsVisited,
+                                                                  SmallValueNumSet&  hashSet)
+{
+    assert((maxDepth == -1) || (maxDepth > 0));
+    assert((phiDefFunc.m_func == VNF_PhiDef) || (phiDefFunc.m_func == VNF_PhiMemoryDef));
+
+    const bool isMemory = (phiDefFunc.m_func == VNF_PhiMemoryDef);
+
+    // VNF_PhiDef(LclNum, SsaNum, VNF_Phi)
+    // VNF_PhiMemoryDef(BB, VNF_Phi)
+    ValueNum nextPhiArg = isMemory ? phiDefFunc.m_args[1] : phiDefFunc.m_args[2];
+    if (!isMemory && (phiDefFunc.m_args[0] == BAD_VAR_NUM))
+    {
+        JITDUMP("We can't process this PhiDef - bail out.\n");
+        return PhiDefWalkResult::InvalidPhiDef;
+    }
+
+    assert(!hashSet.Lookup(phiDefVN));
+    hashSet.Add(m_pComp, phiDefVN);
+
+    // Now we iterate over Phi arguments, e.g.
+    //  2 args: VNF_Phi(SSA1, SSA2)
+    //  3 args: VNF_Phi(SSA1, VNF_Phi(SSA2, SSA3))
+    //  4 args: VNF_Phi(SSA1, VNF_Phi(SSA2, VNF_Phi(SSA3, SSA4)))
+    // etc.
+    //
+    while (nextPhiArg != NoVN)
+    {
+        ValueNum  current = nextPhiArg;
+        VNFuncApp phiArgFuncApp;
+
+        // nextPhiArg is either VNF_Phi or SSA num (last)
+        if (GetVNFunc(nextPhiArg, &phiArgFuncApp))
+        {
+            assert(phiArgFuncApp.m_func == VNF_Phi);
+            current    = phiArgFuncApp.m_args[0];
+            nextPhiArg = phiArgFuncApp.m_args[1];
+        }
+        else
+        {
+            // Ok this was the last Phi arg
+            nextPhiArg = NoVN;
+        }
+
+        // We expect a constant VN here representing the SSA number
+        assert(IsVNConstant(current));
+
+        // Extract the real VN from the Phi arg
+        unsigned phiArgSsaNum = ConstantValue<unsigned>(current);
+
+        ValueNum phiArgVN;
+        if (isMemory)
+        {
+            phiArgVN = m_pComp->GetMemoryPerSsaData(phiArgSsaNum)->m_vnPair.Get(vnKind);
+        }
+        else
+        {
+            // NOTE: it's unfortunate that LclNum is not a VN.
+            unsigned phiDefLclNum = phiDefFunc.m_args[0];
+            phiArgVN              = m_pComp->lvaTable[phiDefLclNum].GetPerSsaData(phiArgSsaNum)->m_vnPair.Get(vnKind);
+        }
+
+        if (hashSet.Lookup(phiArgVN))
+        {
+            // We already visited this VN which means we have a cycle
+            continue;
+        }
+
+        // Now check the actual VN arg, if it's yet another PhiDef, walk it recursively
+        VNFuncApp argFuncApp;
+        if (GetVNFunc(phiArgVN, &argFuncApp) &&
+            ((argFuncApp.m_func == VNF_PhiDef) || (argFuncApp.m_func == VNF_PhiMemoryDef)))
+        {
+            if (maxDepth == 1)
+            {
+                JITDUMP("optWalkPhiVNs: maxDepth reached - bail out.\n");
+                return PhiDefWalkResult::HitLimitations;
+            }
+
+            // MaxDepth being -1 means we don't have depth limits
+            int              newDepth = maxDepth == -1 ? -1 : maxDepth - 1;
+            PhiDefWalkResult result =
+                VNWalkPhisInternal(phiArgVN, argFuncApp, phiArgVisitor, newDepth, vnKind, realVNsVisited, hashSet);
+            if (result != PhiDefWalkResult::Completed)
+            {
+                return result;
+            }
+        }
+        else
+        {
+            if (phiArgVisitor(phiArgVN) == PhiArgWalkResult::Abort)
+            {
+                // Caller realized that it doesn't make sense to continue any further
+                JITDUMP("optWalkPhiVNs: aborted by callback.\n");
+                return PhiDefWalkResult::Aborted;
+            }
+            realVNsVisited++;
+            hashSet.Add(m_pComp, phiArgVN);
+        }
+    }
+
+    return PhiDefWalkResult::Completed;
+}

src/coreclr/jit/valuenum.h

@@ -514,6 +514,36 @@ class ValueNumStore
 
     void PeelOffsets(ValueNum* vn, target_ssize_t* offset);
 
+    enum class PhiArgWalkResult
+    {
+        Continue,
+        Abort,
+    };
+    enum class PhiDefWalkResult
+    {
+        Completed,
+        Aborted,
+        InvalidPhiDef,
+        HitLimitations,
+    };
+    template <typename TPhiArgVisitorFunc>
+    PhiDefWalkResult VNWalkPhis(ValueNum           phiDefVN,
+                                VNFuncApp          phiDefFunc,
+                                TPhiArgVisitorFunc walkPhiVnsFn,
+                                int                maxDepth = -1,
+                                ValueNumKind       vnKind   = VNK_Conservative);
+
+private:
+    template <typename TPhiArgVisitorFunc>
+    PhiDefWalkResult VNWalkPhisInternal(ValueNum                phiDefVN,
+                                        VNFuncApp               phiDefFunc,
+                                        TPhiArgVisitorFunc      walkPhiVnsFn,
+                                        int                     maxDepth,
+                                        ValueNumKind            vnKind,
+                                        int&                    realVNsVisited,
+                                        class SmallValueNumSet& hashSet);
+public:
+
     // And the single constant for an object reference type.
     static ValueNum VNForNull()
     {