Merge pull request #708 from undingen/rewriter_release_scratch

rewriter: release allocated scratch space

src/asm_writing/rewriter.cpp

@@ -551,6 +551,12 @@ void Rewriter::_setAttr(RewriterVar* ptr, int offset, RewriterVar* val) {
     }
 
     ptr->bumpUse();
+
+    // If the value is a scratch allocated memory array we have to make sure we won't release it immediately.
+    // Because this setAttr stored a reference to it inside a field and the rewriter can't currently track this uses and
+    // will think it's unused.
+    if (val->hasScratchAllocation())
+        val->resetHasScratchAllocation();
     val->bumpUse();
 
     assertConsistent();
@@ -891,15 +897,6 @@ void Rewriter::_setupCall(RewriterVar* result, bool has_side_effects, const Rewr
     }
 #endif
 
-    for (RewriterVar* arg : args) {
-        arg->bumpUse();
-    }
-    for (RewriterVar* arg_xmm : args_xmm) {
-        arg_xmm->bumpUse();
-    }
-
-    assertConsistent();
-
     // Spill caller-saved registers:
     for (auto check_reg : caller_save_registers) {
         // check_reg.dump();
@@ -970,6 +967,15 @@ void Rewriter::_call(RewriterVar* result, bool has_side_effects, void* func_addr
 
     _setupCall(result, has_side_effects, args, args_xmm);
 
+    for (RewriterVar* arg : args) {
+        arg->bumpUse();
+    }
+    for (RewriterVar* arg_xmm : args_xmm) {
+        arg_xmm->bumpUse();
+    }
+
+    assertConsistent();
+
     // make sure setupCall doesn't use R11
     assert(vars_by_location.count(assembler::R11) == 0);
 
@@ -1019,6 +1025,17 @@ void RewriterVar::bumpUse() {
         for (Location loc : locations) {
             rewriter->vars_by_location.erase(loc);
         }
+
+        // releases allocated scratch space
+        if (hasScratchAllocation()) {
+            for (int i = 0; i < scratch_allocation.second; ++i) {
+                Location l = Location(Location::Scratch, (scratch_allocation.first + i) * 8);
+                assert(rewriter->vars_by_location[l] == LOCATION_PLACEHOLDER);
+                rewriter->vars_by_location.erase(l);
+            }
+            resetHasScratchAllocation();
+        }
+
         this->locations.clear();
     }
 }
@@ -1033,6 +1050,17 @@ void RewriterVar::releaseIfNoUses() {
         for (Location loc : locations) {
             rewriter->vars_by_location.erase(loc);
         }
+
+        // releases allocated scratch space
+        if (hasScratchAllocation()) {
+            for (int i = 0; i < scratch_allocation.second; ++i) {
+                Location l = Location(Location::Scratch, (scratch_allocation.first + i) * 8);
+                assert(rewriter->vars_by_location[l] == LOCATION_PLACEHOLDER);
+                rewriter->vars_by_location.erase(l);
+            }
+            resetHasScratchAllocation();
+        }
+
         this->locations.clear();
     }
 }
@@ -1326,7 +1354,6 @@ Location Rewriter::allocScratch() {
             return l;
         }
     }
-
     failed = true;
     return Location(Location::None, 0);
 }
@@ -1394,6 +1421,9 @@ int Rewriter::_allocate(RewriterVar* result, int n) {
                     vars_by_location[m] = LOCATION_PLACEHOLDER;
                 }
 
+                assert(result->scratch_allocation == std::make_pair(0, 0));
+                result->scratch_allocation = std::make_pair(a, n);
+
                 assembler::Register r = result->initializeInReg();
 
                 // TODO we could do something like we do for constants and only load

src/asm_writing/rewriter.h

@@ -240,13 +240,16 @@ private:
     void bumpUse();
     void releaseIfNoUses();
     bool isDoneUsing() { return next_use == uses.size(); }
+    bool hasScratchAllocation() const { return scratch_allocation.second > 0; }
+    void resetHasScratchAllocation() { scratch_allocation = std::make_pair(0, 0); }
 
     // Indicates if this variable is an arg, and if so, what location the arg is from.
     bool is_arg;
-    Location arg_loc;
-
     bool is_constant;
+
     uint64_t constant_value;
+    Location arg_loc;
+    std::pair<int /*offset*/, int /*size*/> scratch_allocation;
 
     llvm::SmallSet<std::tuple<int, uint64_t, bool>, 4> attr_guards; // used to detect duplicate guards
 

src/codegen/baseline_jit.cpp

@@ -712,7 +712,6 @@ void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, const
         for (int i = 6; i < args.size(); ++i) {
             assembler::Register reg = args[i]->getInReg(Location::any(), true);
             assembler->mov(reg, assembler::Indirect(assembler::RSP, sizeof(void*) * (i - 6)));
-            args[i]->bumpUse();
         }
         RewriterVar::SmallVector reg_args(args.begin(), args.begin() + 6);
         assert(reg_args.size() == 6);
@@ -747,6 +746,12 @@ void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, const
         pp_scratch_location -= 8;
     }
 
+    for (RewriterVar* arg : args) {
+        arg->bumpUse();
+    }
+
+    assertConsistent();
+
     StackInfo stack_info(pp_scratch_size, pp_scratch_location);
     pp_infos.emplace_back(PPInfo{ func_addr, pp_start, pp_end, std::move(setup_info), stack_info });