bjit: args in the entry block are non null + pass info on to successor block

when possible

src/codegen/ast_interpreter.cpp

@@ -128,7 +128,8 @@ private:
     Value visit_langPrimitive(AST_LangPrimitive* node);
 
     // for doc on 'exit_offset' have a look at JitFragmentWriter::num_bytes_exit and num_bytes_overlapping
-    void startJITing(CFGBlock* block, int exit_offset = 0);
+    void startJITing(CFGBlock* block, int exit_offset = 0,
+                     llvm::DenseSet<int> known_non_null_vregs = llvm::DenseSet<int>());
     void abortJITing();
     void finishJITing(CFGBlock* continue_block = NULL);
     Box* execJITedBlock(CFGBlock* b);
@@ -293,7 +294,7 @@ void ASTInterpreter::initArguments(BoxedClosure* _closure, BoxedGenerator* _gene
     assert(i == param_names.totalParameters());
 }
 
-void ASTInterpreter::startJITing(CFGBlock* block, int exit_offset) {
+void ASTInterpreter::startJITing(CFGBlock* block, int exit_offset, llvm::DenseSet<int> known_non_null_vregs) {
     assert(ENABLE_BASELINEJIT);
     assert(!jit);
 
@@ -308,7 +309,18 @@ void ASTInterpreter::startJITing(CFGBlock* block, int exit_offset) {
         exit_offset = 0;
     }
 
-    jit = code_block->newFragment(block, exit_offset);
+    // small optimization: we know that the passed arguments in the entry block are non zero
+    if (block == block->cfg->getStartingBlock() && block->predecessors.empty()) {
+        auto param_names = getMD()->param_names;
+        for (auto&& arg : param_names.arg_names) {
+            known_non_null_vregs.insert(arg->vreg);
+        }
+        if (param_names.vararg_name)
+            known_non_null_vregs.insert(param_names.vararg_name->vreg);
+        if (param_names.kwarg_name)
+            known_non_null_vregs.insert(param_names.kwarg_name->vreg);
+    }
+    jit = code_block->newFragment(block, exit_offset, std::move(known_non_null_vregs));
 }
 
 void ASTInterpreter::abortJITing() {
@@ -323,10 +335,20 @@ void ASTInterpreter::abortJITing() {
 void ASTInterpreter::finishJITing(CFGBlock* continue_block) {
     if (!jit)
         return;
-    int exit_offset = jit->finishCompilation();
+
+    int exit_offset = 0;
+    llvm::DenseSet<int> known_non_null;
+    std::tie(exit_offset, known_non_null) = jit->finishCompilation();
     jit.reset();
-    if (continue_block && !continue_block->code)
-        startJITing(continue_block, exit_offset);
+    if (continue_block && !continue_block->code) {
+        // check if we can reuse the known non null vreg set
+        if (continue_block->predecessors.size() == 1)
+            assert(current_block == continue_block->predecessors[0]);
+        else
+            known_non_null.clear();
+
+        startJITing(continue_block, exit_offset, std::move(known_non_null));
+    }
 }
 
 Box* ASTInterpreter::execJITedBlock(CFGBlock* b) {

src/codegen/baseline_jit.cpp

@@ -123,7 +123,8 @@ JitCodeBlock::JitCodeBlock(llvm::StringRef name)
     g.func_addr_registry.registerFunction(unique_name, code, code_size, NULL);
 }
 
-std::unique_ptr<JitFragmentWriter> JitCodeBlock::newFragment(CFGBlock* block, int patch_jump_offset) {
+std::unique_ptr<JitFragmentWriter> JitCodeBlock::newFragment(CFGBlock* block, int patch_jump_offset,
+                                                             llvm::DenseSet<int> known_non_null_vregs) {
     if (is_currently_writing || blocks_aborted.count(block))
         return std::unique_ptr<JitFragmentWriter>();
 
@@ -142,8 +143,9 @@ std::unique_ptr<JitFragmentWriter> JitCodeBlock::newFragment(CFGBlock* block, in
                                                std::vector<Location>(), bjit_allocatable_regs));
     std::unique_ptr<ICSlotRewrite> rewrite = ic_info->startRewrite("");
 
-    return std::unique_ptr<JitFragmentWriter>(new JitFragmentWriter(
-        block, std::move(ic_info), std::move(rewrite), fragment_offset, patch_jump_offset, a.getStartAddr(), *this));
+    return std::unique_ptr<JitFragmentWriter>(
+        new JitFragmentWriter(block, std::move(ic_info), std::move(rewrite), fragment_offset, patch_jump_offset,
+                              a.getStartAddr(), *this, std::move(known_non_null_vregs)));
 }
 
 void JitCodeBlock::fragmentAbort(bool not_enough_space) {
@@ -164,7 +166,8 @@ 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)
+                                     void* entry_code, JitCodeBlock& code_block,
+                                     llvm::DenseSet<int> known_non_null_vregs)
     : Rewriter(std::move(rewrite), 0, {}, /* needs_invalidation_support = */ false),
       block(block),
       code_offset(code_offset),
@@ -173,6 +176,7 @@ JitFragmentWriter::JitFragmentWriter(CFGBlock* block, std::unique_ptr<ICInfo> ic
       entry_code(entry_code),
       code_block(code_block),
       interp(0),
+      known_non_null_vregs(std::move(known_non_null_vregs)),
       ic_info(std::move(ic_info)) {
 
     added_changing_action = true;
@@ -766,14 +770,14 @@ void JitFragmentWriter::abortCompilation() {
     abort();
 }
 
-int JitFragmentWriter::finishCompilation() {
+std::pair<int, llvm::DenseSet<int>> JitFragmentWriter::finishCompilation() {
     RELEASE_ASSERT(!assembler->hasFailed(), "");
 
     commit();
     if (failed) {
         blocks_aborted.insert(block);
         code_block.fragmentAbort(false);
-        return 0;
+        return std::make_pair(0, llvm::DenseSet<int>());
     }
 
     if (assembler->hasFailed()) {
@@ -792,7 +796,7 @@ int JitFragmentWriter::finishCompilation() {
             // block for the next attempt.
             code_block.fragmentAbort(true /* not_enough_space */);
         }
-        return 0;
+        return std::make_pair(0, llvm::DenseSet<int>());
     }
 
     block->code = (void*)((uint64_t)entry_code + code_offset);
@@ -865,7 +869,7 @@ int JitFragmentWriter::finishCompilation() {
     registerGCTrackedICInfo(ic_info.release());
 #endif
 
-    return exit_info.num_bytes;
+    return std::make_pair(exit_info.num_bytes, std::move(known_non_null_vregs));
 }
 
 bool JitFragmentWriter::finishAssembly(int continue_offset, bool& should_fill_with_nops, bool& variable_size_slots) {

src/codegen/baseline_jit.h

@@ -181,7 +181,8 @@ private:
 public:
     JitCodeBlock(llvm::StringRef name);
 
-    std::unique_ptr<JitFragmentWriter> newFragment(CFGBlock* block, int patch_jump_offset = 0);
+    std::unique_ptr<JitFragmentWriter> newFragment(CFGBlock* block, int patch_jump_offset,
+                                                   llvm::DenseSet<int> known_non_null_vregs);
     bool shouldCreateNewBlock() const { return asm_failed || a.bytesLeft() < 128; }
     void fragmentAbort(bool not_enough_space);
     void fragmentFinished(int bytes_witten, int num_bytes_overlapping, void* next_fragment_start, ICInfo& ic_info);
@@ -216,7 +217,6 @@ private:
     RewriterVar* vregs_array;
     llvm::DenseMap<InternedString, RewriterVar*> local_syms;
     // keeps track which non block local vregs are known to have a non NULL value
-    // TODO: in the future we could reuse this information between different basic blocks
     llvm::DenseSet<int> known_non_null_vregs;
     std::unique_ptr<ICInfo> ic_info;
     llvm::SmallPtrSet<RewriterVar*, 4> var_is_a_python_bool;
@@ -241,7 +241,8 @@ private:
 
 public:
     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);
+                      int code_offset, int num_bytes_overlapping, void* entry_code, JitCodeBlock& code_block,
+                      llvm::DenseSet<int> known_non_null_vregs);
 
     RewriterVar* getInterp();
     RewriterVar* imm(uint64_t val);
@@ -311,7 +312,8 @@ public:
     void emitUncacheExcInfo();
 
     void abortCompilation();
-    int finishCompilation();
+    // returns pair of the number of bytes for the overwriteable jump and known non null vregs at end of current block
+    std::pair<int, llvm::DenseSet<int>> finishCompilation();
 
     bool finishAssembly(int continue_offset, bool& should_fill_with_nops, bool& variable_size_slots) override;