Skip to content

Commit 835cef9

Browse files
fffonionfffonion
committed
fix(jit): restore region state on interrupts
1 parent 9f987a0 commit 835cef9

2 files changed

Lines changed: 306 additions & 24 deletions

File tree

‎src/vm/jit/native/lower.rs‎

Lines changed: 153 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,9 @@ fn try_compile_ssa_trace(
282282
let call_value_block = call_value_exits
283283
.contains_key(&exit.id)
284284
.then(|| b.create_block());
285+
let interrupt_block = (interrupt_settings.is_some()
286+
&& internal_links.iter().any(|link| link.exit == exit.id))
287+
.then(|| b.create_block());
285288
for value in &inputs {
286289
let Some(repr) = value_reprs.get(value).copied() else {
287290
return Ok(None);
@@ -294,13 +297,17 @@ fn try_compile_ssa_trace(
294297
if let Some(call_value_block) = call_value_block {
295298
b.append_block_param(call_value_block, ty);
296299
}
300+
if let Some(interrupt_block) = interrupt_block {
301+
b.append_block_param(interrupt_block, ty);
302+
}
297303
}
298304
exit_specs.insert(
299305
exit.id,
300306
SsaExitLowering {
301307
trace_exit_block,
302308
halted_block,
303309
call_value_block,
310+
interrupt_block,
304311
inputs,
305312
},
306313
);
@@ -330,6 +337,8 @@ fn try_compile_ssa_trace(
330337
ssa.entry_stack_depth,
331338
entry_local_count,
332339
)?;
340+
let interrupt_ops_to_advance = u32::try_from(trace.op_names.len().max(1))
341+
.map_err(|_| VmError::JitNative("trace op count exceeds u32".to_string()))?;
333342
let lower_ctx = SsaLowerCtx {
334343
vm_ptr,
335344
active_stack_base,
@@ -349,6 +358,8 @@ fn try_compile_ssa_trace(
349358
borrowed_array_gets: &borrowed_array_gets,
350359
value_reprs: &value_reprs,
351360
tagged_constant_addrs: &tagged_constant_addrs,
361+
interrupt_settings,
362+
interrupt_ops_to_advance,
352363
};
353364
let root_ip = b.ins().iconst(
354365
pointer_type,
@@ -376,8 +387,6 @@ fn try_compile_ssa_trace(
376387
b.ins().jump(entry_handle, &entry_args);
377388

378389
let charge_blocks = ssa_interrupt_charge_blocks(ssa);
379-
let ops_to_advance = u32::try_from(trace.op_names.len().max(1))
380-
.map_err(|_| VmError::JitNative("trace op count exceeds u32".to_string()))?;
381390
for block in &ssa.blocks {
382391
let handle = *block_handles
383392
.get(&block.id)
@@ -388,15 +397,16 @@ fn try_compile_ssa_trace(
388397
for (param, lowered) in block.params.iter().zip(block_params.iter().copied()) {
389398
values.insert(param.value.id, lowered);
390399
}
391-
if charge_blocks.contains(&block.id)
400+
if internal_links.is_empty()
401+
&& charge_blocks.contains(&block.id)
392402
&& let Some(settings) = interrupt_settings
393403
{
394404
emit_interrupt_tick_inline(
395405
&mut b,
396406
vm_ptr,
397407
exit_block,
398408
offsets,
399-
ops_to_advance,
409+
lower_ctx.interrupt_ops_to_advance,
400410
settings,
401411
);
402412
}
@@ -446,6 +456,9 @@ fn try_compile_ssa_trace(
446456
},
447457
)?;
448458
}
459+
if spec.interrupt_block.is_some() {
460+
lower_ssa_exit_block(&mut b, lower_ctx, exit, spec, SsaExitAction::InterruptYield)?;
461+
}
449462
}
450463

451464
b.switch_to_block(exit_block);
@@ -500,6 +513,7 @@ struct SsaExitLowering {
500513
trace_exit_block: Block,
501514
halted_block: Block,
502515
call_value_block: Option<Block>,
516+
interrupt_block: Option<Block>,
503517
inputs: Vec<SsaValueId>,
504518
}
505519

@@ -514,6 +528,7 @@ enum SsaExitAction {
514528
call_ip: usize,
515529
resume_ip: usize,
516530
},
531+
InterruptYield,
517532
}
518533

519534
#[derive(Clone, Copy)]
@@ -636,6 +651,8 @@ struct SsaLowerCtx<'a> {
636651
borrowed_array_gets: &'a BTreeSet<SsaValueId>,
637652
value_reprs: &'a HashMap<SsaValueId, SsaValueRepr>,
638653
tagged_constant_addrs: &'a HashMap<SsaValueId, usize>,
654+
interrupt_settings: Option<NativeInterruptSettings>,
655+
interrupt_ops_to_advance: u32,
639656
}
640657

641658
#[derive(Clone, Copy)]
@@ -3209,7 +3226,7 @@ fn lower_ssa_terminator(
32093226
let condition = *values
32103227
.get(condition)
32113228
.ok_or_else(|| VmError::JitNative("SSA branch condition missing".to_string()))?;
3212-
let true_target = ssa_branch_target_block(
3229+
let (true_block, true_args, true_pending) = ssa_branch_target_block(
32133230
b,
32143231
ctx,
32153232
if_true,
@@ -3218,7 +3235,7 @@ fn lower_ssa_terminator(
32183235
exit_specs,
32193236
internal_links,
32203237
)?;
3221-
let false_target = ssa_branch_target_block(
3238+
let (false_block, false_args, false_pending) = ssa_branch_target_block(
32223239
b,
32233240
ctx,
32243241
if_false,
@@ -3227,25 +3244,42 @@ fn lower_ssa_terminator(
32273244
exit_specs,
32283245
internal_links,
32293246
)?;
3230-
let true_args = ssa_block_args(true_target.1);
3231-
let false_args = ssa_block_args(false_target.1);
3232-
b.ins().brif(
3233-
condition,
3234-
true_target.0,
3235-
&true_args,
3236-
false_target.0,
3237-
&false_args,
3238-
);
3247+
let true_args = ssa_block_args(true_args);
3248+
let false_args = ssa_block_args(false_args);
3249+
b.ins()
3250+
.brif(condition, true_block, &true_args, false_block, &false_args);
3251+
if let Some(pending) = true_pending {
3252+
lower_pending_internal_link_gate(b, ctx, pending);
3253+
}
3254+
if let Some(pending) = false_pending {
3255+
lower_pending_internal_link_gate(b, ctx, pending);
3256+
}
32393257
}
32403258
SsaTerminator::Exit { exit } => {
32413259
if let Some((link, slots)) = internal_links.find(*exit) {
32423260
let handle = *block_handles.get(&link.child_entry).ok_or_else(|| {
32433261
VmError::JitNative("fused SSA child entry block missing".to_string())
32443262
})?;
3245-
let args = lower_internal_link_args(b, ctx, link, slots, values)?;
3246-
increment_native_region_edge_count(b, ctx);
3263+
let child_args = lower_internal_link_args(b, ctx, link, slots, values)?;
3264+
let spec = exit_specs.get(exit).ok_or_else(|| {
3265+
VmError::JitNative("SSA internal exit lowering missing".to_string())
3266+
})?;
3267+
let restore_args = spec
3268+
.inputs
3269+
.iter()
3270+
.map(|value| {
3271+
values.get(value).copied().ok_or_else(|| {
3272+
VmError::JitNative("SSA internal restore arg missing".to_string())
3273+
})
3274+
})
3275+
.collect::<VmResult<Vec<_>>>()?;
3276+
let (target, args, pending) =
3277+
lower_internal_link_target(b, ctx, handle, child_args, spec, restore_args)?;
32473278
let args = ssa_block_args(args);
3248-
b.ins().jump(handle, &args);
3279+
b.ins().jump(target, &args);
3280+
if let Some(pending) = pending {
3281+
lower_pending_internal_link_gate(b, ctx, pending);
3282+
}
32493283
return Ok(());
32503284
}
32513285
let spec = exit_specs
@@ -3311,7 +3345,11 @@ fn ssa_branch_target_block(
33113345
block_handles: &HashMap<crate::vm::jit::ir::SsaBlockId, Block>,
33123346
exit_specs: &HashMap<SsaExitId, SsaExitLowering>,
33133347
internal_links: InternalRegionLinks<'_>,
3314-
) -> VmResult<(Block, Vec<cranelift_codegen::ir::Value>)> {
3348+
) -> VmResult<(
3349+
Block,
3350+
Vec<cranelift_codegen::ir::Value>,
3351+
Option<PendingInternalLinkGate>,
3352+
)> {
33153353
match target {
33163354
SsaBranchTarget::Block { target, args } => {
33173355
let handle = *block_handles
@@ -3326,16 +3364,29 @@ fn ssa_branch_target_block(
33263364
.ok_or_else(|| VmError::JitNative("SSA branch arg missing".to_string()))
33273365
})
33283366
.collect::<VmResult<Vec<_>>>()?;
3329-
Ok((handle, lowered_args))
3367+
Ok((handle, lowered_args, None))
33303368
}
33313369
SsaBranchTarget::Exit(exit) => {
33323370
if let Some((link, slots)) = internal_links.find(*exit) {
33333371
let handle = *block_handles.get(&link.child_entry).ok_or_else(|| {
33343372
VmError::JitNative("fused SSA child entry block missing".to_string())
33353373
})?;
3336-
let args = lower_internal_link_args(b, ctx, link, slots, values)?;
3337-
increment_native_region_edge_count(b, ctx);
3338-
return Ok((handle, args));
3374+
let child_args = lower_internal_link_args(b, ctx, link, slots, values)?;
3375+
let spec = exit_specs.get(exit).ok_or_else(|| {
3376+
VmError::JitNative("SSA internal branch exit lowering missing".to_string())
3377+
})?;
3378+
let restore_args = spec
3379+
.inputs
3380+
.iter()
3381+
.map(|value| {
3382+
values.get(value).copied().ok_or_else(|| {
3383+
VmError::JitNative(
3384+
"SSA internal branch restore arg missing".to_string(),
3385+
)
3386+
})
3387+
})
3388+
.collect::<VmResult<Vec<_>>>()?;
3389+
return lower_internal_link_target(b, ctx, handle, child_args, spec, restore_args);
33393390
}
33403391
let spec = exit_specs.get(exit).ok_or_else(|| {
33413392
VmError::JitNative("SSA branch exit lowering missing".to_string())
@@ -3350,11 +3401,85 @@ fn ssa_branch_target_block(
33503401
.ok_or_else(|| VmError::JitNative("SSA exit arg missing".to_string()))
33513402
})
33523403
.collect::<VmResult<Vec<_>>>()?;
3353-
Ok((spec.trace_exit_block, lowered_args))
3404+
Ok((spec.trace_exit_block, lowered_args, None))
33543405
}
33553406
}
33563407
}
33573408

3409+
struct PendingInternalLinkGate {
3410+
gate_block: Block,
3411+
child_handle: Block,
3412+
child_arg_count: usize,
3413+
interrupt_block: Block,
3414+
}
3415+
3416+
fn lower_internal_link_target(
3417+
b: &mut FunctionBuilder,
3418+
ctx: SsaLowerCtx<'_>,
3419+
child_handle: Block,
3420+
child_args: Vec<cranelift_codegen::ir::Value>,
3421+
exit_spec: &SsaExitLowering,
3422+
restore_args: Vec<cranelift_codegen::ir::Value>,
3423+
) -> VmResult<(
3424+
Block,
3425+
Vec<cranelift_codegen::ir::Value>,
3426+
Option<PendingInternalLinkGate>,
3427+
)> {
3428+
let Some(_settings) = ctx.interrupt_settings else {
3429+
increment_native_region_edge_count(b, ctx);
3430+
return Ok((child_handle, child_args, None));
3431+
};
3432+
let gate_block = b.create_block();
3433+
for value in child_args.iter().chain(&restore_args) {
3434+
b.append_block_param(gate_block, b.func.dfg.value_type(*value));
3435+
}
3436+
let child_arg_count = child_args.len();
3437+
let mut gate_args = child_args;
3438+
gate_args.extend(restore_args);
3439+
let interrupt_block = exit_spec
3440+
.interrupt_block
3441+
.ok_or_else(|| VmError::JitNative("SSA internal interrupt block missing".to_string()))?;
3442+
Ok((
3443+
gate_block,
3444+
gate_args,
3445+
Some(PendingInternalLinkGate {
3446+
gate_block,
3447+
child_handle,
3448+
child_arg_count,
3449+
interrupt_block,
3450+
}),
3451+
))
3452+
}
3453+
3454+
fn lower_pending_internal_link_gate(
3455+
b: &mut FunctionBuilder,
3456+
ctx: SsaLowerCtx<'_>,
3457+
pending: PendingInternalLinkGate,
3458+
) {
3459+
let settings = ctx
3460+
.interrupt_settings
3461+
.expect("pending internal link gate requires interrupt settings");
3462+
b.switch_to_block(pending.gate_block);
3463+
let gate_params = b.block_params(pending.gate_block).to_vec();
3464+
let interrupt_status_block = b.create_block();
3465+
b.append_block_param(interrupt_status_block, types::I32);
3466+
emit_interrupt_tick_inline(
3467+
b,
3468+
ctx.vm_ptr,
3469+
interrupt_status_block,
3470+
ctx.offsets,
3471+
ctx.interrupt_ops_to_advance,
3472+
settings,
3473+
);
3474+
increment_native_region_edge_count(b, ctx);
3475+
let child_params = ssa_block_args(gate_params[..pending.child_arg_count].to_vec());
3476+
b.ins().jump(pending.child_handle, &child_params);
3477+
3478+
b.switch_to_block(interrupt_status_block);
3479+
let restore_params = ssa_block_args(gate_params[pending.child_arg_count..].to_vec());
3480+
b.ins().jump(pending.interrupt_block, &restore_params);
3481+
}
3482+
33583483
fn increment_native_region_edge_count(b: &mut FunctionBuilder, ctx: SsaLowerCtx<'_>) {
33593484
let count = b.ins().load(
33603485
types::I64,
@@ -3401,6 +3526,7 @@ fn ssa_exit_action_status(
34013526
action: SsaExitAction,
34023527
) -> VmResult<cranelift_codegen::ir::Value> {
34033528
match action {
3529+
SsaExitAction::InterruptYield => Ok(b.ins().iconst(types::I32, STATUS_OUT_OF_FUEL as i64)),
34043530
SsaExitAction::Return => ssa_leave_frame_status(
34053531
b,
34063532
pointer_type,
@@ -3482,6 +3608,9 @@ fn lower_ssa_exit_block(
34823608
SsaExitAction::CallValue { .. } => spec
34833609
.call_value_block
34843610
.ok_or_else(|| VmError::JitNative("SSA call-value exit block missing".to_string()))?,
3611+
SsaExitAction::InterruptYield => spec
3612+
.interrupt_block
3613+
.ok_or_else(|| VmError::JitNative("SSA interrupt exit block missing".to_string()))?,
34853614
};
34863615
b.switch_to_block(block);
34873616
let block_params = b.block_params(block).to_vec();

0 commit comments

Comments
 (0)