Add compiled_autograd_id to CompileId (#83)

land with pytorch/pytorch#141907

src/lib.rs

@@ -378,14 +378,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             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}") },
-                )
+                .map_or(format!("unknown_{lineno}"), |cid| cid.as_directory_name())
                 .into();
             let parser: Box<dyn StructuredLogParser> =
                 Box::new(crate::parsers::CompilationMetricsParser {
@@ -450,7 +443,10 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             }
             let mut cid = e.compile_id.clone();
             if let Some(c) = cid.as_mut() {
-                c.attempt = 0;
+                if let Some(_frame_id) = c.frame_compile_id {
+                    // data migration for old logs that don't have attempt
+                    c.attempt = Some(0);
+                }
             }
             metrics_index.entry(cid).or_default().push(m.clone());
         }

src/parsers.rs

@@ -56,14 +56,7 @@ fn simple_file_output(
 ) -> anyhow::Result<ParserResults> {
     let compile_id_dir: PathBuf = compile_id
         .as_ref()
-        .map_or(
-            format!("unknown_{lineno}"),
-            |CompileId {
-                 frame_id,
-                 frame_compile_id,
-                 attempt,
-             }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
-        )
+        .map_or(format!("unknown_{lineno}"), |cid| cid.as_directory_name())
         .into();
     let subdir = PathBuf::from(compile_id_dir);
     let f = subdir.join(filename);
@@ -380,7 +373,10 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
                 .map_or("(unknown) ".to_string(), |c| format!("{cid} ", cid = c));
             let mut cid = compile_id.clone();
             if let Some(c) = cid.as_mut() {
-                c.attempt = 0;
+                if let Some(_frame_id) = c.frame_compile_id {
+                    // data migration for old logs that don't have attempt
+                    c.attempt = Some(0);
+                }
             }
             let stack_html = self
                 .stack_index

src/templates.rs

@@ -117,6 +117,11 @@ common cause of recompilation is a graph break in an inlined function call, whic
 and avoid inlining the function in the first place.
 </p>
 <p>
+When compiled autograd is enabled, the compile id will include a prefix signifier <code>[!a/x/y]</code>,
+where a is the <strong>compiled autograd id</strong>. For instance, <code>[!0/-/-]</code> refers 
+to the first graph captured by compiled autograd. It is then traced by torch.compile as <code>[!0/x/y_z]</code>.
+</p>
+<p>
 Here is a high level description of PT2's compilation phases, and the intermediate products each
 phase generates:
 </p>

src/types.rs

@@ -128,21 +128,50 @@ impl StackTrieNode {
 
 #[derive(Eq, PartialEq, Hash, Deserialize, Serialize, Debug, Clone)]
 pub struct CompileId {
-    pub frame_id: u32,
-    pub frame_compile_id: u32,
-    pub attempt: u32,
+    pub compiled_autograd_id: Option<u32>,
+    pub frame_id: Option<u32>,
+    pub frame_compile_id: Option<u32>,
+    pub attempt: Option<u32>,
 }
 
 impl fmt::Display for CompileId {
+    // NOTE: If you want to elide an id e.g. attempt, compiled_autograd_id, you need to ensure
+    // the representation remains unique. One way is to use a unique prefix.
+
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "[{}/{}", self.frame_id, self.frame_compile_id)?;
-        if self.attempt != 0 {
-            write!(f, "_{}", self.attempt)?;
+        write!(f, "[")?;
+        if let Some(compiled_autograd_id) = self.compiled_autograd_id {
+            write!(f, "!{}/", compiled_autograd_id)?;
+        }
+        let frame_id = self.frame_id.map_or("-".to_string(), |v| v.to_string());
+        let frame_compile_id = self
+            .frame_compile_id
+            .map_or("-".to_string(), |v| v.to_string());
+        write!(f, "{}/{}", frame_id, frame_compile_id)?;
+        if let Some(attempt) = self.attempt {
+            if attempt != 0 {
+                write!(f, "_{}", attempt)?;
+            }
         }
         write!(f, "]")
     }
 }
 
+impl CompileId {
+    pub fn as_directory_name(&self) -> String {
+        let compiled_autograd_id_str = self
+            .compiled_autograd_id
+            .map_or("-".to_string(), |v| v.to_string());
+        let frame_id_str = self.frame_id.map_or("-".to_string(), |v| v.to_string());
+        let frame_compile_id_str = self
+            .frame_compile_id
+            .map_or("-".to_string(), |v| v.to_string());
+        let attempt_str = self.attempt.map_or("-".to_string(), |v| v.to_string());
+
+        format!("{compiled_autograd_id_str}_{frame_id_str}_{frame_compile_id_str}_{attempt_str}")
+    }
+}
+
 #[derive(Default, Debug)]
 pub struct Stats {
     pub ok: u64,