diff --git a/cranelift/codegen/src/write.rs b/cranelift/codegen/src/write.rs index 72496a3192..d046da505a 100644 --- a/cranelift/codegen/src/write.rs +++ b/cranelift/codegen/src/write.rs @@ -216,12 +216,18 @@ pub fn write_block_header( block: Block, indent: usize, ) -> fmt::Result { + let cold = if func.layout.is_cold(block) { + " cold" + } else { + "" + }; + // The `indent` is the instruction indentation. block headers are 4 spaces out from that. write!(w, "{1:0$}{2}", indent - 4, "", block)?; let mut args = func.dfg.block_params(block).iter().cloned(); match args.next() { - None => return writeln!(w, ":"), + None => return writeln!(w, "{}:", cold), Some(arg) => { write!(w, "(")?; write_arg(w, func, arg)?; @@ -232,7 +238,7 @@ pub fn write_block_header( write!(w, ", ")?; write_arg(w, func, arg)?; } - writeln!(w, "):") + writeln!(w, "){}:", cold) } fn decorate_block( @@ -666,4 +672,26 @@ mod tests { "function u0:0() fast {\nblock0(v3: i32):\n v0 -> v3\n v2 -> v0\n v4 = iconst.i32 42\n v5 = iadd v0, v0\n v1 -> v5\n v6 = iconst.i32 23\n v7 = iadd v1, v1\n}\n" ); } + + #[test] + fn cold_blocks() { + let mut func = Function::new(); + { + let mut pos = FuncCursor::new(&mut func); + + let block0 = pos.func.dfg.make_block(); + pos.insert_block(block0); + pos.func.layout.set_cold(block0); + + let block1 = pos.func.dfg.make_block(); + pos.insert_block(block1); + pos.func.dfg.append_block_param(block1, types::I32); + pos.func.layout.set_cold(block1); + } + + assert_eq!( + func.to_string(), + "function u0:0() fast {\nblock0 cold:\n\nblock1(v0: i32) cold:\n}\n" + ); + } } diff --git a/cranelift/reader/src/lexer.rs b/cranelift/reader/src/lexer.rs index 45b419f1e3..f8f0272347 100644 --- a/cranelift/reader/src/lexer.rs +++ b/cranelift/reader/src/lexer.rs @@ -34,6 +34,7 @@ pub enum Token<'a> { Type(types::Type), // i32, f32, b32x4, ... Value(Value), // v12, v7 Block(Block), // block3 + Cold, // cold (flag on block) StackSlot(u32), // ss3 GlobalValue(u32), // gv3 Heap(u32), // heap2 @@ -326,6 +327,7 @@ impl<'a> Lexer<'a> { .unwrap_or_else(|| match text { "iflags" => Token::Type(types::IFLAGS), "fflags" => Token::Type(types::FFLAGS), + "cold" => Token::Cold, _ => Token::Identifier(text), }), loc, diff --git a/cranelift/reader/src/parser.rs b/cranelift/reader/src/parser.rs index bc5089d37e..5850c2f974 100644 --- a/cranelift/reader/src/parser.rs +++ b/cranelift/reader/src/parser.rs @@ -447,6 +447,11 @@ impl Context { self.function.layout.append_block(block); Ok(block) } + + /// Set a block as cold. + fn set_cold_block(&mut self, block: Block) { + self.function.layout.set_cold(block); + } } impl<'a> Parser<'a> { @@ -1856,7 +1861,8 @@ impl<'a> Parser<'a> { // Parse a basic block, add contents to `ctx`. // // extended-basic-block ::= * block-header { instruction } - // block-header ::= Block(block) [block-params] ":" + // block-header ::= Block(block) [block-params] [block-flags] ":" + // block-flags ::= [Cold] // fn parse_basic_block(&mut self, ctx: &mut Context) -> ParseResult<()> { // Collect comments for the next block. @@ -1869,12 +1875,16 @@ impl<'a> Parser<'a> { return Err(self.error("too many blocks")); } - if !self.optional(Token::Colon) { - // block-header ::= Block(block) [ * block-params ] ":" + if self.token() == Some(Token::LPar) { self.parse_block_params(ctx, block)?; - self.match_token(Token::Colon, "expected ':' after block parameters")?; } + if self.optional(Token::Cold) { + ctx.set_cold_block(block); + } + + self.match_token(Token::Colon, "expected ':' after block parameters")?; + // Collect any trailing comments. self.token(); self.claim_gathered_comments(block); @@ -3728,4 +3738,23 @@ mod tests { "0x00000003000000020000000100000000" ); } + + #[test] + fn parse_cold_blocks() { + let code = "function %test() { + block0 cold: + return + block1(v0: i32) cold: + return + block2(v1: i32): + return + }"; + + let mut parser = Parser::new(code); + let func = parser.parse_function().unwrap().0; + assert_eq!(func.layout.blocks().count(), 3); + assert!(func.layout.is_cold(Block::from_u32(0))); + assert!(func.layout.is_cold(Block::from_u32(1))); + assert!(!func.layout.is_cold(Block::from_u32(2))); + } }