diff --git a/src/lib.rs b/src/lib.rs index daa03ab..bbdacda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -115,6 +115,16 @@ fn run_parser<'t>( )); *output_count += 1; } + ParserOutput::GlobalFile(filename, out) => { + output.push((filename.clone(), out)); + let filename_str = format!("{}", filename.to_string_lossy()); + compile_directory.push(( + filename_str.clone(), + filename_str, + *output_count, + )); + *output_count += 1; + } ParserOutput::Link(name, url) => { compile_directory.push((url, name, *output_count)); *output_count += 1; @@ -344,12 +354,25 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result = Box::new(crate::parsers::CompilationMetricsParser { tt: &tt, stack_index: &stack_index, symbolic_shape_specialization_index: &symbolic_shape_specialization_index, output_files: &copied_directory, + compile_id_dir: &compile_id_dir, }); run_parser( lineno, @@ -362,18 +385,6 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result { pub stack_index: &'t RefCell, pub symbolic_shape_specialization_index: &'t RefCell, pub output_files: &'t Vec<(String, String, i32)>, + pub compile_id_dir: &'t PathBuf, } impl StructuredLogParser for CompilationMetricsParser<'_> { fn name(&self) -> &'static str { @@ -380,6 +383,18 @@ impl StructuredLogParser for CompilationMetricsParser<'_> { .borrow() .get(&cid) .map_or("".to_string(), format_stack); + let mini_stack_html = if let (Some(name), Some(filename), Some(line)) = + (&m.co_name, &m.co_filename, m.co_firstlineno) + { + format_stack(&Vec::from([FrameSummary { + uninterned_filename: Some(filename.clone()), + filename: u32::MAX, + line: line, + name: name.clone(), + }])) + } else { + "".to_string() + }; let specializations = self .symbolic_shape_specialization_index .borrow_mut() @@ -413,8 +428,10 @@ impl StructuredLogParser for CompilationMetricsParser<'_> { m: &m, compile_id: id, stack_html: stack_html, + mini_stack_html: mini_stack_html, symbolic_shape_specializations: specializations, output_files: &output_files, + compile_id_dir: &self.compile_id_dir, }; let output = self.tt.render(&filename, &context)?; simple_file_output(&filename, lineno, compile_id, &output) @@ -502,6 +519,87 @@ impl StructuredLogParser for BwdCompilationMetricsParser<'_> { } } +pub struct DumpFileParser; +impl StructuredLogParser for DumpFileParser { + fn name(&self) -> &'static str { + "dump_file" + } + fn get_metadata<'e>(&self, e: &'e Envelope) -> Option> { + e.dump_file.as_ref().map(|m| Metadata::DumpFile(m)) + } + fn parse<'e>( + &self, + lineno: usize, + metadata: Metadata<'e>, + _rank: Option, + compile_id: &Option, + payload: &str, + ) -> anyhow::Result { + if let Metadata::DumpFile(metadata) = metadata { + let mb_fx_id = extract_eval_with_key_id(&metadata.name); + let filename = if let Some(fx_id) = mb_fx_id { + format!("eval_with_key_{}.html", fx_id) + } else { + format!("{}.html", metadata.name) + }; + let subdir = PathBuf::from("dump_file"); + let f = subdir.join(filename); + Ok(Vec::from([ParserOutput::GlobalFile( + f, + anchor_source(payload), + )])) + } else { + Err(anyhow::anyhow!("Expected DumpFile metadata")) + } + } +} + +pub fn anchor_source(text: &str) -> String { + let lines: Vec<&str> = text.lines().collect(); + let mut html = String::from( + r#" + + + + + Source Code + + + +
"#,
+    );
+
+    for (i, line) in lines.iter().enumerate() {
+        let line_number = i + 1;
+        html.push_str(&format!(
+            r#"{}"#,
+            line_number,
+            encode_text(line)
+        ));
+    }
+
+    html.push_str("
"); + html +} + pub struct ArtifactParser; impl StructuredLogParser for ArtifactParser { fn name(&self) -> &'static str { @@ -578,6 +676,7 @@ pub fn default_parsers<'t>( Box::new(BwdCompilationMetricsParser { tt }), // TODO: use own tt instances Box::new(LinkParser), Box::new(ArtifactParser), + Box::new(DumpFileParser), ]; result diff --git a/src/templates.rs b/src/templates.rs index dc1d34b..de1e77d 100644 --- a/src/templates.rs +++ b/src/templates.rs @@ -206,13 +206,15 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#" {css} Compilation Metrics +

Compilation Info for {compile_id}

+

{mini_stack_html | format_unescaped}

Output files:

Stack

@@ -249,7 +251,7 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#"

Cache Size: {m.cache_size}

Accumulated Cache Size: {m.accumulated_cache_size}

Graph Metrics

-

Guard Count: {m.guard_count}

+

Guard Count: {m.guard_count}

Shape Env Guards: {m.shape_env_guard_count}

Graph Ops: {m.graph_op_count}

Graph Nodes: {m.graph_node_count}

diff --git a/src/types.rs b/src/types.rs index 8f91e11..439c06a 100644 --- a/src/types.rs +++ b/src/types.rs @@ -21,6 +21,13 @@ pub type SymbolicShapeSpecializationIndex = pub type FxIndexMap = IndexMap>; +pub fn extract_eval_with_key_id(filename: &str) -> Option { + let re = Regex::new(r"\.([0-9]+)").unwrap(); + re.captures(filename) + .and_then(|caps| caps.get(1)) + .and_then(|m| m.as_str().parse::().ok()) +} + pub static INTERN_TABLE: Lazy>> = Lazy::new(|| Mutex::new(FxHashMap::default())); @@ -103,10 +110,10 @@ impl StackTrieNode { // If the node has multiple children, increase the indent and print a hyphen writeln!( f, - "
  • {star}{}
      ", - frame, + "
    • {star}", star = star )?; + writeln!(f, "{}
        ", frame)?; node.fmt_inner(f, mb_metrics_index)?; write!(f, "
    • ")?; } else { @@ -153,6 +160,7 @@ pub struct FrameSummary { pub filename: u32, pub line: i32, pub name: String, + pub uninterned_filename: Option, } pub fn simplify_filename<'a>(filename: &'a str) -> &'a str { @@ -180,16 +188,32 @@ pub fn unintern_str(interned_str: u32) -> String { impl fmt::Display for FrameSummary { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let intern_table = INTERN_TABLE.lock().unwrap(); - let filename = intern_table - .get(&self.filename) - .map_or("(unknown)", |s| s.as_str()); - write!( - f, - "{}:{} in {}", - encode_text(simplify_filename(filename)), - self.line, - encode_text(&self.name) - ) + let filename = if let Some(f) = &self.uninterned_filename { + f.as_str() + } else { + intern_table + .get(&self.filename) + .map_or("(unknown)", |s| s.as_str()) + }; + if let Some(fx_id) = extract_eval_with_key_id(filename) { + write!( + f, + "{filename}:{line} in {name}", + fx_id = fx_id, + filename = encode_text(simplify_filename(filename)), + line = self.line, + name = encode_text(&self.name) + )?; + } else { + write!( + f, + "{}:{} in {}", + encode_text(simplify_filename(filename)), + self.line, + encode_text(&self.name) + )?; + } + Ok(()) } } @@ -254,7 +278,10 @@ pub struct ArtifactMetadata { #[derive(Debug, Deserialize, Serialize, Clone)] pub struct CompilationMetricsMetadata { - // Other information like frame_key, co_name, etc. are already in envelope + // Other information like frame_key are already in envelope + pub co_name: Option, + pub co_filename: Option, + pub co_firstlineno: Option, pub cache_size: Option, pub accumulated_cache_size: Option, pub guard_count: Option, @@ -325,6 +352,8 @@ pub struct CompilationMetricsContext<'e> { pub stack_html: String, pub symbolic_shape_specializations: Vec, pub output_files: &'e Vec<(String, String, i32)>, + pub compile_id_dir: &'e PathBuf, + pub mini_stack_html: String, } #[derive(Debug, Serialize)] @@ -381,6 +410,12 @@ pub enum Metadata<'e> { AOTAutogradBackwardCompilationMetrics(&'e AOTAutogradBackwardCompilationMetricsMetadata), BwdCompilationMetrics(&'e BwdCompilationMetricsMetadata), Artifact(&'e ArtifactMetadata), + DumpFile(&'e DumpFileMetadata), +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct DumpFileMetadata { + pub name: String, } #[derive(Debug, Deserialize)] @@ -416,6 +451,7 @@ pub struct Envelope { pub describe_storage: Option, pub describe_tensor: Option, pub describe_source: Option, + pub dump_file: Option, #[serde(flatten)] pub _other: FxHashMap, }