Changeset 254788 in webkit

Timestamp:

Jan 17, 2020 7:24:31 PM (4 years ago)

Author:

sbarati@apple.com

Message:

Air O0 should have better stack allocation
​https://bugs.webkit.org/show_bug.cgi?id=206436

Reviewed by Tadeu Zagallo.

JSTests:

• wasm/stress/dont-stack-overflow-in-air.js: Added.

Source/JavaScriptCore:

This patch adds a simple stack slot allocator to Air O0 to make code
use smaller stack frames. The huge stack frames from the old stack
allocator were leading to stack overflows in some programs. Before,
each Tmp got its own stack slot. The new allocator works similar to O0's
register allocator. This stack allocator linearizes the program and uses live
range end as an opportunity to place the stack slot on a free list of
available stack slots. This patch also fixes an issue in our linearization code
where the head of a block and the tail of another block would share the
same linearization index. This didn't matter for register allocation, but
does matter for the stack allocator. So "live at head", and "live at tail"
now get their own linearization index.

• b3/air/AirAllocateRegistersAndStackAndGenerateCode.cpp:

(JSC::B3::Air::GenerateAndAllocateRegisters::buildLiveRanges):
(JSC::B3::Air::GenerateAndAllocateRegisters::prepareForGeneration):
(JSC::B3::Air::GenerateAndAllocateRegisters::generate):

• b3/air/AirAllocateRegistersAndStackAndGenerateCode.h:
• b3/air/AirLiveness.h:

Location:

trunk

Files:

• JSTests/ChangeLog (modified) (1 diff)
• JSTests/wasm/stress/dont-stack-overflow-in-air.js (added)
• Source/JavaScriptCore/ChangeLog (modified) (1 diff)
• Source/JavaScriptCore/b3/air/AirAllocateRegistersAndStackAndGenerateCode.cpp (modified) (10 diffs)
• Source/JavaScriptCore/b3/air/AirAllocateRegistersAndStackAndGenerateCode.h (modified) (2 diffs)
• Source/JavaScriptCore/b3/air/AirLiveness.h (modified) (1 diff)

trunk/JSTests/ChangeLog

@@ +1 @@
+2020-01-17  Saam Barati  <sbarati@apple.com>
+
+        Air O0 should have better stack allocation
+        https://bugs.webkit.org/show_bug.cgi?id=206436
+
+        Reviewed by Tadeu Zagallo.
+
+        * wasm/stress/dont-stack-overflow-in-air.js: Added.
+
 2020-01-17  Mark Lam  <mark.lam@apple.com>
 

trunk/Source/JavaScriptCore/ChangeLog

@@ +1 @@
+2020-01-17  Saam Barati  <sbarati@apple.com>
+
+        Air O0 should have better stack allocation
+        https://bugs.webkit.org/show_bug.cgi?id=206436
+
+        Reviewed by Tadeu Zagallo.
+
+        This patch adds a simple stack slot allocator to Air O0 to make code
+        use smaller stack frames. The huge stack frames from the old stack
+        allocator were leading to stack overflows in some programs. Before,
+        each Tmp got its own stack slot. The new allocator works similar to O0's
+        register allocator. This stack allocator linearizes the program and uses live
+        range end as an opportunity to place the stack slot on a free list of
+        available stack slots. This patch also fixes an issue in our linearization code
+        where the head of a block and the tail of another block would share the
+        same linearization index. This didn't matter for register allocation, but
+        does matter for the stack allocator. So "live at head", and "live at tail"
+        now get their own linearization index.
+
+        * b3/air/AirAllocateRegistersAndStackAndGenerateCode.cpp:
+        (JSC::B3::Air::GenerateAndAllocateRegisters::buildLiveRanges):
+        (JSC::B3::Air::GenerateAndAllocateRegisters::prepareForGeneration):
+        (JSC::B3::Air::GenerateAndAllocateRegisters::generate):
+        * b3/air/AirAllocateRegistersAndStackAndGenerateCode.h:
+        * b3/air/AirLiveness.h:
+
 2020-01-17  David Kilzer  <ddkilzer@apple.com>
 

