interpreter: Assign fixed slots (vregs) to symbols with fast or closure scope.

This removes a bottleneck of the interpreter/bjit:
most var accesses introduced a DenseMap lookup, with this change we use a fixed offset per var.
The bjit stores the pointer to the vregs array inside r14 for fast accesses.

src/asm_writing/assembler.cpp

@@ -838,7 +838,6 @@ void Assembler::test(Register reg1, Register reg2) {
         reg1_idx -= 8;
     }
     if (reg2_idx >= 8) {
-        trap();
         rex |= REX_B;
         reg2_idx -= 8;
     }

src/codegen/ast_interpreter.cpp

@@ -66,10 +66,7 @@ extern "C" Box* executeInnerAndSetupFrame(ASTInterpreter& interpreter, CFGBlock*
  */
 class ASTInterpreter {
 public:
-    typedef ContiguousMap<InternedString, Box*, llvm::SmallDenseMap<InternedString, int, 16>,
-                          llvm::SmallVector<Box*, 512>> SymMap;
-
-    ASTInterpreter(CLFunction* clfunc);
+    ASTInterpreter(CLFunction* clfunc, Box** vregs);
 
     void initArguments(int nargs, BoxedClosure* closure, BoxedGenerator* generator, Box* arg1, Box* arg2, Box* arg3,
                        Box** args);
@@ -145,7 +142,7 @@ private:
     ScopeInfo* scope_info;
     PhiAnalysis* phis;
 
-    SymMap sym_table;
+    Box** vregs;
     ExcInfo last_exception;
     BoxedClosure* passed_closure, *created_closure;
     BoxedGenerator* generator;
@@ -158,6 +155,11 @@ private:
     std::unique_ptr<JitFragmentWriter> jit;
 
 public:
+    llvm::DenseMap<InternedString, int>& getSymVRegMap() {
+        assert(source_info->cfg);
+        return source_info->cfg->sym_vreg_map;
+    }
+
     AST_stmt* getCurrentStatement() {
         assert(current_inst);
         return current_inst;
@@ -171,7 +173,7 @@ public:
     CLFunction* getCL() { return clfunc; }
     FrameInfo* getFrameInfo() { return &frame_info; }
     BoxedClosure* getPassedClosure() { return passed_closure; }
-    const SymMap& getSymbolTable() { return sym_table; }
+    Box** getVRegs() { return vregs; }
     const ScopeInfo* getScopeInfo() { return scope_info; }
 
     void addSymbol(InternedString name, Box* value, bool allow_duplicates);
@@ -186,9 +188,10 @@ public:
 };
 
 void ASTInterpreter::addSymbol(InternedString name, Box* value, bool allow_duplicates) {
+    assert(getSymVRegMap().count(name));
     if (!allow_duplicates)
-        assert(sym_table.count(name) == 0);
-    sym_table[name] = value;
+        assert(vregs[getSymVRegMap()[name]] == NULL);
+    vregs[getSymVRegMap()[name]] = value;
 }
 
 void ASTInterpreter::setGenerator(Box* gen) {
@@ -222,13 +225,14 @@ void ASTInterpreter::setGlobals(Box* globals) {
     this->globals = globals;
 }
 
-ASTInterpreter::ASTInterpreter(CLFunction* clfunc)
+ASTInterpreter::ASTInterpreter(CLFunction* clfunc, Box** vregs)
     : current_block(0),
       current_inst(0),
       clfunc(clfunc),
       source_info(clfunc->source.get()),
       scope_info(0),
       phis(NULL),
+      vregs(vregs),
       last_exception(NULL, NULL, NULL),
       passed_closure(0),
       created_closure(0),
@@ -313,7 +317,7 @@ void ASTInterpreter::finishJITing(CFGBlock* continue_block) {
 Box* ASTInterpreter::execJITedBlock(CFGBlock* b) {
     try {
         UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_baseline_jitted_code");
-        std::pair<CFGBlock*, Box*> rtn = b->entry_code(this, b);
+        std::pair<CFGBlock*, Box*> rtn = b->entry_code(this, b, vregs);
         next_block = rtn.first;
         if (!next_block)
             return rtn.second;
@@ -447,6 +451,7 @@ void ASTInterpreter::doStore(AST_Name* node, Value value) {
             jit->emitSetItemName(name.getBox(), value);
         assert(frame_info.boxedLocals != NULL);
         // TODO should probably pre-box the names when it's a scope that usesNameLookup
+        assert(gc::isValidGCObject(value.o));
         setitem(frame_info.boxedLocals, name.getBox(), value.o);
     } else {
         bool closure = vst == ScopeInfo::VarScopeType::CLOSURE;
@@ -454,14 +459,17 @@ void ASTInterpreter::doStore(AST_Name* node, Value value) {
             if (!closure) {
                 bool is_live = source_info->getLiveness()->isLiveAtEnd(name, current_block);
                 if (is_live)
-                    jit->emitSetLocal(name, closure, value);
+                    jit->emitSetLocal(name, node->vreg, closure, value);
                 else
                     jit->emitSetBlockLocal(name, value);
             } else
-                jit->emitSetLocal(name, closure, value);
+                jit->emitSetLocal(name, node->vreg, closure, value);
         }
 
-        sym_table[name] = value.o;
+        assert(getSymVRegMap().count(name));
+        assert(getSymVRegMap()[name] == node->vreg);
+        vregs[node->vreg] = value.o;
+
         if (closure) {
             created_closure->elts[scope_info->getClosureOffset(name)] = value.o;
         }
@@ -626,17 +634,24 @@ Box* ASTInterpreter::doOSR(AST_Jump* node) {
     std::unique_ptr<PhiAnalysis> phis
         = computeRequiredPhis(clfunc->param_names, source_info->cfg, liveness, scope_info);
 
+    llvm::DenseMap<int, InternedString> offset_name_map;
+    for (auto&& v : getSymVRegMap()) {
+        offset_name_map[v.second] = v.first;
+    }
+
     std::vector<InternedString> dead_symbols;
-    for (auto& it : sym_table) {
-        if (!liveness->isLiveAtEnd(it.first, current_block)) {
-            dead_symbols.push_back(it.first);
-        } else if (phis->isRequiredAfter(it.first, current_block)) {
-            assert(scope_info->getScopeTypeOfName(it.first) != ScopeInfo::VarScopeType::GLOBAL);
+    for (int i = 0; i < getSymVRegMap().size(); ++i) {
+        if (!liveness->isLiveAtEnd(offset_name_map[i], current_block)) {
+            dead_symbols.push_back(offset_name_map[i]);
+        } else if (phis->isRequiredAfter(offset_name_map[i], current_block)) {
+            assert(scope_info->getScopeTypeOfName(offset_name_map[i]) != ScopeInfo::VarScopeType::GLOBAL);
         } else {
         }
     }
-    for (auto&& dead : dead_symbols)
-        sym_table.erase(dead);
+    for (auto&& dead : dead_symbols) {
+        assert(getSymVRegMap().count(dead));
+        vregs[getSymVRegMap()[dead]] = NULL;
+    }
 
     const OSREntryDescriptor* found_entry = nullptr;
     for (auto& p : clfunc->osr_versions) {
@@ -652,20 +667,19 @@ Box* ASTInterpreter::doOSR(AST_Jump* node) {
     static Box* const VAL_UNDEFINED = (Box*)-1;
 
     for (auto& name : phis->definedness.getDefinedNamesAtEnd(current_block)) {
-        auto it = sym_table.find(name);
+        assert(getSymVRegMap().count(name));
+        Box* val = vregs[getSymVRegMap()[name]];
         if (!liveness->isLiveAtEnd(name, current_block))
             continue;
 
         if (phis->isPotentiallyUndefinedAfter(name, current_block)) {
-            bool is_defined = it != sym_table.end();
+            bool is_defined = val != NULL;
             // TODO only mangle once
             sorted_symbol_table[getIsDefinedName(name, source_info->getInternedStrings())] = (Box*)is_defined;
-            if (is_defined)
-                assert(sym_table.getMapped(it->second) != NULL);
-            sorted_symbol_table[name] = is_defined ? sym_table.getMapped(it->second) : VAL_UNDEFINED;
+            sorted_symbol_table[name] = is_defined ? val : VAL_UNDEFINED;
         } else {
-            ASSERT(it != sym_table.end(), "%s", name.c_str());
-            Box* v = sorted_symbol_table[it->first] = sym_table.getMapped(it->second);
+            ASSERT(val != NULL, "%s", name.c_str());
+            Box* v = sorted_symbol_table[name] = val;
             assert(gc::isValidGCObject(v));
         }
     }
@@ -1078,8 +1092,10 @@ Value ASTInterpreter::visit_assert(AST_Assert* node) {
 
 Value ASTInterpreter::visit_global(AST_Global* node) {
     abortJITing();
-    for (auto name : node->names)
-        sym_table.erase(name);
+    for (auto name : node->names) {
+        if (getSymVRegMap().count(name))
+            vregs[getSymVRegMap()[name]] = NULL;
+    }
     return Value();
 }
 
@@ -1123,12 +1139,14 @@ Value ASTInterpreter::visit_delete(AST_Delete* node) {
                 } else {
                     assert(vst == ScopeInfo::VarScopeType::FAST);
 
-                    if (sym_table.count(target->id) == 0) {
+                    assert(getSymVRegMap().count(target->id));
+                    assert(getSymVRegMap()[target->id] == target->vreg);
+                    if (vregs[target->vreg] == 0) {
                         assertNameDefined(0, target->id.c_str(), NameError, true /* local_var_msg */);
                         return Value();
                     }
 
-                    sym_table.erase(target->id);
+                    vregs[target->vreg] = NULL;
                 }
                 break;
             }
@@ -1431,19 +1449,30 @@ Value ASTInterpreter::visit_name(AST_Name* node) {
                     is_live = source_info->getLiveness()->isLiveAtEnd(node->id, current_block);
 
                 if (is_live)
-                    v.var = jit->emitGetLocal(node->id);
+                    v.var = jit->emitGetLocal(node->id, node->vreg);
                 else
-                    v.var = jit->emitGetBlockLocal(node->id);
+                    v.var = jit->emitGetBlockLocal(node->id, node->vreg);
             }
 
-            v.o = ASTInterpreterJitInterface::getLocalHelper(this, node->id);
+            assert(node->vreg >= 0);
+            assert(getSymVRegMap().count(node->id));
+            assert(getSymVRegMap()[node->id] == node->vreg);
+            Box* val = vregs[node->vreg];
+            if (val) {
+                assert(gc::isValidGCObject(val));
+                v.o = val;
                 return v;
             }
+
+            assertNameDefined(0, node->id.c_str(), UnboundLocalError, true);
+            RELEASE_ASSERT(0, "should be unreachable");
+        }
         case ScopeInfo::VarScopeType::NAME: {
             Value v;
             if (jit)
                 v.var = jit->emitGetBoxedLocal(node->id.getBox());
             v.o = boxedLocalsGet(frame_info.boxedLocals, node->id.getBox(), globals);
+            assert(gc::isValidGCObject(v.o));
             return v;
         }
         default:
@@ -1537,20 +1566,6 @@ Box* ASTInterpreterJitInterface::doOSRHelper(void* _interpreter, AST_Jump* node)
     return NULL;
 }
 
-Box* ASTInterpreterJitInterface::getLocalHelper(void* _interpreter, InternedString id) {
-    ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
-
-    auto it = interpreter->sym_table.find(id);
-    if (it != interpreter->sym_table.end()) {
-        Box* v = interpreter->sym_table.getMapped(it->second);
-        assert(gc::isValidGCObject(v));
-        return v;
-    }
-
-    assertNameDefined(0, id.c_str(), UnboundLocalError, true);
-    return 0;
-}
-
 Box* ASTInterpreterJitInterface::landingpadHelper(void* _interpreter) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
     ExcInfo& last_exception = interpreter->last_exception;
@@ -1574,24 +1589,16 @@ Box* ASTInterpreterJitInterface::uncacheExcInfoHelper(void* _interpreter) {
     return None;
 }
 
-void ASTInterpreterJitInterface::setLocalClosureHelper(void* _interpreter, InternedString id, Box* v) {
+void ASTInterpreterJitInterface::setLocalClosureHelper(void* _interpreter, long vreg, InternedString id, Box* v) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
 
     assert(gc::isValidGCObject(v));
-    interpreter->sym_table[id] = v;
-    RELEASE_ASSERT(interpreter->sym_table.size() < 512, "sym table to small");
-
+    assert(interpreter->getSymVRegMap().count(id));
+    assert(interpreter->getSymVRegMap()[id] == vreg);
+    interpreter->vregs[vreg] = v;
     interpreter->created_closure->elts[interpreter->scope_info->getClosureOffset(id)] = v;
 }
 
-void ASTInterpreterJitInterface::setLocalHelper(void* _interpreter, InternedString id, Box* v) {
-    ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
-
-    assert(gc::isValidGCObject(v));
-    interpreter->sym_table[id] = v;
-    RELEASE_ASSERT(interpreter->sym_table.size() < 512, "sym table to small");
-}
-
 const void* interpreter_instr_addr = (void*)&executeInnerAndSetupFrame;
 
 // small wrapper around executeInner because we can not directly call the member function from asm.
@@ -1599,6 +1606,24 @@ extern "C" Box* executeInnerFromASM(ASTInterpreter& interpreter, CFGBlock* start
     return ASTInterpreter::executeInner(interpreter, start_block, start_at);
 }
 
+static int calculateNumVRegs(CLFunction* clfunc) {
+    ScopeInfo* scope_info = clfunc->source->getScopeInfo();
+    SourceInfo* source_info = clfunc->source.get();
+
+    // Note: due to some (avoidable) restrictions, this check is pretty constrained in where
+    // it can go, due to the fact that it can throw an exception.
+    // It can't go in the ASTInterpreter constructor, since that will cause the C++ runtime to
+    // delete the partially-constructed memory which we don't currently handle.  It can't go into
+    // executeInner since we want the SyntaxErrors to happen *before* the stack frame is entered.
+    // (For instance, throwing the exception will try to fetch the current statement, but we determine
+    // that by looking at the cfg.)
+    if (!source_info->cfg)
+        source_info->cfg = computeCFG(source_info, source_info->body);
+
+    source_info->cfg->assignVRegs(clfunc->param_names, scope_info);
+    return source_info->cfg->sym_vreg_map.size();
+}
+
 Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* generator, Box* globals, Box* arg1,
                           Box* arg2, Box* arg3, Box** args) {
     UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_interpreter");
@@ -1650,10 +1675,18 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
         return optimized->call(arg1, arg2, arg3, args);
     }
 
+    Box** vregs = NULL;
+    int num_vregs = calculateNumVRegs(clfunc);
+    if (num_vregs > 0) {
+        vregs = (Box**)alloca(sizeof(Box*) * num_vregs);
+        memset(vregs, 0, sizeof(Box*) * num_vregs);
+    }
+
     ++clfunc->times_interpreted;
-    ASTInterpreter interpreter(clfunc);
+    ASTInterpreter interpreter(clfunc, vregs);
 
     ScopeInfo* scope_info = clfunc->source->getScopeInfo();
+
     if (unlikely(scope_info->usesNameLookup())) {
         interpreter.setBoxedLocals(new BoxedDict());
     }
@@ -1673,13 +1706,17 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
 Box* astInterpretFunctionEval(CLFunction* clfunc, Box* globals, Box* boxedLocals) {
     ++clfunc->times_interpreted;
 
-    ASTInterpreter interpreter(clfunc);
+    Box** vregs = NULL;
+    int num_vregs = calculateNumVRegs(clfunc);
+    if (num_vregs > 0) {
+        vregs = (Box**)alloca(sizeof(Box*) * num_vregs);
+        memset(vregs, 0, sizeof(Box*) * num_vregs);
+    }
+
+    ASTInterpreter interpreter(clfunc, vregs);
     interpreter.initArguments(0, NULL, NULL, NULL, NULL, NULL, NULL);
     interpreter.setBoxedLocals(boxedLocals);
 
-    ScopeInfo* scope_info = clfunc->source->getScopeInfo();
-    SourceInfo* source_info = clfunc->source.get();
-
     assert(!clfunc->source->scoping->areGlobalsFromModule());
     assert(globals);
     interpreter.setGlobals(globals);
@@ -1690,16 +1727,24 @@ Box* astInterpretFunctionEval(CLFunction* clfunc, Box* globals, Box* boxedLocals
 
 Box* astInterpretDeopt(CLFunction* clfunc, AST_expr* after_expr, AST_stmt* enclosing_stmt, Box* expr_val,
                        FrameStackState frame_state) {
+    RELEASE_ASSERT(0, "didn't check if this still works");
     assert(clfunc);
     assert(enclosing_stmt);
     assert(frame_state.locals);
     assert(after_expr);
     assert(expr_val);
 
-    ASTInterpreter interpreter(clfunc);
-
-    ScopeInfo* scope_info = clfunc->source->getScopeInfo();
     SourceInfo* source_info = clfunc->source.get();
+
+    Box** vregs = NULL;
+    int num_vregs = calculateNumVRegs(clfunc);
+    if (num_vregs > 0) {
+        vregs = (Box**)alloca(sizeof(Box*) * num_vregs);
+        memset(vregs, 0, sizeof(Box*) * num_vregs);
+    }
+
+    ASTInterpreter interpreter(clfunc, vregs);
+
     assert(clfunc->source->scoping->areGlobalsFromModule());
     interpreter.setGlobals(source_info->parent_module);
 
@@ -1806,11 +1851,15 @@ BoxedDict* localsForInterpretedFrame(void* frame_ptr, bool only_user_visible) {
     ASTInterpreter* interpreter = getInterpreterFromFramePtr(frame_ptr);
     assert(interpreter);
     BoxedDict* rtn = new BoxedDict();
-    for (auto& l : interpreter->getSymbolTable()) {
+    for (auto& l : interpreter->getSymVRegMap()) {
         if (only_user_visible && (l.first.s()[0] == '!' || l.first.s()[0] == '#'))
             continue;
 
-        rtn->d[l.first.getBox()] = interpreter->getSymbolTable().getMapped(l.second);
+        Box* val = interpreter->getVRegs()[l.second];
+        if (val) {
+            assert(gc::isValidGCObject(val));
+            rtn->d[l.first.getBox()] = val;
+        }
     }
 
     return rtn;

src/codegen/ast_interpreter.h

@@ -43,12 +43,10 @@ struct ASTInterpreterJitInterface {
 
     static Box* derefHelper(void* interp, InternedString s);
     static Box* doOSRHelper(void* interp, AST_Jump* node);
-    static Box* getLocalHelper(void* interp, InternedString id);
     static Box* landingpadHelper(void* interp);
     static Box* setExcInfoHelper(void* interp, Box* type, Box* value, Box* traceback);
     static Box* uncacheExcInfoHelper(void* interp);
-    static void setLocalClosureHelper(void* interp, InternedString id, Box* v);
-    static void setLocalHelper(void* interp, InternedString id, Box* v);
+    static void setLocalClosureHelper(void* interp, long vreg, InternedString id, Box* v);
 };
 
 class RewriterVar;

src/codegen/baseline_jit.cpp

@@ -36,17 +36,18 @@ static llvm::DenseMap<CFGBlock*, std::vector<void*>> block_patch_locations;
 //
 // long foo(char* c);
 // void bjit() {
+//   asm volatile ("" ::: "r14");
 //   asm volatile ("" ::: "r12");
 //   char scratch[256+16];
 //   foo(scratch);
 // }
 //
-// It omits the frame pointer but saves R12
+// It omits the frame pointer but saves R12 and R14
 const unsigned char eh_info[]
     = { 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x7a, 0x52, 0x00, 0x01, 0x78, 0x10,
         0x01, 0x1b, 0x0c, 0x07, 0x08, 0x90, 0x01, 0x00, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x1c, 0x00,
-        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x00, 0x42, 0x0e, 0x10, 0x47,
-        0x0e, 0xa0, 0x02, 0x8c, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x42, 0x0e, 0x10, 0x42,
+        0x0e, 0x18, 0x47, 0x0e, 0xb0, 0x02, 0x8c, 0x03, 0x8e, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00 };
 static_assert(JitCodeBlock::num_stack_args == 2, "have to update EH table!");
 static_assert(JitCodeBlock::scratch_size == 256, "have to update EH table!");
 
@@ -63,10 +64,12 @@ JitCodeBlock::JitCodeBlock(llvm::StringRef name)
     num_jit_total_bytes.log(code_size);
 
     // emit prolog
+    a.push(assembler::R14);
     a.push(assembler::R12);
-    static_assert(sp_adjustment % 16 == 0, "stack isn't aligned");
+    static_assert(sp_adjustment % 16 == 8, "stack isn't aligned");
     a.sub(assembler::Immediate(sp_adjustment), assembler::RSP);
     a.mov(assembler::RDI, assembler::R12);                                // interpreter pointer
+    a.mov(assembler::RDX, assembler::R14);                                // vreg array
     a.jmp(assembler::Indirect(assembler::RSI, offsetof(CFGBlock, code))); // jump to block
 
     entry_offset = a.bytesWritten();
@@ -139,6 +142,10 @@ JitFragmentWriter::JitFragmentWriter(CFGBlock* block, std::unique_ptr<ICInfo> ic
     interp = createNewVar();
     addLocationToVar(interp, assembler::R12);
     interp->setAttr(ASTInterpreterJitInterface::getCurrentBlockOffset(), imm(block));
+
+    vregs_array = createNewVar();
+    addLocationToVar(vregs_array, assembler::R14);
+    addAction([=]() { vregs_array->bumpUse(); }, vregs_array, ActionType::NORMAL);
 }
 
 RewriterVar* JitFragmentWriter::imm(uint64_t val) {
@@ -279,10 +286,10 @@ RewriterVar* JitFragmentWriter::emitGetAttr(RewriterVar* obj, BoxedString* s, AS
     return emitPPCall((void*)getattr, { obj, imm(s) }, 2, 512, getTypeRecorderForNode(node));
 }
 
-RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s) {
+RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s, int vreg) {
     auto it = local_syms.find(s);
     if (it == local_syms.end())
-        return emitGetLocal(s);
+        return emitGetLocal(s, vreg);
     return it->second;
 }
 
@@ -308,13 +315,11 @@ RewriterVar* JitFragmentWriter::emitGetItem(RewriterVar* value, RewriterVar* sli
     return emitPPCall((void*)getitem, { value, slice }, 2, 512);
 }
 
-RewriterVar* JitFragmentWriter::emitGetLocal(InternedString s) {
-    return call(false, (void*)ASTInterpreterJitInterface::getLocalHelper, getInterp(),
-#ifndef NDEBUG
-                imm(asUInt(s).first), imm(asUInt(s).second));
-#else
-                imm(asUInt(s)));
-#endif
+RewriterVar* JitFragmentWriter::emitGetLocal(InternedString s, int vreg) {
+    assert(vreg >= 0);
+    RewriterVar* val_var = vregs_array->getAttr(vreg * 8);
+    addAction([=]() { _emitGetLocal(val_var, s.c_str()); }, { val_var }, ActionType::NORMAL);
+    return val_var;
 }
 
 RewriterVar* JitFragmentWriter::emitGetPystonIter(RewriterVar* v) {
@@ -471,17 +476,19 @@ void JitFragmentWriter::emitSetItemName(BoxedString* s, RewriterVar* v) {
     emitSetItem(emitGetBoxedLocals(), imm(s), v);
 }
 
-void JitFragmentWriter::emitSetLocal(InternedString s, bool set_closure, RewriterVar* v) {
-    void* func = set_closure ? (void*)ASTInterpreterJitInterface::setLocalClosureHelper
-                             : (void*)ASTInterpreterJitInterface::setLocalHelper;
-
-    call(false, func, getInterp(),
+void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closure, RewriterVar* v) {
+    assert(vreg >= 0);
+    if (set_closure) {
+        call(false, (void*)ASTInterpreterJitInterface::setLocalClosureHelper, getInterp(), imm(vreg),
 #ifndef NDEBUG
              imm(asUInt(s).first), imm(asUInt(s).second),
 #else
              imm(asUInt(s)),
 #endif
              v);
+    } else {
+        vregs_array->setAttr(8 * vreg, v);
+    }
 }
 
 void JitFragmentWriter::emitSideExit(RewriterVar* v, Box* cmp_value, CFGBlock* next_block) {
@@ -619,6 +626,10 @@ RewriterVar* JitFragmentWriter::emitPPCall(void* func_addr, llvm::ArrayRef<Rewri
 #endif
 }
 
+void JitFragmentWriter::assertNameDefinedHelper(const char* id) {
+    assertNameDefined(0, id, UnboundLocalError, true);
+}
+
 Box* JitFragmentWriter::callattrHelper(Box* obj, BoxedString* attr, CallattrFlags flags, TypeRecorder* type_recorder,
                                        Box** args, std::vector<BoxedString*>* keyword_names) {
     auto arg_tuple = getTupleFromArgsArray(&args[0], flags.argspec.totalPassed());
@@ -683,6 +694,18 @@ Box* JitFragmentWriter::runtimeCallHelper(Box* obj, ArgPassSpec argspec, TypeRec
     return recordType(type_recorder, r);
 }
 
+void JitFragmentWriter::_emitGetLocal(RewriterVar* val_var, const char* name) {
+    assembler::Register var_reg = val_var->getInReg();
+    assembler->test(var_reg, var_reg);
+    val_var->bumpUse();
+
+    {
+        assembler::ForwardJump jnz(*assembler, assembler::COND_NOT_ZERO);
+        assembler->mov(assembler::Immediate((uint64_t)name), assembler::RDI);
+        assembler->mov(assembler::Immediate((void*)assertNameDefinedHelper), assembler::R11);
+        assembler->callq(assembler::R11);
+    }
+}
 
 void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, int& size_of_exit_to_interp) {
     size_of_exit_to_interp = 0;
@@ -698,6 +721,7 @@ void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, int& siz
         block_next->getInReg(assembler::RAX, true);
         assembler->add(assembler::Immediate(JitCodeBlock::sp_adjustment), assembler::RSP);
         assembler->pop(assembler::R12);
+        assembler->pop(assembler::R14);
         assembler->retq();
 
         // make sure we have at least 'min_patch_size' of bytes available.
@@ -724,6 +748,7 @@ void JitFragmentWriter::_emitOSRPoint(RewriterVar* result, RewriterVar* node_var
         assembler->mov(assembler::Immediate(0ul), assembler::RAX); // TODO: use xor
         assembler->add(assembler::Immediate(JitCodeBlock::sp_adjustment), assembler::RSP);
         assembler->pop(assembler::R12);
+        assembler->pop(assembler::R14);
         assembler->retq();
     }
 
@@ -794,6 +819,7 @@ void JitFragmentWriter::_emitReturn(RewriterVar* return_val) {
     assembler->mov(assembler::Immediate(0ul), assembler::RAX); // TODO: use xor
     assembler->add(assembler::Immediate(JitCodeBlock::sp_adjustment), assembler::RSP);
     assembler->pop(assembler::R12);
+    assembler->pop(assembler::R14);
     assembler->retq();
     return_val->bumpUse();
 }

src/codegen/baseline_jit.h

@@ -70,12 +70,16 @@ class JitFragmentWriter;
 // This also means that we are allowed to store a Python variable which only lives in the current CFGBLock* inside a
 // register or stack slot but we aren't if it outlives the block - we have to store it in the interpreter instance.
 //
+// We use the following callee-save regs to speed up the generated code:
+//      r12: pointer to ASTInterpreter instance
+//      r14: pointer to the vregs array
+//
 // To execute a specific CFGBlock one has to call:
 //      CFGBlock* block;
-//      block->entry_code(ast_interpreter_instance, block)
+//      block->entry_code(ast_interpreter_instance, block, ast_interpreter_instance->vregs)
 //
 // Signature of a JitCodeBlock:
-//  std::pair<CFGBlock*, Box*>(*entry_code)(ASTInterpreter* interp, CFGBlock* block)
+//  std::pair<CFGBlock*, Box*>(*entry_code)(ASTInterpreter* interp, CFGBlock* block, Box** vregs)
 //  args:
 //      interp:  instance to the ASTInterpreter
 //      block:   block to execute
@@ -87,10 +91,12 @@ class JitFragmentWriter;
 //
 // Basic layout of generated code block is:
 // entry_code:
+//      push   %r14                 ; save r14
 //      push   %r12                 ; save r12
-//      sub    $0x110,%rsp          ; setup scratch, 0x110 = scratch_size + 16 = space for two func args passed on the
-//                                                                               stack
+//      sub    $0x118,%rsp          ; setup scratch, 0x118 = scratch_size + 16 = space for two func args passed on the
+//                                                                               stack + 8 byte for stack alignment
 //      mov    %rdi,%r12            ; copy the pointer to ASTInterpreter instance into r12
+//      mov    %rdx,%r14            ; copy the pointer to the vregs array into r14
 //      jmpq   *0x8(%rsi)           ; jump to block->code
 //                                      possible values: first_JitFragment, second_JitFragment,...
 //
@@ -101,8 +107,9 @@ class JitFragmentWriter;
 //      cmp    %rax,%rcx            ; rax == True
 //      jne    end_side_exit
 //      movabs $0x215bb60,%rax      ; rax = CFGBlock* to interpret next (rax is the 1. return reg)
-//      add    $0x110,%rsp          ; restore stack pointer
+//      add    $0x118,%rsp          ; restore stack pointer
 //      pop    %r12                 ; restore r12
+//      pop    %r14                 ; restore r14
 //      ret                         ; exit to the interpreter which will interpret the specified CFGBLock*
 //    end_side_exit:
 //      ....
@@ -113,8 +120,9 @@ class JitFragmentWriter;
 //      mov    $0,%rax              ; rax contains the next block to interpret.
 //                                    in this case 0 which means we are finished
 //      movabs $0x1270014108,%rdx   ; rdx must contain the Box* value to return
-//      add    $0x110,%rsp          ; restore stack pointer
+//      add    $0x118,%rsp          ; restore stack pointer
 //      pop    %r12                 ; restore r12
+//      pop    %r14                 ; restore r14
 //      ret
 //
 // nth_JitFragment:
@@ -130,7 +138,7 @@ public:
 
     // scratch size + space for passing additional args on the stack without having to adjust the SP when calling
     // functions with more than 6 args.
-    static constexpr int sp_adjustment = scratch_size + num_stack_args * 8;
+    static constexpr int sp_adjustment = scratch_size + num_stack_args * 8 + 8 /* = alignment */;
 
 private:
     std::unique_ptr<uint8_t[]> code;
@@ -168,6 +176,7 @@ private:
     void* entry_code; // JitCodeBlock start address. Must have an offset of 0 into the code block
     JitCodeBlock& code_block;
     RewriterVar* interp;
+    RewriterVar* vregs_array;
     llvm::DenseMap<InternedString, RewriterVar*> local_syms;
     std::unique_ptr<ICInfo> ic_info;
 
@@ -208,13 +217,13 @@ public:
     RewriterVar* emitDeref(InternedString s);
     RewriterVar* emitExceptionMatches(RewriterVar* v, RewriterVar* cls);
     RewriterVar* emitGetAttr(RewriterVar* obj, BoxedString* s, AST_expr* node);
-    RewriterVar* emitGetBlockLocal(InternedString s);
+    RewriterVar* emitGetBlockLocal(InternedString s, int vreg);
     RewriterVar* emitGetBoxedLocal(BoxedString* s);
     RewriterVar* emitGetBoxedLocals();
     RewriterVar* emitGetClsAttr(RewriterVar* obj, BoxedString* s);
     RewriterVar* emitGetGlobal(Box* global, BoxedString* s);
     RewriterVar* emitGetItem(RewriterVar* value, RewriterVar* slice);
-    RewriterVar* emitGetLocal(InternedString s);
+    RewriterVar* emitGetLocal(InternedString s, int vreg);
     RewriterVar* emitGetPystonIter(RewriterVar* v);
     RewriterVar* emitHasnext(RewriterVar* v);
     RewriterVar* emitLandingpad();
@@ -241,7 +250,7 @@ public:
     void emitSetGlobal(Box* global, BoxedString* s, RewriterVar* v);
     void emitSetItemName(BoxedString* s, RewriterVar* v);
     void emitSetItem(RewriterVar* target, RewriterVar* slice, RewriterVar* value);
-    void emitSetLocal(InternedString s, bool set_closure, RewriterVar* v);
+    void emitSetLocal(InternedString s, int vreg, bool set_closure, RewriterVar* v);
     void emitSideExit(RewriterVar* v, Box* cmp_value, CFGBlock* next_block);
     void emitUncacheExcInfo();
 
@@ -262,6 +271,7 @@ private:
     RewriterVar* emitPPCall(void* func_addr, llvm::ArrayRef<RewriterVar*> args, int num_slots, int slot_size,
                             TypeRecorder* type_recorder = NULL);
 
+    static void assertNameDefinedHelper(const char* id);
     static Box* callattrHelper(Box* obj, BoxedString* attr, CallattrFlags flags, TypeRecorder* type_recorder,
                                Box** args, std::vector<BoxedString*>* keyword_names);
     static Box* createDictHelper(uint64_t num, Box** keys, Box** values);
@@ -275,6 +285,7 @@ private:
     static Box* runtimeCallHelper(Box* obj, ArgPassSpec argspec, TypeRecorder* type_recorder, Box** args,
                                   std::vector<BoxedString*>* keyword_names);
 
+    void _emitGetLocal(RewriterVar* val_var, const char* name);
     void _emitJump(CFGBlock* b, RewriterVar* block_next, int& size_of_exit_to_interp);
     void _emitOSRPoint(RewriterVar* result, RewriterVar* node_var);
     void _emitPPCall(RewriterVar* result, void* func_addr, const RewriterVar::SmallVector& args, int num_slots,

src/core/ast.h

@@ -703,6 +703,10 @@ public:
     // different bytecodes.
     ScopeInfo::VarScopeType lookup_type;
 
+    // The interpreter and baseline JIT store variables with FAST and CLOSURE scopes in an array (vregs) this specifies
+    // the zero based index of this variable inside the vregs array. If uninitialized it's value is -1.
+    int vreg;
+
     virtual void accept(ASTVisitor* v);
     virtual void* accept_expr(ExprVisitor* v);
 
@@ -710,7 +714,8 @@ public:
         : AST_expr(AST_TYPE::Name, lineno, col_offset),
           ctx_type(ctx_type),
           id(id),
-          lookup_type(ScopeInfo::VarScopeType::UNKNOWN) {}
+          lookup_type(ScopeInfo::VarScopeType::UNKNOWN),
+          vreg(-1) {}
 
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Name;
 };

src/core/cfg.cpp

@@ -2471,6 +2471,87 @@ void CFG::print() {
         blocks[i]->print();
 }
 
+class AssignVRegsVisitor : public NoopASTVisitor {
+public:
+    int index = 0;
+    llvm::DenseMap<InternedString, int> sym_vreg_map;
+    ScopeInfo* scope_info;
+
+    AssignVRegsVisitor(ScopeInfo* scope_info) : scope_info(scope_info) {}
+
+    bool visit_arguments(AST_arguments* node) override {
+        for (AST_expr* d : node->defaults)
+            d->accept(this);
+        return true;
+    }
+
+    bool visit_classdef(AST_ClassDef* node) override {
+        for (auto e : node->bases)
+            e->accept(this);
+        for (auto e : node->decorator_list)
+            e->accept(this);
+        return true;
+    }
+
+    bool visit_functiondef(AST_FunctionDef* node) override {
+        for (auto* d : node->decorator_list)
+            d->accept(this);
+        node->args->accept(this);
+        return true;
+    }
+
+    bool visit_lambda(AST_Lambda* node) override {
+        node->args->accept(this);
+        return true;
+    }
+
+    bool visit_name(AST_Name* node) override {
+        if (node->vreg != -1)
+            return true;
+
+        if (node->lookup_type == ScopeInfo::VarScopeType::UNKNOWN)
+            node->lookup_type = scope_info->getScopeTypeOfName(node->id);
+
+        if (node->lookup_type == ScopeInfo::VarScopeType::FAST || node->lookup_type == ScopeInfo::VarScopeType::CLOSURE)
+            node->vreg = assignVReg(node->id);
+        return true;
+    }
+
+    int assignVReg(InternedString id) {
+        auto it = sym_vreg_map.find(id);
+        if (sym_vreg_map.end() == it) {
+            sym_vreg_map[id] = index;
+            return index++;
+        }
+        return it->second;
+    }
+};
+
+void CFG::assignVRegs(const ParamNames& param_names, ScopeInfo* scope_info) {
+    if (has_vregs_assigned)
+        return;
+
+    AssignVRegsVisitor visitor(scope_info);
+    for (CFGBlock* b : blocks) {
+        for (AST_stmt* stmt : b->body) {
+            stmt->accept(&visitor);
+        }
+    }
+
+    for (auto* name : param_names.arg_names) {
+        name->accept(&visitor);
+    }
+
+    if (param_names.vararg_name)
+        param_names.vararg_name->accept(&visitor);
+
+    if (param_names.kwarg_name)
+        param_names.kwarg_name->accept(&visitor);
+
+    sym_vreg_map = std::move(visitor.sym_vreg_map);
+    has_vregs_assigned = true;
+}
+
 CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
     STAT_TIMER(t0, "us_timer_computecfg", 0);
 

src/core/cfg.h

@@ -39,6 +39,9 @@ class AST_stmt;
 class Box;
 
 class CFG;
+class ParamNames;
+class ScopeInfo;
+
 class CFGBlock {
 private:
     CFG* cfg;
@@ -48,7 +51,7 @@ public:
     // contains address to the start of the code of this basic block
     void* code;
     // contains the address of the entry function
-    std::pair<CFGBlock*, Box*>(*entry_code)(void* interpeter, CFGBlock* block);
+    std::pair<CFGBlock*, Box*>(*entry_code)(void* interpeter, CFGBlock* block, Box** vregs);
 
     std::vector<AST_stmt*> body;
     std::vector<CFGBlock*> predecessors, successors;
@@ -70,11 +73,14 @@ public:
 class CFG {
 private:
     int next_idx;
+    bool has_vregs_assigned;
 
 public:
     std::vector<CFGBlock*> blocks;
 
-    CFG() : next_idx(0) {}
+    llvm::DenseMap<InternedString, int> sym_vreg_map;
+
+    CFG() : next_idx(0), has_vregs_assigned(false) {}
 
     CFGBlock* getStartingBlock() { return blocks[0]; }
 
@@ -103,6 +109,8 @@ public:
     }
 
     void print();
+
+    void assignVRegs(const ParamNames& param_names, ScopeInfo* scope_info);
 };
 
 class SourceInfo;