VRegInfo: reduce memory usage

src/codegen/ast_interpreter.cpp

@@ -155,9 +155,12 @@ public:
     ~ASTInterpreter() { Py_XDECREF(this->created_closure); }
 
     const VRegInfo& getVRegInfo() const { return source_info->cfg->getVRegInfo(); }
+
+#ifndef NDEBUG
     const llvm::DenseMap<InternedString, DefaultedInt<-1>>& getSymVRegMap() const {
-        return source_info->cfg->getVRegInfo().getSymVRegMap();
+        return getVRegInfo().getSymVRegMap();
     }
+#endif
 
     AST_stmt* getCurrentStatement() {
         assert(frame_info.stmt);
@@ -755,12 +758,13 @@ Box* ASTInterpreter::doOSR(AST_Jump* node) {
         = computeRequiredPhis(getMD()->param_names, source_info->cfg, liveness, scope_info);
 
     llvm::SmallVector<int, 16> dead_vregs;
-    for (auto&& sym : getSymVRegMap()) {
-        if (!liveness->isLiveAtEnd(sym.second, current_block)) {
-            dead_vregs.push_back(sym.second);
-        } else if (phis->isRequiredAfter(sym.second, current_block)) {
-            assert(scope_info->getScopeTypeOfName(sym.first) != ScopeInfo::VarScopeType::GLOBAL);
-        } else {
+
+    for (int vreg = 0; vreg < getVRegInfo().getTotalNumOfVRegs(); ++vreg) {
+        if (!liveness->isLiveAtEnd(vreg, current_block)) {
+            dead_vregs.push_back(vreg);
+        } else if (phis->isRequiredAfter(vreg, current_block)) {
+            assert(!getVRegInfo().vregHasName(vreg)
+                   || scope_info->getScopeTypeOfName(getVRegInfo().getName(vreg)) != ScopeInfo::VarScopeType::GLOBAL);
         }
     }
     for (auto&& vreg_num : dead_vregs) {

src/codegen/irgen/irgenerator.cpp

@@ -2663,18 +2663,6 @@ private:
         }
     }
 
-    void loadArgument(InternedString name, ConcreteCompilerType* t, llvm::Value* v, const UnwindInfo& unw_info) {
-        assert(name.s() != FRAME_INFO_PTR_NAME);
-        CompilerVariable* var = unboxVar(t, v);
-        auto cfg = irstate->getSourceInfo()->cfg;
-        auto vst = irstate->getScopeInfo()->getScopeTypeOfName(name);
-        int vreg = -1;
-        if (vst == ScopeInfo::VarScopeType::FAST || vst == ScopeInfo::VarScopeType::CLOSURE) {
-            vreg = cfg->getVRegInfo().getVReg(name);
-        }
-        _doSet(vreg, name, vst, var, unw_info);
-    }
-
     void loadArgument(AST_expr* name, ConcreteCompilerType* t, llvm::Value* v, const UnwindInfo& unw_info) {
         CompilerVariable* var = unboxVar(t, v);
         _doSet(name, var, unw_info);
@@ -2882,6 +2870,7 @@ public:
         setDefinedVar(vreg, val);
     }
 
+#ifndef NDEBUG
     void giveLocalSymbol(InternedString name, CompilerVariable* var) override {
         assert(name.s() != "None");
         assert(name.s() != FRAME_INFO_PTR_NAME);
@@ -2895,6 +2884,7 @@ public:
         int vreg = irstate->getSourceInfo()->cfg->getVRegInfo().getVReg(name);
         giveLocalSymbol(vreg, var);
     }
+#endif
 
     void giveLocalSymbol(int vreg, CompilerVariable* var) override {
         assert(var->getType()->isUsable());
@@ -2994,13 +2984,11 @@ public:
 
         int i = 0;
         for (; i < param_names.args.size(); i++) {
-            loadArgument(internString(param_names.args[i]), arg_types[i], python_parameters[i],
-                         UnwindInfo::cantUnwind());
+            loadArgument(param_names.arg_names[i], arg_types[i], python_parameters[i], UnwindInfo::cantUnwind());
         }
 
         if (param_names.vararg.size()) {
-            loadArgument(internString(param_names.vararg), arg_types[i], python_parameters[i],
-                         UnwindInfo::cantUnwind());
+            loadArgument(param_names.vararg_name, arg_types[i], python_parameters[i], UnwindInfo::cantUnwind());
             i++;
         }
 
@@ -3030,7 +3018,7 @@ public:
             emitter.refConsumed(passed_dict, null_check);
             emitter.refConsumed(created_dict, isnull_terminator);
 
-            loadArgument(internString(param_names.kwarg), arg_types[i], phi, UnwindInfo::cantUnwind());
+            loadArgument(param_names.kwarg_name, arg_types[i], phi, UnwindInfo::cantUnwind());
             i++;
         }
 

src/codegen/irgen/irgenerator.h

@@ -179,7 +179,9 @@ public:
     virtual void doFunctionEntry(const ParamNames& param_names, const std::vector<ConcreteCompilerType*>& arg_types)
         = 0;
 
+#ifndef NDEBUG
     virtual void giveLocalSymbol(InternedString name, CompilerVariable* var) = 0;
+#endif
     virtual void giveLocalSymbol(int vreg, CompilerVariable* var) = 0;
     virtual void giveDefinednessVar(int vreg, llvm::Value* val) = 0;
     virtual void copySymbolsFrom(SymbolTable* st) = 0;

src/codegen/unwinding.cpp

@@ -890,15 +890,16 @@ 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->getVRegInfo().getUserVisibleSymVRegMap()) {
-                if (is_undefined.count(p.second))
+            int num_vregs_user_visible = cf->md->source->cfg->getVRegInfo().getNumOfUserVisibleVRegs();
+            for (int vreg = 0; vreg < num_vregs_user_visible; ++vreg) {
+                if (is_undefined.count(vreg))
                     assert(0);
 
-                Box* v = vregs[p.second];
+                Box* v = vregs[vreg];
                 if (!v)
                     continue;
 
-                d->d[boxInt(p.second)] = incref(v);
+                d->d[boxInt(vreg)] = incref(v);
             }
 
             for (const auto& p : cf->location_map->vars) {
@@ -930,12 +931,14 @@ BORROWED(Box*) fastLocalsToBoxedLocals() {
 
 static BoxedDict* localsForFrame(Box** vregs, CFG* cfg) {
     BoxedDict* rtn = new BoxedDict();
-    rtn->d.grow(cfg->getVRegInfo().getNumOfUserVisibleVRegs());
-    for (auto& l : cfg->getVRegInfo().getUserVisibleSymVRegMap()) {
-        Box* val = vregs[l.second];
+    auto vregs_sym_map = cfg->getVRegInfo().getVRegSymUserVisibleMap();
+    int num_user_visible_vregs = vregs_sym_map.size();
+    rtn->d.grow(num_user_visible_vregs);
+    for (int vreg = 0; vreg < num_user_visible_vregs; ++vreg) {
+        Box* val = vregs[vreg];
         if (val) {
-            assert(!rtn->d.count(l.first.getBox()));
-            rtn->d[incref(l.first.getBox())] = incref(val);
+            assert(!rtn->d.count(vregs_sym_map[vreg].getBox()));
+            rtn->d[incref(vregs_sym_map[vreg].getBox())] = incref(val);
         }
     }
     return rtn;

src/core/cfg.cpp

@@ -2847,12 +2847,17 @@ void VRegInfo::assignVRegs(CFG* cfg, const ParamNames& param_names, ScopeInfo* s
                 num_vregs = std::max(num_vregs, visitor.next_vreg);
         }
 
-        if (step == AssignVRegsVisitor::UserVisible)
+        if (step == AssignVRegsVisitor::UserVisible) {
+            num_vregs_user_visible = visitor.sym_vreg_map.size();
+#ifndef NDEBUG
             sym_vreg_map_user_visible = visitor.sym_vreg_map;
-        else if (step == AssignVRegsVisitor::CrossBlock)
+#endif
+        } else if (step == AssignVRegsVisitor::CrossBlock)
             num_vregs = num_vregs_cross_block = visitor.next_vreg;
     }
+#ifndef NDEBUG
     sym_vreg_map = std::move(visitor.sym_vreg_map);
+#endif
     vreg_sym_map = std::move(visitor.vreg_sym_map);
     assert(hasVRegsAssigned());
 #if REUSE_VREGS

src/core/cfg.h

@@ -29,6 +29,7 @@
 
 #include <vector>
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 
 #include "core/ast.h"
@@ -107,21 +108,26 @@ public:
 // llvm jit    : [user visible]
 class VRegInfo {
 private:
+#ifndef NDEBUG
+    // this maps use too much memory, we only use them in the debug build for asserts
     llvm::DenseMap<InternedString, DefaultedInt<-1>> sym_vreg_map_user_visible;
     llvm::DenseMap<InternedString, DefaultedInt<-1>> sym_vreg_map;
+#endif
 
     // Reverse map, from vreg->symbol name.
-    // Entries won't exist for all vregs
+    // Entries won't exist for all vregs (=no entries for reused vregs)
     std::vector<InternedString> vreg_sym_map;
 
     int num_vregs_cross_block = -1;
+    int num_vregs_user_visible = -1;
     int num_vregs = -1;
 
 public:
+#ifndef NDEBUG
     // 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, DefaultedInt<-1>>& getSymVRegMap() { return sym_vreg_map; }
-    const llvm::DenseMap<InternedString, DefaultedInt<-1>>& getUserVisibleSymVRegMap() {
+    const llvm::DenseMap<InternedString, DefaultedInt<-1>>& getSymVRegMap() const { return sym_vreg_map; }
+    const llvm::DenseMap<InternedString, DefaultedInt<-1>>& getUserVisibleSymVRegMap() const {
         return sym_vreg_map_user_visible;
     }
 
@@ -133,6 +139,11 @@ public:
         assert(it->second != -1);
         return it->second;
     }
+#endif
+
+    llvm::ArrayRef<InternedString> getVRegSymUserVisibleMap() const {
+        return llvm::makeArrayRef(vreg_sym_map).slice(0, num_vregs_user_visible);
+    }
 
     // Not all vregs correspond to a name; many are our compiler-generated variables.
     bool vregHasName(int vreg) const { return vreg < num_vregs_cross_block; }
@@ -149,12 +160,12 @@ public:
         return vreg_sym_map[vreg];
     }
 
-    bool isUserVisibleVReg(int vreg) const { return vreg < sym_vreg_map_user_visible.size(); }
+    bool isUserVisibleVReg(int vreg) const { return vreg < num_vregs_user_visible; }
     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 getNumOfUserVisibleVRegs() const { return num_vregs_user_visible; }
     int getNumOfCrossBlockVRegs() const { return num_vregs_cross_block; }
 
     bool hasVRegsAssigned() const { return num_vregs != -1; }

test/unittests/analysis.cpp

@@ -24,6 +24,8 @@ protected:
     }
 };
 
+// this test use functions (VRegInfo::getVReg) which are only available in a debug build
+#ifndef NDEBUG
 TEST_F(AnalysisTest, augassign) {
     const std::string fn("test/unittests/analysis_listcomp.py");
     AST_Module* module = caching_parse_file(fn.c_str(), 0);
@@ -146,3 +148,4 @@ TEST_F(AnalysisTest, osr1) {
 TEST_F(AnalysisTest, osr2) {
     doOsrTest(true, true);
 }
+#endif