Merge pull request #961 from ltratt/pass_traceiterator

Pass `[AOT]TraceIterator` to compiler backends
ykjit · Feb 8, 2024 · 41f598f · 41f598f
2 parents 8b6ccad + 1ea1014
commit 41f598f
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 15 deletions.
diff --git a/ykrt/src/compile/jitc_llvm.rs b/ykrt/src/compile/jitc_llvm.rs
@@ -5,7 +5,7 @@ use crate::{
     compile::{CompilationError, CompiledTrace, Compiler},
     location::HotLocation,
     mt::{SideTraceInfo, MT},
-    trace::TracedAOTBlock,
+    trace::{AOTTraceIterator, TracedAOTBlock},
 };
 use object::{Object, ObjectSection};
 use parking_lot::Mutex;
@@ -32,10 +32,11 @@ impl Compiler for JITCLLVM {
     fn compile(
         &self,
         mt: Arc<MT>,
-        irtrace: Vec<TracedAOTBlock>,
+        aottrace_iter: Box<dyn AOTTraceIterator>,
         sti: Option<SideTraceInfo>,
         hl: Arc<Mutex<HotLocation>>,
     ) -> Result<CompiledTrace, CompilationError> {
+        let irtrace = aottrace_iter.collect::<Vec<_>>();
         let (func_names, bbs, trace_len) = self.encode_trace(&irtrace);
 
         let llvmbc = llvmbc_section();

diff --git a/ykrt/src/compile/jitc_yk/mod.rs b/ykrt/src/compile/jitc_yk/mod.rs
@@ -5,7 +5,7 @@ use crate::{
     compile::{CompiledTrace, Compiler},
     location::HotLocation,
     mt::{SideTraceInfo, MT},
-    trace::TracedAOTBlock,
+    trace::AOTTraceIterator,
 };
 use parking_lot::Mutex;
 use std::{
@@ -60,7 +60,7 @@ impl Compiler for JITCYk {
     fn compile(
         &self,
         _mt: Arc<MT>,
-        mtrace: Vec<TracedAOTBlock>,
+        aottrace_iter: Box<dyn AOTTraceIterator>,
         sti: Option<SideTraceInfo>,
         _hl: Arc<Mutex<HotLocation>>,
     ) -> Result<CompiledTrace, CompilationError> {
@@ -77,6 +77,7 @@ impl Compiler for JITCYk {
             eprintln!("--- End aot ---");
         }
 
+        let mtrace = aottrace_iter.collect::<Vec<_>>();
         let jit_mod = trace_builder::build(&aot_mod, &mtrace)?;
 
         if PHASES_TO_PRINT.contains(&IRPhase::PreOpt) {

diff --git a/ykrt/src/compile/mod.rs b/ykrt/src/compile/mod.rs
@@ -1,7 +1,7 @@
 use crate::{
     location::HotLocation,
     mt::{SideTraceInfo, MT},
-    trace::TracedAOTBlock,
+    trace::AOTTraceIterator,
 };
 use libc::c_void;
 use parking_lot::Mutex;
@@ -41,7 +41,7 @@ pub(crate) trait Compiler: Send + Sync {
     fn compile(
         &self,
         mt: Arc<MT>,
-        irtrace: Vec<TracedAOTBlock>,
+        aottrace_iter: Box<dyn AOTTraceIterator>,
         sti: Option<SideTraceInfo>,
         hl: Arc<Mutex<HotLocation>>,
     ) -> Result<CompiledTrace, CompilationError>;

diff --git a/ykrt/src/mt.rs b/ykrt/src/mt.rs
@@ -27,7 +27,7 @@ use crate::print_jit_state;
 use crate::{
     compile::{default_compiler, CompilationError, CompiledTrace, Compiler, GuardId},
     location::{HotLocation, HotLocationKind, Location, TraceFailed},
-    trace::{default_tracer, TraceCollector, TraceIterator, Tracer},
+    trace::{default_tracer, AOTTraceIterator, TraceCollector, Tracer},
     ykstats::{TimingState, YkStats},
 };
 use yktracec::promote;
@@ -476,15 +476,14 @@ impl MT {
     ///     in the `hl_arc`.
     fn queue_compile_job(
         self: &Arc<Self>,
-        trace_iter: Box<dyn TraceIterator>,
+        trace_iter: Box<dyn AOTTraceIterator>,
         hl_arc: Arc<Mutex<HotLocation>>,
         sidetrace: Option<(SideTraceInfo, Arc<CompiledTrace>)>,
     ) {
         self.stats.trace_collected_ok();
         let mt = Arc::clone(self);
         let do_compile = move || {
             mt.stats.timing_state(TimingState::TraceMapping);
-            let irtrace = trace_iter.collect::<Vec<_>>();
             debug_assert!(
                 sidetrace.is_none() || matches!(hl_arc.lock().kind, HotLocationKind::Compiled(_))
             );
@@ -497,7 +496,7 @@ impl MT {
             let guardid = sidetrace.as_ref().map(|x| x.0.guardid);
             match compiler.compile(
                 Arc::clone(&mt),
-                irtrace,
+                trace_iter,
                 sidetrace.as_ref().map(|x| x.0),
                 Arc::clone(&hl_arc),
             ) {

diff --git a/ykrt/src/trace/hwt/mod.rs b/ykrt/src/trace/hwt/mod.rs
@@ -1,6 +1,6 @@
 //! Hardware tracing via hwtracer.
 
-use super::{errors::InvalidTraceError, TraceCollector, TraceIterator, TracedAOTBlock};
+use super::{errors::InvalidTraceError, AOTTraceIterator, TraceCollector, TracedAOTBlock};
 use std::{error::Error, sync::Arc};
 
 pub(crate) mod mapper;
@@ -33,7 +33,7 @@ struct HWTTraceCollector {
 }
 
 impl TraceCollector for HWTTraceCollector {
-    fn stop_collector(self: Box<Self>) -> Result<Box<dyn TraceIterator>, InvalidTraceError> {
+    fn stop_collector(self: Box<Self>) -> Result<Box<dyn AOTTraceIterator>, InvalidTraceError> {
         let tr = self.thread_tracer.stop_collector().unwrap();
         let mut mt = HWTMapper::new();
         let mapped = mt
@@ -53,7 +53,7 @@ struct HWTTraceIterator {
     trace: std::vec::IntoIter<TracedAOTBlock>,
 }
 
-impl TraceIterator for HWTTraceIterator {}
+impl AOTTraceIterator for HWTTraceIterator {}
 
 impl Iterator for HWTTraceIterator {
     type Item = TracedAOTBlock;

diff --git a/ykrt/src/trace/mod.rs b/ykrt/src/trace/mod.rs
@@ -41,11 +41,11 @@ pub(crate) fn default_tracer() -> Result<Arc<dyn Tracer>, Box<dyn Error>> {
 pub(crate) trait TraceCollector {
     /// Stop collecting a trace of the current thread and return an iterator which successively
     /// produces the traced blocks.
-    fn stop_collector(self: Box<Self>) -> Result<Box<dyn TraceIterator>, InvalidTraceError>;
+    fn stop_collector(self: Box<Self>) -> Result<Box<dyn AOTTraceIterator>, InvalidTraceError>;
 }
 
 /// An iterator which takes an underlying raw trace and successively produces [TracedAOTBlock]s.
-pub(crate) trait TraceIterator: Iterator<Item = TracedAOTBlock> + Send {}
+pub(crate) trait AOTTraceIterator: Iterator<Item = TracedAOTBlock> + Send {}
 
 /// An AOT LLVM IR block that has been traced at JIT time.
 #[derive(Debug, Eq, PartialEq)]