[PATCH v3 1/2] rust: std_vendor: add `likely()`, `unlikely()` and `cold_path()`

[PATCH v3 1/2] rust: std_vendor: add `likely()`, `unlikely()` and `cold_path()`

[Miguel Ojeda]

`{un,}likely()` are commonly used macros on the C side -- there are
more than 25k occurrences in the tree. Rust is likely going to provide
functions (not macros) for the same purpose: `core::hint::{un,}likely()`.
Currently they are unstable under `feature(likely_unlikely)` [1].

For instance, given:

    fn f1(a: i32) -> i32 {
        if a < 0 {
            return 123;
        }

        42
    }

    fn f2(a: i32) -> i32 {
        if likely(a < 0) {
            return 124;
        }

        42
    }

    fn f3(a: i32) -> i32 {
        if unlikely(a < 0) {
            return 125;
        }

        42
    }

LLVM emits code equivalent to the one it would for similar C functions:

    f1:
            test   %edi,%edi
            mov    $0x7b,%ecx
            mov    $0x2a,%eax
            cmovs  %ecx,%eax
            ret

    f2:
            mov    $0x7c,%eax
            test   %edi,%edi
        /-- jns    <f2+0xa>
        |   ret
        \-> mov    $0x2a,%eax
            ret

    f3:
            mov    $0x2a,%eax
            test   %edi,%edi
        /-- js     <f3+0xa>
        |   ret
        \-> mov    $0x7d,%eax
            ret

