From 77b803cfc833233a49b497a9342f16c9237833e6 Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Wed, 1 Apr 2026 16:56:09 -0400 Subject: [PATCH 1/4] ZJIT: Clean up branching in polymorphic opt_send_without_block HIR construction (#16634) We have global regalloc now and don't need this manual threading. --- zjit/src/hir.rs | 108 +++++++-------------- zjit/src/hir/opt_tests.rs | 196 +++++++++++++++++++------------------- 2 files changed, 134 insertions(+), 170 deletions(-) diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index bc7be7498a3c81..2758682fbcd181 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -7793,81 +7793,45 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { break; } - { - fn new_branch_block( - fun: &mut Function, - cd: *const rb_call_data, - argc: usize, - opcode: u32, - new_type: Type, - insn_idx: u32, - exit_state: &FrameState, - locals_count: usize, - stack_count: usize, - join_block: BlockId, - ) -> BlockId { - let block = fun.new_block(insn_idx); - let self_param = fun.push_insn(block, Insn::Param); - let mut state = exit_state.clone(); - state.locals.clear(); - state.stack.clear(); - state.locals.extend((0..locals_count).map(|_| fun.push_insn(block, Insn::Param))); - state.stack.extend((0..stack_count).map(|_| fun.push_insn(block, Insn::Param))); - let snapshot = fun.push_insn(block, Insn::Snapshot { state: state.clone() }); - let args = state.stack_pop_n(argc).unwrap(); - let recv = state.stack_pop().unwrap(); - let refined_recv = fun.push_insn(block, Insn::RefineType { val: recv, new_type }); - state.replace(recv, refined_recv); - let send = fun.push_insn(block, Insn::Send { recv: refined_recv, cd, block: None, args, state: snapshot, reason: Uncategorized(opcode) }); - state.stack_push(send); - fun.push_insn(block, Insn::Jump(BranchEdge { target: join_block, args: state.as_args(self_param) })); - block - } - let branch_insn_idx = exit_state.insn_idx as u32; - let locals_count = state.locals.len(); - let stack_count = state.stack.len(); - let recv = state.stack_topn(argc as usize)?; // args are on top - let entry_args = state.as_args(self_param); - if let Some(summary) = fun.polymorphic_summary(&profiles, recv, exit_state.insn_idx) { - let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx)); - // Dedup by expected type so immediate/heap variants - // under the same Ruby class can still get separate branches. - let mut seen_types = Vec::with_capacity(summary.buckets().len()); - for &profiled_type in summary.buckets() { - if profiled_type.is_empty() { break; } - let expected = Type::from_profiled_type(profiled_type); - if seen_types.iter().any(|ty: &Type| ty.bit_equal(expected)) { - continue; - } - seen_types.push(expected); - let has_type = fun.push_insn(block, Insn::HasType { val: recv, expected }); - let iftrue_block = - new_branch_block(&mut fun, cd, argc as usize, opcode, expected, branch_insn_idx, &exit_state, locals_count, stack_count, join_block); - let target = BranchEdge { target: iftrue_block, args: entry_args.clone() }; - fun.push_insn(block, Insn::IfTrue { val: has_type, target }); + let branch_insn_idx = exit_state.insn_idx as u32; + let args = state.stack_pop_n(argc as usize)?; + let recv = state.stack_pop()?; + + if let Some(summary) = fun.polymorphic_summary(&profiles, recv, exit_state.insn_idx) { + let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx)); + let join_param = fun.push_insn(join_block, Insn::Param); + // Dedup by expected type so immediate/heap variants + // under the same Ruby class can still get separate branches. + let mut seen_types = Vec::with_capacity(summary.buckets().len()); + for &profiled_type in summary.buckets() { + if profiled_type.is_empty() { break; } + let expected = Type::from_profiled_type(profiled_type); + if seen_types.iter().any(|ty: &Type| ty.bit_equal(expected)) { + continue; } - // Continue compilation from the join block at the next instruction. - // Make a copy of the current state without the args (pop the receiver - // and push the result) because we just use the locals/stack sizes to - // make the right number of Params - let mut join_state = state.clone(); - join_state.stack_pop_n(argc as usize)?; - queue.push_back((join_state, join_block, insn_idx, local_inval)); - // In the fallthrough case, do a generic interpreter send and then join. - let args = state.stack_pop_n(argc as usize)?; - let recv = state.stack_pop()?; - let reason = SendWithoutBlockPolymorphicFallback; - let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason }); - state.stack_push(send); - fun.push_insn(block, Insn::Jump(BranchEdge { target: join_block, args: state.as_args(self_param) })); - break; // End the block + seen_types.push(expected); + let has_type = fun.push_insn(block, Insn::HasType { val: recv, expected }); + let iftrue_block = fun.new_block(branch_insn_idx); + let target = BranchEdge { target: iftrue_block, args: vec![] }; + fun.push_insn(block, Insn::IfTrue { val: has_type, target }); + + let refined_recv = fun.push_insn(iftrue_block, Insn::RefineType { val: recv, new_type: expected }); + let send = fun.push_insn(iftrue_block, Insn::Send { recv: refined_recv, cd, block: None, args: args.clone(), state: exit_id, reason: Uncategorized(opcode) }); + fun.push_insn(iftrue_block, Insn::Jump(BranchEdge { target: join_block, args: vec![send] })); } - } + // In the fallthrough case, do a generic interpreter send and then join. + let reason = SendWithoutBlockPolymorphicFallback; + let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason }); + fun.push_insn(block, Insn::Jump(BranchEdge { target: join_block, args: vec![send] })); - let args = state.stack_pop_n(argc as usize)?; - let recv = state.stack_pop()?; - let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason: Uncategorized(opcode) }); - state.stack_push(send); + // Continue compilation in the join_block + block = join_block; + state.stack_push(join_param); + } else { + // Maybe monomorphic; handled in type_specialize + let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason: Uncategorized(opcode) }); + state.stack_push(send); + } } YARVINSN_send => { let cd: *const rb_call_data = get_arg(pc, 0).as_ptr(); diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs index fde0ab0583d279..abdb59ce486097 100644 --- a/zjit/src/hir/opt_tests.rs +++ b/zjit/src/hir/opt_tests.rs @@ -7861,19 +7861,19 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v15, bb5(v9, v10, v10) - v24:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v24) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): - v20:ObjectSubclass[class_exact:C] = RefineType v18, ObjectSubclass[class_exact:C] + v16:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v16, bb5() + v21:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v21) + bb5(): + v18:ObjectSubclass[class_exact:C] = RefineType v10, ObjectSubclass[class_exact:C] PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, foo@0x1010, cme:0x1018) - v37:BasicObject = GetIvar v20, :@foo - Jump bb4(v16, v17, v37) - bb4(v26:BasicObject, v27:BasicObject, v28:BasicObject): + v30:BasicObject = GetIvar v18, :@foo + Jump bb4(v30) + bb4(v15:BasicObject): CheckInterrupts - Return v28 + Return v15 "); } @@ -14453,29 +14453,29 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v15, bb5(v9, v10, v10) - v24:CBool = HasType v10, ObjectSubclass[class_exact:D] - IfTrue v24, bb6(v9, v10, v10) - v33:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v33) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v16:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v16, bb5() + v21:CBool = HasType v10, ObjectSubclass[class_exact:D] + IfTrue v21, bb6() + v26:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v26) + bb5(): PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, foo@0x1010, cme:0x1018) - v55:Fixnum[3] = Const Value(3) - Jump bb4(v16, v17, v55) - bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): + v44:Fixnum[3] = Const Value(3) + Jump bb4(v44) + bb6(): PatchPoint NoSingletonClass(D@0x1040) PatchPoint MethodRedefined(D@0x1040, foo@0x1010, cme:0x1048) - v56:Fixnum[4] = Const Value(4) - Jump bb4(v25, v26, v56) - bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): - v40:Fixnum[2] = Const Value(2) + v45:Fixnum[4] = Const Value(4) + Jump bb4(v45) + bb4(v15:BasicObject): + v29:Fixnum[2] = Const Value(2) PatchPoint MethodRedefined(Integer@0x1070, +@0x1078, cme:0x1080) - v59:Fixnum = GuardType v37, Fixnum - v60:Fixnum = FixnumAdd v59, v40 + v48:Fixnum = GuardType v15, Fixnum + v49:Fixnum = FixnumAdd v48, v29 CheckInterrupts - Return v60 + Return v49 "); } @@ -14505,24 +14505,24 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v15, bb5(v9, v10, v10) - v24:CBool = HasType v10, Fixnum - IfTrue v24, bb6(v9, v10, v10) - v33:BasicObject = Send v10, :itself # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v33) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): - v20:ObjectSubclass[class_exact:C] = RefineType v18, ObjectSubclass[class_exact:C] + v16:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v16, bb5() + v21:CBool = HasType v10, Fixnum + IfTrue v21, bb6() + v26:BasicObject = Send v10, :itself # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v26) + bb5(): + v18:ObjectSubclass[class_exact:C] = RefineType v10, ObjectSubclass[class_exact:C] PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, itself@0x1010, cme:0x1018) - Jump bb4(v16, v17, v20) - bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): - v29:Fixnum = RefineType v27, Fixnum + Jump bb4(v18) + bb6(): + v23:Fixnum = RefineType v10, Fixnum PatchPoint MethodRedefined(Integer@0x1040, itself@0x1010, cme:0x1018) - Jump bb4(v25, v26, v29) - bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): + Jump bb4(v23) + bb4(v15:BasicObject): CheckInterrupts - Return v37 + Return v15 "); } @@ -14558,25 +14558,25 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :x@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, Fixnum - IfTrue v15, bb5(v9, v10, v10) - v24:CBool = HasType v10, Bignum - IfTrue v24, bb6(v9, v10, v10) - v33:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v33) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): - v20:Fixnum = RefineType v18, Fixnum + v16:CBool = HasType v10, Fixnum + IfTrue v16, bb5() + v21:CBool = HasType v10, Bignum + IfTrue v21, bb6() + v26:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v26) + bb5(): + v18:Fixnum = RefineType v10, Fixnum PatchPoint MethodRedefined(Integer@0x1008, to_s@0x1010, cme:0x1018) - v46:StringExact = CCallVariadic v20, :Integer#to_s@0x1040 - Jump bb4(v16, v17, v46) - bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): - v29:Bignum = RefineType v27, Bignum + v35:StringExact = CCallVariadic v18, :Integer#to_s@0x1040 + Jump bb4(v35) + bb6(): + v23:Bignum = RefineType v10, Bignum PatchPoint MethodRedefined(Integer@0x1008, to_s@0x1010, cme:0x1018) - v49:StringExact = CCallVariadic v29, :Integer#to_s@0x1040 - Jump bb4(v25, v26, v49) - bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): + v38:StringExact = CCallVariadic v23, :Integer#to_s@0x1040 + Jump bb4(v38) + bb4(v15:BasicObject): CheckInterrupts - Return v37 + Return v15 "); } @@ -14609,25 +14609,25 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :x@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, Flonum - IfTrue v15, bb5(v9, v10, v10) - v24:CBool = HasType v10, HeapFloat - IfTrue v24, bb6(v9, v10, v10) - v33:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v33) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): - v20:Flonum = RefineType v18, Flonum + v16:CBool = HasType v10, Flonum + IfTrue v16, bb5() + v21:CBool = HasType v10, HeapFloat + IfTrue v21, bb6() + v26:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v26) + bb5(): + v18:Flonum = RefineType v10, Flonum PatchPoint MethodRedefined(Float@0x1008, to_s@0x1010, cme:0x1018) - v46:BasicObject = CCallWithFrame v20, :Float#to_s@0x1040 - Jump bb4(v16, v17, v46) - bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): - v29:HeapFloat = RefineType v27, HeapFloat + v35:BasicObject = CCallWithFrame v18, :Float#to_s@0x1040 + Jump bb4(v35) + bb6(): + v23:HeapFloat = RefineType v10, HeapFloat PatchPoint MethodRedefined(Float@0x1008, to_s@0x1010, cme:0x1018) - v49:BasicObject = CCallWithFrame v29, :Float#to_s@0x1040 - Jump bb4(v25, v26, v49) - bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): + v38:BasicObject = CCallWithFrame v23, :Float#to_s@0x1040 + Jump bb4(v38) + bb4(v15:BasicObject): CheckInterrupts - Return v37 + Return v15 "); } @@ -14660,25 +14660,25 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :x@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, StaticSymbol - IfTrue v15, bb5(v9, v10, v10) - v24:CBool = HasType v10, DynamicSymbol - IfTrue v24, bb6(v9, v10, v10) - v33:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v33) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): - v20:StaticSymbol = RefineType v18, StaticSymbol + v16:CBool = HasType v10, StaticSymbol + IfTrue v16, bb5() + v21:CBool = HasType v10, DynamicSymbol + IfTrue v21, bb6() + v26:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v26) + bb5(): + v18:StaticSymbol = RefineType v10, StaticSymbol PatchPoint MethodRedefined(Symbol@0x1008, to_s@0x1010, cme:0x1018) - v48:StringExact = InvokeBuiltin leaf , v20 - Jump bb4(v16, v17, v48) - bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): - v29:DynamicSymbol = RefineType v27, DynamicSymbol + v37:StringExact = InvokeBuiltin leaf , v18 + Jump bb4(v37) + bb6(): + v23:DynamicSymbol = RefineType v10, DynamicSymbol PatchPoint MethodRedefined(Symbol@0x1008, to_s@0x1010, cme:0x1018) - v49:StringExact = InvokeBuiltin leaf , v29 - Jump bb4(v25, v26, v49) - bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): + v38:StringExact = InvokeBuiltin leaf , v23 + Jump bb4(v38) + bb4(v15:BasicObject): CheckInterrupts - Return v37 + Return v15 "); } @@ -14715,18 +14715,18 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v15:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v15, bb5(v9, v10, v10) - v24:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v9, v10, v24) - bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v16:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v16, bb5() + v21:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v21) + bb5(): PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, foo@0x1010, cme:0x1018) - v38:Fixnum[3] = Const Value(3) - Jump bb4(v16, v17, v38) - bb4(v26:BasicObject, v27:BasicObject, v28:BasicObject): + v31:Fixnum[3] = Const Value(3) + Jump bb4(v31) + bb4(v15:BasicObject): CheckInterrupts - Return v28 + Return v15 "); } From c34139963e708546a9822463cf906dda3a813c64 Mon Sep 17 00:00:00 2001 From: Peter Zhu Date: Tue, 31 Mar 2026 20:55:56 -0400 Subject: [PATCH 2/4] Fix memory leak in pm_parse_process pm_parse_process initializes the index_lookup_table but nothing seems to use it after it has been allocated. However, pm_compile_scope_node will overwrite the index_lookup_table and cause it to leak memory. This can be seen during bootup with the following memory leaks reported by ASAN: #0 0x60dba31b7af3 in malloc #1 0x60dba32e0718 in rb_gc_impl_malloc gc/default/default.c:8287:5 #2 0x60dba32c7aa7 in ruby_xmalloc_body gc.c:5373:12 #3 0x60dba32c4a54 in ruby_xmalloc gc.c:5355:34 #4 0x60dba3260314 in pm_index_lookup_table_init_heap prism_compile.h:89:29 #5 0x60dba3209388 in pm_parse_process prism_compile.c:11366:5 --- prism_compile.c | 6 ------ 1 file changed, 6 deletions(-) diff --git a/prism_compile.c b/prism_compile.c index 85e8a2cdfa40d7..d30785bb883dc2 100644 --- a/prism_compile.c +++ b/prism_compile.c @@ -11362,12 +11362,6 @@ pm_parse_process(pm_parse_result_t *result, pm_node_t *node, VALUE *script_lines pm_intern_constants_ctx_t intern_ctx = { .constants = scope_node->constants, .encoding = scope_node->encoding, .index = 0 }; pm_parser_constants_each(parser, pm_intern_constants_callback, &intern_ctx); - pm_constant_id_list_t *locals = &scope_node->locals; - pm_index_lookup_table_init_heap(&scope_node->index_lookup_table, (int) constants_size); - for (size_t index = 0; index < locals->size; index++) { - pm_index_lookup_table_insert(&scope_node->index_lookup_table, locals->ids[index], (int) index); - } - // If we got here, this is a success and we can return Qnil to indicate that // no error should be raised. result->parsed = true; From 25ca647db5ea7ee7b092d0501ace5134050eba30 Mon Sep 17 00:00:00 2001 From: Takashi Kokubun Date: Wed, 1 Apr 2026 16:29:27 -0700 Subject: [PATCH 3/4] ZJIT: Hook into GC W^X callbacks for compaction (#16635) --- gc.c | 6 ++++++ zjit.h | 2 ++ zjit/src/asm/mod.rs | 4 ++++ zjit/src/gc.rs | 30 ++++++++++++++++++++++++------ zjit/src/virtualmem.rs | 16 ++++++++++++++++ 5 files changed, 52 insertions(+), 6 deletions(-) diff --git a/gc.c b/gc.c index 50868ea390b001..1f441b60e56f0c 100644 --- a/gc.c +++ b/gc.c @@ -2475,6 +2475,9 @@ rb_gc_before_updating_jit_code(void) #if USE_YJIT rb_yjit_mark_all_writeable(); #endif +#if USE_ZJIT + rb_zjit_mark_all_writable(); +#endif } /* @@ -2488,6 +2491,9 @@ rb_gc_after_updating_jit_code(void) #if USE_YJIT rb_yjit_mark_all_executable(); #endif +#if USE_ZJIT + rb_zjit_mark_all_executable(); +#endif } static void diff --git a/zjit.h b/zjit.h index e96caa257c6b91..d67c8b82f21efe 100644 --- a/zjit.h +++ b/zjit.h @@ -40,6 +40,8 @@ void rb_zjit_invalidate_no_ep_escape(const rb_iseq_t *iseq); void rb_zjit_constant_state_changed(ID id); void rb_zjit_iseq_mark(void *payload); void rb_zjit_iseq_update_references(void *payload); +void rb_zjit_mark_all_writable(void); +void rb_zjit_mark_all_executable(void); void rb_zjit_iseq_free(const rb_iseq_t *iseq); void rb_zjit_before_ractor_spawn(void); void rb_zjit_tracing_invalidate_all(void); diff --git a/zjit/src/asm/mod.rs b/zjit/src/asm/mod.rs index 1e8f3414ec4f81..6583476594a3bf 100644 --- a/zjit/src/asm/mod.rs +++ b/zjit/src/asm/mod.rs @@ -294,6 +294,10 @@ impl CodeBlock { } /// Make all the code in the region executable. Call this at the end of a write session. + pub fn mark_all_writable(&mut self) { + self.mem_block.borrow_mut().mark_all_writable(); + } + pub fn mark_all_executable(&mut self) { self.mem_block.borrow_mut().mark_all_executable(); } diff --git a/zjit/src/gc.rs b/zjit/src/gc.rs index 239b71d5f48754..7f5bc7891f20cd 100644 --- a/zjit/src/gc.rs +++ b/zjit/src/gc.rs @@ -158,7 +158,9 @@ fn iseq_version_update_references(mut version: IseqVersionRef) { } } - // Move objects baked in JIT code + // Move objects baked in JIT code. + // The code region is already writable because rb_zjit_mark_all_writable() was called + // before the GC update_references phase. We write directly to avoid per-page mprotect calls. let cb = ZJITState::get_code_block(); for &offset in unsafe { version.as_ref() }.gc_offsets.iter() { let value_ptr: *const u8 = offset.raw_ptr(cb); @@ -170,13 +172,10 @@ fn iseq_version_update_references(mut version: IseqVersionRef) { // Only write when the VALUE moves, to be copy-on-write friendly. if new_addr != object { - for (byte_idx, &byte) in new_addr.as_u64().to_le_bytes().iter().enumerate() { - let byte_code_ptr = offset.add_bytes(byte_idx); - cb.write_mem(byte_code_ptr, byte).expect("patching existing code should be within bounds"); - } + let value_ptr = value_ptr as *mut VALUE; + unsafe { value_ptr.write_unaligned(new_addr) }; } } - cb.mark_all_executable(); } /// Append a set of gc_offsets to the iseq's payload @@ -211,6 +210,25 @@ fn ranges_overlap(left: &Range, right: &Range) -> bool where T: Partial left.start < right.end && right.start < left.end } +/// GC callback for making all JIT code writable before updating references in bulk. +/// This avoids toggling W^X permissions per-page during GC compaction. +#[unsafe(no_mangle)] +pub extern "C" fn rb_zjit_mark_all_writable() { + if !ZJITState::has_instance() { + return; + } + ZJITState::get_code_block().mark_all_writable(); +} + +/// GC callback for making all JIT code executable after updating references in bulk. +#[unsafe(no_mangle)] +pub extern "C" fn rb_zjit_mark_all_executable() { + if !ZJITState::has_instance() { + return; + } + ZJITState::get_code_block().mark_all_executable(); +} + /// Callback for marking GC objects inside [crate::invariants::Invariants]. #[unsafe(no_mangle)] pub extern "C" fn rb_zjit_root_mark() { diff --git a/zjit/src/virtualmem.rs b/zjit/src/virtualmem.rs index 9741a7b13867d5..c2a1e13a5dee83 100644 --- a/zjit/src/virtualmem.rs +++ b/zjit/src/virtualmem.rs @@ -265,6 +265,22 @@ impl VirtualMemory { memory_usage_bytes + self.page_size_bytes < memory_limit_bytes } + /// Make all the code in the region writable. Call this before bulk writes (e.g. GC + /// reference updates). See [Self] for usual usage flow. + pub fn mark_all_writable(&mut self) { + self.current_write_page = None; + + let region_start = self.region_start; + let mapped_region_bytes: u32 = self.mapped_region_bytes.try_into().unwrap(); + + // Make mapped region writable + if mapped_region_bytes > 0 { + if !self.allocator.mark_writable(region_start.as_ptr(), mapped_region_bytes) { + panic!("Cannot make JIT memory region writable"); + } + } + } + /// Make all the code in the region executable. Call this at the end of a write session. /// See [Self] for usual usage flow. pub fn mark_all_executable(&mut self) { From f0c73671af8376640eb9ae1d0f4647d3b4f222a8 Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Wed, 1 Apr 2026 19:48:34 -0400 Subject: [PATCH 4/4] Revert "ZJIT: Clean up branching in polymorphic opt_send_without_block HIR construction (#16634)" This reverts commit 77b803cfc833233a49b497a9342f16c9237833e6. There is a new NULL dereference in tool/test/testunit/test_sorting.rb that seem to have started with this commit. --- zjit/src/hir.rs | 108 ++++++++++++++------- zjit/src/hir/opt_tests.rs | 196 +++++++++++++++++++------------------- 2 files changed, 170 insertions(+), 134 deletions(-) diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 2758682fbcd181..bc7be7498a3c81 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -7793,45 +7793,81 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { break; } - let branch_insn_idx = exit_state.insn_idx as u32; - let args = state.stack_pop_n(argc as usize)?; - let recv = state.stack_pop()?; - - if let Some(summary) = fun.polymorphic_summary(&profiles, recv, exit_state.insn_idx) { - let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx)); - let join_param = fun.push_insn(join_block, Insn::Param); - // Dedup by expected type so immediate/heap variants - // under the same Ruby class can still get separate branches. - let mut seen_types = Vec::with_capacity(summary.buckets().len()); - for &profiled_type in summary.buckets() { - if profiled_type.is_empty() { break; } - let expected = Type::from_profiled_type(profiled_type); - if seen_types.iter().any(|ty: &Type| ty.bit_equal(expected)) { - continue; + { + fn new_branch_block( + fun: &mut Function, + cd: *const rb_call_data, + argc: usize, + opcode: u32, + new_type: Type, + insn_idx: u32, + exit_state: &FrameState, + locals_count: usize, + stack_count: usize, + join_block: BlockId, + ) -> BlockId { + let block = fun.new_block(insn_idx); + let self_param = fun.push_insn(block, Insn::Param); + let mut state = exit_state.clone(); + state.locals.clear(); + state.stack.clear(); + state.locals.extend((0..locals_count).map(|_| fun.push_insn(block, Insn::Param))); + state.stack.extend((0..stack_count).map(|_| fun.push_insn(block, Insn::Param))); + let snapshot = fun.push_insn(block, Insn::Snapshot { state: state.clone() }); + let args = state.stack_pop_n(argc).unwrap(); + let recv = state.stack_pop().unwrap(); + let refined_recv = fun.push_insn(block, Insn::RefineType { val: recv, new_type }); + state.replace(recv, refined_recv); + let send = fun.push_insn(block, Insn::Send { recv: refined_recv, cd, block: None, args, state: snapshot, reason: Uncategorized(opcode) }); + state.stack_push(send); + fun.push_insn(block, Insn::Jump(BranchEdge { target: join_block, args: state.as_args(self_param) })); + block + } + let branch_insn_idx = exit_state.insn_idx as u32; + let locals_count = state.locals.len(); + let stack_count = state.stack.len(); + let recv = state.stack_topn(argc as usize)?; // args are on top + let entry_args = state.as_args(self_param); + if let Some(summary) = fun.polymorphic_summary(&profiles, recv, exit_state.insn_idx) { + let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx)); + // Dedup by expected type so immediate/heap variants + // under the same Ruby class can still get separate branches. + let mut seen_types = Vec::with_capacity(summary.buckets().len()); + for &profiled_type in summary.buckets() { + if profiled_type.is_empty() { break; } + let expected = Type::from_profiled_type(profiled_type); + if seen_types.iter().any(|ty: &Type| ty.bit_equal(expected)) { + continue; + } + seen_types.push(expected); + let has_type = fun.push_insn(block, Insn::HasType { val: recv, expected }); + let iftrue_block = + new_branch_block(&mut fun, cd, argc as usize, opcode, expected, branch_insn_idx, &exit_state, locals_count, stack_count, join_block); + let target = BranchEdge { target: iftrue_block, args: entry_args.clone() }; + fun.push_insn(block, Insn::IfTrue { val: has_type, target }); } - seen_types.push(expected); - let has_type = fun.push_insn(block, Insn::HasType { val: recv, expected }); - let iftrue_block = fun.new_block(branch_insn_idx); - let target = BranchEdge { target: iftrue_block, args: vec![] }; - fun.push_insn(block, Insn::IfTrue { val: has_type, target }); - - let refined_recv = fun.push_insn(iftrue_block, Insn::RefineType { val: recv, new_type: expected }); - let send = fun.push_insn(iftrue_block, Insn::Send { recv: refined_recv, cd, block: None, args: args.clone(), state: exit_id, reason: Uncategorized(opcode) }); - fun.push_insn(iftrue_block, Insn::Jump(BranchEdge { target: join_block, args: vec![send] })); + // Continue compilation from the join block at the next instruction. + // Make a copy of the current state without the args (pop the receiver + // and push the result) because we just use the locals/stack sizes to + // make the right number of Params + let mut join_state = state.clone(); + join_state.stack_pop_n(argc as usize)?; + queue.push_back((join_state, join_block, insn_idx, local_inval)); + // In the fallthrough case, do a generic interpreter send and then join. + let args = state.stack_pop_n(argc as usize)?; + let recv = state.stack_pop()?; + let reason = SendWithoutBlockPolymorphicFallback; + let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason }); + state.stack_push(send); + fun.push_insn(block, Insn::Jump(BranchEdge { target: join_block, args: state.as_args(self_param) })); + break; // End the block } - // In the fallthrough case, do a generic interpreter send and then join. - let reason = SendWithoutBlockPolymorphicFallback; - let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason }); - fun.push_insn(block, Insn::Jump(BranchEdge { target: join_block, args: vec![send] })); - - // Continue compilation in the join_block - block = join_block; - state.stack_push(join_param); - } else { - // Maybe monomorphic; handled in type_specialize - let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason: Uncategorized(opcode) }); - state.stack_push(send); } + + let args = state.stack_pop_n(argc as usize)?; + let recv = state.stack_pop()?; + let send = fun.push_insn(block, Insn::Send { recv, cd, block: None, args, state: exit_id, reason: Uncategorized(opcode) }); + state.stack_push(send); } YARVINSN_send => { let cd: *const rb_call_data = get_arg(pc, 0).as_ptr(); diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs index abdb59ce486097..fde0ab0583d279 100644 --- a/zjit/src/hir/opt_tests.rs +++ b/zjit/src/hir/opt_tests.rs @@ -7861,19 +7861,19 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v16, bb5() - v21:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v21) - bb5(): - v18:ObjectSubclass[class_exact:C] = RefineType v10, ObjectSubclass[class_exact:C] + v15:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v15, bb5(v9, v10, v10) + v24:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v24) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v20:ObjectSubclass[class_exact:C] = RefineType v18, ObjectSubclass[class_exact:C] PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, foo@0x1010, cme:0x1018) - v30:BasicObject = GetIvar v18, :@foo - Jump bb4(v30) - bb4(v15:BasicObject): + v37:BasicObject = GetIvar v20, :@foo + Jump bb4(v16, v17, v37) + bb4(v26:BasicObject, v27:BasicObject, v28:BasicObject): CheckInterrupts - Return v15 + Return v28 "); } @@ -14453,29 +14453,29 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v16, bb5() - v21:CBool = HasType v10, ObjectSubclass[class_exact:D] - IfTrue v21, bb6() - v26:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v26) - bb5(): + v15:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v15, bb5(v9, v10, v10) + v24:CBool = HasType v10, ObjectSubclass[class_exact:D] + IfTrue v24, bb6(v9, v10, v10) + v33:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v33) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, foo@0x1010, cme:0x1018) - v44:Fixnum[3] = Const Value(3) - Jump bb4(v44) - bb6(): + v55:Fixnum[3] = Const Value(3) + Jump bb4(v16, v17, v55) + bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): PatchPoint NoSingletonClass(D@0x1040) PatchPoint MethodRedefined(D@0x1040, foo@0x1010, cme:0x1048) - v45:Fixnum[4] = Const Value(4) - Jump bb4(v45) - bb4(v15:BasicObject): - v29:Fixnum[2] = Const Value(2) + v56:Fixnum[4] = Const Value(4) + Jump bb4(v25, v26, v56) + bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): + v40:Fixnum[2] = Const Value(2) PatchPoint MethodRedefined(Integer@0x1070, +@0x1078, cme:0x1080) - v48:Fixnum = GuardType v15, Fixnum - v49:Fixnum = FixnumAdd v48, v29 + v59:Fixnum = GuardType v37, Fixnum + v60:Fixnum = FixnumAdd v59, v40 CheckInterrupts - Return v49 + Return v60 "); } @@ -14505,24 +14505,24 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v16, bb5() - v21:CBool = HasType v10, Fixnum - IfTrue v21, bb6() - v26:BasicObject = Send v10, :itself # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v26) - bb5(): - v18:ObjectSubclass[class_exact:C] = RefineType v10, ObjectSubclass[class_exact:C] + v15:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v15, bb5(v9, v10, v10) + v24:CBool = HasType v10, Fixnum + IfTrue v24, bb6(v9, v10, v10) + v33:BasicObject = Send v10, :itself # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v33) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v20:ObjectSubclass[class_exact:C] = RefineType v18, ObjectSubclass[class_exact:C] PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, itself@0x1010, cme:0x1018) - Jump bb4(v18) - bb6(): - v23:Fixnum = RefineType v10, Fixnum + Jump bb4(v16, v17, v20) + bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): + v29:Fixnum = RefineType v27, Fixnum PatchPoint MethodRedefined(Integer@0x1040, itself@0x1010, cme:0x1018) - Jump bb4(v23) - bb4(v15:BasicObject): + Jump bb4(v25, v26, v29) + bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): CheckInterrupts - Return v15 + Return v37 "); } @@ -14558,25 +14558,25 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :x@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, Fixnum - IfTrue v16, bb5() - v21:CBool = HasType v10, Bignum - IfTrue v21, bb6() - v26:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v26) - bb5(): - v18:Fixnum = RefineType v10, Fixnum + v15:CBool = HasType v10, Fixnum + IfTrue v15, bb5(v9, v10, v10) + v24:CBool = HasType v10, Bignum + IfTrue v24, bb6(v9, v10, v10) + v33:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v33) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v20:Fixnum = RefineType v18, Fixnum PatchPoint MethodRedefined(Integer@0x1008, to_s@0x1010, cme:0x1018) - v35:StringExact = CCallVariadic v18, :Integer#to_s@0x1040 - Jump bb4(v35) - bb6(): - v23:Bignum = RefineType v10, Bignum + v46:StringExact = CCallVariadic v20, :Integer#to_s@0x1040 + Jump bb4(v16, v17, v46) + bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): + v29:Bignum = RefineType v27, Bignum PatchPoint MethodRedefined(Integer@0x1008, to_s@0x1010, cme:0x1018) - v38:StringExact = CCallVariadic v23, :Integer#to_s@0x1040 - Jump bb4(v38) - bb4(v15:BasicObject): + v49:StringExact = CCallVariadic v29, :Integer#to_s@0x1040 + Jump bb4(v25, v26, v49) + bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): CheckInterrupts - Return v15 + Return v37 "); } @@ -14609,25 +14609,25 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :x@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, Flonum - IfTrue v16, bb5() - v21:CBool = HasType v10, HeapFloat - IfTrue v21, bb6() - v26:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v26) - bb5(): - v18:Flonum = RefineType v10, Flonum + v15:CBool = HasType v10, Flonum + IfTrue v15, bb5(v9, v10, v10) + v24:CBool = HasType v10, HeapFloat + IfTrue v24, bb6(v9, v10, v10) + v33:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v33) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v20:Flonum = RefineType v18, Flonum PatchPoint MethodRedefined(Float@0x1008, to_s@0x1010, cme:0x1018) - v35:BasicObject = CCallWithFrame v18, :Float#to_s@0x1040 - Jump bb4(v35) - bb6(): - v23:HeapFloat = RefineType v10, HeapFloat + v46:BasicObject = CCallWithFrame v20, :Float#to_s@0x1040 + Jump bb4(v16, v17, v46) + bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): + v29:HeapFloat = RefineType v27, HeapFloat PatchPoint MethodRedefined(Float@0x1008, to_s@0x1010, cme:0x1018) - v38:BasicObject = CCallWithFrame v23, :Float#to_s@0x1040 - Jump bb4(v38) - bb4(v15:BasicObject): + v49:BasicObject = CCallWithFrame v29, :Float#to_s@0x1040 + Jump bb4(v25, v26, v49) + bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): CheckInterrupts - Return v15 + Return v37 "); } @@ -14660,25 +14660,25 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :x@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, StaticSymbol - IfTrue v16, bb5() - v21:CBool = HasType v10, DynamicSymbol - IfTrue v21, bb6() - v26:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v26) - bb5(): - v18:StaticSymbol = RefineType v10, StaticSymbol + v15:CBool = HasType v10, StaticSymbol + IfTrue v15, bb5(v9, v10, v10) + v24:CBool = HasType v10, DynamicSymbol + IfTrue v24, bb6(v9, v10, v10) + v33:BasicObject = Send v10, :to_s # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v33) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): + v20:StaticSymbol = RefineType v18, StaticSymbol PatchPoint MethodRedefined(Symbol@0x1008, to_s@0x1010, cme:0x1018) - v37:StringExact = InvokeBuiltin leaf , v18 - Jump bb4(v37) - bb6(): - v23:DynamicSymbol = RefineType v10, DynamicSymbol + v48:StringExact = InvokeBuiltin leaf , v20 + Jump bb4(v16, v17, v48) + bb6(v25:BasicObject, v26:BasicObject, v27:BasicObject): + v29:DynamicSymbol = RefineType v27, DynamicSymbol PatchPoint MethodRedefined(Symbol@0x1008, to_s@0x1010, cme:0x1018) - v38:StringExact = InvokeBuiltin leaf , v23 - Jump bb4(v38) - bb4(v15:BasicObject): + v49:StringExact = InvokeBuiltin leaf , v29 + Jump bb4(v25, v26, v49) + bb4(v35:BasicObject, v36:BasicObject, v37:BasicObject): CheckInterrupts - Return v15 + Return v37 "); } @@ -14715,18 +14715,18 @@ mod hir_opt_tests { v7:BasicObject = LoadArg :o@1 Jump bb3(v6, v7) bb3(v9:BasicObject, v10:BasicObject): - v16:CBool = HasType v10, ObjectSubclass[class_exact:C] - IfTrue v16, bb5() - v21:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback - Jump bb4(v21) - bb5(): + v15:CBool = HasType v10, ObjectSubclass[class_exact:C] + IfTrue v15, bb5(v9, v10, v10) + v24:BasicObject = Send v10, :foo # SendFallbackReason: SendWithoutBlock: polymorphic fallback + Jump bb4(v9, v10, v24) + bb5(v16:BasicObject, v17:BasicObject, v18:BasicObject): PatchPoint NoSingletonClass(C@0x1008) PatchPoint MethodRedefined(C@0x1008, foo@0x1010, cme:0x1018) - v31:Fixnum[3] = Const Value(3) - Jump bb4(v31) - bb4(v15:BasicObject): + v38:Fixnum[3] = Const Value(3) + Jump bb4(v16, v17, v38) + bb4(v26:BasicObject, v27:BasicObject, v28:BasicObject): CheckInterrupts - Return v15 + Return v28 "); }