Enable dumping modified bytecode (#60)

src/parsers.rs

@@ -197,6 +197,33 @@ impl StructuredLogParser for DynamoOutputGraphParser {
     }
 }
 
+pub struct BytecodeDumpParser;
+impl StructuredLogParser for BytecodeDumpParser {
+    fn name(&self) -> &'static str {
+        "bytecode_dump"
+    }
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.bytecode_dump
+            .as_ref()
+            .map(|m| Metadata::BytecodeDump(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::BytecodeDump(metadata) = metadata {
+            let filename = format!("{}_{}.txt", metadata.co_name, metadata.bytecode_type);
+            payload_file_output(&filename, lineno, compile_id)
+        } else {
+            Err(anyhow::anyhow!("Expected Bytecode metadata"))
+        }
+    }
+}
+
 pub struct DynamoGuardParser<'t> {
     tt: &'t TinyTemplate<'t>,
 }
@@ -1175,6 +1202,7 @@ pub fn default_parsers<'t>(
         })),
         Box::new(GraphDumpParser),
         Box::new(DynamoOutputGraphParser),
+        Box::new(BytecodeDumpParser),
         Box::new(DynamoGuardParser { tt }),
         Box::new(InductorOutputCodeParser::new(parser_config)),
         Box::new(OptimizeDdpSplitChildParser),

src/types.rs

@@ -425,6 +425,12 @@ pub struct DynamoOutputGraphMetadata {
     _sizes: Option<FxHashMap<String, Vec<SymInt>>>,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct BytecodeDumpMetadata {
+    pub co_name: String,
+    pub bytecode_type: String,
+}
+
 #[derive(Debug, Deserialize)]
 pub struct DynamoStartMetadata {
     pub stack: Option<StackSummary>,
@@ -692,6 +698,7 @@ pub enum Metadata<'e> {
     GraphDump(&'e GraphDumpMetadata),
     DynamoOutputGraph(&'e DynamoOutputGraphMetadata),
     #[allow(dead_code)]
+    BytecodeDump(&'e BytecodeDumpMetadata),
     DynamoStart(&'e DynamoStartMetadata),
     InductorOutputCode(&'e InductorOutputCodeMetadata),
     OptimizeDdpSplitChild(&'e OptimizeDdpSplitChildMetadata),
@@ -763,6 +770,7 @@ pub struct Envelope {
     pub exported_program: Option<EmptyMetadata>,
     #[serde(flatten)]
     pub _other: FxHashMap<String, Value>,
+    pub bytecode_dump: Option<BytecodeDumpMetadata>,
 }
 
 type MetaTensorId = u64;