Auto merge of rust-lang#106428 - saethlin:inline-diverging-functions, r=cjgillot

Permit the MIR inliner to inline diverging functions

This heuristic prevents inlining of `hint::unreachable_unchecked`, which in turn makes `Option/Result::unwrap_unchecked` a bad inlining candidate. I looked through the changes to `core`, `alloc`, `std`, and `hashbrown` by hand and they all seem reasonable. Let's see how this looks in perf...

---

Based on rustc-perf it looks like this regresses ctfe-stress, and the cachegrind diff indicates that this regression is in `InterpCx::statement`. I don't know how to do any deeper analysis because that function is _enormous_ in the try toolchain, which has no debuginfo in it. And a local build produces significantly different codegen for that function, even with LTO.

Author: bors

compiler/rustc_mir_transform/src/inline.rs

@@ -424,13 +424,6 @@ impl<'tcx> Inliner<'tcx> {
         debug!("    final inline threshold = {}", threshold);
 
         // FIXME: Give a bonus to functions with only a single caller
-        let diverges = matches!(
-            callee_body.basic_blocks[START_BLOCK].terminator().kind,
-            TerminatorKind::Unreachable | TerminatorKind::Call { target: None, .. }
-        );
-        if diverges && !matches!(callee_attrs.inline, InlineAttr::Always) {
-            return Err("callee diverges unconditionally");
-        }
 
         let mut checker = CostChecker {
             tcx: self.tcx,

compiler/rustc_mir_transform/src/instcombine.rs

@@ -4,7 +4,7 @@ use crate::MirPass;
 use rustc_hir::Mutability;
 use rustc_middle::mir::{
     BinOp, Body, Constant, ConstantKind, LocalDecls, Operand, Place, ProjectionElem, Rvalue,
-    SourceInfo, Statement, StatementKind, Terminator, TerminatorKind, UnOp,
+    SourceInfo, Statement, StatementKind, SwitchTargets, Terminator, TerminatorKind, UnOp,
 };
 use rustc_middle::ty::layout::ValidityRequirement;
 use rustc_middle::ty::{self, ParamEnv, SubstsRef, Ty, TyCtxt};
@@ -44,6 +44,7 @@ impl<'tcx> MirPass<'tcx> for InstCombine {
                 &mut block.terminator.as_mut().unwrap(),
                 &mut block.statements,
             );
+            ctx.combine_duplicate_switch_targets(&mut block.terminator.as_mut().unwrap());
         }
     }
 }
@@ -217,6 +218,19 @@ impl<'tcx> InstCombineContext<'tcx, '_> {
         terminator.kind = TerminatorKind::Goto { target: destination_block };
     }
 
+    fn combine_duplicate_switch_targets(&self, terminator: &mut Terminator<'tcx>) {
+        let TerminatorKind::SwitchInt { targets, .. } = &mut terminator.kind
+        else { return };
+
+        let otherwise = targets.otherwise();
+        if targets.iter().any(|t| t.1 == otherwise) {
+            *targets = SwitchTargets::new(
+                targets.iter().filter(|t| t.1 != otherwise),
+                targets.otherwise(),
+            );
+        }
+    }
+
     fn combine_intrinsic_assert(
         &self,
         terminator: &mut Terminator<'tcx>,

compiler/rustc_mir_transform/src/simplify.rs

@@ -28,7 +28,7 @@
 //! return.
 
 use crate::MirPass;
-use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::fx::{FxHashSet, FxIndexSet};
 use rustc_index::vec::{Idx, IndexVec};
 use rustc_middle::mir::coverage::*;
 use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
@@ -48,6 +48,7 @@ impl SimplifyCfg {
 
 pub fn simplify_cfg<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     CfgSimplifier::new(body).simplify();
+    remove_duplicate_unreachable_blocks(tcx, body);
     remove_dead_blocks(tcx, body);
 
     // FIXME: Should probably be moved into some kind of pass manager
@@ -259,6 +260,49 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
     }
 }
 
