Starting to enable refcounting in the bjit

Current issue is that this is triggering the "mark IC as something
being inside", which trips up the bjit's "overwrite the last few
bytes of the previous block, assuming that they are just a jump".

src/asm_writing/rewriter.cpp

@@ -477,6 +477,73 @@ void Rewriter::_getAttrFloat(RewriterVar* result, RewriterVar* ptr, int offset,
     assertConsistent();
 }
 
+class Helper {
+public:
+    static void incref(Box* b) {
+        Py_INCREF(b);
+    }
+    static void decref(Box* b) {
+        Py_DECREF(b);
+    }
+    static void xdecref(Box* b) {
+        Py_XDECREF(b);
+    }
+};
+
+void RewriterVar::incref() {
+    rewriter->addAction([=]() {
+        rewriter->_incref(this);
+        this->bumpUse();
+    }, { this }, ActionType::MUTATION);
+}
+
+void RewriterVar::decref() {
+    rewriter->addAction([=]() {
+        rewriter->_decref(this);
+        this->bumpUse();
+    }, { this }, ActionType::MUTATION);
+}
+
+void RewriterVar::xdecref() {
+    rewriter->addAction([=]() {
+        rewriter->_xdecref(this);
+        this->bumpUse();
+    }, { this }, ActionType::MUTATION);
+}
+
+void Rewriter::_incref(RewriterVar* var) {
+    //assembler->trap();
+    //auto reg = var->getInReg();
+    //assembler->incl(assembler::Indirect(reg, offsetof(Box, ob_refcnt)));
+
+    //this->_trap();
+    this->_call(NULL, true, (void*)Helper::incref, llvm::ArrayRef<RewriterVar*>(&var, 1),
+                llvm::ArrayRef<RewriterVar*>());
+
+    // Doesn't call bumpUse, since this function is designed to be callable from other emitting functions.
+    // (ie the caller should call bumpUse)
+}
+
+void Rewriter::_decref(RewriterVar* var) {
+    //assembler->trap();
+
+    this->_call(NULL, true, (void*)Helper::decref, llvm::ArrayRef<RewriterVar*>(&var, 1),
+                llvm::ArrayRef<RewriterVar*>(NULL, (int)0));
+
+    // Doesn't call bumpUse, since this function is designed to be callable from other emitting functions.
+    // (ie the caller should call bumpUse)
+}
+
+void Rewriter::_xdecref(RewriterVar* var) {
+    //assembler->trap();
+
+    this->_call(NULL, true, (void*)Helper::xdecref, llvm::ArrayRef<RewriterVar*>(&var, 1),
+                llvm::ArrayRef<RewriterVar*>(NULL, (int)0));
+
+    // Doesn't call bumpUse, since this function is designed to be callable from other emitting functions.
+    // (ie the caller should call bumpUse)
+}
+
 RewriterVar* RewriterVar::cmp(AST_TYPE::AST_TYPE cmp_type, RewriterVar* other, Location dest) {
     STAT_TIMER(t0, "us_timer_rewriter", 10);
 
@@ -846,6 +913,10 @@ RewriterVar* Rewriter::call(bool has_side_effects, void* func_addr, const Rewrit
     addAction([args_size, xmm_args_size, has_side_effects, this, result, func_addr, _args, _args_xmm]() {
         this->_call(result, has_side_effects, func_addr, llvm::ArrayRef<RewriterVar*>(_args, args_size),
                     llvm::ArrayRef<RewriterVar*>(_args_xmm, xmm_args_size));
+        for (int i = 0; i < args_size; i++)
+            _args[i]->bumpUse();
+        for (int i = 0; i < xmm_args_size; i++)
+            _args_xmm[i]->bumpUse();
     }, uses, type);
 
     return result;
@@ -869,7 +940,7 @@ void Rewriter::_setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> ar
         assert(assembler->bytesWritten() >= IC_INVALDITION_HEADER_SIZE);
     }
 
-    if (has_side_effects) {
+    if (has_side_effects && needs_invalidation_support) {
         if (!marked_inside_ic) {
             uintptr_t counter_addr = (uintptr_t)(&picked_slot->num_inside);
             if (isLargeConstant(counter_addr)) {
@@ -1011,12 +1082,15 @@ void Rewriter::_call(RewriterVar* result, bool has_side_effects, void* func_addr
 
     _setupCall(has_side_effects, args, args_xmm, assembler::R11);
 
+    // This duty is now on the caller:
+    /*
     for (RewriterVar* arg : args) {
         arg->bumpUse();
     }
     for (RewriterVar* arg_xmm : args_xmm) {
         arg_xmm->bumpUse();
     }
+    */
 
     assertConsistent();
 
@@ -1036,12 +1110,14 @@ void Rewriter::_call(RewriterVar* result, bool has_side_effects, void* func_addr
 
     if (!failed) {
         assert(vars_by_location.count(assembler::RAX) == 0);
-        result->initializeInReg(assembler::RAX);
 
+        if (result)
+            result->initializeInReg(assembler::RAX);
         assertConsistent();
     }
 
-    result->releaseIfNoUses();
+    if (result)
+        result->releaseIfNoUses();
 }
 
 void Rewriter::abort() {
@@ -1205,6 +1281,8 @@ void Rewriter::commit() {
     if (marked_inside_ic) {
         assembler->comment("mark inside ic");
 
+        ASSERT(this->needs_invalidation_support, "why did we mark ourselves as inside this?");
+
         uintptr_t counter_addr = (uintptr_t)(&picked_slot->num_inside);
         if (isLargeConstant(counter_addr)) {
             assembler::Register reg = allocReg(Location::any(), getReturnDestination());
@@ -1824,13 +1902,14 @@ TypeRecorder* Rewriter::getTypeRecorder() {
     return rewrite->getTypeRecorder();
 }
 
-Rewriter::Rewriter(std::unique_ptr<ICSlotRewrite> rewrite, int num_args, const LiveOutSet& live_outs)
+Rewriter::Rewriter(std::unique_ptr<ICSlotRewrite> rewrite, int num_args, const LiveOutSet& live_outs, bool needs_invalidation_support)
     : rewrite(std::move(rewrite)),
       assembler(this->rewrite->getAssembler()),
       picked_slot(NULL),
       const_loader(this),
       return_location(this->rewrite->returnRegister()),
       failed(false),
+      needs_invalidation_support(needs_invalidation_support),
       added_changing_action(false),
       marked_inside_ic(false),
       done_guarding(false),

src/asm_writing/rewriter.h

@@ -216,6 +216,9 @@ public:
     void setAttr(int offset, RewriterVar* other);
     RewriterVar* cmp(AST_TYPE::AST_TYPE cmp_type, RewriterVar* other, Location loc = Location::any());
     RewriterVar* toBool(Location loc = Location::any());
+    void incref();
+    void decref();
+    void xdecref();
 
     template <typename Src, typename Dst> inline RewriterVar* getAttrCast(int offset, Location loc = Location::any());
 
@@ -397,8 +400,9 @@ protected:
 
     bool failed;   // if we tried to generate an invalid rewrite.
     bool finished; // committed or aborted
-#ifndef NDEBUG
+    const bool needs_invalidation_support;
 
+#ifndef NDEBUG
     bool phase_emitting;
     void initPhaseCollecting() { phase_emitting = false; }
     void initPhaseEmitting() { phase_emitting = true; }
@@ -417,7 +421,11 @@ protected:
     llvm::SmallVector<RewriterVar*, 8> args;
     llvm::SmallVector<RewriterVar*, 8> live_outs;
 
-    Rewriter(std::unique_ptr<ICSlotRewrite> rewrite, int num_args, const LiveOutSet& live_outs);
+    // needs_invalidation_support: whether we do some extra work to make sure that the code that this Rewriter
+    // produces will support invalidation.  Normally we want this, but the baseline jit needs to turn
+    // this off (for the non-IC assembly it generates).
+    Rewriter(std::unique_ptr<ICSlotRewrite> rewrite, int num_args, const LiveOutSet& live_outs,
+             bool needs_invalidation_support = true);
 
     std::deque<RewriterAction, RegionAllocatorAdaptor<RewriterAction>> actions;
     template <typename F> void addAction(F&& action, llvm::ArrayRef<RewriterVar*> vars, ActionType type) {
@@ -490,6 +498,7 @@ protected:
     void _loadConst(RewriterVar* result, int64_t val);
     void _setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> args, llvm::ArrayRef<RewriterVar*> args_xmm,
                     Location preserve = Location::any());
+    // _call does not call bumpUse on its arguments:
     void _call(RewriterVar* result, bool has_side_effects, void* func_addr, llvm::ArrayRef<RewriterVar*> args,
                llvm::ArrayRef<RewriterVar*> args_xmm);
     void _add(RewriterVar* result, RewriterVar* a, int64_t b, Location dest);
@@ -511,6 +520,11 @@ protected:
               Location loc = Location::any());
     void _toBool(RewriterVar* result, RewriterVar* var, Location loc = Location::any());
 
+    // These do not call bumpUse on their arguments:
+    void _incref(RewriterVar* var);
+    void _decref(RewriterVar* var);
+    void _xdecref(RewriterVar* var);
+
     void assertConsistent() {
 #ifndef NDEBUG
         for (RewriterVar& var : vars) {

src/codegen/ast_interpreter.cpp

@@ -300,8 +300,6 @@ void ASTInterpreter::startJITing(CFGBlock* block, int exit_offset) {
     assert(ENABLE_BASELINEJIT);
     assert(!jit);
 
-    assert(0 && "refcounting not set up");
-
     auto& code_blocks = md->code_blocks;
     JitCodeBlock* code_block = NULL;
     if (!code_blocks.empty())
@@ -474,7 +472,6 @@ void ASTInterpreter::doStore(AST_Name* node, STOLEN(Value) value) {
     } else {
         bool closure = vst == ScopeInfo::VarScopeType::CLOSURE;
         if (jit) {
-            assert(0 && "check refcounting");
             if (!closure) {
                 bool is_live = source_info->getLiveness()->isLiveAtEnd(name, current_block);
                 if (is_live)
@@ -1550,7 +1547,6 @@ Value ASTInterpreter::visit_name(AST_Name* node) {
         case ScopeInfo::VarScopeType::CLOSURE: {
             Value v;
             if (jit) {
-                assert(0 && "check refcounting");
                 bool is_live = false;
                 if (node->lookup_type == ScopeInfo::VarScopeType::FAST)
                     is_live = source_info->getLiveness()->isLiveAtEnd(node->id, current_block);

src/codegen/baseline_jit.cpp

@@ -132,10 +132,10 @@ void JitCodeBlock::fragmentFinished(int bytes_written, int num_bytes_overlapping
 JitFragmentWriter::JitFragmentWriter(CFGBlock* block, std::unique_ptr<ICInfo> ic_info,
                                      std::unique_ptr<ICSlotRewrite> rewrite, int code_offset, int num_bytes_overlapping,
                                      void* entry_code, JitCodeBlock& code_block)
-    : Rewriter(std::move(rewrite), 0, {}),
+    : Rewriter(std::move(rewrite), 0, {}, /* needs_invalidation_support = */ false),
       block(block),
       code_offset(code_offset),
-      num_bytes_exit(0),
+      exit_info(),
       num_bytes_overlapping(num_bytes_overlapping),
       entry_code(entry_code),
       code_block(code_block),
@@ -296,6 +296,8 @@ RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s, int vreg) {
     auto it = local_syms.find(s);
     if (it == local_syms.end())
         return emitGetLocal(s, vreg);
+    // TODO xincref?
+    it->second->incref();
     return it->second;
 }
 
@@ -462,7 +464,7 @@ void JitFragmentWriter::emitExec(RewriterVar* code, RewriterVar* globals, Rewrit
 
 void JitFragmentWriter::emitJump(CFGBlock* b) {
     RewriterVar* next = imm(b);
-    addAction([=]() { _emitJump(b, next, num_bytes_exit); }, { next }, ActionType::NORMAL);
+    addAction([=]() { _emitJump(b, next, exit_info); }, { next }, ActionType::NORMAL);
 }
 
 void JitFragmentWriter::emitOSRPoint(AST_Jump* node) {
@@ -496,7 +498,12 @@ void JitFragmentWriter::emitSetAttr(AST_expr* node, RewriterVar* obj, BoxedStrin
 }
 
 void JitFragmentWriter::emitSetBlockLocal(InternedString s, RewriterVar* v) {
+    RewriterVar* prev = local_syms[s];
     local_syms[s] = v;
+    if (prev) {
+        // TODO: xdecref?
+        prev->decref();
+    }
 }
 
 void JitFragmentWriter::emitSetCurrentInst(AST_stmt* node) {
@@ -522,6 +529,7 @@ void JitFragmentWriter::emitSetItemName(BoxedString* s, RewriterVar* v) {
 void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closure, RewriterVar* v) {
     assert(vreg >= 0);
     if (set_closure) {
+        assert(0 && "check refcounting");
         call(false, (void*)ASTInterpreterJitInterface::setLocalClosureHelper, getInterp(), imm(vreg),
 #ifndef NDEBUG
              imm(asUInt(s).first), imm(asUInt(s).second),
@@ -530,7 +538,9 @@ void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closur
 #endif
              v);
     } else {
+        RewriterVar* prev = vregs_array->getAttr(8 * vreg);
         vregs_array->setAttr(8 * vreg, v);
+        prev->xdecref();
     }
 }
 
@@ -624,6 +634,10 @@ int JitFragmentWriter::finishCompilation() {
     }
 
     void* next_fragment_start = (uint8_t*)block->code + assembler->bytesWritten();
+    if (exit_info.num_bytes)
+        ASSERT(assembler->curInstPointer() == (uint8_t*)exit_info.exit_start + exit_info.num_bytes,
+               "Error! wrote more bytes out after the 'retq' that we thought was going to be the end of the assembly.  "
+               "We will end up overwriting those instructions.");
     code_block.fragmentFinished(assembler->bytesWritten(), num_bytes_overlapping, next_fragment_start);
 
 #if MOVING_GC
@@ -634,7 +648,7 @@ int JitFragmentWriter::finishCompilation() {
     registerGCTrackedICInfo(ic_info.release());
 #endif
 
-    return num_bytes_exit;
+    return exit_info.num_bytes;
 }
 
 bool JitFragmentWriter::finishAssembly(int continue_offset) {
@@ -777,10 +791,13 @@ void JitFragmentWriter::_emitGetLocal(RewriterVar* val_var, const char* name) {
         assembler->mov(assembler::Immediate((void*)assertNameDefinedHelper), assembler::R11);
         assembler->callq(assembler::R11);
     }
+
+    _incref(val_var);
 }
 
-void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, int& size_of_exit_to_interp) {
-    size_of_exit_to_interp = 0;
+void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, ExitInfo& exit_info) {
+    assert(exit_info.num_bytes == 0);
+    assert(exit_info.exit_start == NULL);
     if (b->code) {
         int64_t offset = (uint64_t)b->code - ((uint64_t)entry_code + code_offset);
         if (isLargeConstant(offset)) {
@@ -790,6 +807,7 @@ void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, int& siz
             assembler->jmp(assembler::JumpDestination::fromStart(offset));
     } else {
         int num_bytes = assembler->bytesWritten();
+        exit_info.exit_start = assembler->curInstPointer();
         block_next->getInReg(assembler::RAX, true);
         assembler->add(assembler::Immediate(JitCodeBlock::sp_adjustment), assembler::RSP);
         assembler->pop(assembler::R12);
@@ -800,8 +818,8 @@ void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, int& siz
         for (int i = assembler->bytesWritten() - num_bytes; i < min_patch_size; ++i)
             assembler->trap(); // we could use nops but traps may help if something goes wrong
 
-        size_of_exit_to_interp = assembler->bytesWritten() - num_bytes;
-        assert(assembler->hasFailed() || size_of_exit_to_interp >= min_patch_size);
+        exit_info.num_bytes = assembler->bytesWritten() - num_bytes;
+        assert(assembler->hasFailed() || exit_info.num_bytes >= min_patch_size);
     }
     block_next->bumpUse();
 }
@@ -813,6 +831,8 @@ void JitFragmentWriter::_emitOSRPoint(RewriterVar* result, RewriterVar* node_var
     _call(result, false, (void*)ASTInterpreterJitInterface::doOSRHelper, args, RewriterVar::SmallVector());
     auto result_reg = result->getInReg(assembler::RDX);
     result->bumpUse();
+    for (auto a : args)
+        a->bumpUse();
 
     assembler->test(result_reg, result_reg);
     {
@@ -935,12 +955,13 @@ void JitFragmentWriter::_emitSideExit(RewriterVar* var, RewriterVar* val_constan
 
     {
         assembler::ForwardJump jne(*assembler, assembler::COND_EQUAL);
-        int exit_size = 0;
-        _emitJump(next_block, next_block_var, exit_size);
-        if (exit_size) {
+        ExitInfo exit_info;
+        _emitJump(next_block, next_block_var, exit_info);
+        if (exit_info.num_bytes) {
+            assert(assembler->curInstPointer() == (uint8_t*)exit_info.exit_start + exit_info.num_bytes);
             RELEASE_ASSERT(!side_exit_patch_location.first,
                            "if we start to emit more than one side exit we should make this a vector");
-            side_exit_patch_location = std::make_pair(next_block, assembler->bytesWritten() - exit_size);
+            side_exit_patch_location = std::make_pair(next_block, assembler->bytesWritten() - exit_info.num_bytes);
         }
     }
 

src/codegen/baseline_jit.h

@@ -158,6 +158,13 @@ public:
 
 class JitFragmentWriter : public Rewriter {
 private:
+    struct ExitInfo {
+        int num_bytes; // the number of bytes for the overwriteable jump
+        void* exit_start; // where that jump starts
+
+        ExitInfo() : num_bytes(0), exit_start(NULL) {}
+    };
+
     static constexpr int min_patch_size = 13;
 
     CFGBlock* block;
@@ -169,7 +176,7 @@ private:
     // value which will make the fragment start at the instruction where the last block is exiting to the interpreter to
     // interpret the new block -> we overwrite the exit with the code of the new block.
     // If there is nothing to overwrite this field will be 0.
-    int num_bytes_exit;
+    ExitInfo exit_info;
     int num_bytes_overlapping; // num of bytes this block overlaps with the prev. used to patch unessary jumps
 
     void* entry_code; // JitCodeBlock start address. Must have an offset of 0 into the code block
@@ -292,7 +299,7 @@ private:
                                   std::vector<BoxedString*>* keyword_names);
 
     void _emitGetLocal(RewriterVar* val_var, const char* name);
-    void _emitJump(CFGBlock* b, RewriterVar* block_next, int& size_of_exit_to_interp);
+    void _emitJump(CFGBlock* b, RewriterVar* block_next, ExitInfo& exit_info);
     void _emitOSRPoint(RewriterVar* result, RewriterVar* node_var);
     void _emitPPCall(RewriterVar* result, void* func_addr, llvm::ArrayRef<RewriterVar*> args, int num_slots,
                      int slot_size, AST* ast_node);

src/core/options.cpp

@@ -34,7 +34,7 @@ bool DUMPJIT = false;
 bool TRAP = false;
 bool USE_STRIPPED_STDLIB = true; // always true
 bool ENABLE_INTERPRETER = true;
-bool ENABLE_BASELINEJIT = 0; // XXX
+bool ENABLE_BASELINEJIT = 1;
 bool ENABLE_PYPA_PARSER = true;
 bool ENABLE_CPYTHON_PARSER = false;
 bool USE_REGALLOC_BASIC = true;