trunk/Source/JavaScriptCore/b3/air/AirAllocateRegistersAndStackAndGenerateCode.cpp

@@ -85 +85 @@
                 m_liveRangeEnd[tmp] = m_globalInstIndex;
         }
+        ++m_globalInstIndex;
         for (Inst& inst : *block) {
             inst.forEachTmpFast([&] (Tmp tmp) {
@@ -96 +97 @@
                 m_liveRangeEnd[tmp] = m_globalInstIndex;
         }
+        ++m_globalInstIndex;
     }
 }
@@ -281 +283 @@
     allocateEscapedStackSlots(m_code);
 
-    // Each Tmp gets its own stack slot.
-    auto createStackSlot = [&] (const Tmp& tmp) {
-        TmpData data;
-        data.spillSlot = m_code.addStackSlot(8, StackSlotKind::Spill);
-        data.reg = Reg();
-        m_map[tmp] = data;
+    insertBlocksForFlushAfterTerminalPatchpoints();
+
 #if ASSERT_ENABLED
-        m_allTmps[tmp.bank()].append(tmp);
-#endif
-    };
-
     m_code.forEachTmp([&] (Tmp tmp) {
         ASSERT(!tmp.isReg());
-        createStackSlot(tmp);
+        m_allTmps[tmp.bank()].append(tmp);
     });
+#endif
+
+    m_liveness = makeUnique<UnifiedTmpLiveness>(m_code);
+
+    {
+        buildLiveRanges(*m_liveness);
+
+        Vector<StackSlot*, 16> freeSlots;
+        Vector<StackSlot*, 4> toFree;
+        m_globalInstIndex = 0;
+        for (BasicBlock* block : m_code) {
+            auto assignStackSlotToTmp = [&] (Tmp tmp) {
+                if (tmp.isReg())
+                    return;
+
+                TmpData& data = m_map[tmp];
+                if (data.spillSlot) {
+                    if (m_liveRangeEnd[tmp] == m_globalInstIndex)
+                        toFree.append(data.spillSlot);
+                    return;
+                }
+
+                if (freeSlots.size())
+                    data.spillSlot = freeSlots.takeLast();
+                else
+                    data.spillSlot = m_code.addStackSlot(8, StackSlotKind::Spill);
+                data.reg = Reg();
+            };
+
+            auto flushToFreeList = [&] {
+                for (auto* stackSlot : toFree)
+                    freeSlots.append(stackSlot);
+                toFree.clear();
+            };
+
+            for (Tmp tmp : m_liveness->liveAtHead(block))
+                assignStackSlotToTmp(tmp);
+            flushToFreeList();
+
+            ++m_globalInstIndex;
+
+            for (Inst& inst : *block) {
+                Vector<Tmp, 4> seenTmps;
+                inst.forEachTmpFast([&] (Tmp tmp) {
+                    if (seenTmps.contains(tmp))
+                        return;
+                    seenTmps.append(tmp);
+                    assignStackSlotToTmp(tmp);
+                });
+
+                flushToFreeList();
+                ++m_globalInstIndex;
+            }
+
+            for (Tmp tmp : m_liveness->liveAtTail(block))
+                assignStackSlotToTmp(tmp);
+            flushToFreeList();
+
+            ++m_globalInstIndex;
+        }
+    } 
 
     m_allowedRegisters = RegisterSet();
@@ -304 +359 @@
         for (Reg reg : m_registers[bank]) {
             m_allowedRegisters.set(reg);
-            createStackSlot(Tmp(reg));
+            TmpData& data = m_map[Tmp(reg)];
+            data.spillSlot = m_code.addStackSlot(8, StackSlotKind::Spill);
+            data.reg = Reg();
         }
     });
@@ -324 +381 @@
     lowerStackArgs(m_code);
 
