Add dump_file parsing support

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<ParseOu
 
         if let Some(ref m) = e.compilation_metrics {
             let copied_directory = compile_directory.clone();
+            let compile_id_dir: PathBuf = e
+                .compile_id
+                .as_ref()
+                .map_or(
+                    format!("unknown_{lineno}"),
+                    |CompileId {
+                         frame_id,
+                         frame_compile_id,
+                         attempt,
+                     }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
+                )
+                .into();
             let parser: Box<dyn StructuredLogParser> =
                 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<ParseOu
                 &multi,
                 &mut stats,
             );
-            let compile_id_dir: PathBuf = e
-                .compile_id
-                .as_ref()
-                .map_or(
-                    format!("unknown_{lineno}"),
-                    |CompileId {
-                         frame_id,
-                         frame_compile_id,
-                         attempt,
-                     }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
-                )
-                .into();
 
             // compilation metrics is always the last output, since it just ran
             let metrics_filename = format!(

src/parsers.rs

@@ -1,4 +1,5 @@
 use crate::{types::*, ParseConfig};
+use html_escape::encode_text;
 use std::cell::RefCell;
 use std::ffi::{OsStr, OsString};
 use std::path::Path;
@@ -11,7 +12,8 @@ use syntect::highlighting::ThemeSet;
 use syntect::parsing::SyntaxSet;
 
 pub enum ParserOutput {
-    File(PathBuf, String), // File to be saved on disk
+    File(PathBuf, String),       // File to be saved on disk
+    GlobalFile(PathBuf, String), // Like file, but don't give a unique suffix
     Link(String, String), // External href to (name, url) (linked in compile_directory, not returned)
 }
 
@@ -348,6 +350,7 @@ pub struct CompilationMetricsParser<'t> {
     pub stack_index: &'t RefCell<StackIndex>,
     pub symbolic_shape_specialization_index: &'t RefCell<SymbolicShapeSpecializationIndex>,
     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<Metadata<'e>> {
+        e.dump_file.as_ref().map(|m| Metadata::DumpFile(m))
+    }
+    fn parse<'e>(
+        &self,
+        lineno: usize,
+        metadata: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        payload: &str,
+    ) -> anyhow::Result<ParserResults> {
+        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#"<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Source Code</title>
+    <style>
+        pre {
+            counter-reset: line;
+        }
+        pre span {
+            display: block;
+        }
+        pre span:before {
+            counter-increment: line;
+            content: counter(line);
+            display: inline-block;
+            padding: 0 .5em;
+            margin-right: .5em;
+            color: #888;
+        }
+        pre span:target {
+            background-color: #ffff00;
+        }
+    </style>
+</head>
+<body>
+    <pre>"#,
+    );
+
+    for (i, line) in lines.iter().enumerate() {
+        let line_number = i + 1;
+        html.push_str(&format!(
+            r#"<span id="L{}">{}</span>"#,
+            line_number,
+            encode_text(line)
+        ));
+    }
+
+    html.push_str("</pre></body></html>");
+    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

src/templates.rs

@@ -206,13 +206,15 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#"
     {css}
     </style>
     <title>Compilation Metrics</title>
+    <base href="..">
 </head>
 <body>
     <h1>Compilation Info for {compile_id}</h1>
+    <p>{mini_stack_html | format_unescaped}</p>
     <h2>Output files:</h2>
     <ul>
         {{ for path_idx in output_files }}
-            <li><a href="{path_idx.0}">{path_idx.1}</a> ({path_idx.2})</li>
+            <li><a href="{compile_id_dir}/{path_idx.0}">{path_idx.1}</a> ({path_idx.2})</li>
         {{ endfor }}
     </ul>
     <h2>Stack</h2>
@@ -249,7 +251,7 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#"
     <p>Cache Size: {m.cache_size}</p>
     <p>Accumulated Cache Size: {m.accumulated_cache_size}</p>
     <h2>Graph Metrics</h2>
-    <p><a href='dynamo_guards.html'>Guard</a> Count: {m.guard_count}</p>
+    <p>Guard Count: {m.guard_count}</p>
     <p>Shape Env Guards: {m.shape_env_guard_count}</p>
     <p>Graph Ops: {m.graph_op_count}</p>
     <p>Graph Nodes: {m.graph_node_count}</p>

src/types.rs

@@ -21,6 +21,13 @@ pub type SymbolicShapeSpecializationIndex =
 
 pub type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<FxHasher>>;
 
+pub fn extract_eval_with_key_id(filename: &str) -> Option<u64> {
+    let re = Regex::new(r"<eval_with_key>\.([0-9]+)").unwrap();
+    re.captures(filename)
+        .and_then(|caps| caps.get(1))
+        .and_then(|m| m.as_str().parse::<u64>().ok())
+}
+
 pub static INTERN_TABLE: Lazy<Mutex<FxHashMap<u32, String>>> =
     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,
-                    "<li><span onclick='toggleList(this)' class='marker'></span>{star}{}<ul>",
-                    frame,
+                    "<li><span onclick='toggleList(this)' class='marker'></span>{star}",
                     star = star
                 )?;
+                writeln!(f, "{}<ul>", frame)?;
                 node.fmt_inner(f, mb_metrics_index)?;
                 write!(f, "</ul></li>")?;
             } else {
@@ -153,6 +160,7 @@ pub struct FrameSummary {
     pub filename: u32,
     pub line: i32,
     pub name: String,
+    pub uninterned_filename: Option<String>,
 }
 
 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,
+                "<a href='dump_file/eval_with_key_{fx_id}.html#L{line}'>{filename}:{line}</a> 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<String>,
+    pub co_filename: Option<String>,
+    pub co_firstlineno: Option<i32>,
     pub cache_size: Option<u64>,
     pub accumulated_cache_size: Option<u64>,
     pub guard_count: Option<u64>,
@@ -325,6 +352,8 @@ pub struct CompilationMetricsContext<'e> {
     pub stack_html: String,
     pub symbolic_shape_specializations: Vec<SymbolicShapeSpecializationContext>,
     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<StorageDesc>,
     pub describe_tensor: Option<TensorDesc>,
     pub describe_source: Option<SourceDesc>,
+    pub dump_file: Option<DumpFileMetadata>,
     #[serde(flatten)]
     pub _other: FxHashMap<String, Value>,
 }