Skip to content

Commit 735d1ed

Browse files
[release/8.0] JIT: Fix exponential blowup of memory dependency arrays in VNForMapSelectWork (#93388)
* JIT: Fix exponential blowup of memory dependency arrays in VNForMapSelectWork Switch to sets when accumulating the memory dependencies found in VNForMapSelectWork. Otherwise we might duplicate each memory dependency for every path back to the store that induced it, and there can be an exponential number of those. Fix #93342 * Add test * Use SmallValueNumSet in MapSelectWorkCacheEntry * Revert "Use SmallValueNumSet in MapSelectWorkCacheEntry" This reverts commit b4292d7. * Fix function header --------- Co-authored-by: Jakob Botsch Nielsen <jakob.botsch.nielsen@gmail.com>
1 parent 7a84b55 commit 735d1ed

File tree

4 files changed

+239
-49
lines changed

4 files changed

+239
-49
lines changed

src/coreclr/jit/valuenum.cpp

+118-40
Original file line numberDiff line numberDiff line change
@@ -2860,6 +2860,81 @@ ValueNum ValueNumStore::VNForMapPhysicalSelect(
28602860
return result;
28612861
}
28622862

2863+
typedef JitHashTable<ValueNum, JitSmallPrimitiveKeyFuncs<ValueNum>, bool> ValueNumSet;
2864+
2865+
class SmallValueNumSet
2866+
{
2867+
union {
2868+
ValueNum m_inlineElements[4];
2869+
ValueNumSet* m_set;
2870+
};
2871+
unsigned m_numElements = 0;
2872+
2873+
public:
2874+
unsigned Count()
2875+
{
2876+
return m_numElements;
2877+
}
2878+
2879+
template <typename Func>
2880+
void ForEach(Func func)
2881+
{
2882+
if (m_numElements <= ArrLen(m_inlineElements))
2883+
{
2884+
for (unsigned i = 0; i < m_numElements; i++)
2885+
{
2886+
func(m_inlineElements[i]);
2887+
}
2888+
}
2889+
else
2890+
{
2891+
for (ValueNum vn : ValueNumSet::KeyIteration(m_set))
2892+
{
2893+
func(vn);
2894+
}
2895+
}
2896+
}
2897+
2898+
void Add(Compiler* comp, ValueNum vn)
2899+
{
2900+
if (m_numElements <= ArrLen(m_inlineElements))
2901+
{
2902+
for (unsigned i = 0; i < m_numElements; i++)
2903+
{
2904+
if (m_inlineElements[i] == vn)
2905+
{
2906+
return;
2907+
}
2908+
}
2909+
2910+
if (m_numElements < ArrLen(m_inlineElements))
2911+
{
2912+
m_inlineElements[m_numElements] = vn;
2913+
m_numElements++;
2914+
}
2915+
else
2916+
{
2917+
ValueNumSet* set = new (comp, CMK_ValueNumber) ValueNumSet(comp->getAllocator(CMK_ValueNumber));
2918+
for (ValueNum oldVn : m_inlineElements)
2919+
{
2920+
set->Set(oldVn, true);
2921+
}
2922+
2923+
set->Set(vn, true);
2924+
2925+
m_set = set;
2926+
m_numElements++;
2927+
assert(m_numElements == set->GetCount());
2928+
}
2929+
}
2930+
else
2931+
{
2932+
m_set->Set(vn, true, ValueNumSet::SetKind::Overwrite);
2933+
m_numElements = m_set->GetCount();
2934+
}
2935+
}
2936+
};
2937+
28632938
//------------------------------------------------------------------------------
28642939
// VNForMapSelectInner: Select value from a map and record loop memory dependencies.
28652940
//
@@ -2874,10 +2949,10 @@ ValueNum ValueNumStore::VNForMapPhysicalSelect(
28742949
//
28752950
ValueNum ValueNumStore::VNForMapSelectInner(ValueNumKind vnk, var_types type, ValueNum map, ValueNum index)
28762951
{
2877-
int budget = m_mapSelectBudget;
2878-
bool usedRecursiveVN = false;
2879-
ArrayStack<ValueNum> memoryDependencies(m_alloc);
2880-
ValueNum result = VNForMapSelectWork(vnk, type, map, index, &budget, &usedRecursiveVN, &memoryDependencies);
2952+
int budget = m_mapSelectBudget;
2953+
bool usedRecursiveVN = false;
2954+
SmallValueNumSet memoryDependencies;
2955+
ValueNum result = VNForMapSelectWork(vnk, type, map, index, &budget, &usedRecursiveVN, memoryDependencies);
28812956

28822957
// The remaining budget should always be between [0..m_mapSelectBudget]
28832958
assert((budget >= 0) && (budget <= m_mapSelectBudget));
@@ -2888,11 +2963,9 @@ ValueNum ValueNumStore::VNForMapSelectInner(ValueNumKind vnk, var_types type, Va
28882963
if ((m_pComp->compCurBB != nullptr) && (m_pComp->compCurTree != nullptr) &&
28892964
m_pComp->compCurBB->bbNatLoopNum != BasicBlock::NOT_IN_LOOP)
28902965
{
2891-
for (int i = 0; i < memoryDependencies.Height(); i++)
2892-
{
2893-
m_pComp->optRecordLoopMemoryDependence(m_pComp->compCurTree, m_pComp->compCurBB,
2894-
memoryDependencies.Bottom(i));
2895-
}
2966+
memoryDependencies.ForEach([this](ValueNum vn) {
2967+
m_pComp->optRecordLoopMemoryDependence(m_pComp->compCurTree, m_pComp->compCurBB, vn);
2968+
});
28962969
}
28972970

28982971
return result;
@@ -2903,19 +2976,16 @@ ValueNum ValueNumStore::VNForMapSelectInner(ValueNumKind vnk, var_types type, Va
29032976
// cache entry.
29042977
//
29052978
// Arguments:
2906-
// alloc - Allocator to use if memory is required.
2907-
// deps - Array stack containing the memory dependencies.
2908-
// startIndex - Start index into 'deps' of memory dependencies.
2979+
// comp - Compiler instance
2980+
// set - Set of memory dependencies to store in the entry.
29092981
//
2910-
void ValueNumStore::MapSelectWorkCacheEntry::SetMemoryDependencies(CompAllocator alloc,
2911-
ArrayStack<ValueNum>& deps,
2912-
unsigned startIndex)
2982+
void ValueNumStore::MapSelectWorkCacheEntry::SetMemoryDependencies(Compiler* comp, SmallValueNumSet& set)
29132983
{
2914-
m_numMemoryDependencies = deps.Height() - startIndex;
2984+
m_numMemoryDependencies = set.Count();
29152985
ValueNum* arr;
29162986
if (m_numMemoryDependencies > ArrLen(m_inlineMemoryDependencies))
29172987
{
2918-
m_memoryDependencies = new (alloc) ValueNum[m_numMemoryDependencies];
2988+
m_memoryDependencies = new (comp, CMK_ValueNumber) ValueNum[m_numMemoryDependencies];
29192989

29202990
arr = m_memoryDependencies;
29212991
}
@@ -2924,27 +2994,29 @@ void ValueNumStore::MapSelectWorkCacheEntry::SetMemoryDependencies(CompAllocator
29242994
arr = m_inlineMemoryDependencies;
29252995
}
29262996

2927-
for (unsigned i = 0; i < m_numMemoryDependencies; i++)
2928-
{
2929-
arr[i] = deps.Bottom(startIndex + i);
2930-
}
2997+
size_t i = 0;
2998+
set.ForEach([&i, arr](ValueNum vn) {
2999+
arr[i] = vn;
3000+
i++;
3001+
});
29313002
}
29323003

29333004
//------------------------------------------------------------------------------
29343005
// GetMemoryDependencies: Push all of the memory dependencies cached in this
2935-
// entry into the specified array stack.
3006+
// entry into the specified set.
29363007
//
29373008
// Arguments:
2938-
// result - Array stack to push memory dependencies into.
3009+
// comp - Compiler instance
3010+
// result - Set to add memory dependencies to.
29393011
//
2940-
void ValueNumStore::MapSelectWorkCacheEntry::GetMemoryDependencies(ArrayStack<ValueNum>& result)
3012+
void ValueNumStore::MapSelectWorkCacheEntry::GetMemoryDependencies(Compiler* comp, SmallValueNumSet& result)
29413013
{
29423014
ValueNum* arr = m_numMemoryDependencies <= ArrLen(m_inlineMemoryDependencies) ? m_inlineMemoryDependencies
29433015
: m_memoryDependencies;
29443016

29453017
for (unsigned i = 0; i < m_numMemoryDependencies; i++)
29463018
{
2947-
result.Push(arr[i]);
3019+
result.Add(comp, arr[i]);
29483020
}
29493021
}
29503022

@@ -2959,7 +3031,7 @@ void ValueNumStore::MapSelectWorkCacheEntry::GetMemoryDependencies(ArrayStack<Va
29593031
// pBudget - Remaining budget for the outer evaluation
29603032
// pUsedRecursiveVN - Out-parameter that is set to true iff RecursiveVN was returned from this method
29613033
// or from a method called during one of recursive invocations.
2962-
// memoryDependencies - Array stack that records VNs of memories that the result is dependent upon.
3034+
// memoryDependencies - Set that records VNs of memories that the result is dependent upon.
29633035
//
29643036
// Return Value:
29653037
// Value number for the result of the evaluation.
@@ -2969,13 +3041,13 @@ void ValueNumStore::MapSelectWorkCacheEntry::GetMemoryDependencies(ArrayStack<Va
29693041
// "select(m1, ind)", ..., "select(mk, ind)" to see if they agree. It needs to know which kind of value number
29703042
// (liberal/conservative) to read from the SSA def referenced in the phi argument.
29713043
//
2972-
ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
2973-
var_types type,
2974-
ValueNum map,
2975-
ValueNum index,
2976-
int* pBudget,
2977-
bool* pUsedRecursiveVN,
2978-
ArrayStack<ValueNum>* memoryDependencies)
3044+
ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
3045+
var_types type,
3046+
ValueNum map,
3047+
ValueNum index,
3048+
int* pBudget,
3049+
bool* pUsedRecursiveVN,
3050+
SmallValueNumSet& memoryDependencies)
29793051
{
29803052
TailCall:
29813053
// This label allows us to directly implement a tail call by setting up the arguments, and doing a goto to here.
@@ -2997,13 +3069,12 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
29973069
assert(selLim == 0 || m_numMapSels < selLim);
29983070
#endif
29993071

3000-
int firstMemoryDependency = memoryDependencies->Height();
30013072
MapSelectWorkCacheEntry entry;
30023073

30033074
VNDefFuncApp<2> fstruct(VNF_MapSelect, map, index);
30043075
if (GetMapSelectWorkCache()->Lookup(fstruct, &entry))
30053076
{
3006-
entry.GetMemoryDependencies(*memoryDependencies);
3077+
entry.GetMemoryDependencies(m_pComp, memoryDependencies);
30073078
return entry.Result;
30083079
}
30093080

@@ -3029,6 +3100,8 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
30293100
return RecursiveVN;
30303101
}
30313102

3103+
SmallValueNumSet recMemoryDependencies;
3104+
30323105
VNFuncApp funcApp;
30333106
if (GetVNFunc(map, &funcApp))
30343107
{
@@ -3047,7 +3120,7 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
30473120
funcApp.m_args[0], map, funcApp.m_args[1], funcApp.m_args[2], index, funcApp.m_args[2]);
30483121
#endif
30493122

3050-
memoryDependencies->Push(funcApp.m_args[0]);
3123+
memoryDependencies.Add(m_pComp, funcApp.m_args[0]);
30513124

30523125
return funcApp.m_args[2];
30533126
}
@@ -3191,7 +3264,7 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
31913264
bool allSame = true;
31923265
ValueNum argRest = phiFuncApp.m_args[1];
31933266
ValueNum sameSelResult = VNForMapSelectWork(vnk, type, phiArgVN, index, pBudget,
3194-
pUsedRecursiveVN, memoryDependencies);
3267+
pUsedRecursiveVN, recMemoryDependencies);
31953268

31963269
// It is possible that we just now exceeded our budget, if so we need to force an early exit
31973270
// and stop calling VNForMapSelectWork
@@ -3233,7 +3306,7 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
32333306
{
32343307
bool usedRecursiveVN = false;
32353308
ValueNum curResult = VNForMapSelectWork(vnk, type, phiArgVN, index, pBudget,
3236-
&usedRecursiveVN, memoryDependencies);
3309+
&usedRecursiveVN, recMemoryDependencies);
32373310

32383311
*pUsedRecursiveVN |= usedRecursiveVN;
32393312
if (sameSelResult == ValueNumStore::RecursiveVN)
@@ -3261,11 +3334,14 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
32613334
if (!*pUsedRecursiveVN)
32623335
{
32633336
entry.Result = sameSelResult;
3264-
entry.SetMemoryDependencies(m_alloc, *memoryDependencies, firstMemoryDependency);
3337+
entry.SetMemoryDependencies(m_pComp, recMemoryDependencies);
32653338

32663339
GetMapSelectWorkCache()->Set(fstruct, entry);
32673340
}
32683341

3342+
recMemoryDependencies.ForEach(
3343+
[this, &memoryDependencies](ValueNum vn) { memoryDependencies.Add(m_pComp, vn); });
3344+
32693345
return sameSelResult;
32703346
}
32713347
// Otherwise, fall through to creating the select(phi(m1, m2), x) function application.
@@ -3294,11 +3370,13 @@ ValueNum ValueNumStore::VNForMapSelectWork(ValueNumKind vnk,
32943370
fapp->m_args[1] = fstruct.m_args[1];
32953371

32963372
entry.Result = c->m_baseVN + offsetWithinChunk;
3297-
entry.SetMemoryDependencies(m_alloc, *memoryDependencies, firstMemoryDependency);
3373+
entry.SetMemoryDependencies(m_pComp, recMemoryDependencies);
32983374

32993375
GetMapSelectWorkCache()->Set(fstruct, entry);
33003376
}
33013377

3378+
recMemoryDependencies.ForEach([this, &memoryDependencies](ValueNum vn) { memoryDependencies.Add(m_pComp, vn); });
3379+
33023380
return entry.Result;
33033381
}
33043382

src/coreclr/jit/valuenum.h

+9-9
Original file line numberDiff line numberDiff line change
@@ -684,13 +684,13 @@ class ValueNumStore
684684
ValueNum VNForMapSelectInner(ValueNumKind vnk, var_types type, ValueNum map, ValueNum index);
685685

686686
// A method that does the work for VNForMapSelect and may call itself recursively.
687-
ValueNum VNForMapSelectWork(ValueNumKind vnk,
688-
var_types type,
689-
ValueNum map,
690-
ValueNum index,
691-
int* pBudget,
692-
bool* pUsedRecursiveVN,
693-
ArrayStack<ValueNum>* loopMemoryDependencies);
687+
ValueNum VNForMapSelectWork(ValueNumKind vnk,
688+
var_types type,
689+
ValueNum map,
690+
ValueNum index,
691+
int* pBudget,
692+
bool* pUsedRecursiveVN,
693+
class SmallValueNumSet& loopMemoryDependencies);
694694

695695
// A specialized version of VNForFunc that is used for VNF_MapStore and provides some logging when verbose is set
696696
ValueNum VNForMapStore(ValueNum map, ValueNum index, ValueNum value);
@@ -1821,8 +1821,8 @@ class ValueNumStore
18211821
public:
18221822
ValueNum Result;
18231823

1824-
void SetMemoryDependencies(CompAllocator alloc, ArrayStack<ValueNum>& deps, unsigned startIndex);
1825-
void GetMemoryDependencies(ArrayStack<ValueNum>& deps);
1824+
void SetMemoryDependencies(Compiler* comp, class SmallValueNumSet& deps);
1825+
void GetMemoryDependencies(Compiler* comp, class SmallValueNumSet& deps);
18261826
};
18271827

18281828
typedef JitHashTable<VNDefFuncApp<2>, VNDefFuncAppKeyFuncs<2>, MapSelectWorkCacheEntry> MapSelectWorkCache;

0 commit comments

Comments
 (0)