+pub fn remove_duplicate_unreachable_blocks<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    struct OptApplier<'tcx> {
+        tcx: TyCtxt<'tcx>,
+        duplicates: FxIndexSet<BasicBlock>,
+    }
+
+    impl<'tcx> MutVisitor<'tcx> for OptApplier<'tcx> {
+        fn tcx(&self) -> TyCtxt<'tcx> {
+            self.tcx
+        }
+
+        fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
+            for target in terminator.successors_mut() {
+                // We don't have to check whether `target` is a cleanup block, because have
+                // entirely excluded cleanup blocks in building the set of duplicates.
+                if self.duplicates.contains(target) {
+                    *target = self.duplicates[0];
+                }
+            }
+
+            self.super_terminator(terminator, location);
+        }
+    }
+
+    let unreachable_blocks = body
+        .basic_blocks
+        .iter_enumerated()
+        .filter(|(_, bb)| {
+            // CfgSimplifier::simplify leaves behind some unreachable basic blocks without a
+            // terminator. Those blocks will be deleted by remove_dead_blocks, but we run just
+            // before then so we need to handle missing terminators.
+            // We also need to prevent confusing cleanup and non-cleanup blocks. In practice we
+            // don't emit empty unreachable cleanup blocks, so this simple check suffices.
+            bb.terminator.is_some() && bb.is_empty_unreachable() && !bb.is_cleanup
+        })
+        .map(|(block, _)| block)
+        .collect::<FxIndexSet<_>>();
+
+    if unreachable_blocks.len() > 1 {
+        OptApplier { tcx, duplicates: unreachable_blocks }.visit_body(body);
+    }
+}
+
 pub fn remove_dead_blocks<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let reachable = traversal::reachable_as_bitset(body);
     let num_blocks = body.basic_blocks.len();

src/bootstrap/builder.rs

@@ -1941,6 +1941,12 @@ impl<'a> Builder<'a> {
                 rustflags.arg("-Zvalidate-mir");
                 rustflags.arg(&format!("-Zmir-opt-level={}", mir_opt_level));
             }
+            // Always enable inlining MIR when building the standard library.
+            // Without this flag, MIR inlining is disabled when incremental compilation is enabled.
+            // That causes some mir-opt tests which inline functions from the standard library to
+            // break when incremental compilation is enabled. So this overrides the "no inlining
+            // during incremental builds" heuristic for the standard library.
+            rustflags.arg("-Zinline-mir");
         }
 
         Cargo { command: cargo, rustflags, rustdocflags, allow_features }

tests/mir-opt/building/async_await.b-{closure#0}.generator_resume.0.mir

@@ -90,7 +90,7 @@ fn b::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:14:18: 17:2]>,
 
     bb0: {
         _39 = discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2]))); // scope 0 at $DIR/async_await.rs:+0:18: +3:2
-        switchInt(move _39) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb30]; // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        switchInt(move _39) -> [0: bb1, 1: bb28, 3: bb26, 4: bb27, otherwise: bb29]; // scope 0 at $DIR/async_await.rs:+0:18: +3:2
     }
 
     bb1: {
@@ -263,7 +263,7 @@ fn b::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:14:18: 17:2]>,
         StorageDead(_29);                // scope 5 at $DIR/async_await.rs:+2:13: +2:14
         StorageDead(_26);                // scope 5 at $DIR/async_await.rs:+2:13: +2:14
         _32 = discriminant(_25);         // scope 4 at $DIR/async_await.rs:+2:8: +2:14
