Merge pull request #1239 from undingen/bjit_opt2_with_vreg_reuse

vregs: reuse block local vregs

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();