Use is_val_statically_known to optimize pow

In the dynamic exponent case, it's preferred to not increase code size, so use solely the loop-based implementation there. This shows about 4% penalty in the variable exponent benchmarks on x86_64.
rust-lang · Jul 12, 2024 · 2f23534 · 2f23534
1 parent 1faa101
commit 2f23534
Show file tree

Hide file tree

Showing 3 changed files with 125 additions and 92 deletions.
diff --git a/library/core/src/lib.rs b/library/core/src/lib.rs
@@ -170,6 +170,7 @@
 #![feature(internal_impls_macro)]
 #![feature(ip)]
 #![feature(is_ascii_octdigit)]
+#![feature(is_val_statically_known)]
 #![feature(isqrt)]
 #![feature(link_cfg)]
 #![feature(offset_of_enum)]

diff --git a/library/core/src/num/int_macros.rs b/library/core/src/num/int_macros.rs
@@ -2172,35 +2172,43 @@ macro_rules! int_impl {
         #[must_use = "this returns the result of the operation, \
                       without modifying the original"]
         #[inline]
+        #[rustc_allow_const_fn_unstable(is_val_statically_known)]
         pub const fn wrapping_pow(self, mut exp: u32) -> Self {
             let mut base = self;
 
-            // Unroll multiplications for small exponent values.
-            // This gives the optimizer a way to efficiently inline call sites
-            // for the most common use cases with constant exponents.
-            // Currently, LLVM is unable to unroll the loop below.
-            match exp {
-                0 => return 1,
-                1 => return base,
-                2 => return base.wrapping_mul(base),
-                3 => {
-                    let squared = base.wrapping_mul(base);
-                    return squared.wrapping_mul(base);
-                }
-                4 => {
-                    let squared = base.wrapping_mul(base);
-                    return squared.wrapping_mul(squared);
-                }
-                5 => {
-                    let squared = base.wrapping_mul(base);
-                    return squared.wrapping_mul(squared).wrapping_mul(base);
+            if intrinsics::is_val_statically_known(exp) {
+                // Unroll multiplications for small exponent values.
+                // This gives the optimizer a way to efficiently inline call sites
+                // for the most common use cases with constant exponents.
+                // Currently, LLVM is unable to unroll the loop below.
+                match exp {
+                    0 => return 1,
+                    1 => return base,
+                    2 => return base.wrapping_mul(base),
+                    3 => {
+                        let squared = base.wrapping_mul(base);
+                        return squared.wrapping_mul(base);
+                    }
+                    4 => {
+                        let squared = base.wrapping_mul(base);
+                        return squared.wrapping_mul(squared);
+                    }
+                    5 => {
+                        let squared = base.wrapping_mul(base);
+                        return squared.wrapping_mul(squared).wrapping_mul(base);
+                    }
+                    6 => {
+                        let cubed = base.wrapping_mul(base).wrapping_mul(base);
+                        return cubed.wrapping_mul(cubed);
+                    }
+                    _ => {}
                 }
-                6 => {
-                    let cubed = base.wrapping_mul(base).wrapping_mul(base);
-                    return cubed.wrapping_mul(cubed);
+            } else {
+                if exp == 0 {
+                    return 1;
                 }
-                _ => {}
             }
+            debug_assert!(exp != 0);
 
             let mut acc: Self = 1;
 
@@ -2743,35 +2751,43 @@ macro_rules! int_impl {
                       without modifying the original"]
         #[inline]
         #[rustc_inherit_overflow_checks]
+        #[rustc_allow_const_fn_unstable(is_val_statically_known)]
         pub const fn pow(self, mut exp: u32) -> Self {
             let mut base = self;
 
-            // Unroll multiplications for small exponent values.
-            // This gives the optimizer a way to efficiently inline call sites
-            // for the most common use cases with constant exponents.
-            // Currently, LLVM is unable to unroll the loop below.
-            match exp {
-                0 => return 1,
-                1 => return base,
-                2 => return base * base,
-                3 => {
-                    let squared = base * base;
-                    return squared * base;
-                }
-                4 => {
-                    let squared = base * base;
-                    return squared * squared;
-                }
-                5 => {
-                    let squared = base * base;
-                    return squared * squared * base;
+            if intrinsics::is_val_statically_known(exp) {
+                // Unroll multiplications for small exponent values.
+                // This gives the optimizer a way to efficiently inline call sites
+                // for the most common use cases with constant exponents.
+                // Currently, LLVM is unable to unroll the loop below.
+                match exp {
+                    0 => return 1,
+                    1 => return base,
+                    2 => return base * base,
+                    3 => {
+                        let squared = base * base;
+                        return squared * base;
+                    }
+                    4 => {
+                        let squared = base * base;
+                        return squared * squared;
+                    }
+                    5 => {
+                        let squared = base * base;
+                        return squared * squared * base;
+                    }
+                    6 => {
+                        let cubed = base * base * base;
+                        return cubed * cubed;
+                    }
+                    _ => {}
                 }
-                6 => {
-                    let cubed = base * base * base;
-                    return cubed * cubed;
+            } else {
+                if exp == 0 {
+                    return 1;
                 }
-                _ => {}
             }
+            debug_assert!(exp != 0);
 
             let mut acc = 1;
 

diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs
@@ -2048,35 +2048,43 @@ macro_rules! uint_impl {
         #[must_use = "this returns the result of the operation, \
                       without modifying the original"]
         #[inline]
+        #[rustc_allow_const_fn_unstable(is_val_statically_known)]
         pub const fn wrapping_pow(self, mut exp: u32) -> Self {
             let mut base = self;
 
-            // Unroll multiplications for small exponent values.
-            // This gives the optimizer a way to efficiently inline call sites
-            // for the most common use cases with constant exponents.
-            // Currently, LLVM is unable to unroll the loop below.
-            match exp {
-                0 => return 1,
-                1 => return base,
-                2 => return base.wrapping_mul(base),
-                3 => {
-                    let squared = base.wrapping_mul(base);
-                    return squared.wrapping_mul(base);
-                }
-                4 => {
-                    let squared = base.wrapping_mul(base);
-                    return squared.wrapping_mul(squared);
-                }
-                5 => {
-                    let squared = base.wrapping_mul(base);
-                    return squared.wrapping_mul(squared).wrapping_mul(base);
+            if intrinsics::is_val_statically_known(exp) {
+                // Unroll multiplications for small exponent values.
+                // This gives the optimizer a way to efficiently inline call sites
+                // for the most common use cases with constant exponents.
+                // Currently, LLVM is unable to unroll the loop below.
+                match exp {
+                    0 => return 1,
+                    1 => return base,
+                    2 => return base.wrapping_mul(base),
+                    3 => {
+                        let squared = base.wrapping_mul(base);
+                        return squared.wrapping_mul(base);
+                    }
+                    4 => {
+                        let squared = base.wrapping_mul(base);
+                        return squared.wrapping_mul(squared);
+                    }
+                    5 => {
+                        let squared = base.wrapping_mul(base);
+                        return squared.wrapping_mul(squared).wrapping_mul(base);
+                    }
+                    6 => {
+                        let cubed = base.wrapping_mul(base).wrapping_mul(base);
+                        return cubed.wrapping_mul(cubed);
+                    }
+                    _ => {}
                 }
-                6 => {
-                    let cubed = base.wrapping_mul(base).wrapping_mul(base);
-                    return cubed.wrapping_mul(cubed);
+            } else {
+                if exp == 0 {
+                    return 1;
                 }
-                _ => {}
             }
+            debug_assert!(exp != 0);
 
             let mut acc: Self = 1;
 
@@ -2568,35 +2576,43 @@ macro_rules! uint_impl {
                       without modifying the original"]
         #[inline]
         #[rustc_inherit_overflow_checks]
+        #[rustc_allow_const_fn_unstable(is_val_statically_known)]
         pub const fn pow(self, mut exp: u32) -> Self {
             let mut base = self;
 
-            // Unroll multiplications for small exponent values.
-            // This gives the optimizer a way to efficiently inline call sites
-            // for the most common use cases with constant exponents.
-            // Currently, LLVM is unable to unroll the loop below.
-            match exp {
-                0 => return 1,
-                1 => return base,
-                2 => return base * base,
-                3 => {
-                    let squared = base * base;
-                    return squared * base;
-                }
-                4 => {
-                    let squared = base * base;
-                    return squared * squared;
-                }
-                5 => {
-                    let squared = base * base;
-                    return squared * squared * base;
+            if intrinsics::is_val_statically_known(exp) {
+                // Unroll multiplications for small exponent values.
+                // This gives the optimizer a way to efficiently inline call sites
+                // for the most common use cases with constant exponents.
+                // Currently, LLVM is unable to unroll the loop below.
+                match exp {
+                    0 => return 1,
+                    1 => return base,
+                    2 => return base * base,
+                    3 => {
+                        let squared = base * base;
+                        return squared * base;
+                    }
+                    4 => {
+                        let squared = base * base;
+                        return squared * squared;
+                    }
+                    5 => {
+                        let squared = base * base;
+                        return squared * squared * base;
+                    }
+                    6 => {
+                        let cubed = base * base * base;
+                        return cubed * cubed;
+                    }
+                    _ => {}
                 }
-                6 => {
-                    let cubed = base * base * base;
-                    return cubed * cubed;
+            } else {
+                if exp == 0 {
+                    return 1;
                 }
-                _ => {}
             }
+            debug_assert!(exp != 0);
 
             let mut acc = 1;