rewriter: directly embed additional ref uses and bump uses early when emitting calls

src/asm_writing/rewriter.cpp

@@ -917,18 +917,9 @@ RewriterVar* Rewriter::loadConst(int64_t val, Location dest) {
 }
 
 RewriterVar* Rewriter::call(bool has_side_effects, void* func_addr, llvm::ArrayRef<RewriterVar*> args,
-                            llvm::ArrayRef<RewriterVar*> args_xmm) {
+                            llvm::ArrayRef<RewriterVar*> args_xmm, llvm::ArrayRef<RewriterVar*> additional_uses) {
     STAT_TIMER(t0, "us_timer_rewriter", 10);
     RewriterVar* result = createNewVar();
-    RewriterVar::SmallVector uses;
-    for (RewriterVar* v : args) {
-        assert(v != NULL);
-        uses.push_back(v);
-    }
-    for (RewriterVar* v : args_xmm) {
-        assert(v != NULL);
-        uses.push_back(v);
-    }
 
     ActionType type;
     if (has_side_effects)
@@ -936,36 +927,65 @@ RewriterVar* Rewriter::call(bool has_side_effects, void* func_addr, llvm::ArrayR
     else
         type = ActionType::NORMAL;
 
-    // It's not nice to pass llvm::SmallVectors through a closure, especially with our SmallFunction
-    // optimization, so just regionAlloc them and copy the data in:
-    RewriterVar** _args = (RewriterVar**)this->regionAlloc(sizeof(RewriterVar*) * args.size());
-    memcpy(_args, args.begin(), sizeof(RewriterVar*) * args.size());
-    RewriterVar** _args_xmm = (RewriterVar**)this->regionAlloc(sizeof(RewriterVar*) * args_xmm.size());
-    memcpy(_args_xmm, args_xmm.begin(), sizeof(RewriterVar*) * args_xmm.size());
-
-    int args_size = args.size();
-    assert(args_xmm.size() <= 0x7fff);
-    // Hack: pack this into a short to make sure it fits in the closure
-    short xmm_args_size = args_xmm.size();
 
     // TODO: we don't need to generate the decref info for calls which can't throw
     bool can_throw = true;
+    auto args_array_ref = regionAllocArgs(args, args_xmm, additional_uses);
 
     // Hack: explicitly order the closure arguments so they pad nicer
-    addAction([args_size, xmm_args_size, has_side_effects, can_throw, this, result, func_addr, _args, _args_xmm]() {
-        this->_call(result, has_side_effects, can_throw, 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);
+    struct LambdaClosure {
+        RewriterVar** args_array;
+
+        struct {
+            unsigned int has_side_effects : 1;
+            unsigned int can_throw : 1;
+            unsigned int num_args : 16;
+            unsigned int num_args_xmm : 16;
+            unsigned int num_additional_uses : 16;
+        };
+
+        llvm::ArrayRef<RewriterVar*> allArgs() const {
+            return llvm::makeArrayRef(args_array, num_args + num_args_xmm + num_additional_uses);
+        }
+
+        llvm::ArrayRef<RewriterVar*> args() const { return allArgs().slice(0, num_args); }
+
+        llvm::ArrayRef<RewriterVar*> argsXmm() const { return allArgs().slice(num_args, num_args_xmm); }
+
+        llvm::ArrayRef<RewriterVar*> additionalUses() const {
+            return allArgs().slice((int)num_args + (int)num_args_xmm, num_additional_uses);
+        }
+
+        LambdaClosure(llvm::MutableArrayRef<RewriterVar*> args_array_ref, llvm::ArrayRef<RewriterVar*> _args,
+                      llvm::ArrayRef<RewriterVar*> _args_xmm, llvm::ArrayRef<RewriterVar*> _addition_uses,
+                      bool has_side_effects, bool can_throw)
+            : args_array(args_array_ref.data()),
+              has_side_effects(has_side_effects),
+              can_throw(can_throw),
+              num_args(_args.size()),
+              num_args_xmm(_args_xmm.size()),
+              num_additional_uses(_addition_uses.size()) {
+            assert(_args.size() < 1 << 16);
+            assert(_args_xmm.size() < 1 << 16);
+            assert(_addition_uses.size() < 1 << 16);
+        }
+
+    } lambda_closure(args_array_ref, args, args_xmm, additional_uses, has_side_effects, can_throw);
+    assert(lambda_closure.args().size() == args.size());
+    assert(lambda_closure.argsXmm().size() == args_xmm.size());
+    assert(lambda_closure.additionalUses().size() == additional_uses.size());
+
+    addAction([this, result, func_addr, lambda_closure]() {
+        this->_call(result, lambda_closure.has_side_effects, lambda_closure.can_throw, func_addr, lambda_closure.args(),
+                    lambda_closure.argsXmm(), lambda_closure.allArgs());
+    }, lambda_closure.allArgs(), type);
 
     return result;
 }
 
 void Rewriter::_setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> args,
-                          llvm::ArrayRef<RewriterVar*> args_xmm, Location preserve) {
+                          llvm::ArrayRef<RewriterVar*> args_xmm, Location preserve,
+                          llvm::ArrayRef<RewriterVar*> bump_if_possible) {
     if (has_side_effects)
         assert(done_guarding);
 
@@ -1052,6 +1072,10 @@ void Rewriter::_setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> ar
     }
 #endif
 
+    for (auto&& use : bump_if_possible) {
+        use->bumpUseEarlyIfPossible();
+    }
+
     // Spill caller-saved registers:
     for (auto check_reg : caller_save_registers) {
         // check_reg.dump();
@@ -1114,7 +1138,8 @@ void Rewriter::_setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> ar
 }
 
 void Rewriter::_call(RewriterVar* result, bool has_side_effects, bool can_throw, void* func_addr,
-                     llvm::ArrayRef<RewriterVar*> args, llvm::ArrayRef<RewriterVar*> args_xmm) {
+                     llvm::ArrayRef<RewriterVar*> args, llvm::ArrayRef<RewriterVar*> args_xmm,
+                     llvm::ArrayRef<RewriterVar*> vars_to_bump) {
     if (LOG_IC_ASSEMBLY)
         assembler->comment("_call");
 
@@ -1123,17 +1148,7 @@ void Rewriter::_call(RewriterVar* result, bool has_side_effects, bool can_throw,
     if (failed)
         return;
 
-    _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();
-    }
-    */
+    _setupCall(has_side_effects, args, args_xmm, assembler::R11, vars_to_bump);
 
     assertConsistent();
 
@@ -1164,6 +1179,10 @@ void Rewriter::_call(RewriterVar* result, bool has_side_effects, bool can_throw,
 
     if (result)
         result->releaseIfNoUses();
+
+    for (RewriterVar* var : vars_to_bump) {
+        var->bumpUseLateIfNecessary();
+    }
 }
 
 std::vector<Location> Rewriter::getDecrefLocations() {
@@ -1286,12 +1305,7 @@ void RewriterVar::refConsumed(RewriterAction* action) {
     last_refconsumed_numuses = uses.size();
     if (!action)
         action = rewriter->getLastAction();
-    action->consumed_refs.emplace_back(this);
-}
-
-void RewriterVar::refUsed() {
-    // TODO: This is a pretty silly implementation that might prevent other optimizations?
-    rewriter->addAction([=]() { this->bumpUse(); }, { this }, ActionType::NORMAL);
+    action->consumed_refs.push_front(this);
 }
 
 bool RewriterVar::needsDecref(int current_action_index) {
@@ -1892,7 +1906,7 @@ void Rewriter::_checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val, asse
     } else
         assembler->cmp(var_reg, assembler::Immediate(exc_val), type);
 
-    _setupCall(false, RewriterVar::SmallVector(), RewriterVar::SmallVector());
+    _setupCall(false, {});
     {
         assembler::ForwardJump jnz(*assembler, assembler::COND_NOT_ZERO);
         assembler->mov(assembler::Immediate((void*)throwCAPIException), assembler::R11);

src/asm_writing/rewriter.h

@@ -16,6 +16,7 @@
 #define PYSTON_ASMWRITING_REWRITER_H
 
 #include <deque>
+#include <forward_list>
 #include <list>
 #include <map>
 #include <memory>
@@ -192,8 +193,6 @@ public:
     // if no action is specified it will assume the last action consumed the reference
     void refConsumed(RewriterAction* action = NULL);
 
-    void refUsed();
-
     // registerOwnedAttr tells the refcounter that a certain memory location holds a pointer
     // to an owned reference.  This must be paired with a call to deregisterOwnedAttr
     // Call these right before emitting the store (for register) or decref (for deregister).
@@ -237,11 +236,11 @@ private:
     // /* some code */
     // bumpUseLateIfNecessary();
     void bumpUseEarlyIfPossible() {
-        if (reftype != RefType::OWNED)
+        if (reftype != RefType::OWNED && !hasScratchAllocation())
             bumpUse();
     }
     void bumpUseLateIfNecessary() {
-        if (reftype == RefType::OWNED)
+        if (reftype == RefType::OWNED || hasScratchAllocation())
             bumpUse();
     }
 
@@ -339,8 +338,9 @@ public:
 
 class RewriterAction {
 public:
-    SmallFunction<56> action;
-    std::vector<RewriterVar*> consumed_refs;
+    SmallFunction<48> action;
+    std::forward_list<RewriterVar*> consumed_refs;
+
 
     template <typename F> RewriterAction(F&& action) : action(std::forward<F>(action)) {}
 
@@ -367,7 +367,33 @@ private:
 
 protected:
     // Allocates `bytes` bytes of data.  The allocation will get freed when the rewriter gets freed.
-    void* regionAlloc(size_t bytes) { return allocator.Allocate(bytes, 16 /* alignment */); }
+    void* regionAlloc(size_t bytes, int alignment = 16) { return allocator.Allocate(bytes, alignment); }
+    template <typename T> llvm::MutableArrayRef<T> regionAlloc(size_t num_elements) {
+        return llvm::MutableArrayRef<T>(allocator.Allocate<T>(num_elements), num_elements);
+    }
+
+    // This takes a variable number of llvm::ArrayRef<RewriterVar*> and copies in all elements into a single contiguous
+    // memory location.
+    template <typename... Args>
+    llvm::MutableArrayRef<RewriterVar*> regionAllocArgs(llvm::ArrayRef<RewriterVar*> arg1, Args... args) {
+        size_t num_total_args = 0;
+        for (auto&& array : { arg1, args... }) {
+            num_total_args += array.size();
+        }
+        if (num_total_args == 0)
+            return llvm::MutableArrayRef<RewriterVar*>();
+
+        auto args_array_ref = regionAlloc<RewriterVar*>(num_total_args);
+        auto insert_point = args_array_ref;
+        for (auto&& array : { arg1, args... }) {
+            if (!array.empty()) {
+                memcpy(insert_point.data(), array.data(), array.size() * sizeof(RewriterVar*));
+                insert_point = insert_point.slice(array.size());
+            }
+        }
+        assert(insert_point.size() == 0);
+        return args_array_ref;
+    }
 
     // Helps generating the best code for loading a const integer value.
     // By keeping track of the last known value of every register and reusing it.
@@ -511,11 +537,13 @@ protected:
     void _slowpathJump(bool condition_eq);
     void _trap();
     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());
+    void _setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> args = {},
+                    llvm::ArrayRef<RewriterVar*> args_xmm = {}, Location preserve = Location::any(),
+                    llvm::ArrayRef<RewriterVar*> bump_if_possible = {});
     // _call does not call bumpUse on its arguments:
     void _call(RewriterVar* result, bool has_side_effects, bool can_throw, void* func_addr,
-               llvm::ArrayRef<RewriterVar*> args, llvm::ArrayRef<RewriterVar*> args_xmm = {});
+               llvm::ArrayRef<RewriterVar*> args, llvm::ArrayRef<RewriterVar*> args_xmm = {},
+               llvm::ArrayRef<RewriterVar*> vars_to_bump = {});
     void _add(RewriterVar* result, RewriterVar* a, int64_t b, Location dest);
     int _allocate(RewriterVar* result, int n);
     void _allocateAndCopy(RewriterVar* result, RewriterVar* array, int n);
@@ -613,7 +641,7 @@ public:
     // inline cache.  (Extra allocations don't count even though they're potentially visible if you look
     // hard enough.)
     RewriterVar* call(bool has_side_effects, void* func_addr, llvm::ArrayRef<RewriterVar*> args = {},
-                      llvm::ArrayRef<RewriterVar*> args_xmm = {});
+                      llvm::ArrayRef<RewriterVar*> args_xmm = {}, llvm::ArrayRef<RewriterVar*> additional_uses = {});
     template <typename... Args>
     RewriterVar* call(bool has_side_effects, void* func_addr, RewriterVar* arg1, Args... args) {
         return call(has_side_effects, func_addr, llvm::ArrayRef<RewriterVar*>({ arg1, args... }), {});

src/codegen/ast_interpreter.cpp

@@ -1172,12 +1172,9 @@ Value ASTInterpreter::createFunction(AST* node, AST_arguments* args, const std::
             closure_var = jit->imm(0ul);
         if (!passed_globals_var)
             passed_globals_var = jit->imm(0ul);
-        rtn.var = jit->call(false, (void*)createFunctionFromMetadata, jit->imm(md), closure_var, passed_globals_var,
-                            defaults_var, jit->imm(args->defaults.size()))->setType(RefType::OWNED);
-
-        for (auto d_var : defaults_vars) {
-            d_var->refUsed();
-        }
+        rtn.var = jit->call(false, (void*)createFunctionFromMetadata, { jit->imm(md), closure_var, passed_globals_var,
+                                                                        defaults_var, jit->imm(args->defaults.size()) },
+                            {}, defaults_vars)->setType(RefType::OWNED);
     }
 
     rtn.o = createFunctionFromMetadata(md, closure, passed_globals, u.il);

src/codegen/baseline_jit.cpp

@@ -222,10 +222,10 @@ RewriterVar* JitFragmentWriter::emitCallattr(AST_expr* node, RewriterVar* obj, B
     call_args.push_back(args.size() > 1 ? args[1] : imm(0ul));
     call_args.push_back(args.size() > 2 ? args[2] : imm(0ul));
 
+    llvm::ArrayRef<RewriterVar*> additional_uses;
     if (args.size() > 3) {
-        RewriterVar* scratch = allocate(args.size() - 3);
-        for (int i = 0; i < args.size() - 3; ++i)
-            scratch->setAttr(i * sizeof(void*), args[i + 3], RewriterVar::SetattrType::REF_USED);
+        additional_uses = args.slice(3);
+        RewriterVar* scratch = allocArgs(additional_uses, RewriterVar::SetattrType::REF_USED);
         call_args.push_back(scratch);
     } else if (keyword_names) {
         call_args.push_back(imm(0ul));
@@ -234,11 +234,8 @@ RewriterVar* JitFragmentWriter::emitCallattr(AST_expr* node, RewriterVar* obj, B
     if (keyword_names)
         call_args.push_back(imm(keyword_names));
 
-    auto r = emitPPCall((void*)callattr, call_args, 2, 640, node, type_recorder).first->setType(RefType::OWNED);
-    for (int i = 3; i < args.size(); i++) {
-        args[i]->refUsed();
-    }
-    return r;
+    return emitPPCall((void*)callattr, call_args, 2, 640, node, type_recorder, additional_uses)
+        .first->setType(RefType::OWNED);
 #else
     // We could make this faster but for now: keep it simple, stupid...
     RewriterVar* attr_var = imm(attr);
@@ -285,44 +282,9 @@ void JitFragmentWriter::emitDictSet(RewriterVar* dict, RewriterVar* k, RewriterV
     v->refConsumed();
 }
 
-// TODO: merge this function's functionality with refUsed
 RewriterVar* JitFragmentWriter::emitCallWithAllocatedArgs(void* func_addr, const llvm::ArrayRef<RewriterVar*> args,
                                                           const llvm::ArrayRef<RewriterVar*> additional_uses) {
-    RewriterVar* result = createNewVar();
-    RewriterVar::SmallVector uses;
-    for (RewriterVar* v : args) {
-        assert(v != NULL);
-        uses.push_back(v);
-    }
-    for (RewriterVar* v : additional_uses) {
-        assert(v != NULL);
-        uses.push_back(v);
-    }
-
-    // It's not nice to pass llvm::SmallVectors through a closure, especially with our SmallFunction
-    // optimization, so just regionAlloc them and copy the data in:
-    RewriterVar** _args = (RewriterVar**)this->regionAlloc(sizeof(RewriterVar*) * args.size());
-    memcpy(_args, args.begin(), sizeof(RewriterVar*) * args.size());
-    RewriterVar** _args_additional = (RewriterVar**)this->regionAlloc(sizeof(RewriterVar*) * additional_uses.size());
-    memcpy(_args_additional, additional_uses.begin(), sizeof(RewriterVar*) * additional_uses.size());
-
-    int args_size = args.size();
-    assert(additional_uses.size() <= 0x7fff);
-    // Hack: pack this into a short to make sure it fits in the closure
-    short args_additional_size = additional_uses.size();
-
-    bool can_throw = true;
-
-    // Hack: explicitly order the closure arguments so they pad nicer
-    addAction([args_size, args_additional_size, can_throw, func_addr, this, result, _args, _args_additional]() {
-        this->_call(result, false, can_throw, func_addr, llvm::ArrayRef<RewriterVar*>(_args, args_size),
-                    llvm::ArrayRef<RewriterVar*>());
-        for (int i = 0; i < args_size; i++)
-            _args[i]->bumpUse();
-        for (int i = 0; i < args_additional_size; i++)
-            _args_additional[i]->bumpUse();
-    }, uses, ActionType::NORMAL);
-    return result;
+    return call(false, func_addr, args, {}, additional_uses);
 }
 
 RewriterVar* JitFragmentWriter::emitCreateList(const llvm::ArrayRef<RewriterVar*> values) {
@@ -504,21 +466,18 @@ RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj
     call_args.push_back(args.size() > 1 ? args[1] : imm(0ul));
     call_args.push_back(args.size() > 2 ? args[2] : imm(0ul));
 
+    llvm::ArrayRef<RewriterVar*> additional_uses;
     if (args.size() > 3) {
-        RewriterVar* scratch = allocate(args.size() - 3);
-        for (int i = 0; i < args.size() - 3; ++i)
-            scratch->setAttr(i * sizeof(void*), args[i + 3], RewriterVar::SetattrType::REF_USED);
+        additional_uses = args.slice(3);
+        RewriterVar* scratch = allocArgs(additional_uses, RewriterVar::SetattrType::REF_USED);
         call_args.push_back(scratch);
     } else
         call_args.push_back(imm(0ul));
     if (keyword_names)
         call_args.push_back(imm(keyword_names));
 
-    auto r = emitPPCall((void*)runtimeCall, call_args, 2, 640, node, type_recorder).first->setType(RefType::OWNED);
-    for (int i = 3; i < args.size(); i++) {
-        args[i]->refUsed();
-    }
-    return r;
+    return emitPPCall((void*)runtimeCall, call_args, 2, 640, node, type_recorder, additional_uses)
+        .first->setType(RefType::OWNED);
 #else
     RewriterVar* argspec_var = imm(argspec.asInt());
     RewriterVar* keyword_names_var = keyword_names ? imm(keyword_names) : nullptr;
@@ -576,11 +535,9 @@ RewriterVar* JitFragmentWriter::emitYield(RewriterVar* v) {
 
     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);
+    auto rtn = call(false, (void*)yield, { generator, v, args, imm(local_args.size()) }, {}, local_args)
+                   ->setType(RefType::OWNED);
     v->refConsumed();
-    for (auto&& var : local_args) {
-        var->refUsed();
-    }
     return rtn;
 }
 
@@ -906,31 +863,38 @@ uint64_t JitFragmentWriter::asUInt(InternedString s) {
 }
 #endif
 
-std::pair<RewriterVar*, RewriterAction*> JitFragmentWriter::emitPPCall(void* func_addr,
-                                                                       llvm::ArrayRef<RewriterVar*> args, int num_slots,
-                                                                       int slot_size, AST* ast_node,
-                                                                       TypeRecorder* type_recorder) {
+std::pair<RewriterVar*, RewriterAction*>
+JitFragmentWriter::emitPPCall(void* func_addr, llvm::ArrayRef<RewriterVar*> args, unsigned char num_slots,
+                              unsigned short slot_size, AST* ast_node, TypeRecorder* type_recorder,
+                              llvm::ArrayRef<RewriterVar*> additional_uses) {
     if (LOG_BJIT_ASSEMBLY)
         comment("BJIT: emitPPCall() start");
     RewriterVar::SmallVector args_vec(args.begin(), args.end());
 #if ENABLE_BASELINEJIT_ICS
     RewriterVar* result = createNewVar();
 
+    auto args_array_ref = regionAllocArgs(args, additional_uses);
+
+    assert(args.size() < 1ul << 32);
     int args_size = args.size();
-    RewriterVar** _args = (RewriterVar**)regionAlloc(sizeof(RewriterVar*) * args_size);
-    memcpy(_args, args.begin(), sizeof(RewriterVar*) * args_size);
+    assert(additional_uses.size() < 1 << 8);
+    unsigned char num_additional = additional_uses.size();
+    RewriterVar** args_array = args_array_ref.data();
 
-    RewriterAction* call_action = addAction([=]() {
-        this->_emitPPCall(result, func_addr, llvm::ArrayRef<RewriterVar*>(_args, args_size), num_slots, slot_size,
-                          ast_node);
-    }, args, ActionType::NORMAL);
+    RewriterAction* call_action
+        = addAction([this, result, func_addr, ast_node, args_array, args_size, slot_size, num_slots, num_additional]() {
+            auto all_args = llvm::makeArrayRef(args_array, args_size + num_additional);
+            auto args = all_args.slice(0, args_size);
+            this->_emitPPCall(result, func_addr, args, num_slots, slot_size, ast_node, all_args);
+        }, args_array_ref, ActionType::NORMAL);
 
     if (type_recorder) {
         RewriterVar* type_recorder_var = imm(type_recorder);
         RewriterVar* obj_cls_var = result->getAttr(offsetof(Box, cls));
-        addAction([=]() { _emitRecordType(type_recorder_var, obj_cls_var); }, { type_recorder_var, obj_cls_var },
-                  ActionType::NORMAL);
-        result->refUsed();
+        addAction([=]() {
+            _emitRecordType(type_recorder_var, obj_cls_var);
+            result->bumpUse();
+        }, { type_recorder_var, obj_cls_var, result }, ActionType::NORMAL);
 
         emitPendingCallsCheck();
         return std::make_pair(result, call_action);
@@ -1025,7 +989,7 @@ void JitFragmentWriter::_emitGetLocal(RewriterVar* val_var, const char* name) {
     assembler::Register var_reg = val_var->getInReg();
     assembler->test(var_reg, var_reg);
 
-    _setupCall(false, RewriterVar::SmallVector(), RewriterVar::SmallVector());
+    _setupCall(false, {});
     {
         assembler::ForwardJump jnz(*assembler, assembler::COND_NOT_ZERO);
         assembler->mov(assembler::Immediate((uint64_t)name), assembler::RDI);
@@ -1100,7 +1064,8 @@ void JitFragmentWriter::_emitOSRPoint() {
 }
 
 void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, llvm::ArrayRef<RewriterVar*> args,
-                                    int num_slots, int slot_size, AST* ast_node) {
+                                    int num_slots, int slot_size, AST* ast_node,
+                                    llvm::ArrayRef<RewriterVar*> vars_to_bump) {
     assembler::Register r = allocReg(assembler::R11);
 
     if (args.size() > 6) { // only 6 args can get passed in registers.
@@ -1109,11 +1074,9 @@ void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, llvm::
             assembler::Register reg = args[i]->getInReg(Location::any(), true);
             assembler->mov(reg, assembler::Indirect(assembler::RSP, sizeof(void*) * (i - 6)));
         }
-        RewriterVar::SmallVector reg_args(args.begin(), args.begin() + 6);
-        assert(reg_args.size() == 6);
-        _setupCall(false, reg_args, RewriterVar::SmallVector());
+        _setupCall(false, args.slice(0, 6), {}, Location::any(), vars_to_bump);
     } else
-        _setupCall(false, args, RewriterVar::SmallVector());
+        _setupCall(false, args, {}, Location::any(), vars_to_bump);
 
     if (failed)
         return;
@@ -1158,15 +1121,8 @@ void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, llvm::
 
     result->releaseIfNoUses();
 
-    // TODO: it would be nice to be able to bumpUse on these earlier so that we could potentially avoid spilling
-    // the args across the call if we don't need to.
-    // This had to get moved to the very end of this function due to the fact that bumpUse can cause refcounting
-    // operations to happen.
-    // I'm not sure though that just moving this earlier is good enough though -- we also might need to teach setupCall
-    // that the args might not be needed afterwards?
-    // Anyway this feels like micro-optimizations at the moment and we can figure it out later.
-    for (RewriterVar* arg : args) {
-        arg->bumpUse();
+    for (RewriterVar* use : vars_to_bump) {
+        use->bumpUseLateIfNecessary();
     }
 }
 

src/codegen/baseline_jit.h

@@ -316,8 +316,9 @@ private:
     RewriterVar* emitCallWithAllocatedArgs(void* func_addr, const llvm::ArrayRef<RewriterVar*> args,
                                            const llvm::ArrayRef<RewriterVar*> additional_uses);
     std::pair<RewriterVar*, RewriterAction*> emitPPCall(void* func_addr, llvm::ArrayRef<RewriterVar*> args,
-                                                        int num_slots, int slot_size, AST* ast_node = NULL,
-                                                        TypeRecorder* type_recorder = NULL);
+                                                        unsigned char num_slots, unsigned short slot_size,
+                                                        AST* ast_node = NULL, TypeRecorder* type_recorder = NULL,
+                                                        llvm::ArrayRef<RewriterVar*> additional_uses = {});
 
     static void assertNameDefinedHelper(const char* id);
     static Box* callattrHelper(Box* obj, BoxedString* attr, CallattrFlags flags, TypeRecorder* type_recorder,
@@ -337,7 +338,7 @@ private:
     void _emitJump(CFGBlock* b, RewriterVar* block_next, ExitInfo& exit_info);
     void _emitOSRPoint();
     void _emitPPCall(RewriterVar* result, void* func_addr, llvm::ArrayRef<RewriterVar*> args, int num_slots,
-                     int slot_size, AST* ast_node);
+                     int slot_size, AST* ast_node, llvm::ArrayRef<RewriterVar*> vars_to_bump);
     void _emitRecordType(RewriterVar* type_recorder_var, RewriterVar* obj_cls_var);
     void _emitReturn(RewriterVar* v);
     void _emitSideExit(STOLEN(RewriterVar*) var, RewriterVar* val_constant, CFGBlock* next_block,

src/runtime/objmodel.cpp

@@ -4870,32 +4870,32 @@ Box* callCLFunc(FunctionMetadata* md, CallRewriteArgs* rewrite_args, int num_out
             } else {
                 // Hacky workaround: the rewriter can only pass arguments in registers, so use this helper function
                 // to unpack some of the additional arguments:
+                llvm::SmallVector<RewriterVar*, 4> additional_uses;
                 RewriterVar* arg_array = rewrite_args->rewriter->allocate(4);
                 arg_vec.push_back(arg_array);
-                if (num_output_args >= 1)
+                if (num_output_args >= 1) {
                     arg_array->setAttr(0, rewrite_args->arg1, RewriterVar::SetattrType::REF_USED);
-                if (num_output_args >= 2)
+                    additional_uses.push_back(rewrite_args->arg1);
+                }
+                if (num_output_args >= 2) {
                     arg_array->setAttr(8, rewrite_args->arg2, RewriterVar::SetattrType::REF_USED);
-                if (num_output_args >= 3)
+                    additional_uses.push_back(rewrite_args->arg2);
+                }
+                if (num_output_args >= 3) {
                     arg_array->setAttr(16, rewrite_args->arg3, RewriterVar::SetattrType::REF_USED);
-                if (num_output_args >= 4)
+                    additional_uses.push_back(rewrite_args->arg3);
+                }
+                if (num_output_args >= 4) {
                     arg_array->setAttr(24, rewrite_args->args, RewriterVar::SetattrType::REF_USED);
+                    additional_uses.push_back(rewrite_args->args);
+                }
 
                 if (S == CXX)
-                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec)
-                                                ->setType(RefType::OWNED);
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec, {},
+                                                                         additional_uses)->setType(RefType::OWNED);
                 else
-                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec)
-                                                ->setType(RefType::OWNED);
-
-                if (num_output_args >= 1)
-                    rewrite_args->arg1->refUsed();
-                if (num_output_args >= 2)
-                    rewrite_args->arg2->refUsed();
-                if (num_output_args >= 3)
-                    rewrite_args->arg3->refUsed();
-                if (num_output_args >= 4)
-                    rewrite_args->args->refUsed();
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec,
+                                                                         {}, additional_uses)->setType(RefType::OWNED);
             }
 
             rewrite_args->out_success = true;