-        switchInt(move _32) -> [0: bb22, 1: bb20, otherwise: bb21]; // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        switchInt(move _32) -> [0: bb21, 1: bb20, otherwise: bb9]; // scope 4 at $DIR/async_await.rs:+2:8: +2:14
     }
 
     bb20: {
@@ -281,10 +281,6 @@ fn b::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:14:18: 17:2]>,
     }
 
     bb21: {
-        unreachable;                     // scope 4 at $DIR/async_await.rs:+2:8: +2:14
-    }
-
-    bb22: {
         StorageLive(_33);                // scope 4 at $DIR/async_await.rs:+2:5: +2:14
         _33 = ((_25 as Ready).0: ());    // scope 4 at $DIR/async_await.rs:+2:5: +2:14
         _37 = _33;                       // scope 6 at $DIR/async_await.rs:+2:5: +2:14
@@ -293,34 +289,34 @@ fn b::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:14:18: 17:2]>,
         StorageDead(_28);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
         StorageDead(_25);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
         StorageDead(_24);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
-        goto -> bb24;                    // scope 0 at $DIR/async_await.rs:+2:13: +2:14
+        goto -> bb23;                    // scope 0 at $DIR/async_await.rs:+2:13: +2:14
     }
 
-    bb23: {
+    bb22: {
         StorageDead(_36);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
         _38 = move _35;                  // scope 4 at $DIR/async_await.rs:+2:8: +2:14
         StorageDead(_35);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
         _7 = const ();                   // scope 4 at $DIR/async_await.rs:+2:8: +2:14
         goto -> bb16;                    // scope 4 at $DIR/async_await.rs:+2:8: +2:14
     }
 
-    bb24: {
+    bb23: {
         nop;                             // scope 0 at $DIR/async_await.rs:+2:13: +2:14
-        goto -> bb25;                    // scope 0 at $DIR/async_await.rs:+3:1: +3:2
+        goto -> bb24;                    // scope 0 at $DIR/async_await.rs:+3:1: +3:2
     }
 
-    bb25: {
+    bb24: {
         StorageDead(_21);                // scope 0 at $DIR/async_await.rs:+3:1: +3:2
-        goto -> bb26;                    // scope 0 at $DIR/async_await.rs:+3:1: +3:2
+        goto -> bb25;                    // scope 0 at $DIR/async_await.rs:+3:1: +3:2
     }
 
-    bb26: {
+    bb25: {
         _0 = Poll::<()>::Ready(move _37); // scope 0 at $DIR/async_await.rs:+3:2: +3:2
         discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2]))) = 1; // scope 0 at $DIR/async_await.rs:+3:2: +3:2
         return;                          // scope 0 at $DIR/async_await.rs:+3:2: +3:2
     }
 
-    bb27: {
+    bb26: {
         StorageLive(_3);                 // scope 0 at $DIR/async_await.rs:+0:18: +3:2
         StorageLive(_4);                 // scope 0 at $DIR/async_await.rs:+0:18: +3:2
         StorageLive(_19);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
@@ -329,19 +325,19 @@ fn b::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:14:18: 17:2]>,
         goto -> bb11;                    // scope 0 at $DIR/async_await.rs:+0:18: +3:2
     }
 
