wip some fixes for block local opt

Author: undingen

src/asm_writing/rewriter.cpp

@@ -1231,7 +1231,6 @@ std::vector<Location> Rewriter::getDecrefLocations() {
     for (RewriterVar& var : vars) {
         if (var.locations.size() && var.needsDecref()) {
             bool found_location = false;
-
             for (Location l : var.locations) {
                 if (l.type == Location::Scratch) {
                     // convert to stack based location because later on we may not know the offset of the scratch area
@@ -1249,7 +1248,10 @@ std::vector<Location> Rewriter::getDecrefLocations() {
                 } else
                     RELEASE_ASSERT(0, "not implemented");
             }
-            RELEASE_ASSERT(found_location, "");
+            if (!found_location) {
+                assert(0);
+                failed = true;
+            }
         }
     }
 
@@ -1353,7 +1355,22 @@ void RewriterVar::refUsed() {
 }
 
 bool RewriterVar::needsDecref() {
-    return reftype == RefType::OWNED && !this->refHandedOff();
+    if (reftype == RefType::OWNED && !this->refHandedOff())
+        return true;
+    if (reftype == RefType::OWNED && this->num_refs_consumed > 0) {
+        if (uses[0] == rewriter->current_action_idx)
+            return false;
+        if (uses[this->last_refconsumed_numuses - 1] > rewriter->current_action_idx) {
+            return true;
+        }
+        return false;
+    }
+    if (reftype == RefType::OWNED) {
+        if (uses[0] == rewriter->current_action_idx)
+            return false;
+        return true;
+    }
+    return false;
 }
 
 void RewriterVar::registerOwnedAttr(int byte_offset) {
@@ -1506,6 +1523,7 @@ void Rewriter::commit() {
             _incref(var, 1);
         }
 
+        current_action_idx = i;
         actions[i].action();
 
         if (failed) {

src/asm_writing/rewriter.h

@@ -566,6 +566,7 @@ class Rewriter : public ICSlotRewrite::CommitHook {
     }
 
     llvm::ArrayRef<assembler::Register> allocatable_regs;
+    int current_action_idx = -1;
 
 public:
     // This should be called exactly once for each argument

src/codegen/ast_interpreter.cpp

@@ -477,7 +477,7 @@ void ASTInterpreter::doStore(AST_Name* node, STOLEN(Value) value) {
             if (is_live)
                 jit->emitSetLocal(name, node->vreg, closure, value);
             else
-                jit->emitSetBlockLocal(name, value);
+                jit->emitSetBlockLocal(name, node->vreg, value);
         }
 
         if (closure) {
@@ -1781,6 +1781,10 @@ int ASTInterpreterJitInterface::getBoxedLocalsOffset() {
     return offsetof(ASTInterpreter, frame_info.boxedLocals);
 }
 
+int ASTInterpreterJitInterface::getCreatedClosureOffset() {
+    return offsetof(ASTInterpreter, created_closure);
+}
+
 int ASTInterpreterJitInterface::getCurrentBlockOffset() {
     return offsetof(ASTInterpreter, current_block);
 }

src/codegen/ast_interpreter.h

@@ -39,6 +39,7 @@ struct ASTInterpreterJitInterface {
     static constexpr uint64_t osr_dummy_value = -1;
 
     static int getBoxedLocalsOffset();
+    static int getCreatedClosureOffset();
     static int getCurrentBlockOffset();
     static int getCurrentInstOffset();
     static int getEdgeCountOffset();

src/codegen/baseline_jit.cpp

@@ -393,8 +393,7 @@ RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s, int vreg) {
     if (it == local_syms.end()) {
         auto r = emitGetLocal(s, vreg);
         assert(r->reftype == RefType::OWNED);
-        emitSetLocal(s, vreg, false, imm(nullptr)); // clear out the vreg
-        emitSetBlockLocal(s, r);
+        emitSetBlockLocal(s, vreg, r);
         return r;
     }
     return it->second;
@@ -403,6 +402,8 @@ RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s, int vreg) {
 void JitFragmentWriter::emitKillTemporary(InternedString s, int vreg) {
     if (!local_syms.count(s))
         emitSetLocal(s, vreg, false, imm(nullptr));
+    else
+        local_syms[s]->refUsed();
 }
 
 RewriterVar* JitFragmentWriter::emitGetBoxedLocal(BoxedString* s) {
@@ -558,8 +559,19 @@ std::vector<RewriterVar*> JitFragmentWriter::emitUnpackIntoArray(RewriterVar* v,
 }
 
 RewriterVar* JitFragmentWriter::emitYield(RewriterVar* v) {
-    auto rtn = call(false, (void*)ASTInterpreterJitInterface::yieldHelper, getInterp(), v)->setType(RefType::OWNED);
+    llvm::SmallVector<RewriterVar*, 8> local_args;
+    local_args.push_back(interp->getAttr(ASTInterpreterJitInterface::getCreatedClosureOffset()));
+    for (auto&& sym : local_syms) {
+        if (sym.second->reftype == RefType::OWNED)
+            local_args.push_back(sym.second);
+    }
+    auto&& args = allocArgs(local_args, RewriterVar::SetattrType::REF_USED);
+    RewriterVar* generator = interp->getAttr(ASTInterpreterJitInterface::getGeneratorOffset());
+    auto rtn = call(false, (void*)yield, generator, v, args, imm(local_args.size()))->setType(RefType::OWNED);
     v->refConsumed();
+    for (auto&& var : local_args) {
+        var->refUsed();
+    }
     return rtn;
 }
 
@@ -650,10 +662,13 @@ void JitFragmentWriter::emitSetAttr(AST_expr* node, RewriterVar* obj, BoxedStrin
     attr->refConsumed(rtn.second);
 }
 
-void JitFragmentWriter::emitSetBlockLocal(InternedString s, STOLEN(RewriterVar*) v) {
+void JitFragmentWriter::emitSetBlockLocal(InternedString s, int vreg, STOLEN(RewriterVar*) v) {
     if (LOG_BJIT_ASSEMBLY)
         comment("BJIT: emitSetBlockLocal() start");
     RewriterVar* prev = local_syms[s];
+    if (!prev) {
+        emitSetLocal(s, vreg, false, imm(nullptr)); // clear out the vreg
+    }
     local_syms[s] = v;
     if (LOG_BJIT_ASSEMBLY)
         comment("BJIT: emitSetBlockLocal() end");

src/codegen/baseline_jit.h

@@ -278,7 +278,7 @@ class JitFragmentWriter : public Rewriter {
     void emitRaise3(RewriterVar* arg0, RewriterVar* arg1, RewriterVar* arg2);
     void emitReturn(RewriterVar* v);
     void emitSetAttr(AST_expr* node, RewriterVar* obj, BoxedString* s, STOLEN(RewriterVar*) attr);
-    void emitSetBlockLocal(InternedString s, STOLEN(RewriterVar*) v);
+    void emitSetBlockLocal(InternedString s, int vreg, STOLEN(RewriterVar*) v);
     void emitSetCurrentInst(AST_stmt* node);
     void emitSetExcInfo(RewriterVar* type, RewriterVar* value, RewriterVar* traceback);
     void emitSetGlobal(Box* global, BoxedString* s, STOLEN(RewriterVar*) v);