vregs: split them in three parts and reuse them in some cases

user visible: used for all non compiler generated names, name could be used in a single block or multiple
              all frames contain atleast this vregs in order to do frame introspection
cross block : used for compiler generated names which get used in several blocks or which have closure scope
single block: used by compiler created names which are only used in a single block.
              get reused for different names

we assign the lowest numbers to the user visible ones, followed by the cross block ones and finally the single block
ones. we do this because not all tiers use all of the vregs and it still makes it fast to switch between tiers.

src/codegen/ast_interpreter.cpp

@@ -66,7 +66,7 @@ extern "C" Box* executeInnerAndSetupFrame(ASTInterpreter& interpreter, CFGBlock*
  */
 class ASTInterpreter {
 public:
-    ASTInterpreter(FunctionMetadata* md, Box** vregs, int num_vregs, FrameInfo* deopt_frame_info = NULL);
+    ASTInterpreter(FunctionMetadata* md, Box** vregs, FrameInfo* deopt_frame_info = NULL);
 
     void initArguments(BoxedClosure* closure, BoxedGenerator* generator, Box* arg1, Box* arg2, Box* arg3, Box** args);
 
@@ -154,9 +154,9 @@ private:
 public:
     ~ASTInterpreter() { Py_XDECREF(this->created_closure); }
 
-    llvm::DenseMap<InternedString, int>& getSymVRegMap() {
-        assert(source_info->cfg);
-        return source_info->cfg->sym_vreg_map;
+    const VRegInfo& getVRegInfo() const { return source_info->cfg->getVRegInfo(); }
+    const llvm::DenseMap<InternedString, int>& getSymVRegMap() const {
+        return source_info->cfg->getVRegInfo().getSymVRegMap();
     }
 
     AST_stmt* getCurrentStatement() {
@@ -183,8 +183,7 @@ public:
 };
 
 void ASTInterpreter::addSymbol(InternedString name, Box* new_value, bool allow_duplicates) {
-    assert(getSymVRegMap().count(name));
-    Box*& value = vregs[getSymVRegMap()[name]];
+    Box*& value = vregs[getVRegInfo().getVReg(name)];
     Box* old_value = value;
     value = incref(new_value);
     if (allow_duplicates)
@@ -222,7 +221,7 @@ void ASTInterpreter::setGlobals(Box* globals) {
     this->frame_info.globals = incref(globals);
 }
 
-ASTInterpreter::ASTInterpreter(FunctionMetadata* md, Box** vregs, int num_vregs, FrameInfo* deopt_frame_info)
+ASTInterpreter::ASTInterpreter(FunctionMetadata* md, Box** vregs, FrameInfo* deopt_frame_info)
     : current_block(0),
       frame_info(ExcInfo(NULL, NULL, NULL)),
       edgecount(0),
@@ -255,7 +254,7 @@ ASTInterpreter::ASTInterpreter(FunctionMetadata* md, Box** vregs, int num_vregs,
 
     frame_info.vregs = vregs;
     frame_info.md = md;
-    frame_info.num_vregs = num_vregs;
+    frame_info.num_vregs = getVRegInfo().getNumOfCrossBlockVRegs();
 
     assert(scope_info);
 }
@@ -482,8 +481,8 @@ void ASTInterpreter::doStore(AST_Name* node, STOLEN(Value) value) {
         if (closure) {
             ASTInterpreterJitInterface::setLocalClosureHelper(this, node->vreg, name, value.o);
         } else {
-            assert(getSymVRegMap().count(name));
-            assert(getSymVRegMap()[name] == node->vreg);
+            assert(getVRegInfo().getVReg(node->id) == node->vreg);
+            frame_info.num_vregs = std::max(frame_info.num_vregs, node->vreg + 1);
             Box* prev = vregs[node->vreg];
             vregs[node->vreg] = value.o;
             Py_XDECREF(prev);
@@ -712,23 +711,16 @@ Box* ASTInterpreter::doOSR(AST_Jump* node) {
     std::unique_ptr<PhiAnalysis> phis
         = computeRequiredPhis(getMD()->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 (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);
+    llvm::SmallVector<int, 16> dead_vregs;
+    for (auto&& sym : getSymVRegMap()) {
+        if (!liveness->isLiveAtEnd(sym.first, current_block)) {
+            dead_vregs.push_back(sym.second);
+        } else if (phis->isRequiredAfter(sym.first, current_block)) {
+            assert(scope_info->getScopeTypeOfName(sym.first) != ScopeInfo::VarScopeType::GLOBAL);
         } else {
         }
     }
-    for (auto&& dead : dead_symbols) {
-        assert(getSymVRegMap().count(dead));
-        int vreg_num = getSymVRegMap()[dead];
+    for (auto&& vreg_num : dead_vregs) {
         Py_CLEAR(vregs[vreg_num]);
     }
 
@@ -747,8 +739,7 @@ Box* ASTInterpreter::doOSR(AST_Jump* node) {
     static Box* const VAL_UNDEFINED = (Box*)None;
 
     for (auto& name : phis->definedness.getDefinedNamesAtEnd(current_block)) {
-        assert(getSymVRegMap().count(name));
-        Box* val = vregs[getSymVRegMap()[name]];
+        Box* val = vregs[getVRegInfo().getVReg(name)];
         if (!liveness->isLiveAtEnd(name, current_block))
             continue;
 
@@ -1343,8 +1334,7 @@ Value ASTInterpreter::visit_delete(AST_Delete* node) {
                 } else {
                     assert(vst == ScopeInfo::VarScopeType::FAST);
 
-                    assert(getSymVRegMap().count(target->id));
-                    assert(getSymVRegMap()[target->id] == target->vreg);
+                    assert(getVRegInfo().getVReg(target->id) == target->vreg);
 
                     if (target->id.s()[0] == '#') {
                         assert(vregs[target->vreg] != NULL);
@@ -1358,6 +1348,7 @@ Value ASTInterpreter::visit_delete(AST_Delete* node) {
                         }
                     }
 
+                    frame_info.num_vregs = std::max(frame_info.num_vregs, target->vreg + 1);
                     Py_DECREF(vregs[target->vreg]);
                     vregs[target->vreg] = NULL;
                 }
@@ -1690,8 +1681,8 @@ Value ASTInterpreter::visit_name(AST_Name* node) {
             }
 
             assert(node->vreg >= 0);
-            assert(getSymVRegMap().count(node->id));
-            assert(getSymVRegMap()[node->id] == node->vreg);
+            assert(getVRegInfo().getVReg(node->id) == node->vreg);
+            frame_info.num_vregs = std::max(frame_info.num_vregs, node->vreg + 1);
             Box* val = vregs[node->vreg];
 
             if (val) {
@@ -1867,8 +1858,8 @@ void ASTInterpreterJitInterface::setExcInfoHelper(void* _interpreter, STOLEN(Box
 void ASTInterpreterJitInterface::setLocalClosureHelper(void* _interpreter, long vreg, InternedString id, Box* v) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
 
-    assert(interpreter->getSymVRegMap().count(id));
-    assert(interpreter->getSymVRegMap()[id] == vreg);
+    interpreter->frame_info.num_vregs = std::max(interpreter->frame_info.num_vregs, (int)vreg + 1);
+    assert(interpreter->getVRegInfo().getVReg(id) == vreg);
     Box* prev = interpreter->vregs[vreg];
     interpreter->vregs[vreg] = v;
     auto closure_offset = interpreter->scope_info->getClosureOffset(id);
@@ -1980,17 +1971,17 @@ Box* astInterpretFunction(FunctionMetadata* md, Box* closure, Box* generator, Bo
     // (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 = computeCFG(source_info, source_info->body, md->param_names);
 
     Box** vregs = NULL;
-    int num_vregs = md->calculateNumVRegs();
+    int num_vregs = source_info->cfg->getVRegInfo().getTotalNumOfVRegs();
     if (num_vregs > 0) {
         vregs = (Box**)alloca(sizeof(Box*) * num_vregs);
         memset(vregs, 0, sizeof(Box*) * num_vregs);
     }
 
     ++md->times_interpreted;
-    ASTInterpreter interpreter(md, vregs, num_vregs);
+    ASTInterpreter interpreter(md, vregs);
 
     ScopeInfo* scope_info = md->source->getScopeInfo();
 
@@ -2022,16 +2013,16 @@ Box* astInterpretFunctionEval(FunctionMetadata* md, Box* globals, Box* boxedLoca
     // that by looking at the cfg.)
     SourceInfo* source_info = md->source.get();
     if (!source_info->cfg)
-        source_info->cfg = computeCFG(source_info, source_info->body);
+        source_info->cfg = computeCFG(source_info, source_info->body, md->param_names);
 
     Box** vregs = NULL;
-    int num_vregs = md->calculateNumVRegs();
+    int num_vregs = source_info->cfg->getVRegInfo().getTotalNumOfVRegs();
     if (num_vregs > 0) {
         vregs = (Box**)alloca(sizeof(Box*) * num_vregs);
         memset(vregs, 0, sizeof(Box*) * num_vregs);
     }
 
-    ASTInterpreter interpreter(md, vregs, num_vregs);
+    ASTInterpreter interpreter(md, vregs);
     interpreter.initArguments(NULL, NULL, NULL, NULL, NULL, NULL);
     interpreter.setBoxedLocals(incref(boxedLocals));
 
@@ -2060,7 +2051,7 @@ extern "C" Box* astInterpretDeoptFromASM(FunctionMetadata* md, AST_expr* after_e
     // there wouldn't be enough space for the compiler generated ones which the interpreter (+bjit) stores inside the
     // vreg array.
     Box** vregs = NULL;
-    int num_vregs = md->calculateNumVRegs();
+    int num_vregs = source_info->cfg->getVRegInfo().getTotalNumOfVRegs();
     if (num_vregs > 0) {
         vregs = (Box**)alloca(sizeof(Box*) * num_vregs);
         memset(vregs, 0, sizeof(Box*) * num_vregs);
@@ -2071,7 +2062,7 @@ extern "C" Box* astInterpretDeoptFromASM(FunctionMetadata* md, AST_expr* after_e
     RELEASE_ASSERT(cur_thread_state.frame_info == frame_state.frame_info, "");
     cur_thread_state.frame_info = frame_state.frame_info->back;
 
-    ASTInterpreter interpreter(md, vregs, num_vregs, frame_state.frame_info);
+    ASTInterpreter interpreter(md, vregs, frame_state.frame_info);
 
     for (const auto& p : *frame_state.locals) {
         assert(p.first->cls == str_cls);

src/codegen/baseline_jit.cpp

@@ -634,7 +634,9 @@ void JitFragmentWriter::emitSetBlockLocal(InternedString s, int vreg, STOLEN(Rew
     if (LOG_BJIT_ASSEMBLY)
         comment("BJIT: emitSetBlockLocal() start");
     RewriterVar* prev = local_syms[s];
-    if (!prev)
+    // if we never set this sym before in this BB and the symbol gets accessed in several blocks clear it because it
+    // could have been set in a previous block.
+    if (!prev && !block->cfg->getVRegInfo().isBlockLocalVReg(vreg))
         emitSetLocal(s, vreg, false, imm(nullptr)); // clear out the vreg
     local_syms[s] = v;
     if (LOG_BJIT_ASSEMBLY)
@@ -690,6 +692,7 @@ void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closur
         // but I suspect is not that big a deal as long as the llvm jit implements this kind of optimization.
         bool prev_nullable = true;
 
+        assert(!block->cfg->getVRegInfo().isBlockLocalVReg(vreg));
         vregs_array->replaceAttr(8 * vreg, v, prev_nullable);
     }
     if (LOG_BJIT_ASSEMBLY)

src/codegen/codegen.cpp

@@ -73,34 +73,6 @@ BORROWED(BoxedCode*) FunctionMetadata::getCode() {
     return code_obj;
 }
 
-int FunctionMetadata::calculateNumVRegs() {
-    SourceInfo* source_info = source.get();
-
-    CFG* cfg = source_info->cfg;
-    assert(cfg && "We don't calculate the CFG inside this function because it can raise an exception and its "
-                  "therefore not safe to call at every point");
-
-    if (!cfg->hasVregsAssigned()) {
-        ScopeInfo* scope_info = source->getScopeInfo();
-        cfg->assignVRegs(param_names, scope_info);
-    }
-    return cfg->sym_vreg_map.size();
-}
-
-int FunctionMetadata::calculateNumUserVisibleVRegs() {
-    SourceInfo* source_info = source.get();
-
-    CFG* cfg = source_info->cfg;
-    assert(cfg && "We don't calculate the CFG inside this function because it can raise an exception and its "
-                  "therefore not safe to call at every point");
-
-    if (!cfg->hasVregsAssigned()) {
-        ScopeInfo* scope_info = source->getScopeInfo();
-        cfg->assignVRegs(param_names, scope_info);
-    }
-    return cfg->sym_vreg_map_user_visible.size();
-}
-
 void FunctionMetadata::addVersion(CompiledFunction* compiled) {
     assert(compiled);
     assert((compiled->spec != NULL) + (compiled->entry_descriptor != NULL) == 1);

src/codegen/irgen.cpp

@@ -1009,8 +1009,6 @@ CompiledFunction* doCompile(FunctionMetadata* md, SourceInfo* source, ParamNames
 
     assert((entry_descriptor != NULL) + (spec != NULL) == 1);
 
-    md->calculateNumVRegs();
-
     if (VERBOSITY("irgen") >= 2)
         source->cfg->print();
 

src/codegen/irgen/hooks.cpp

@@ -271,7 +271,7 @@ CompiledFunction* compileFunction(FunctionMetadata* f, FunctionSpecialization* s
 
     // Do the analysis now if we had deferred it earlier:
     if (source->cfg == NULL) {
-        source->cfg = computeCFG(source, source->body);
+        source->cfg = computeCFG(source, source->body, f->param_names);
     }
 
 

src/codegen/irgen/irgenerator.cpp

@@ -244,7 +244,7 @@ void IRGenState::setupFrameInfoVar(llvm::Value* passed_closure, llvm::Value* pas
         assert(al->isStaticAlloca());
 
         assert(!vregs);
-        int num_user_visible_vregs = getMD()->calculateNumUserVisibleVRegs();
+        int num_user_visible_vregs = getSourceInfo()->cfg->getVRegInfo().getNumOfUserVisibleVRegs();
         if (num_user_visible_vregs > 0) {
             auto* vregs_alloca
                 = builder.CreateAlloca(g.llvm_value_type_ptr, getConstantInt(num_user_visible_vregs), "vregs");
@@ -1896,7 +1896,7 @@ private:
         auto cfg = irstate->getSourceInfo()->cfg;
         assert(vreg >= 0);
 
-        if (vreg < cfg->sym_vreg_map_user_visible.size()) {
+        if (cfg->getVRegInfo().isUserVisibleVReg(vreg)) {
             // looks like this store don't have to be volatile because llvm knows that the vregs are visible thru the
             // FrameInfo which escapes.
             auto* gep = emitter.getBuilder()->CreateConstInBoundsGEP1_64(irstate->getVRegsVar(), vreg);
@@ -2617,8 +2617,7 @@ private:
         auto vst = irstate->getScopeInfo()->getScopeTypeOfName(name);
         int vreg = -1;
         if (vst == ScopeInfo::VarScopeType::FAST || vst == ScopeInfo::VarScopeType::CLOSURE) {
-            assert(cfg->sym_vreg_map.count(name));
-            vreg = cfg->sym_vreg_map[name];
+            vreg = cfg->getVRegInfo().getVReg(name);
         }
         _doSet(vreg, name, vst, var, unw_info);
     }

src/codegen/unwinding.cpp

@@ -892,10 +892,9 @@ DeoptState getDeoptState() {
             // and assigning them to the new vregs array...
             // But deopts are so rare it's not really worth it.
             Box** vregs = frame_iter->getFrameInfo()->vregs;
-            for (const auto& p : cf->md->source->cfg->sym_vreg_map_user_visible) {
+            for (const auto& p : cf->md->source->cfg->getVRegInfo().getUserVisibleSymVRegMap()) {
                 if (is_undefined.count(p.first.s()))
                     continue;
-                assert(p.second >= 0 && p.second < cf->md->source->cfg->sym_vreg_map_user_visible.size());
 
                 Box* v = vregs[p.second];
                 if (!v)
@@ -950,8 +949,8 @@ BORROWED(Box*) fastLocalsToBoxedLocals() {
 
 static BoxedDict* localsForFrame(Box** vregs, CFG* cfg) {
     BoxedDict* rtn = new BoxedDict();
-    rtn->d.grow(cfg->sym_vreg_map_user_visible.size());
-    for (auto& l : cfg->sym_vreg_map_user_visible) {
+    rtn->d.grow(cfg->getVRegInfo().getNumOfUserVisibleVRegs());
+    for (auto& l : cfg->getVRegInfo().getUserVisibleSymVRegMap()) {
         Box* val = vregs[l.second];
         if (val) {
             assert(!rtn->d.count(l.first.getBox()));

src/core/cfg.cpp

@@ -181,7 +181,7 @@ private:
 
     unsigned int next_var_index = 0;
 
-    friend CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body);
+    friend CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body, const ParamNames& param_names);
 
 public:
     CFGVisitor(SourceInfo* source, AST_TYPE::AST_TYPE root_type, FutureFlags future_flags,
@@ -2556,13 +2556,15 @@ void CFG::print(llvm::raw_ostream& stream) {
 
 class AssignVRegsVisitor : public NoopASTVisitor {
 public:
-    int index = 0;
-    bool only_user_visible;
-    llvm::DenseMap<InternedString, int> sym_vreg_map;
     ScopeInfo* scope_info;
+    CFGBlock* current_block;
+    int next_vreg;
+    llvm::DenseMap<InternedString, int> sym_vreg_map;
+    llvm::DenseMap<InternedString, std::unordered_set<CFGBlock*>> sym_blocks_map;
+
+    enum Step { TrackBlockUsage = 0, UserVisible, CrossBlock, SingleBlockUse } step;
 
-    AssignVRegsVisitor(ScopeInfo* scope_info, bool only_user_visible)
-        : only_user_visible(only_user_visible), scope_info(scope_info) {}
+    AssignVRegsVisitor(ScopeInfo* scope_info) : scope_info(scope_info), current_block(0), next_vreg(0) {}
 
     bool visit_arguments(AST_arguments* node) override {
         for (AST_expr* d : node->defaults)
@@ -2590,66 +2592,94 @@ public:
         return true;
     }
 
+
+    bool isNameUsedInSingleBlock(InternedString id) {
+        assert(step != TrackBlockUsage);
+        assert(sym_blocks_map.count(id));
+        return sym_blocks_map[id].size() == 1;
+    }
+
     bool visit_name(AST_Name* node) override {
         if (node->vreg != -1)
             return true;
 
-        if (only_user_visible && node->id.isCompilerCreatedName())
-            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);
+        if (node->lookup_type != ScopeInfo::VarScopeType::FAST && node->lookup_type != ScopeInfo::VarScopeType::CLOSURE)
+            return true;
+
+        if (step == TrackBlockUsage) {
+            sym_blocks_map[node->id].insert(current_block);
+            return true;
+        } else if (step == UserVisible) {
+            if (node->id.isCompilerCreatedName())
+                return true;
+        } else {
+            bool is_block_local = node->lookup_type == ScopeInfo::VarScopeType::FAST
+                                  && isNameUsedInSingleBlock(node->id);
+            if (step == CrossBlock && is_block_local)
+                return true;
+            if (step == SingleBlockUse && !is_block_local)
+                return true;
+        }
+        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++;
+            sym_vreg_map[id] = next_vreg;
+            return next_vreg++;
         }
         return it->second;
     }
 };
 
-void CFG::assignVRegs(const ParamNames& param_names, ScopeInfo* scope_info) {
-    if (has_vregs_assigned)
-        return;
+void VRegInfo::assignVRegs(CFG* cfg, const ParamNames& param_names, ScopeInfo* scope_info) {
+    assert(!hasVRegsAssigned());
+
+    // warning: don't rearrange the steps, they need to be run in this exact order!
+    AssignVRegsVisitor visitor(scope_info);
+    for (auto step : { AssignVRegsVisitor::TrackBlockUsage, AssignVRegsVisitor::UserVisible,
+                       AssignVRegsVisitor::CrossBlock, AssignVRegsVisitor::SingleBlockUse }) {
+        visitor.step = step;
+        for (CFGBlock* b : cfg->blocks) {
+            visitor.current_block = b;
+
+            if (step == AssignVRegsVisitor::SingleBlockUse)
+                visitor.next_vreg = num_vregs_cross_block;
+
+            if (b == cfg->getStartingBlock()) {
+                for (auto* name : param_names.arg_names) {
+                    name->accept(&visitor);
+                }
+                if (param_names.vararg_name)
+                    param_names.vararg_name->accept(&visitor);
 
-    AssignVRegsVisitor visitor(scope_info, true);
+                if (param_names.kwarg_name)
+                    param_names.kwarg_name->accept(&visitor);
+            }
 
-    // we need todo two passes: first we assign the user visible vars a vreg and then the compiler created get there
-    // value.
-    for (int i = 0; i < 2; ++i) {
-        for (CFGBlock* b : blocks) {
             for (AST_stmt* stmt : b->body) {
                 stmt->accept(&visitor);
             }
-        }
 
-        for (auto* name : param_names.arg_names) {
-            name->accept(&visitor);
+            if (step == AssignVRegsVisitor::SingleBlockUse)
+                num_vregs = std::max(num_vregs, visitor.next_vreg);
         }
 
-        if (param_names.vararg_name)
-            param_names.vararg_name->accept(&visitor);
-
-        if (param_names.kwarg_name)
-            param_names.kwarg_name->accept(&visitor);
-
-        if (visitor.only_user_visible) {
-            visitor.only_user_visible = false;
+        if (step == AssignVRegsVisitor::UserVisible)
             sym_vreg_map_user_visible = visitor.sym_vreg_map;
-        }
+        else if (step == AssignVRegsVisitor::CrossBlock)
+            num_vregs = num_vregs_cross_block = visitor.next_vreg;
     }
     sym_vreg_map = std::move(visitor.sym_vreg_map);
-    has_vregs_assigned = true;
+    assert(hasVRegsAssigned());
 }
 
-CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
+CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body, const ParamNames& param_names) {
     STAT_TIMER(t0, "us_timer_computecfg", 0);
 
     CFG* rtn = new CFG();
@@ -2890,6 +2920,7 @@ CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
         rtn->print();
     }
 
+    rtn->getVRegInfo().assignVRegs(rtn, param_names, source->getScopeInfo());
 
     return rtn;
 }

src/core/cfg.h

@@ -43,10 +43,9 @@ class ParamNames;
 class ScopeInfo;
 
 class CFGBlock {
-private:
+public:
     CFG* cfg;
 
-public:
     // Baseline JIT helper fields:
     // contains address to the start of the code of this basic block
     void* code;
@@ -70,23 +69,68 @@ public:
     void _print() { print(); }
 };
 
+// the vregs are split into three parts.
+// user visible: used for all non compiler generated names, name could be used in a single block or multiple
+//               all frames contain atleast this vregs in order to do frame introspection
+// cross block : used for compiler generated names which get used in several blocks or which have closure scope
+// single block: used by compiler created names which are only used in a single block.
+//               get reused for different names
+//
+// we assign the lowest numbers to the user visible ones, followed by the cross block ones and finally the single block
+// ones. we do this because not all tiers use all of the vregs and it still makes it fast to switch between tiers.
+//
+// usage by our different tiers:
+// interpreter : [user visible] [cross block] [single block]
+// baseline jit: [user visible] [cross block]
+// llvm jit    : [user visible]
+class VRegInfo {
+private:
+    llvm::DenseMap<InternedString, int> sym_vreg_map_user_visible;
+    llvm::DenseMap<InternedString, int> sym_vreg_map;
+    int num_vregs_cross_block = -1;
+    int num_vregs = -1;
+
+public:
+    // map of all assigned names. if the name is block local the vreg number is not unique because this vregs get reused
+    // between blocks.
+    const llvm::DenseMap<InternedString, int>& getSymVRegMap() { return sym_vreg_map; }
+    const llvm::DenseMap<InternedString, int>& getUserVisibleSymVRegMap() { return sym_vreg_map_user_visible; }
+
+    int getVReg(InternedString name) const {
+        assert(hasVRegsAssigned());
+        assert(sym_vreg_map.count(name));
+        auto it = sym_vreg_map.find(name);
+        assert(it != sym_vreg_map.end());
+        assert(it->second != -1);
+        return it->second;
+    }
+
+    bool isUserVisibleVReg(int vreg) const { return vreg < sym_vreg_map_user_visible.size(); }
+    bool isCrossBlockVReg(int vreg) const { return !isUserVisibleVReg(vreg) && vreg < num_vregs_cross_block; }
+    bool isBlockLocalVReg(int vreg) const { return vreg >= num_vregs_cross_block; }
+
+    int getTotalNumOfVRegs() const { return num_vregs; }
+    int getNumOfUserVisibleVRegs() const { return sym_vreg_map_user_visible.size(); }
+    int getNumOfCrossBlockVRegs() const { return num_vregs_cross_block; }
+
+    bool hasVRegsAssigned() const { return num_vregs != -1; }
+    void assignVRegs(CFG* cfg, const ParamNames& param_names, ScopeInfo* scope_info);
+};
+
 // Control Flow Graph
 class CFG {
 private:
     int next_idx;
-    bool has_vregs_assigned;
+    VRegInfo vreg_info;
 
 public:
     std::vector<CFGBlock*> blocks;
 
-    // Contains the vreg assignment for every name including the user visible ones
-    // (which will have lower ids than the compiler generated ones).
-    llvm::DenseMap<InternedString, int> sym_vreg_map;
-    llvm::DenseMap<InternedString, int> sym_vreg_map_user_visible;
-
-    CFG() : next_idx(0), has_vregs_assigned(false) {}
+public:
+    CFG() : next_idx(0) {}
 
     CFGBlock* getStartingBlock() { return blocks[0]; }
+    VRegInfo& getVRegInfo() { return vreg_info; }
 
     CFGBlock* addBlock() {
         int idx = next_idx;
@@ -113,13 +157,10 @@ public:
     }
 
     void print(llvm::raw_ostream& stream = llvm::outs());
-
-    bool hasVregsAssigned() { return has_vregs_assigned; }
-    void assignVRegs(const ParamNames& param_names, ScopeInfo* scope_info);
 };
 
 class SourceInfo;
-CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body);
+CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body, const ParamNames& param_names);
 void printCFG(CFG* cfg);
 }
 

src/core/types.h

@@ -508,9 +508,6 @@ public:
     void addVersion(void* f, ConcreteCompilerType* rtn_type, const std::vector<ConcreteCompilerType*>& arg_types,
                     ExceptionStyle exception_style = CXX);
 
-    int calculateNumVRegs();
-    int calculateNumUserVisibleVRegs();
-
     // Helper function, meant for the C++ runtime, which allocates a FunctionMetadata object and calls addVersion
     // once to it.
     static FunctionMetadata* create(void* f, ConcreteCompilerType* rtn_type, int nargs, bool takes_varargs,

test/unittests/analysis.cpp

@@ -43,7 +43,8 @@ TEST_F(AnalysisTest, augassign) {
     SourceInfo* si = new SourceInfo(createModule(boxString("augassign"), fn.c_str()), scoping, future_flags, func,
             func->body, boxString(fn));
 
-    CFG* cfg = computeCFG(si, func->body);
+    ParamNames param_names(si->ast, si->getInternedStrings());
+    CFG* cfg = computeCFG(si, func->body, param_names);
     std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg);
 
     //cfg->print();
@@ -74,7 +75,7 @@ void doOsrTest(bool is_osr, bool i_maybe_undefined) {
                     fn.c_str()), scoping, future_flags, func, func->body, boxString(fn)));
     FunctionMetadata* clfunc = new FunctionMetadata(0, false, false, std::move(si));
 
-    CFG* cfg = computeCFG(clfunc->source.get(), func->body);
+    CFG* cfg = computeCFG(clfunc->source.get(), func->body, clfunc->param_names);
     std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg);
 
     // cfg->print();