+#if ASSERT_ENABLED
     // Verify none of these passes add any tmps.
-#if ASSERT_ENABLED
     forEachBank([&] (Bank bank) {
-        ASSERT(m_allTmps[bank].size() - m_registers[bank].size() == m_code.numTmps(bank));
+        ASSERT(m_allTmps[bank].size() == m_code.numTmps(bank));
     });
+
+    {
+        // Verify that lowerStackArgs didn't change Tmp liveness at the boundaries for the Tmps and Registers we model.
+        UnifiedTmpLiveness liveness(m_code);
+        for (BasicBlock* block : m_code) {
+            auto assertLivenessAreEqual = [&] (auto a, auto b) {
+                HashSet<Tmp> livenessA;
+                HashSet<Tmp> livenessB;
+                for (Tmp tmp : a) {
+                    if (tmp.isReg() && isDisallowedRegister(tmp.reg()))
+                        continue;
+                    livenessA.add(tmp);
+                }
+                for (Tmp tmp : b) {
+                    if (tmp.isReg() && isDisallowedRegister(tmp.reg()))
+                        continue;
+                    livenessB.add(tmp);
+                }
+
+                ASSERT(livenessA == livenessB);
+            };
+
+            assertLivenessAreEqual(m_liveness->liveAtHead(block), liveness.liveAtHead(block));
+            assertLivenessAreEqual(m_liveness->liveAtTail(block), liveness.liveAtTail(block));
+        }
+    }
 #endif
 }
@@ -338 +421 @@
     TimingScope timingScope("Air::generateAndAllocateRegisters");
 
-    insertBlocksForFlushAfterTerminalPatchpoints();
-
     DisallowMacroScratchRegisterUsage disallowScratch(*m_jit);
 
-    UnifiedTmpLiveness liveness(m_code);
-    buildLiveRanges(liveness);
+    buildLiveRanges(*m_liveness);
 
     IndexMap<BasicBlock*, IndexMap<Reg, Tmp>> currentAllocationMap(m_code.size());
@@ -356 +436 @@
         for (unsigned i = m_code.numEntrypoints(); i--;) {
             BasicBlock* entrypoint = m_code.entrypoint(i).block();
-            for (Tmp tmp : liveness.liveAtHead(entrypoint)) {
+            for (Tmp tmp : m_liveness->liveAtHead(entrypoint)) {
                 if (tmp.isReg())
                     currentAllocationMap[entrypoint][tmp.reg()] = tmp;
@@ -443 +523 @@
             m_availableRegs[tmp.bank()].clear(reg);
         }
+
+        ++m_globalInstIndex;
 
         bool isReplayingSameInst = false;
@@ -674 +756 @@
                 }
                 if (!everySuccessorGetsOurRegisterState) {
-                    for (Tmp tmp : liveness.liveAtTail(block)) {
+                    for (Tmp tmp : m_liveness->liveAtTail(block)) {
                         if (tmp.isReg() && isDisallowedRegister(tmp.reg()))
                             continue;
@@ -753 +835 @@
                 m_map[tmp].reg = Reg();
         }
+
+        ++m_globalInstIndex;
     }
 

trunk/Source/JavaScriptCore/b3/air/AirAllocateRegistersAndStackAndGenerateCode.h

@@ -44 +44 @@
 
     struct TmpData {
-        StackSlot* spillSlot;
+        StackSlot* spillSlot { nullptr };
         Reg reg;
     };
@@ -85 +85 @@
     RegisterSet m_namedDefdRegs;
     RegisterSet m_allowedRegisters;
+    std::unique_ptr<UnifiedTmpLiveness> m_liveness;
 
     struct PatchSpillData {

trunk/Source/JavaScriptCore/b3/air/AirLiveness.h

@@ -37 +37 @@
 template<typename Adapter>
 class Liveness : public WTF::Liveness<Adapter> {
+    WTF_MAKE_FAST_ALLOCATED;
 public:
     Liveness(Code& code)