diff --git a/src/parsers.rs b/src/parsers.rs index c3f3771..095b892 100644 --- a/src/parsers.rs +++ b/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> { + e.bytecode_dump + .as_ref() + .map(|m| Metadata::BytecodeDump(m)) + } + fn parse<'e>( + &self, + lineno: usize, + metadata: Metadata<'e>, + _rank: Option, + compile_id: &Option, + _payload: &str, + ) -> anyhow::Result { + 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), diff --git a/src/types.rs b/src/types.rs index b229c0a..acfba6f 100644 --- a/src/types.rs +++ b/src/types.rs @@ -425,6 +425,12 @@ pub struct DynamoOutputGraphMetadata { _sizes: Option>>, } +#[derive(Debug, Deserialize)] +pub struct BytecodeDumpMetadata { + pub co_name: String, + pub bytecode_type: String, +} + #[derive(Debug, Deserialize)] pub struct DynamoStartMetadata { pub stack: Option, @@ -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, #[serde(flatten)] pub _other: FxHashMap, + pub bytecode_dump: Option, } type MetaTensorId = u64;