Add inductor provenance tracking higlighter

src/cli.rs

@@ -44,6 +44,9 @@ pub struct Cli {
     /// For export specific logs
     #[arg(short, long)]
     export: bool,
+    /// For inductor provenance tracking highlighter
+    #[arg(short, long)]
+    inductor_provenance: bool,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -93,6 +96,7 @@ fn main() -> anyhow::Result<()> {
         verbose: cli.verbose,
         plain_text: cli.plain_text,
         export: cli.export,
+        inductor_provenance: cli.inductor_provenance,
     };
 
     let output = parse_path(&path, config)?;

src/lib.rs

@@ -29,6 +29,7 @@ pub struct ParseConfig {
     pub verbose: bool,
     pub plain_text: bool,
     pub export: bool,
+    pub inductor_provenance: bool,
 }
 
 impl Default for ParseConfig {
@@ -41,6 +42,7 @@ impl Default for ParseConfig {
             verbose: false,
             plain_text: false,
             export: false,
+            inductor_provenance: false,
         }
     }
 }
@@ -248,6 +250,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS,
         )?;
     }
+    tt.add_template("provenance_tracking.html", TEMPLATE_PROVENANCE_TRACKING)?;
 
     let mut unknown_fields: FxHashSet<String> = FxHashSet::default();
 
@@ -668,6 +671,14 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
 
     let has_unknown_compile_id = directory.contains_key(&None);
 
+    let directory_names: Vec<String> = directory
+        .iter()
+        .map(|(x, _)| {
+            x.as_ref()
+                .map_or("(unknown)".to_string(), |e| e.as_directory_name())
+        })
+        .collect();
+
     let index_context = IndexContext {
         css: CSS,
         javascript: JAVASCRIPT,
@@ -686,6 +697,8 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
         num_breaks: breaks.failures.len(),
         has_chromium_events: !chromium_events.is_empty(),
         qps: TEMPLATE_QUERY_PARAM_SCRIPT,
+        has_inductor_provenance: config.inductor_provenance,
+        directory_names: directory_names.clone(),
     };
     output.push((
         PathBuf::from("index.html"),
@@ -713,5 +726,55 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
         return Err(anyhow!("Some log entries did not have compile id"));
     }
 
+    if config.inductor_provenance {
+        // Helper function to get file content for a specific directory name
+        fn get_file_content(
+            output: &[(PathBuf, String)],
+            filename_pattern: &str,
+            directory_name: &str,
+        ) -> String {
+            output
+                .iter()
+                .find(|(path, _)| {
+                    path.to_string_lossy()
+                        .contains(&format!("{}/{}", directory_name, filename_pattern))
+                })
+                .map(|(_, content)| content.clone())
+                .unwrap_or_default()
+        }
+
+        // Generate HTML for each directory name
+        for directory_name in &directory_names {
+            let pre_grad_graph_content =
+                get_file_content(&output, "inductor_pre_grad_graph", directory_name);
+            let post_grad_graph_content =
+                get_file_content(&output, "inductor_post_grad_graph", directory_name);
+            let output_code_content =
+                get_file_content(&output, "inductor_output_code", directory_name);
+            let aot_code_content = get_file_content(&output, "inductor_aot_code", directory_name);
+            let node_mappings_content = get_file_content(
+                &output,
+                "inductor_provenance_tracking_node_mappings",
+                directory_name,
+            );
+
+            output.push((
+                PathBuf::from(format!("provenance_tracking_{}.html", directory_name)),
+                tt.render(
+                    "provenance_tracking.html",
+                    &ProvenanceContext {
+                        css: PROVENANCE_CSS,
+                        js: PROVENANCE_JS,
+                        pre_grad_graph_content,
+                        post_grad_graph_content,
+                        output_code_content,
+                        aot_code_content,
+                        node_mappings_content,
+                    },
+                )?,
+            ));
+        }
+    }
+
     Ok(output)
 }

src/provenance.css

