Merge pull request ruby#67 from Shopify/yjit-2d-blocks

YJIT: Use 2D array to group block versions

Author: maximecb

darray.h

@@ -60,6 +60,11 @@
 #define rb_darray_foreach(ary, idx_name, elem_ptr_var) \
     for (int idx_name = 0; idx_name < rb_darray_size(ary) && ((elem_ptr_var) = rb_darray_ref(ary, idx_name)); ++idx_name)
 
+// Iterate over valid indicies in the array in a for loop
+//
+#define rb_darray_for(ary, idx_name) \
+    for (int idx_name = 0; idx_name < rb_darray_size(ary); ++idx_name)
+
 // Make a dynamic array of a certain size. All bytes backing the elements are set to zero.
 // Return 1 on success and 0 on failure.
 //

ujit_codegen.c

@@ -578,13 +578,27 @@ gen_getinstancevariable(jitstate_t* jit, ctx_t* ctx)
 
 
 
+    /*
+    num_versions = count_block_versions(this_instruction);
 
+    if (num_versions > N)
+        return JIT_CANT_COMPILE;
 
-    /*
+    
     if (defer_compilation(this_instruction, ctx))
         return JIT_END_BLOCK;
 
+
     VALUE top_val = jit_peek_at_stack();
+    
+
+
+
+    class = get_ruby_class(top_val);
+
+
+
+    guard_object_class(class, current_instr);
     */
 
 
@@ -595,6 +609,9 @@ gen_getinstancevariable(jitstate_t* jit, ctx_t* ctx)
 
 
 
+
+
+
     // If the class uses the default allocator, instances should all be T_OBJECT
     // NOTE: This assumes nobody changes the allocator of the class after allocation.
     //       Eventually, we can encode whether an object is T_OBJECT or not

ujit_core.c

@@ -144,17 +144,26 @@ int ctx_diff(const ctx_t* src, const ctx_t* dst)
     return diff;
 }
 
-static block_t *
-get_first_version(const rb_iseq_t *iseq, unsigned idx)
+// Get all blocks for a particular place in an iseq.
+static rb_ujit_block_array_t
+get_version_array(const rb_iseq_t *iseq, unsigned idx)
 {
     struct rb_iseq_constant_body *body = iseq->body;
+
     if (rb_darray_size(body->ujit_blocks) == 0) {
         return NULL;
     }
+
     RUBY_ASSERT((unsigned)rb_darray_size(body->ujit_blocks) == body->iseq_size);
     return rb_darray_get(body->ujit_blocks, idx);
 }
 
+// Count the number of block versions matching a given blockid
+static size_t get_num_versions(blockid_t blockid)
+{
+    return rb_darray_size(get_version_array(blockid.iseq, blockid.idx));
+}
+
 // Keep track of a block version. Block should be fully constructed.
 static void
 add_block_version(blockid_t blockid, block_t* block)
@@ -181,19 +190,14 @@ add_block_version(blockid_t blockid, block_t* block)
 #endif
     }
 
-    block_t *first_version = get_first_version(iseq, blockid.idx);
+    RUBY_ASSERT((int32_t)blockid.idx < rb_darray_size(body->ujit_blocks));
+    rb_ujit_block_array_t *block_array_ref = rb_darray_ref(body->ujit_blocks, blockid.idx);
 
-    // If there exists a version for this block id
-    if (first_version != NULL) {
-        // Link to the next version in a linked list
-        RUBY_ASSERT(block->next == NULL);
-        block->next = first_version;
+    // Add the new block
+    if (!rb_darray_append(block_array_ref, block)) {
+        rb_bug("allocation failed");
     }
 
-    // Make new block the first version
-    rb_darray_set(body->ujit_blocks, blockid.idx, block);
-    RUBY_ASSERT(find_block_version(blockid, &block->ctx) != NULL);
-
     {
         // By writing the new block to the iseq, the iseq now
         // contains new references to Ruby objects. Run write barriers.
@@ -214,50 +218,28 @@ add_block_version(blockid_t blockid, block_t* block)
     }
 }
 
