OSR: enable type specialization

src/codegen/ast_interpreter.cpp

@@ -662,24 +662,13 @@ Box* ASTInterpreter::doOSR(BST_Jump* node) {
     LivenessAnalysis* liveness = getLiveness();
     std::unique_ptr<PhiAnalysis> phis = computeRequiredPhis(getCode()->param_names, source_info->cfg, liveness);
 
-    llvm::SmallVector<int, 16> dead_vregs;
-
     for (int vreg = 0; vreg < getVRegInfo().getTotalNumOfVRegs(); ++vreg) {
         if (!liveness->isLiveAtEnd(vreg, current_block)) {
-            dead_vregs.push_back(vreg);
+            Py_CLEAR(vregs[vreg]);
         }
     }
-    for (auto&& vreg_num : dead_vregs) {
-        Py_CLEAR(vregs[vreg_num]);
-    }
 
-    const OSREntryDescriptor* found_entry = nullptr;
-    for (auto& p : getCode()->osr_versions) {
-        if (p.first->backedge != node)
-            continue;
 
-        found_entry = p.first;
-    }
 
     int num_vregs = source_info->cfg->getVRegInfo().getTotalNumOfVRegs();
     VRegMap<Box*> sorted_symbol_table(num_vregs);
@@ -718,12 +707,30 @@ Box* ASTInterpreter::doOSR(BST_Jump* node) {
         return nullptr;
     }
 
+    const OSREntryDescriptor* found_entry = nullptr;
+    for (auto& f : getCode()->osr_versions) {
+        if (f->entry_descriptor->backedge != node)
+            continue;
+
+        bool same_types = true;
+        for (auto&& p : sorted_symbol_table) {
+            same_types = f->entry_descriptor->args[p.first] == UNKNOWN
+                         || f->entry_descriptor->args[p.first] == typeFromClass(p.second->cls);
+            if (!same_types)
+                break;
+        }
+        if (same_types)
+            found_entry = f->entry_descriptor;
+    }
+
     if (found_entry == nullptr) {
         OSREntryDescriptor* entry = OSREntryDescriptor::create(getCode(), node, CXX);
 
-        // TODO can we just get rid of this?
+        // for the first OSR we try to create a type specialized version.
+        bool create_type_specialized_osr = getCode()->osr_versions.empty();
+
         for (auto&& p : sorted_symbol_table) {
-            entry->args[p.first] = UNKNOWN;
+            entry->args[p.first] = create_type_specialized_osr ? typeFromClass(p.second->cls) : UNKNOWN;
         }
 
         entry->potentially_undefined = potentially_undefined;
@@ -733,7 +740,7 @@ Box* ASTInterpreter::doOSR(BST_Jump* node) {
 
     OSRExit exit(found_entry);
 
-    std::vector<Box*> arg_array;
+    llvm::SmallVector<Box*, 8> arg_array;
     arg_array.reserve(sorted_symbol_table.numSet() + potentially_undefined.numSet());
     for (auto&& p : sorted_symbol_table) {
         arg_array.push_back(p.second);
@@ -746,8 +753,8 @@ Box* ASTInterpreter::doOSR(BST_Jump* node) {
     UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_jitted_code");
     CompiledFunction* partial_func = compilePartialFuncInternal(&exit);
 
-    // generated is only borrowed in order to not introduce cycles
-    Box* r = partial_func->call_osr(generator, created_closure, &frame_info, &arg_array[0]);
+    // generator is only borrowed in order to not introduce cycles
+    Box* r = partial_func->call_osr(generator, created_closure, &frame_info, arg_array.data());
 
     if (partial_func->exception_style == CXX) {
         assert(r);

src/codegen/codegen.cpp

@@ -54,7 +54,7 @@ void BoxedCode::addVersion(CompiledFunction* compiled) {
         assert(compiled->spec->arg_types.size() == numReceivedArgs());
         versions.push_back(compiled);
     } else {
-        osr_versions.emplace_front(compiled->entry_descriptor, compiled);
+        osr_versions.push_back(compiled);
     }
 }
 

src/codegen/irgen.cpp

@@ -493,7 +493,6 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
                 // good to go
                 v = from_arg;
             } else if (p.second->canConvertTo(phi_type)) {
-                RELEASE_ASSERT(0, "this hasn't been hit in a very long time -- check refcounting");
                 // not sure if/when this happens, but if there's a type mismatch but one we know
                 // can be handled (such as casting from a subclass to a superclass), handle it:
                 ConcreteCompilerVariable* converted = var->makeConverted(*unbox_emitter, phi_type);

src/codegen/irgen/hooks.cpp

@@ -539,14 +539,12 @@ void CompiledFunction::speculationFailed() {
         }
 
         if (!found) {
-            code->osr_versions.remove_if([&](const std::pair<const OSREntryDescriptor*, CompiledFunction*>& e) {
-                if (e.second == this) {
-                    this->dependent_callsites.invalidateAll();
-                    found = true;
-                    return true;
-                }
-                return false;
-            });
+            auto it = std::find(code->osr_versions.begin(), code->osr_versions.end(), this);
+            if (it != code->osr_versions.end()) {
+                code->osr_versions.erase(it);
+                this->dependent_callsites.invalidateAll();
+                found = true;
+            }
         }
 
         if (!found) {
@@ -558,7 +556,6 @@ void CompiledFunction::speculationFailed() {
     }
 }
 
-std::unordered_set<CompiledFunction*> all_compiled_functions;
 CompiledFunction::CompiledFunction(BoxedCode* code_obj, FunctionSpecialization* spec, void* code, EffortLevel effort,
                                    ExceptionStyle exception_style, const OSREntryDescriptor* entry_descriptor)
     : code_obj(code_obj),
@@ -629,16 +626,15 @@ CompiledFunction* compilePartialFuncInternal(OSRExit* exit) {
     BoxedCode* code = exit->entry->code;
     assert(code);
     for (auto&& osr_functions : code->osr_versions) {
-        if (osr_functions.first == exit->entry)
-            return osr_functions.second;
+        if (osr_functions->entry_descriptor == exit->entry)
+            return osr_functions;
     }
 
     EffortLevel new_effort = EffortLevel::MAXIMAL;
     CompiledFunction* compiled
         = compileFunction(code, NULL, new_effort, exit->entry, true, exit->entry->exception_style);
     stat_osr_compiles.log();
-    assert(std::find(code->osr_versions.begin(), code->osr_versions.end(), std::make_pair(exit->entry, compiled))
-           != code->osr_versions.end());
+    assert(std::find(code->osr_versions.begin(), code->osr_versions.end(), compiled) != code->osr_versions.end());
     return compiled;
 }
 

src/runtime/types.h

@@ -1143,7 +1143,7 @@ class BoxedCode : public Box {
         versions; // any compiled versions along with their type parameters; in order from most preferred to least
     ExceptionSwitchable<CompiledFunction*>
         always_use_version; // if this version is set, always use it (for unboxed cases)
-    std::forward_list<std::pair<const OSREntryDescriptor*, CompiledFunction*>> osr_versions;
+    llvm::TinyPtrVector<CompiledFunction*> osr_versions;
 
     // Profiling counter:
     int propagated_cxx_exceptions = 0;