rewriter: directly embed additional ref uses

Author: undingen

src/asm_writing/rewriter.cpp

@@ -917,55 +917,75 @@ 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)
         type = ActionType::MUTATION;
     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*> stuff_to_bump) {
     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 : stuff_to_bump) {
+        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*> stuff_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, stuff_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* arg : stuff_to_bump) {
+        arg->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() {
@@ -1893,7 +1907,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 @@ class RewriterVar {
     // 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 @@ class RewriterVar {
     // /* 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 @@ template <int N = 24> class SmallFunction {
 
 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 @@ class Rewriter : public ICSlotRewrite::CommitHook {
 
 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.
@@ -509,11 +535,13 @@ class Rewriter : public ICSlotRewrite::CommitHook {
     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*> stuff_to_bump = {});
     // _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*> stuff_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);
@@ -612,7 +640,7 @@ class Rewriter : public ICSlotRewrite::CommitHook {
     // 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);