@@ -0,0 +1,152 @@
+body {
+    display: flex;
+    flex-direction: column;
+    font-family: Arial, sans-serif;
+    margin: 0;
+    padding: 0;
+    height: 100vh;
+    overflow: hidden;
+}
+
+.url-inputs {
+    display: flex;
+    padding: 10px;
+    background-color: #f0f0f0;
+}
+
+.url-input {
+    flex: 1;
+    margin-right: 10px;
+    display: flex;
+    align-items: center;
+    gap: 10px;
+}
+
+.url-input input {
+    width: 100%;
+    padding: 5px;
+}
+
+.editor-container {
+    display: flex;
+    flex: 1;
+    overflow: hidden;
+}
+
+.editor {
+    height: 100%;
+    overflow-y: auto;
+    border: 1px solid #ddd;
+    padding: 10px;
+    box-sizing: border-box;
+    flex: 1;
+    font-family: monospace;
+}
+
+.line {
+    padding: 2px 5px;
+    cursor: pointer;
+    white-space: nowrap;
+    transition: background-color 0.2s ease;
+}
+
+.highlight {
+    background-color: yellow;
+    transition: background-color 0.2s ease;
+}
+
+.mapped-line {
+    font-weight: bold;
+}
+
+.line-number {
+    color: #888;
+    display: inline-block;
+    width: 30px;
+    text-align: right;
+    margin-right: 10px;
+}
+
+.divider {
+    width: 10px;
+    background-color: #ccc;
+    cursor: col-resize;
+}
+
+.line-content {
+    white-space: pre;
+    display: inline;
+}
+
+.json-popup {
+    display: none;
+    position: fixed;
+    background: white;
+    border: 1px solid #ccc;
+    padding: 10px;
+    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
+    max-width: 400px;
+    max-height: 300px;
+    overflow: auto;
+    z-index: 1000;
+    white-space: pre-wrap;
+    font-family: monospace;
+    font-size: 12px;
+}
+
+.json-popup .close-button {
+    position: sticky;
+    float: right;
+    top: 0;
+    right: 0;
+    cursor: pointer;
+    font-size: 20px;
+    background: white;
+    padding: 0 5px;
+    margin-left: 10px;
+}
+
+#jsonContent {
+    margin-top: 10px;
+}
+
+.toggle-container {
+    margin-left: 20px;
+    display: inline-flex;
+    align-items: center;
+    cursor: pointer;
+    min-width: 200px;
+    white-space: nowrap;
+}
+
+.toggle-container input[type="checkbox"] {
+    margin-right: 5px;
+}
+
+.file-input {
+    display: none;
+}
+
+.url-input {
+    flex: 1;
+    margin-right: 10px;
+    display: flex;
+    align-items: center;
+    gap: 10px;
+}
+
+.file-label {
+    background-color: #f0f0f0;
+    padding: 5px 10px;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+    cursor: pointer;
+}
+
+.file-label:hover {
+    background-color: #e0e0e0;
+}
+
+.has-match {
+    font-weight: bold;
+} 

src/provenance.html

@@ -0,0 +1,35 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Inductor Provenance Tracking Highlighter</title>
+    <style>
+        {css | format_unescaped}
+    </style>
+</head>
+
+<body>
+    <div class="editor-container">
+        <div id="preGradGraph" class="editor">
+            <pre>{pre_grad_graph_content}</pre>
+        </div>
+        <div id="divider" class="divider"></div>
+        <div id="postGradGraph" class="editor">
+            <pre>{post_grad_graph_content}</pre>
+        </div>
+        <div id="divider" class="divider"></div>
+        <div id="generatedCode" class="editor">
+            <pre>{output_code_content | format_unescaped}{aot_code_content}</pre>
+        </div>
+    </div>
+
+    <script>
+        // Make node mappings available to JavaScript
+        const nodeMappings = {node_mappings_content | format_unescaped};
+        {js | format_unescaped}
+    </script>
+</body>
+
+</html>