Thus vendor (essentially) the `core` library code to provide the same
API (so that we may use `core`'s directly in the future).

We do so by implementing `{un,}likely()` like `core`'s `intrinsics` do
(to save a layer) and based on calling our own `cold_path()` instead of
using `core`'s (which is stable since 1.95.0 [2][3]).

It turns out that simply using `#[cold]` to emulate `cold_path()`,
suggested by Gary [4], works in the cases I tested across all supported
versions so far, so we can avoid the complexity of having one more Kconfig
and conditionally enabling the feature in callers (since `cold_path()`
was added in Rust 1.86.0 [5]).

We could optionally call `cold_path()` for Rust >= 1.86.0 within our own
`cold_path()`, but it seems that it is not needed, and anyway since Rust
1.98.0 `cold_path()` is itself `#[cold]` so that it works even if the
MIR inliner can't inline it [6].

Even if we find that it is not optimal in some cases, it can be improved,
and meanwhile we will have the means to annotate these cases already.

Apart from that, the only differences with `core`'s version are the
usual changes (`pr_*!` calls and their `\n`, imports, attributes, lint
cleanup), plus a fix for a polarity typo (I sent a fix upstream [7])
and removal of unneeded code documentation.

In addition, add them to the prelude (inlined, so that the documentation
is rendered at least there, until we can use `core`'s).

Link: https://github.com/rust-lang/rust/issues/151619 [1]
Link: https://github.com/rust-lang/rust/issues/136873 [2]
Link: https://github.com/rust-lang/rust/pull/151576 [3]
Link: https://lore.kernel.org/rust-for-linux/DGA6GUR58QJ7.1XZZ0P4VZNW86@garyguo.net/ [4]
Link: https://github.com/rust-lang/rust/pull/133695 [5]
Link: https://github.com/rust-lang/rust/pull/156863 [6]
Link: https://github.com/rust-lang/rust/pull/164374 [7]
Assisted-by: LLM
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>
---
v3: I used Gary's `#[cold]` suggestion, and thus removed the Kconfig.
I rebased and made sure the code is up to date with `core`, tested a
bunch of things, compared again with C and the assembly generated, found
a strange behavior on Clang I will report separately, a polarity typo in
`core` (fixed here) and inlined the docs in the prelude to get it
rendered somewhere until we bump the MSRV. I think the potential issue
I found in v2's discussion may have been with `-Os`, but it happens the
same way with `core`'s, so it does not matter.

Alice, Gary: I removed your Reviewed-by's just in case. It is pretty
similar to last time, but the commit message is reworded, and the commit
is simpler.

 rust/kernel/prelude.rs    |   7 ++
 rust/kernel/std_vendor.rs | 136 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 143 insertions(+)

diff --git a/rust/kernel/prelude.rs b/rust/kernel/prelude.rs
index edaf1e4c174b..de5a5aad16a1 100644
--- a/rust/kernel/prelude.rs
+++ b/rust/kernel/prelude.rs
@@ -133,3 +133,10 @@
 // `super::std_vendor` is hidden, which makes the macro inline for some reason.
 #[doc(no_inline)]
 pub use super::dbg;
+
+#[doc(inline)]
+pub use super::std_vendor::{
+    cold_path,
+    likely,
+    unlikely, //
+};
diff --git a/rust/kernel/std_vendor.rs b/rust/kernel/std_vendor.rs
index abbab5050cc5..8738c387fb4a 100644
--- a/rust/kernel/std_vendor.rs
+++ b/rust/kernel/std_vendor.rs
@@ -162,3 +162,139 @@ macro_rules! dbg {
         ($($crate::dbg!($val)),+,)
     };
 }
+
+/// Hints to the compiler that a branch condition is likely to be true.
+/// Returns the value passed to it.
+///
+/// It can be used with `if` or boolean `match` expressions.
+///
+/// When used outside of a branch condition, it may still influence a nearby branch, but
+/// probably will not have any effect.
+///
+/// It can also be applied to parts of expressions, such as `likely(a) && unlikely(b)`, or to
+/// compound expressions, such as `likely(a && b)`. When applied to compound expressions, it has
+/// the following effect:
+/// ```text
+///     likely(!a) => !unlikely(a)
+///     likely(a && b) => likely(a) && likely(b)
+///     likely(a || b) => a || likely(b)
+/// ```
+///
+/// See also the function [`cold_path()`] which may be more appropriate for idiomatic Rust code.
+///
+/// # Examples
+///
+/// ```
+/// fn foo(x: i32) {
+///     if likely(x > 0) {
+///         pr_info!("this branch is likely to be taken\n");
+///     } else {
+///         pr_info!("this branch is unlikely to be taken\n");
+///     }
+///
+///     match likely(x > 0) {
+///         true => pr_info!("this branch is likely to be taken\n"),
+///         false => pr_info!("this branch is unlikely to be taken\n"),
+///     }
+///
+///     // Use outside of a branch condition may still influence a nearby branch
+///     let cond = likely(x != 0);
+///     if cond {
+///         pr_info!("this branch is likely to be taken\n");
+///     }
+/// }
+/// ```
+#[inline(always)]
+pub const fn likely(b: bool) -> bool {
+    if b {
+        true
+    } else {
+        cold_path();
+        false
+    }
+}
+
+/// Hints to the compiler that a branch condition is unlikely to be true.
+/// Returns the value passed to it.
+///
+/// It can be used with `if` or boolean `match` expressions.
+///
+/// When used outside of a branch condition, it may still influence a nearby branch, but
+/// probably will not have any effect.
+///
+/// It can also be applied to parts of expressions, such as `likely(a) && unlikely(b)`, or to
+/// compound expressions, such as `unlikely(a && b)`. When applied to compound expressions, it has
+/// the following effect:
+/// ```text
+///     unlikely(!a) => !likely(a)
+///     unlikely(a && b) => a && unlikely(b)
+///     unlikely(a || b) => unlikely(a) || unlikely(b)
+/// ```
+///
+/// See also the function [`cold_path()`] which may be more appropriate for idiomatic Rust code.
+///
+/// # Examples
+///
+/// ```
+/// fn foo(x: i32) {
+///     if unlikely(x > 0) {
+///         pr_info!("this branch is unlikely to be taken\n");
+///     } else {
+///         pr_info!("this branch is likely to be taken\n");
+///     }
+///
+///     match unlikely(x > 0) {
+///         true => pr_info!("this branch is unlikely to be taken\n"),
+///         false => pr_info!("this branch is likely to be taken\n"),
+///     }
+///
+///     // Use outside of a branch condition may still influence a nearby branch
+///     let cond = unlikely(x != 0);
+///     if cond {
+///         pr_info!("this branch is unlikely to be taken\n");
+///     }
+/// }
+/// ```
+#[inline(always)]
+pub const fn unlikely(b: bool) -> bool {
+    if b {
+        cold_path();
+        true
+    } else {
+        false
+    }
+}
+
+/// Hints to the compiler that given path is cold, i.e., unlikely to be taken. The compiler may
+/// choose to optimize paths that are not cold at the expense of paths that are cold.
+///
+/// Note that like all hints, the exact effect to codegen is not guaranteed. Using `cold_path`
+/// can actually *decrease* performance if the branch is called more than expected. It is advisable
+/// to perform benchmarks to tell if this function is useful.
+///
+/// # Examples
+///
+/// ```
+/// fn foo(x: &[i32]) {
+///     if let Some(first) = x.first() {
+///         // this is the fast path
+///     } else {
+///         // this path is unlikely
+///         cold_path();
+///     }
+/// }
+///
+/// fn bar(x: i32) -> i32 {
+///     match x {
+///         1 => 10,
+///         2 => 100,
+///         3 => { cold_path(); 1000 }, // this branch is unlikely
+///         _ => { cold_path(); 10000 }, // this is also unlikely
+///     }
+/// }
+/// ```
+///
+/// See also the [`likely()`] and [`unlikely()`] functions, which are similar to the C side macros.
+#[inline(always)]
+#[cold]
+pub const fn cold_path() {}

base-commit: 7c306f00835a237a70024a8f968dbc3bc8db8fcf
--
2.56.0

[PATCH v3 2/2] rust: bug: hint `warn_on!` conditions as unlikely

[Miguel Ojeda]

The C side `WARN_ON()` macro hints the compiler that its condition is
unlikely to be true, both for the branch that reports the warning and
for the value the macro evaluates to, so that callers such as:

    if (WARN_ON(x))
        return -EINVAL;

get the warning path laid out of line.

Now that `unlikely()` is available, do the same in `warn_on!`.

Suggested-by: Gary Guo <gary@garyguo.net>
Link: https://lore.kernel.org/rust-for-linux/DJUY1ZYYW5K6.39XJ3ZV7DGZQW@garyguo.net/
Assisted-by: LLM
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>
---
 rust/kernel/bug.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/rust/kernel/bug.rs b/rust/kernel/bug.rs
index 3b40b0bd37b5..f1209f0dbccb 100644
--- a/rust/kernel/bug.rs
+++ b/rust/kernel/bug.rs
@@ -147,12 +147,12 @@ macro_rules! warn_on {
         #[cfg(not(CONFIG_DEBUG_BUGVERBOSE_DETAILED))]
         const COND_STR: &str = file!();
 
-        if cond {
+        if $crate::prelude::unlikely(cond) {
             const WARN_ON_FLAGS: u32 = $crate::bug::bugflag_taint($crate::bindings::TAINT_WARN);
 
             $crate::warn_flags!(COND_STR, WARN_ON_FLAGS);
         }
-        cond
+        $crate::prelude::unlikely(cond)
     }};
 }
 
-- 
2.56.0