-// Count the number of block versions matching a given blockid
-static size_t count_block_versions(blockid_t blockid)
-{
-    size_t count = 0;
-    block_t *first_version = get_first_version(blockid.iseq, blockid.idx);
-
-    // For each version matching the blockid
-    for (block_t *version = first_version; version != NULL; version = version->next)
-    {
-        count += 1;
-    }
-
-    return count;
-}
-
 // Retrieve a basic block version for an (iseq, idx) tuple
 block_t* find_block_version(blockid_t blockid, const ctx_t* ctx)
 {
-    block_t *first_version = get_first_version(blockid.iseq, blockid.idx);
-
-    // If there exists a version for this block id
-    if (!first_version) return NULL;
+    rb_ujit_block_array_t versions = get_version_array(blockid.iseq, blockid.idx);
 
     // Best match found
     block_t* best_version = NULL;
     int best_diff = INT_MAX;
 
     // For each version matching the blockid
-    for (block_t* version = first_version; version != NULL; version = version->next)
-    {
+    rb_darray_for(versions, idx) {
+        block_t *version = rb_darray_get(versions, idx);
         int diff = ctx_diff(ctx, &version->ctx);
 
-        if (diff < best_diff)
-        {
+        // Note that we always prefer the first matching
+        // version because of inline-cache chains
+        if (diff < best_diff) {
             best_version = version;
             best_diff = diff;
         }
     }
 
-    if (best_version == NULL)
-    {
-        return NULL;
-    }
-
     return best_version;
 }
 
@@ -393,7 +375,7 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_
     ctx_t generic_ctx = DEFAULT_CTX;
     generic_ctx.stack_size = target_ctx->stack_size;
     generic_ctx.sp_offset = target_ctx->sp_offset;
-    if (count_block_versions(target) >= MAX_VERSIONS - 1)
+    if (get_num_versions(target) >= MAX_VERSIONS - 1)
     {
         fprintf(stderr, "version limit hit in branch_stub_hit\n");
         target_ctx = &generic_ctx;
@@ -559,7 +541,7 @@ void gen_direct_jump(
     ctx_t generic_ctx = DEFAULT_CTX;
     generic_ctx.stack_size = ctx->stack_size;
     generic_ctx.sp_offset = ctx->sp_offset;
-    if (count_block_versions(target0) >= MAX_VERSIONS - 1)
+    if (get_num_versions(target0) >= MAX_VERSIONS - 1)
     {
         fprintf(stderr, "version limit hit in gen_direct_jump\n");
         ctx = &generic_ctx;
@@ -658,6 +640,26 @@ ujit_free_block(block_t *block)
     free(block);
 }
 
+// Remove a block version without reordering the version array
+static bool
+block_array_remove(rb_ujit_block_array_t block_array, block_t *block)
+{
+    bool after_target = false;
+    block_t **element;
+    rb_darray_foreach(block_array, idx, element) {
+        if (after_target) {
+            rb_darray_set(block_array, idx - 1, *element);
+        }
+        else if (*element == block) {
+            after_target = true;
+        }
+    }
+
+    if (after_target) rb_darray_pop_back(block_array);
+
+    return after_target;
+}
+
 // Invalidate one specific block version
 void
 invalidate_block_version(block_t* block)
@@ -667,25 +669,11 @@ invalidate_block_version(block_t* block)
     // fprintf(stderr, "invalidating block (%p, %d)\n", block->blockid.iseq, block->blockid.idx);
     // fprintf(stderr, "block=%p\n", block);
 
-    block_t *first_block = get_first_version(iseq, block->blockid.idx);
-    RUBY_ASSERT(first_block != NULL);
-
-    // Remove references to this block
-    if (first_block == block) {
-        // Make the next block the new first version
-        rb_darray_set(iseq->body->ujit_blocks, block->blockid.idx, block->next);
-    }
-    else {
-        bool deleted = false;
-        for (block_t* cur = first_block; cur != NULL; cur = cur->next) {
-            if (cur->next == block) {
-                cur->next = cur->next->next;
-                deleted = true;
-                break;
-            }
-        }
-        RUBY_ASSERT(deleted);
-    }
+    // Remove this block from the version array
+    rb_ujit_block_array_t versions = get_version_array(iseq, block->blockid.idx);
+    RB_UNUSED_VAR(bool removed);
+    removed = block_array_remove(versions, block);
+    RUBY_ASSERT(removed);
 
     // Get a pointer to the generated code for this block
     uint8_t* code_ptr = cb_get_ptr(cb, block->start_pos);
@@ -741,7 +729,7 @@ invalidate_block_version(block_t* block)
     // Should check how it's used in exit and side-exit
     const void * const *handler_table = rb_vm_get_insns_address_table();
     void* handler_addr = (void*)handler_table[entry_opcode];
-    iseq->body->iseq_encoded[idx] = (VALUE)handler_addr;    
+    iseq->body->iseq_encoded[idx] = (VALUE)handler_addr;
 
     // TODO:
     // May want to recompile a new entry point (for interpreter entry blocks)

ujit_core.h

@@ -124,9 +124,6 @@ typedef struct ujit_block_version
     // Offsets for GC managed objects in the mainline code block
     int32_array_t gc_object_offsets;
 
-    // Next block version for this blockid (singly-linked list)
-    struct ujit_block_version *next;
-
     // GC managed objects that this block depend on
     struct {
         VALUE cc;

ujit_iface.c

@@ -379,12 +379,15 @@ ujit_blocks_for(VALUE mod, VALUE rb_iseq)
     }
 
     const rb_iseq_t *iseq = rb_iseqw_to_iseq(rb_iseq);
-    block_t **element;
+
     VALUE all_versions = rb_ary_new();
+    rb_darray_for(iseq->body->ujit_blocks, version_array_idx) {
+        rb_ujit_block_array_t versions = rb_darray_get(iseq->body->ujit_blocks, version_array_idx);
+
+        rb_darray_for(versions, block_idx) {
+            block_t *block = rb_darray_get(versions, block_idx);
 
-    rb_darray_foreach(iseq->body->ujit_blocks, idx, element) {
-        for (block_t *version = *element; version; version = version->next) {
-            VALUE rb_block = TypedData_Wrap_Struct(cUjitBlock, &ujit_block_type, version);
+            VALUE rb_block = TypedData_Wrap_Struct(cUjitBlock, &ujit_block_type, block);
             rb_ary_push(all_versions, rb_block);
         }
     }
@@ -677,11 +680,13 @@ print_ujit_stats(void)
 void
 rb_ujit_iseq_mark(const struct rb_iseq_constant_body *body)
 {
-    block_t **element;
-    rb_darray_foreach(body->ujit_blocks, idx, element) {
-        for (block_t *block = *element; block; block = block->next) {
-            rb_gc_mark_movable((VALUE)block->blockid.iseq);
+    rb_darray_for(body->ujit_blocks, version_array_idx) {
+        rb_ujit_block_array_t version_array = rb_darray_get(body->ujit_blocks, version_array_idx);
 
+        rb_darray_for(version_array, block_idx) {
+            block_t *block = rb_darray_get(version_array, block_idx);
+
+            rb_gc_mark_movable((VALUE)block->blockid.iseq);
             rb_gc_mark_movable(block->dependencies.cc);
             rb_gc_mark_movable(block->dependencies.cme);
             rb_gc_mark_movable(block->dependencies.iseq);
@@ -703,9 +708,12 @@ rb_ujit_iseq_mark(const struct rb_iseq_constant_body *body)
 void
 rb_ujit_iseq_update_references(const struct rb_iseq_constant_body *body)
 {
-    block_t **element;
-    rb_darray_foreach(body->ujit_blocks, idx, element) {
-        for (block_t *block = *element; block; block = block->next) {
+    rb_darray_for(body->ujit_blocks, version_array_idx) {
+        rb_ujit_block_array_t version_array = rb_darray_get(body->ujit_blocks, version_array_idx);
+
+        rb_darray_for(version_array, block_idx) {
+            block_t *block = rb_darray_get(version_array, block_idx);
+
             block->blockid.iseq = (const rb_iseq_t *)rb_gc_location((VALUE)block->blockid.iseq);
 
             block->dependencies.cc = rb_gc_location(block->dependencies.cc);
@@ -734,14 +742,15 @@ rb_ujit_iseq_update_references(const struct rb_iseq_constant_body *body)
 void
 rb_ujit_iseq_free(const struct rb_iseq_constant_body *body)
 {
-    block_t **element;
-    rb_darray_foreach(body->ujit_blocks, idx, element) {
-        block_t *block = *element;
-        while (block) {
-            block_t *next = block->next;
+    rb_darray_for(body->ujit_blocks, version_array_idx) {
+        rb_ujit_block_array_t version_array = rb_darray_get(body->ujit_blocks, version_array_idx);
+
+        rb_darray_for(version_array, block_idx) {
+            block_t *block = rb_darray_get(version_array, block_idx);
             ujit_free_block(block);
-            block = next;
         }
+
+        rb_darray_free(version_array);
     }
 
     rb_darray_free(body->ujit_blocks);

vm_core.h

@@ -303,7 +303,9 @@ pathobj_realpath(VALUE pathobj)
 /* Forward declarations */
 struct rb_mjit_unit;
 
+// List of YJIT block versions
 typedef rb_darray(struct ujit_block_version *) rb_ujit_block_array_t;
+typedef rb_darray(rb_ujit_block_array_t) rb_ujit_block_array_array_t;
 
 struct rb_iseq_constant_body {
     enum iseq_type {
@@ -441,7 +443,7 @@ struct rb_iseq_constant_body {
     struct rb_mjit_unit *jit_unit;
 #endif
 
-    rb_ujit_block_array_t ujit_blocks; // empty, or has a size equal to iseq_size
+    rb_ujit_block_array_array_t ujit_blocks; // empty, or has a size equal to iseq_size
 };
 
 /* T_IMEMO/iseq */