Re: [PATCH v5 7/8] rust: percpu: Add pin-hole optimizations for numerics

Yury Norov
2026-04-11T03:06:36Z

On Fri, Apr 10, 2026 at 02:35:37PM -0700, Mitchell Levy wrote:
> The C implementations of `this_cpu_add`, `this_cpu_sub`, etc., are
> optimized to save an instruction by avoiding having to compute
> `this_cpu_ptr(&x)` for some per-CPU variable `x`. For example, rather
> than
>
> u64 *x_ptr = this_cpu_ptr(&x);
> *x_ptr += 5;
>
> the implementation of `this_cpu_add` is clever enough to make use of the
> fact that per-CPU variables are implemented on x86 via segment
> registers, and so we can use only a single instruction (where we assume
> `&x` is already in `rax`)
>
> add gs:[rax], 5
>
> Add this optimization via a `PerCpuNumeric` type to enable code-reuse
> between `DynamicPerCpu` and `StaticPerCpu`.
>
> Signed-off-by: Mitchell Levy <levymitchell0@gmail.com>
> ---
> rust/kernel/percpu.rs | 1 +
> rust/kernel/percpu/dynamic.rs | 10 ++-
> rust/kernel/percpu/numeric.rs | 138 ++++++++++++++++++++++++++++++++++++++++++
> samples/rust/rust_percpu.rs | 36 +++++++++++
> 4 files changed, 184 insertions(+), 1 deletion(-)
>
> diff --git a/rust/kernel/percpu.rs b/rust/kernel/percpu.rs
> index 72c83fef68ee..ff04607ee047 100644
> --- a/rust/kernel/percpu.rs
> +++ b/rust/kernel/percpu.rs
> @@ -6,6 +6,7 @@
>
> pub mod cpu_guard;
> mod dynamic;
> +pub mod numeric;
> mod static_;
>
> #[doc(inline)]
> diff --git a/rust/kernel/percpu/dynamic.rs b/rust/kernel/percpu/dynamic.rs
> index 40514704b3d0..a717138b93dc 100644
> --- a/rust/kernel/percpu/dynamic.rs
> +++ b/rust/kernel/percpu/dynamic.rs
> @@ -28,7 +28,7 @@
> /// the memory location on any particular CPU has been initialized. This means that it cannot tell
> /// whether it should drop the *contents* of the allocation when it is dropped. It is up to the
> /// user to do this via something like [`core::ptr::drop_in_place`].
> -pub struct PerCpuAllocation<T>(PerCpuPtr<T>);
> +pub struct PerCpuAllocation<T>(pub(super) PerCpuPtr<T>);
>
> impl<T: Zeroable> PerCpuAllocation<T> {
> /// Dynamically allocates a space in the per-CPU area suitably sized and aligned to hold a `T`,
> @@ -162,6 +162,14 @@ pub fn new_from(mut initer: impl FnMut(CpuId) -> T, flags: Flags) -> Option<Self
> }
> }
>
> +impl<T> DynamicPerCpu<T> {
> + /// Gets the allocation backing this per-CPU variable.
> + pub(crate) fn alloc(&self) -> &Arc<PerCpuAllocation<T>> {
> + // SAFETY: This type's invariant ensures that `self.alloc` is `Some`.
> + unsafe { self.alloc.as_ref().unwrap_unchecked() }
> + }
> +}
> +
> impl<T> PerCpu<T> for DynamicPerCpu<T> {
> unsafe fn get_mut(&mut self, guard: CpuGuard) -> PerCpuToken<'_, T> {
> // SAFETY:
> diff --git a/rust/kernel/percpu/numeric.rs b/rust/kernel/percpu/numeric.rs
> new file mode 100644
> index 000000000000..13b4ab4a794d
> --- /dev/null
> +++ b/rust/kernel/percpu/numeric.rs
> @@ -0,0 +1,138 @@
> +// SPDX-License-Identifier: GPL-2.0
> +//! Pin-hole optimizations for [`PerCpu<T>`] where T is a numeric type.
> +
> +use super::*;
> +use core::arch::asm;
> +
> +/// Represents a per-CPU variable that can be manipulated with machine-intrinsic numeric
> +/// operations.
> +pub struct PerCpuNumeric<'a, T> {
> + // INVARIANT: `ptr.0` is a valid offset into the per-CPU area and is initialized on all CPUs
> + // (since we don't have a CPU guard, we have to be pessimistic and assume we could be on any
> + // CPU).
> + ptr: &'a PerCpuPtr<T>,
> +}
> +
> +macro_rules! impl_ops {
> + ($ty:ty, $reg:tt) => {
> + impl DynamicPerCpu<$ty> {
> + /// Returns a [`PerCpuNumeric`] that can be used to manipulate the underlying per-CPU
> + /// variable.
> + #[inline]
> + pub fn num(&mut self) -> PerCpuNumeric<'_, $ty> {
> + // The invariant is satisfied because `DynamicPerCpu`'s invariant guarantees that
> + // this pointer is valid and initialized on all CPUs.
> + PerCpuNumeric { ptr: &self.alloc().0 }
> + }
> + }
> + impl StaticPerCpu<$ty> {
> + /// Returns a [`PerCpuNumeric`] that can be used to manipulate the underlying per-CPU
> + /// variable.
> + #[inline]
> + pub fn num(&mut self) -> PerCpuNumeric<'_, $ty> {
> + // The invariant is satisfied because `StaticPerCpu`'s invariant guarantees that
> + // this pointer is valid and initialized on all CPUs.
> + PerCpuNumeric { ptr: &self.0 }
> + }
> + }
> +
> + impl PerCpuNumeric<'_, $ty> {
> + /// Adds `rhs` to the per-CPU variable.
> + #[inline]
> + pub fn add(&mut self, rhs: $ty) {
> + // SAFETY: `self.ptr.0` is a valid offset into the per-CPU area (i.e., valid as a
> + // pointer relative to the `gs` segment register) by the invariants of this type.
> + unsafe {
> + asm!(
> + concat!("add gs:[{off}], {val:", $reg, "}"),
> + off = in(reg) self.ptr.0.cast::<$ty>(),
> + val = in(reg) rhs,

So, every user of .add() now will be only compilable against x86_64?
I don't think it's right. Can you make it in a more convenient way:
implement a generic version, and then an x86_64-optimized.

How bad the generic x86_64 version looks comparing to the optimized
one?

Thanks,
Yury

> + );
> + }
> + }
> + }
> + impl PerCpuNumeric<'_, $ty> {
> + /// Subtracts `rhs` from the per-CPU variable.
> + #[inline]
> + pub fn sub(&mut self, rhs: $ty) {
> + // SAFETY: `self.ptr.0` is a valid offset into the per-CPU area (i.e., valid as a
> + // pointer relative to the `gs` segment register) by the invariants of this type.
> + unsafe {
> + asm!(
> + concat!("sub gs:[{off}], {val:", $reg, "}"),
> + off = in(reg) self.ptr.0.cast::<$ty>(),
> + val = in(reg) rhs,
> + );
> + }
> + }
> + }
> + };
> +}
> +
> +macro_rules! impl_ops_byte {
> + ($ty:ty) => {
> + impl DynamicPerCpu<$ty> {
> + /// Returns a [`PerCpuNumeric`] that can be used to manipulate the underlying per-CPU
> + /// variable.
> + #[inline]
> + pub fn num(&mut self) -> PerCpuNumeric<'_, $ty> {
> + // The invariant is satisfied because `DynamicPerCpu`'s invariant guarantees that
> + // this pointer is valid and initialized on all CPUs.
> + PerCpuNumeric { ptr: &self.alloc().0 }
> + }
> + }
> + impl StaticPerCpu<$ty> {
> + /// Returns a [`PerCpuNumeric`] that can be used to manipulate the underlying per-CPU
> + /// variable.
> + #[inline]
> + pub fn num(&mut self) -> PerCpuNumeric<'_, $ty> {
> + // The invariant is satisfied because `StaticPerCpu`'s invariant guarantees that
> + // this pointer is valid and initialized on all CPUs.
> + PerCpuNumeric { ptr: &self.0 }
> + }
> + }
> +
> + impl PerCpuNumeric<'_, $ty> {
> + /// Adds `rhs` to the per-CPU variable.
> + #[inline]
> + pub fn add(&mut self, rhs: $ty) {
> + // SAFETY: `self.ptr.0` is a valid offset into the per-CPU area (i.e., valid as a
> + // pointer relative to the `gs` segment register) by the invariants of this type.
> + unsafe {
> + asm!(
> + "add gs:[{off}], {val}",
> + off = in(reg) self.ptr.0.cast::<$ty>(),
> + val = in(reg_byte) rhs,
> + );
> + }
> + }
> + }
> + impl PerCpuNumeric<'_, $ty> {
> + /// Subtracts `rhs` from the per-CPU variable.
> + #[inline]
> + pub fn sub(&mut self, rhs: $ty) {
> + // SAFETY: `self.ptr.0` is a valid offset into the per-CPU area (i.e., valid as a
> + // pointer relative to the `gs` segment register) by the invariants of this type.
> + unsafe {
> + asm!(
> + "sub gs:[{off}], {val}",
> + off = in(reg) self.ptr.0.cast::<$ty>(),
> + val = in(reg_byte) rhs,
> + );
> + }
> + }
> + }
> + };
> +}
> +
> +impl_ops_byte!(i8);
> +impl_ops!(i16, "x");
> +impl_ops!(i32, "e");
> +impl_ops!(i64, "r");
> +impl_ops!(isize, "r");
> +
> +impl_ops_byte!(u8);
> +impl_ops!(u16, "x");
> +impl_ops!(u32, "e");
> +impl_ops!(u64, "r");
> +impl_ops!(usize, "r");
> diff --git a/samples/rust/rust_percpu.rs b/samples/rust/rust_percpu.rs
> index 5adb30509bd4..90f5debd3c7a 100644
> --- a/samples/rust/rust_percpu.rs
> +++ b/samples/rust/rust_percpu.rs
> @@ -28,6 +28,26 @@
> define_per_cpu!(UPERCPU: u64 = 0);
> define_per_cpu!(CHECKED: RefCell<u64> = RefCell::new(0));
>
> +macro_rules! make_optimization_test {
> + ($ty:ty) => {
> + let mut test: DynamicPerCpu<$ty> = DynamicPerCpu::new_zero(GFP_KERNEL).unwrap();
> + {
> + let _guard = CpuGuard::new();
> + // SAFETY: No other usage of `test`
> + unsafe { test.get_mut(CpuGuard::new()) }.with(|val: &mut $ty| *val = 10);
> + test.num().add(1);
> + // SAFETY: No other usage of `test`
> + unsafe { test.get_mut(CpuGuard::new()) }.with(|val: &mut $ty| assert_eq!(*val, 11));
> + test.num().add(10);
> + // SAFETY: No other usage of `test`
> + unsafe { test.get_mut(CpuGuard::new()) }.with(|val: &mut $ty| assert_eq!(*val, 21));
> + test.num().sub(5);
> + // SAFETY: No other usage of `test`
> + unsafe { test.get_mut(CpuGuard::new()) }.with(|val: &mut $ty| assert_eq!(*val, 16));
> + }
> + };
> +}
> +
> impl kernel::Module for PerCpuMod {
> fn init(_module: &'static ThisModule) -> Result<Self, Error> {
> pr_info!("rust percpu test start\n");
> @@ -228,6 +248,22 @@ fn init(_module: &'static ThisModule) -> Result<Self, Error> {
>
> pr_info!("rust dynamic percpu test done\n");
>
> + pr_info!("rust numeric optimizations test start\n");
> +
> + make_optimization_test!(u8);
> + make_optimization_test!(u16);
> + make_optimization_test!(u32);
> + make_optimization_test!(u64);
> + make_optimization_test!(usize);
> +
> + make_optimization_test!(i8);
> + make_optimization_test!(i16);
> + make_optimization_test!(i32);
> + make_optimization_test!(i64);
> + make_optimization_test!(isize);
> +
> + pr_info!("rust numeric optimizations test done\n");
> +
> // Return Err to unload the module
> Result::Err(EINVAL)
> }
>
> --
> 2.34.1