-    bb28: {
+    bb27: {
         StorageLive(_21);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
         StorageLive(_35);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
         StorageLive(_36);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
         _35 = move _2;                   // scope 0 at $DIR/async_await.rs:+0:18: +3:2
-        goto -> bb23;                    // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        goto -> bb22;                    // scope 0 at $DIR/async_await.rs:+0:18: +3:2
     }
 
-    bb29: {
-        assert(const false, "`async fn` resumed after completion") -> bb29; // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    bb28: {
+        assert(const false, "`async fn` resumed after completion") -> bb28; // scope 0 at $DIR/async_await.rs:+0:18: +3:2
     }
 
-    bb30: {
+    bb29: {
         unreachable;                     // scope 0 at $DIR/async_await.rs:+0:18: +3:2
     }
 }

tests/mir-opt/inline/unchecked_shifts.rs

@@ -0,0 +1,17 @@
+#![crate_type = "lib"]
+#![feature(unchecked_math)]
+
+// ignore-debug: the debug assertions prevent the inlining we are testing for
+// compile-flags: -Zmir-opt-level=2 -Zinline-mir
+
+// EMIT_MIR unchecked_shifts.unchecked_shl_unsigned_smaller.Inline.diff
+// EMIT_MIR unchecked_shifts.unchecked_shl_unsigned_smaller.PreCodegen.after.mir
+pub unsafe fn unchecked_shl_unsigned_smaller(a: u16, b: u32) -> u16 {
+    a.unchecked_shl(b)
+}
+
+// EMIT_MIR unchecked_shifts.unchecked_shr_signed_smaller.Inline.diff
+// EMIT_MIR unchecked_shifts.unchecked_shr_signed_smaller.PreCodegen.after.mir
+pub unsafe fn unchecked_shr_signed_smaller(a: i16, b: u32) -> i16 {
+    a.unchecked_shr(b)
+}

tests/mir-opt/inline/unchecked_shifts.unchecked_shl_unsigned_smaller.Inline.diff

@@ -0,0 +1,115 @@
+- // MIR for `unchecked_shl_unsigned_smaller` before Inline
++ // MIR for `unchecked_shl_unsigned_smaller` after Inline
+  
+  fn unchecked_shl_unsigned_smaller(_1: u16, _2: u32) -> u16 {
+      debug a => _1;                       // in scope 0 at $DIR/unchecked_shifts.rs:+0:46: +0:47
+      debug b => _2;                       // in scope 0 at $DIR/unchecked_shifts.rs:+0:54: +0:55
+      let mut _0: u16;                     // return place in scope 0 at $DIR/unchecked_shifts.rs:+0:65: +0:68
+      let mut _3: u16;                     // in scope 0 at $DIR/unchecked_shifts.rs:+1:5: +1:6
+      let mut _4: u32;                     // in scope 0 at $DIR/unchecked_shifts.rs:+1:21: +1:22
++     scope 1 (inlined core::num::<impl u16>::unchecked_shl) { // at $DIR/unchecked_shifts.rs:10:7: 10:23
++         debug self => _3;                // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         debug rhs => _4;                 // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         let mut _5: u16;                 // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         let mut _6: std::option::Option<u16>; // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         let mut _7: std::result::Result<u16, std::num::TryFromIntError>; // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         scope 2 {
++             scope 3 (inlined Result::<u16, TryFromIntError>::ok) { // at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++                 debug self => _7;        // in scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++                 let mut _8: isize;       // in scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++                 let _9: u16;             // in scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++                 scope 4 {
++                     debug x => _9;       // in scope 4 at $SRC_DIR/core/src/result.rs:LL:COL
++                 }
++                 scope 5 {
++                     scope 6 {
++                         debug x => const TryFromIntError(()); // in scope 6 at $SRC_DIR/core/src/result.rs:LL:COL
++                     }
++                 }
++             }
++             scope 7 (inlined #[track_caller] Option::<u16>::unwrap_unchecked) { // at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++                 debug self => _6;        // in scope 7 at $SRC_DIR/core/src/option.rs:LL:COL
++                 let mut _10: &std::option::Option<u16>; // in scope 7 at $SRC_DIR/core/src/option.rs:LL:COL
++                 let mut _11: isize;      // in scope 7 at $SRC_DIR/core/src/option.rs:LL:COL
++                 scope 8 {
++                     debug val => _5;     // in scope 8 at $SRC_DIR/core/src/option.rs:LL:COL
++                 }
++                 scope 9 {
++                     scope 11 (inlined unreachable_unchecked) { // at $SRC_DIR/core/src/option.rs:LL:COL
++                         scope 12 {
++                             scope 13 (inlined unreachable_unchecked::runtime) { // at $SRC_DIR/core/src/intrinsics.rs:LL:COL
++                             }
++                         }
++                     }
++                 }
++                 scope 10 (inlined Option::<u16>::is_some) { // at $SRC_DIR/core/src/option.rs:LL:COL
++                     debug self => _10;   // in scope 10 at $SRC_DIR/core/src/option.rs:LL:COL
++                 }
++             }
++         }
++     }
+  
+      bb0: {
+          StorageLive(_3);                 // scope 0 at $DIR/unchecked_shifts.rs:+1:5: +1:6
+          _3 = _1;                         // scope 0 at $DIR/unchecked_shifts.rs:+1:5: +1:6
+          StorageLive(_4);                 // scope 0 at $DIR/unchecked_shifts.rs:+1:21: +1:22
+          _4 = _2;                         // scope 0 at $DIR/unchecked_shifts.rs:+1:21: +1:22
+-         _0 = core::num::<impl u16>::unchecked_shl(move _3, move _4) -> bb1; // scope 0 at $DIR/unchecked_shifts.rs:+1:5: +1:23
++         StorageLive(_5);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         StorageLive(_6);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         StorageLive(_7);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         _7 = <u32 as TryInto<u16>>::try_into(_4) -> bb1; // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
+                                           // mir::Constant
+-                                          // + span: $DIR/unchecked_shifts.rs:10:7: 10:20
+-                                          // + literal: Const { ty: unsafe fn(u16, u32) -> u16 {core::num::<impl u16>::unchecked_shl}, val: Value(<ZST>) }
++                                          // + span: $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++                                          // + literal: Const { ty: fn(u32) -> Result<u16, <u32 as TryInto<u16>>::Error> {<u32 as TryInto<u16>>::try_into}, val: Value(<ZST>) }
+      }
+  
+      bb1: {
++         StorageLive(_9);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         _8 = discriminant(_7);           // scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++         switchInt(move _8) -> [0: bb6, 1: bb4, otherwise: bb5]; // scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++     }
++ 
++     bb2: {
++         StorageDead(_9);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         StorageDead(_7);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         StorageLive(_10);                // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         _11 = discriminant(_6);          // scope 7 at $SRC_DIR/core/src/option.rs:LL:COL
++         switchInt(move _11) -> [1: bb7, otherwise: bb5]; // scope 7 at $SRC_DIR/core/src/option.rs:LL:COL
++     }
++ 
++     bb3: {
++         StorageDead(_5);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
+          StorageDead(_4);                 // scope 0 at $DIR/unchecked_shifts.rs:+1:22: +1:23
+          StorageDead(_3);                 // scope 0 at $DIR/unchecked_shifts.rs:+1:22: +1:23
+          return;                          // scope 0 at $DIR/unchecked_shifts.rs:+2:2: +2:2
++     }
++ 
++     bb4: {
++         _6 = Option::<u16>::None;        // scope 6 at $SRC_DIR/core/src/result.rs:LL:COL
++         goto -> bb2;                     // scope 5 at $SRC_DIR/core/src/result.rs:LL:COL
++     }
++ 
++     bb5: {
++         unreachable;                     // scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++     }
++ 
++     bb6: {
++         _9 = move ((_7 as Ok).0: u16);   // scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++         _6 = Option::<u16>::Some(move _9); // scope 4 at $SRC_DIR/core/src/result.rs:LL:COL
++         goto -> bb2;                     // scope 3 at $SRC_DIR/core/src/result.rs:LL:COL
++     }
++ 
++     bb7: {
++         _5 = move ((_6 as Some).0: u16); // scope 7 at $SRC_DIR/core/src/option.rs:LL:COL
++         StorageDead(_10);                // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         StorageDead(_6);                 // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++         _0 = unchecked_shl::<u16>(_3, move _5) -> bb3; // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++                                          // mir::Constant
++                                          // + span: $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
++                                          // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(u16, u16) -> u16 {unchecked_shl::<u16>}, val: Value(<ZST>) }
+      }
+  }
+