From f489b83835803523844908b9beacdfee312f9365 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Tue, 18 Jan 2022 17:19:08 -0800
Subject: [PATCH] Cranelift: add support for cold blocks.

This PR adds a flag to each block that can be set via the frontend/builder
interface that indicates that the block will not be frequently
executed. As such, the compiler backend should place the block "out of
line" in the final machine code, so that the ordinary, more frequent
execution path that excludes the block does not have to jump around it.

This is useful for adding handlers for exceptional conditions
(slow-paths, guard violations) in a way that minimizes performance cost.

Fixes #2747.
---
 cranelift/codegen/src/ir/layout.rs           | 14 ++++++++++++++
 cranelift/codegen/src/machinst/blockorder.rs | 15 ++++++++++++++-
 cranelift/frontend/src/frontend.rs           |  8 ++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/cranelift/codegen/src/ir/layout.rs b/cranelift/codegen/src/ir/layout.rs
index 37842b59be..47ec6b1fda 100644
--- a/cranelift/codegen/src/ir/layout.rs
+++ b/cranelift/codegen/src/ir/layout.rs
@@ -469,6 +469,19 @@ impl Layout {
     pub fn next_block(&self, block: Block) -> Option<Block> {
         self.blocks[block].next.expand()
     }
+
+    /// Mark a block as "cold".
+    ///
+    /// This will try to move it out of the ordinary path of execution
+    /// when lowered to machine code.
+    pub fn set_cold(&mut self, block: Block) {
+        self.blocks[block].cold = true;
+    }
+
+    /// Is the given block cold?
+    pub fn is_cold(&self, block: Block) -> bool {
+        self.blocks[block].cold
+    }
 }
 
 #[derive(Clone, Debug, Default)]
@@ -478,6 +491,7 @@ struct BlockNode {
     first_inst: PackedOption<Inst>,
     last_inst: PackedOption<Inst>,
     seq: SequenceNumber,
+    cold: bool,
 }
 
 /// Iterate over blocks in layout order. See [crate::ir::layout::Layout::blocks].
diff --git a/cranelift/codegen/src/machinst/blockorder.rs b/cranelift/codegen/src/machinst/blockorder.rs
index 73c5ea2453..3cdf1228ca 100644
--- a/cranelift/codegen/src/machinst/blockorder.rs
+++ b/cranelift/codegen/src/machinst/blockorder.rs
@@ -378,11 +378,24 @@ impl BlockLoweringOrder {
 
         postorder.reverse();
         let mut rpo = postorder;
+
+        // Step 3: sink any cold blocks to the end of the
+        // function. Put the "deferred last" block truly at the end;
+        // this is a correctness requirement (for fallthrough
+        // returns).
+        rpo.sort_by_key(|block| {
+            block
+                .0
+                .orig_block()
+                .map(|block| f.layout.is_cold(block))
+                .unwrap_or(false)
+        });
+
         if let Some(d) = deferred_last {
             rpo.push(d);
         }
 
-        // Step 3: now that we have RPO, build the BlockIndex/BB fwd/rev maps.
+        // Step 4: now that we have RPO, build the BlockIndex/BB fwd/rev maps.
         let mut lowered_order = vec![];
         let mut lowered_succ_ranges = vec![];
         let mut lb_to_bindex = FxHashMap::default();
diff --git a/cranelift/frontend/src/frontend.rs b/cranelift/frontend/src/frontend.rs
index 7611e69c0a..ee16400e17 100644
--- a/cranelift/frontend/src/frontend.rs
+++ b/cranelift/frontend/src/frontend.rs
@@ -229,6 +229,14 @@ impl<'a> FunctionBuilder<'a> {
         block
     }
 
+    /// Mark a block as "cold".
+    ///
+    /// This will try to move it out of the ordinary path of execution
+    /// when lowered to machine code.
+    pub fn set_cold_block(&mut self, block: Block) {
+        self.func.layout.set_cold(block);
+    }
+
     /// Insert `block` in the layout *after* the existing block `after`.
     pub fn insert_block_after(&mut self, block: Block, after: Block) {
         self.func.layout.insert_block_after(block, after);