diff --git a/ykrt/src/compile/jitc_yk/trace_builder.rs b/ykrt/src/compile/jitc_yk/trace_builder.rs index de187a031..ae8435342 100644 --- a/ykrt/src/compile/jitc_yk/trace_builder.rs +++ b/ykrt/src/compile/jitc_yk/trace_builder.rs @@ -138,6 +138,29 @@ impl TraceBuilder { Ok(()) } + /// Process promotions inside an otherwise outlined block. + fn process_promotions_only(&mut self, bid: &aot_ir::BBlockId) -> Result<(), CompilationError> { + let blk = self.aot_mod.bblock(bid); + + for inst in blk.insts.iter() { + match inst { + aot_ir::Inst::Promote { val, tyidx, .. } => match val { + aot_ir::Operand::LocalVariable(_) => { + let width_bits = match self.aot_mod.type_(*tyidx) { + aot_ir::Ty::Integer(it) => it.num_bits(), + _ => unreachable!(), + }; + let width_bytes = usize::try_from(width_bits).unwrap() / 8; + self.promote_idx += width_bytes; + } + _ => (), + }, + _ => (), + } + } + Ok(()) + } + /// Walk over a traced AOT block, translating the constituent instructions into the JIT module. fn process_block( &mut self, @@ -1258,7 +1281,32 @@ impl TraceBuilder { // blocks normally. self.outline_target_blk = None; } else { - // We are outlining so just skip this block. + // We are outlining so just skip this block. However, we still need to + // process promoted values to make sure we've processed all promotion + // data and haven't messed up the mapping. + #[cfg(tracer_hwt)] + { + if self.aot_mod.bblock(&bid).is_return() { + last_blk_is_return = true; + } else { + last_blk_is_return = false; + } + // Due to hardware tracing we see the same block twice whenever + // there is a call. We only need to process one of them. We can + // skip the block if: + // a) The previous block had a return. + // b) The previous block is unmappable and the current block isn't + // an entry block. + if last_blk_is_return { + prev_bid = Some(bid); + continue; + } + if prev_bid.is_none() && !bid.is_entry() { + prev_bid = Some(bid); + continue; + } + } + self.process_promotions_only(&bid)?; prev_bid = Some(bid); continue; } @@ -1306,12 +1354,17 @@ impl TraceBuilder { prev_bid = Some(bid); } None => { - // UnmappableBBlock block + // Unmappable block + #[cfg(tracer_hwt)] + { + last_blk_is_return = false; + } prev_bid = None; } } } + debug_assert_eq!(self.promote_idx, self.promotions.len()); let blk = self.aot_mod.bblock(self.cp_block.as_ref().unwrap()); let cpcall = blk.insts.iter().rev().nth(1).unwrap(); debug_assert!(cpcall.is_control_point(self.aot_mod));