[PATCH 0/3] rust: add `num` module

[PATCH 0/3] rust: add `num` module

[Alexandre Courbot]

This patchset was part of the Nova FWSEC-FRTS submission [1], but has
been extracted as it was not very visible there, and seems to require
more discussion that the Nova code. I hope that presenting it separately
will help make it converge towards something that can be merged.

The motivation for this new module is to provide numerical functions not
supported (or not optimally supported for the kernel use-cases) by Rust
core. One example is how kernel code often aligns addresses or length to
some power-of-two. While Rust has `next_multiple_of` that provides the
same result, it also supports non-power-of-two values and is thus not as
efficient as the simple substract-and-mask operation that can be done
with a power of two.

Another missing operation is `last_set_bit` (`fls` in the C code), which
is again commonly used in the kernel but has no equivalent yet in Rust
core.

This patchset introduces these two features available:

- A `PowerOfTwo` newtype for unsigned integers with `align_down` and
  `align_up` operations optimized for the kernel. The newtype ensures
  that these operations cannot be used with a non-power-of-two, which
  would yield an incorrect result.
- A set of `last_set_bit` const functions for unsigned integers types.

The last patch makes use of these features in the Nova driver, as was
originally done in the patchset. It requires the v6 of the patch series
[2] to be applied.

[1] https://lore.kernel.org/rust-for-linux/DANP9ATT1T5W.1KP4992E26FTP@nvidia.com/
[2] https://lore.kernel.org/rust-for-linux/20250619-nova-frts-v6-0-ecf41ef99252@nvidia.com/

Changes since split from the nova-core series:

- Rename `fls` to `last_set_bit`,
- Generate per-type doctests,
- Add invariants section to `PowerOfTwo`.
- Do not use reference to `self` in `PowerOfTwo` methods since it
  implements `Copy`,
- Use #[derive] where possible instead of implementing traits manually,
- Remove `Deref` and `Borrow` implementations.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
Alexandre Courbot (3):
      rust: add `num` module with `PowerOfTwo` type
      rust: num: add the `last_set_bit` operation
      nova-core: use `num` module

 Documentation/gpu/nova/core/todo.rst      |  15 ---
 drivers/gpu/nova-core/falcon/hal/ga102.rs |   4 +-
 drivers/gpu/nova-core/fb.rs               |   6 +-
 drivers/gpu/nova-core/firmware/fwsec.rs   |   7 +-
 drivers/gpu/nova-core/vbios.rs            |   4 +-
 rust/kernel/lib.rs                        |   1 +
 rust/kernel/num.rs                        | 201 ++++++++++++++++++++++++++++++
 7 files changed, 211 insertions(+), 27 deletions(-)
---
base-commit: c7864f7ab73417e352d8d00e46d3e9e6a228c5ab
change-id: 20250620-num-9420281c02c7

Best regards,
-- 
Alexandre Courbot <acourbot@nvidia.com>

[PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

Introduce the `num` module, featuring the `PowerOfTwo` unsigned wrapper
that guarantees (at build-time or runtime) that a value is a power of
two.

Such a property is often useful to maintain. In the context of the
kernel, powers of two are often used to align addresses or sizes up and
down, or to create masks. These operations are provided by this type.

It is introduced to be first used by the nova-core driver.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
 rust/kernel/lib.rs |   1 +
 rust/kernel/num.rs | 163 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 164 insertions(+)

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 6b4774b2b1c37f4da1866e993be6230bc6715841..2955f65da1278dd4cba1e4272ff178b8211a892c 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -89,6 +89,7 @@
 pub mod mm;
 #[cfg(CONFIG_NET)]
 pub mod net;
+pub mod num;
 pub mod of;
 #[cfg(CONFIG_PM_OPP)]
 pub mod opp;
diff --git a/rust/kernel/num.rs b/rust/kernel/num.rs
new file mode 100644
index 0000000000000000000000000000000000000000..6ecff037893dd25420a6369ea0cd6adbe3460b29
--- /dev/null
+++ b/rust/kernel/num.rs
@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Numerical and binary utilities for primitive types.
+
+use crate::build_assert;
+use core::fmt::Debug;
+use core::hash::Hash;
+
+/// An unsigned integer which is guaranteed to be a power of 2.
+///
+/// # Invariants
+///
+/// The stored value is guaranteed to be a power of two.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[repr(transparent)]
+pub struct PowerOfTwo<T>(T);
+
+macro_rules! power_of_two_impl {
+    ($($t:ty),+) => {
+        $(
+            impl PowerOfTwo<$t> {
+                /// Validates that `v` is a power of two at build-time, and returns it wrapped into
+                /// [`PowerOfTwo`].
+                ///
+                /// A build error is triggered if `v` cannot be asserted to be a power of two.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
+                /// assert_eq!(v.value(), 16);
+                /// ```
+                #[inline(always)]
+                pub const fn new(v: $t) -> Self {
+                    build_assert!(v.count_ones() == 1);
+                    Self(v)
+                }
+
+                /// Validates that `v` is a power of two at runtime, and returns it wrapped into
+                /// [`PowerOfTwo`].
+                ///
+                /// [`None`] is returned if `v` was not a power of two.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::try_new(16), Some(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(16)));"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::try_new(15), None);"
+                )]
+                /// ```
+                #[inline(always)]
+                pub const fn try_new(v: $t) -> Option<Self> {
+                    match v.count_ones() {
+                        1 => Some(Self(v)),
+                        _ => None,
+                    }
+                }
+
+                /// Returns the value of this instance.
+                ///
+                /// It is guaranteed to be a power of two.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
+                /// assert_eq!(v.value(), 16);
+                /// ```
+                #[inline(always)]
+                pub const fn value(self) -> $t {
+                    self.0
+                }
+
+                /// Returns the mask corresponding to `self.value() - 1`.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(0x10);")]
+                /// assert_eq!(v.mask(), 0xf);
+                /// ```
+                #[inline(always)]
+                pub const fn mask(self) -> $t {
+                    self.0.wrapping_sub(1)
+                }
+
+                /// Aligns `self` down to `alignment`.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_down(0x2f), 0x20);"
+                )]
+                /// ```
+                #[inline(always)]
+                pub const fn align_down(self, value: $t) -> $t {
+                    value & !self.mask()
+                }
+
+                /// Aligns `value` up to `self`.
+                ///
+                /// Wraps around to `0` if the requested alignment pushes the result above the
+                /// type's limits.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(0x4f), 0x50);"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(0x40), 0x40);"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(0x0), 0x0);"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(",
+                    stringify!($t), "::MAX), 0x0);"
+                )]
+                /// ```
+                #[inline(always)]
+                pub const fn align_up(self, value: $t) -> $t {
+                    self.align_down(value.wrapping_add(self.mask()))
+                }
+            }
+        )+
+    };
+}
+
+power_of_two_impl!(usize, u8, u16, u32, u64, u128);

-- 
2.49.0

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Miguel Ojeda]

On Fri, Jun 20, 2025 at 3:15 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>
> Introduce the `num` module, featuring the `PowerOfTwo` unsigned wrapper
> that guarantees (at build-time or runtime) that a value is a power of
> two.
>
> Such a property is often useful to maintain. In the context of the
> kernel, powers of two are often used to align addresses or sizes up and
> down, or to create masks. These operations are provided by this type.

Before I forget: the other day in a call we discussed powers of two
and I mentioned that there is `Alignment` in the standard library:

    https://doc.rust-lang.org/std/ptr/struct.Alignment.html

    "A type storing a `usize` which is a power of two"

So it would be nice to ask upstream the following if they have plans
to stabilize it, and whether they have considered a generic
`PowerOfTwo<T>` type like this one, rather than one just for alignment
purposes (possibly with an alias or newtype for `Alignment` if
needed).

Similarly, if they stabilize the `Alignment` one (only) and we end up
only using our `PowerOfTwo<T>` for `usize` and those use cases, then
we should consider using the upstream one (and adding any/all methods
that we need).

So I will ask them the next time we meet. I have added
`ptr_alignment_type` to our list (in the "nice to have" section).

(Apologies if this was already discussed!)

Cheers,
Miguel

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

On Fri Jun 20, 2025 at 10:35 PM JST, Miguel Ojeda wrote:
> On Fri, Jun 20, 2025 at 3:15 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>>
>> Introduce the `num` module, featuring the `PowerOfTwo` unsigned wrapper
>> that guarantees (at build-time or runtime) that a value is a power of
>> two.
>>
>> Such a property is often useful to maintain. In the context of the
>> kernel, powers of two are often used to align addresses or sizes up and
>> down, or to create masks. These operations are provided by this type.
>
> Before I forget: the other day in a call we discussed powers of two
> and I mentioned that there is `Alignment` in the standard library:
>
>     https://doc.rust-lang.org/std/ptr/struct.Alignment.html
>
>     "A type storing a `usize` which is a power of two"

Haha, I wasn't aware of this effort, and am quite amazed by how close
its API is to my own design. This is reassuring; maybe I am finally
starting to grok Rust after all. ;)

>
> So it would be nice to ask upstream the following if they have plans
> to stabilize it, and whether they have considered a generic
> `PowerOfTwo<T>` type like this one, rather than one just for alignment
> purposes (possibly with an alias or newtype for `Alignment` if
> needed).

Mmm indeed I don't quite see the fundamental difference between
`Alignment` and `PowerOfTwo`, although `Alignment` might better capture
what we are doing with our type anyway.

>
> Similarly, if they stabilize the `Alignment` one (only) and we end up
> only using our `PowerOfTwo<T>` for `usize` and those use cases, then
> we should consider using the upstream one (and adding any/all methods
> that we need).

`Alignment` is very close to what we need, so I don't see a reason to
not adopt the same name at the very least.

This reminds me that I should also check whether upstream Rust would be
interested in `prev_multiple_of` and `last_set_bit`. The docs I've read
for contributing looked a bit intimidating, with RFCs to write and all.
Would you have a pointer for where I should start? Maybe a Zulip thread?

>
> So I will ask them the next time we meet. I have added
> `ptr_alignment_type` to our list (in the "nice to have" section).
>
> (Apologies if this was already discussed!)

I wasn't aware of this, so thanks for bringing it up!

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alice Ryhl]

On Fri, Jun 20, 2025 at 3:59 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
> > Similarly, if they stabilize the `Alignment` one (only) and we end up
> > only using our `PowerOfTwo<T>` for `usize` and those use cases, then
> > we should consider using the upstream one (and adding any/all methods
> > that we need).
>
> `Alignment` is very close to what we need, so I don't see a reason to
> not adopt the same name at the very least.
>
> This reminds me that I should also check whether upstream Rust would be
> interested in `prev_multiple_of` and `last_set_bit`. The docs I've read
> for contributing looked a bit intimidating, with RFCs to write and all.
> Would you have a pointer for where I should start? Maybe a Zulip thread?

If you want to add a new library function, the correct procedure would
be opening an ACP, which is more light-weight than the RFC process:
https://std-dev-guide.rust-lang.org/development/feature-lifecycle.html

RFCs are mainly for much bigger changes.

Alice

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

On Fri Jun 20, 2025 at 11:02 PM JST, Alice Ryhl wrote:
> On Fri, Jun 20, 2025 at 3:59 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>> > Similarly, if they stabilize the `Alignment` one (only) and we end up
>> > only using our `PowerOfTwo<T>` for `usize` and those use cases, then
>> > we should consider using the upstream one (and adding any/all methods
>> > that we need).
>>
>> `Alignment` is very close to what we need, so I don't see a reason to
>> not adopt the same name at the very least.
>>
>> This reminds me that I should also check whether upstream Rust would be
>> interested in `prev_multiple_of` and `last_set_bit`. The docs I've read
>> for contributing looked a bit intimidating, with RFCs to write and all.
>> Would you have a pointer for where I should start? Maybe a Zulip thread?
>
> If you want to add a new library function, the correct procedure would
> be opening an ACP, which is more light-weight than the RFC process:
> https://std-dev-guide.rust-lang.org/development/feature-lifecycle.html
>
> RFCs are mainly for much bigger changes.

Belated thanks for the suggestion; I have finally opened an ACP for
`last_set_bit` (and `first_set_bit` while we are at it):
https://github.com/rust-lang/libs-team/issues/631

I am still entangled with how to best leverage `Alignment` for our
purposes, but think I am getting close to a v2 of this patchset. 

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Miguel Ojeda]

On Sat, Aug 2, 2025 at 4:02 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>
> Belated thanks for the suggestion; I have finally opened an ACP for
> `last_set_bit` (and `first_set_bit` while we are at it):
> https://github.com/rust-lang/libs-team/issues/631
>
> I am still entangled with how to best leverage `Alignment` for our
> purposes, but think I am getting close to a v2 of this patchset.

Thanks for filling that one -- linked now from our usual lists :)

    https://github.com/Rust-for-Linux/linux/issues/514

Cheers,
Miguel

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

On Sat Aug 2, 2025 at 11:18 PM JST, Miguel Ojeda wrote:
> On Sat, Aug 2, 2025 at 4:02 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>>
>> Belated thanks for the suggestion; I have finally opened an ACP for
>> `last_set_bit` (and `first_set_bit` while we are at it):
>> https://github.com/rust-lang/libs-team/issues/631
>>
>> I am still entangled with how to best leverage `Alignment` for our
>> purposes, but think I am getting close to a v2 of this patchset.
>
> Thanks for filling that one -- linked now from our usual lists :)
>
>     https://github.com/Rust-for-Linux/linux/issues/514

We got some interesting feedback on the ACP already. I have been pointed
to `checked_ilog2` as an equivalent of `last_set_bit`, and it *does*
indeed work well as a replacement - with the caveat that the name is
not very natural to me (or anyone familiar with the C interface). Is
this something we can live with? If we decide to go with the existing
standard library method, how can we make sure that folks looking for an
equivalent of `fls` find `checked_ilog2`?

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Miguel Ojeda]

On Sun, Aug 3, 2025 at 3:13 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>
> We got some interesting feedback on the ACP already. I have been pointed
> to `checked_ilog2` as an equivalent of `last_set_bit`, and it *does*
> indeed work well as a replacement - with the caveat that the name is
> not very natural to me (or anyone familiar with the C interface). Is
> this something we can live with? If we decide to go with the existing
> standard library method, how can we make sure that folks looking for an
> equivalent of `fls` find `checked_ilog2`?

One option is using the `doc(alias = ...)` attribute, which makes it
appear in the search in the rendered docs, and would show easily in
greps too.

Another option is simply wrapping it in an `inline(always)`, I guess,
but I think we can just use the upstream ones, unless we want slightly
different semantics.

Cheers,
Miguel

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

On Mon Aug 4, 2025 at 12:15 AM JST, Miguel Ojeda wrote:
> On Sun, Aug 3, 2025 at 3:13 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>>
>> We got some interesting feedback on the ACP already. I have been pointed
>> to `checked_ilog2` as an equivalent of `last_set_bit`, and it *does*
>> indeed work well as a replacement - with the caveat that the name is
>> not very natural to me (or anyone familiar with the C interface). Is
>> this something we can live with? If we decide to go with the existing
>> standard library method, how can we make sure that folks looking for an
>> equivalent of `fls` find `checked_ilog2`?
>
> One option is using the `doc(alias = ...)` attribute, which makes it
> appear in the search in the rendered docs, and would show easily in
> greps too.
>
> Another option is simply wrapping it in an `inline(always)`, I guess,
> but I think we can just use the upstream ones, unless we want slightly
> different semantics.

That would be useful - let's see what the Rust lib folks say, as you
brought up that question on the ACP as well.

In any case, since we have reasonable alternatives for both `fls`
(`checked_ilog2`) and `ffs` (`NonZero::trailing_zeros`), I guess this
means we want to use these directly in the kernel and can drop patch
2 of this series?

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

On Mon Aug 4, 2025 at 4:32 PM JST, Alexandre Courbot wrote:
> On Mon Aug 4, 2025 at 12:15 AM JST, Miguel Ojeda wrote:
>> On Sun, Aug 3, 2025 at 3:13 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>>>
>>> We got some interesting feedback on the ACP already. I have been pointed
>>> to `checked_ilog2` as an equivalent of `last_set_bit`, and it *does*
>>> indeed work well as a replacement - with the caveat that the name is
>>> not very natural to me (or anyone familiar with the C interface). Is
>>> this something we can live with? If we decide to go with the existing
>>> standard library method, how can we make sure that folks looking for an
>>> equivalent of `fls` find `checked_ilog2`?
>>
>> One option is using the `doc(alias = ...)` attribute, which makes it
>> appear in the search in the rendered docs, and would show easily in
>> greps too.
>>
>> Another option is simply wrapping it in an `inline(always)`, I guess,
>> but I think we can just use the upstream ones, unless we want slightly
>> different semantics.
>
> That would be useful - let's see what the Rust lib folks say, as you
> brought up that question on the ACP as well.
>
> In any case, since we have reasonable alternatives for both `fls`
> (`checked_ilog2`) and `ffs` (`NonZero::trailing_zeros`), I guess this
> means we want to use these directly in the kernel and can drop patch
> 2 of this series?

I didn't expect that, but it looks like the Rust folks want these
methods after all:

https://github.com/rust-lang/libs-team/issues/631#issuecomment-3156000663

I'll proceed with sending a PR, and I guess we can have temporary
implementations in the kernel as well?

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Miguel Ojeda]

On Fri, Jun 20, 2025 at 3:59 PM Alexandre Courbot <acourbot@nvidia.com> wrote:
>
> This reminds me that I should also check whether upstream Rust would be
> interested in `prev_multiple_of` and `last_set_bit`. The docs I've read
> for contributing looked a bit intimidating, with RFCs to write and all.
> Would you have a pointer for where I should start? Maybe a Zulip thread?

I would say it is actually quite straightforward compared to, say, WG14/21.

The first thing I added was a method too, to see how things worked,
and I started just by sending the PR.

So I would say, go for it! :)

Cheers,
Miguel

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Benno Lossin]

On Fri Jun 20, 2025 at 3:14 PM CEST, Alexandre Courbot wrote:
> +/// An unsigned integer which is guaranteed to be a power of 2.
> +///
> +/// # Invariants
> +///
> +/// The stored value is guaranteed to be a power of two.
> +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
> +#[repr(transparent)]
> +pub struct PowerOfTwo<T>(T);
> +
> +macro_rules! power_of_two_impl {
> +    ($($t:ty),+) => {
> +        $(
> +            impl PowerOfTwo<$t> {

I tried to use this type in a doctest like this:

    use kernel::num::PowerOfTwo;
   
    fn new(x: usize) -> PowerOfTwo<usize> {
        PowerOfTwo::new(1 << x)
    }

And it doesn't compile :(

    error[E0034]: multiple applicable items in scope
        --> rust/doctests_kernel_generated.rs:4930:17
         |
    4930 |     PowerOfTwo::new(1 << x)
         |                 ^^^ multiple `new` found
         |
         = note: candidate #1 is defined in an impl for the type `PowerOfTwo<u128>`
         = note: candidate #2 is defined in an impl for the type `PowerOfTwo<u16>`
         = note: candidate #3 is defined in an impl for the type `PowerOfTwo<u32>`
         = note: candidate #4 is defined in an impl for the type `PowerOfTwo<u64>`
         = note: and 2 others
    
    error: aborting due to 1 previous error

The problem is that the function `new` exists 6 times for each of the
integer types. You can write `PowerOfTwo::<usize>::new()` instead, but
that's annoying...

We probably need an `Integer` trait and then do

    impl<I: Integer> PowerOfTwo<I> {
        pub const fn new(value: I) -> Self;
    }

> +                /// Validates that `v` is a power of two at build-time, and returns it wrapped into
> +                /// [`PowerOfTwo`].
> +                ///
> +                /// A build error is triggered if `v` cannot be asserted to be a power of two.
> +                ///
> +                /// # Examples
> +                ///
> +                /// ```
> +                /// use kernel::num::PowerOfTwo;
> +                ///
> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
> +                /// assert_eq!(v.value(), 16);
> +                /// ```
> +                #[inline(always)]
> +                pub const fn new(v: $t) -> Self {
> +                    build_assert!(v.count_ones() == 1);

Why not `v.is_power_of_two()`?

> +                    Self(v)

Missing `// INVARIANT` comment.

> +                }
> +
> +                /// Validates that `v` is a power of two at runtime, and returns it wrapped into
> +                /// [`PowerOfTwo`].
> +                ///
> +                /// [`None`] is returned if `v` was not a power of two.
> +                ///
> +                /// # Examples
> +                ///
> +                /// ```
> +                /// use kernel::num::PowerOfTwo;
> +                ///
> +                #[doc = concat!(
> +                    "assert_eq!(PowerOfTwo::<",
> +                    stringify!($t),
> +                    ">::try_new(16), Some(PowerOfTwo::<",
> +                    stringify!($t),
> +                    ">::new(16)));"
> +                )]
> +                #[doc = concat!(
> +                    "assert_eq!(PowerOfTwo::<",
> +                    stringify!($t),
> +                    ">::try_new(15), None);"
> +                )]
> +                /// ```
> +                #[inline(always)]
> +                pub const fn try_new(v: $t) -> Option<Self> {

Maybe `new_checked` is a better name, since it doesn't return a result?

> +                    match v.count_ones() {

Why not `is_power_of_two()`?

> +                        1 => Some(Self(v)),

Missing `// INVARIANT` comment.

> +                        _ => None,
> +                    }
> +                }
> +
> +                /// Returns the value of this instance.
> +                ///
> +                /// It is guaranteed to be a power of two.
> +                ///
> +                /// # Examples
> +                ///
> +                /// ```
> +                /// use kernel::num::PowerOfTwo;
> +                ///
> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
> +                /// assert_eq!(v.value(), 16);
> +                /// ```
> +                #[inline(always)]
> +                pub const fn value(self) -> $t {
> +                    self.0

Let's add:

    if !self.0.is_power_of_two() {
        core::hint::unreachable_unchecked()
    }
    self.0

> +                }
> +
> +                /// Returns the mask corresponding to `self.value() - 1`.
> +                ///
> +                /// # Examples
> +                ///
> +                /// ```
> +                /// use kernel::num::PowerOfTwo;
> +                ///
> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(0x10);")]
> +                /// assert_eq!(v.mask(), 0xf);
> +                /// ```
> +                #[inline(always)]
> +                pub const fn mask(self) -> $t {
> +                    self.0.wrapping_sub(1)

Then use `self.value().wrapping_sub(1)` here instead to also propagate
the information.

---
Cheers,
Benno

> +                }

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot]

Hi Benno,

Sorry, took some time to come back to this!

On Sun Jun 22, 2025 at 5:11 PM JST, Benno Lossin wrote:
> On Fri Jun 20, 2025 at 3:14 PM CEST, Alexandre Courbot wrote:
>> +/// An unsigned integer which is guaranteed to be a power of 2.
>> +///
>> +/// # Invariants
>> +///
>> +/// The stored value is guaranteed to be a power of two.
>> +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
>> +#[repr(transparent)]
>> +pub struct PowerOfTwo<T>(T);
>> +
>> +macro_rules! power_of_two_impl {
>> +    ($($t:ty),+) => {
>> +        $(
>> +            impl PowerOfTwo<$t> {
>
> I tried to use this type in a doctest like this:
>
>     use kernel::num::PowerOfTwo;
>    
>     fn new(x: usize) -> PowerOfTwo<usize> {
>         PowerOfTwo::new(1 << x)
>     }
>
> And it doesn't compile :(
>
>     error[E0034]: multiple applicable items in scope
>         --> rust/doctests_kernel_generated.rs:4930:17
>          |
>     4930 |     PowerOfTwo::new(1 << x)
>          |                 ^^^ multiple `new` found
>          |
>          = note: candidate #1 is defined in an impl for the type `PowerOfTwo<u128>`
>          = note: candidate #2 is defined in an impl for the type `PowerOfTwo<u16>`
>          = note: candidate #3 is defined in an impl for the type `PowerOfTwo<u32>`
>          = note: candidate #4 is defined in an impl for the type `PowerOfTwo<u64>`
>          = note: and 2 others
>     
>     error: aborting due to 1 previous error
>
> The problem is that the function `new` exists 6 times for each of the
> integer types. You can write `PowerOfTwo::<usize>::new()` instead, but
> that's annoying...

This should go away as we switch to the non-generic `Alignment` type
thankfully.

>
> We probably need an `Integer` trait and then do
>
>     impl<I: Integer> PowerOfTwo<I> {
>         pub const fn new(value: I) -> Self;
>     }
>
>> +                /// Validates that `v` is a power of two at build-time, and returns it wrapped into
>> +                /// [`PowerOfTwo`].
>> +                ///
>> +                /// A build error is triggered if `v` cannot be asserted to be a power of two.
>> +                ///
>> +                /// # Examples
>> +                ///
>> +                /// ```
>> +                /// use kernel::num::PowerOfTwo;
>> +                ///
>> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
>> +                /// assert_eq!(v.value(), 16);
>> +                /// ```
>> +                #[inline(always)]
>> +                pub const fn new(v: $t) -> Self {
>> +                    build_assert!(v.count_ones() == 1);
>
> Why not `v.is_power_of_two()`?

Why not indeed. :) Fixed.

>
>> +                    Self(v)
>
> Missing `// INVARIANT` comment.

Added (and in other places as well).

>
>> +                }
>> +
>> +                /// Validates that `v` is a power of two at runtime, and returns it wrapped into
>> +                /// [`PowerOfTwo`].
>> +                ///
>> +                /// [`None`] is returned if `v` was not a power of two.
>> +                ///
>> +                /// # Examples
>> +                ///
>> +                /// ```
>> +                /// use kernel::num::PowerOfTwo;
>> +                ///
>> +                #[doc = concat!(
>> +                    "assert_eq!(PowerOfTwo::<",
>> +                    stringify!($t),
>> +                    ">::try_new(16), Some(PowerOfTwo::<",
>> +                    stringify!($t),
>> +                    ">::new(16)));"
>> +                )]
>> +                #[doc = concat!(
>> +                    "assert_eq!(PowerOfTwo::<",
>> +                    stringify!($t),
>> +                    ">::try_new(15), None);"
>> +                )]
>> +                /// ```
>> +                #[inline(always)]
>> +                pub const fn try_new(v: $t) -> Option<Self> {
>
> Maybe `new_checked` is a better name, since it doesn't return a result?

Definitely.

>
>> +                    match v.count_ones() {
>
> Why not `is_power_of_two()`?

Fixed, thanks.

>
>> +                        1 => Some(Self(v)),
>
> Missing `// INVARIANT` comment.
>
>> +                        _ => None,
>> +                    }
>> +                }
>> +
>> +                /// Returns the value of this instance.
>> +                ///
>> +                /// It is guaranteed to be a power of two.
>> +                ///
>> +                /// # Examples
>> +                ///
>> +                /// ```
>> +                /// use kernel::num::PowerOfTwo;
>> +                ///
>> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
>> +                /// assert_eq!(v.value(), 16);
>> +                /// ```
>> +                #[inline(always)]
>> +                pub const fn value(self) -> $t {
>> +                    self.0
>
> Let's add:
>
>     if !self.0.is_power_of_two() {
>         core::hint::unreachable_unchecked()
>     }
>     self.0

Sure. Is it to enable compiler optimizations by making assumptions about
the returned value?

>
>> +                }
>> +
>> +                /// Returns the mask corresponding to `self.value() - 1`.
>> +                ///
>> +                /// # Examples
>> +                ///
>> +                /// ```
>> +                /// use kernel::num::PowerOfTwo;
>> +                ///
>> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(0x10);")]
>> +                /// assert_eq!(v.mask(), 0xf);
>> +                /// ```
>> +                #[inline(always)]
>> +                pub const fn mask(self) -> $t {
>> +                    self.0.wrapping_sub(1)
>
> Then use `self.value().wrapping_sub(1)` here instead to also propagate
> the information.

Ack.

Re: [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Benno Lossin]

On Fri Jul 25, 2025 at 5:38 AM CEST, Alexandre Courbot wrote:
> Hi Benno,
>
> Sorry, took some time to come back to this!

No worries!

> On Sun Jun 22, 2025 at 5:11 PM JST, Benno Lossin wrote:
>> On Fri Jun 20, 2025 at 3:14 PM CEST, Alexandre Courbot wrote:
>>> +                /// Returns the value of this instance.
>>> +                ///
>>> +                /// It is guaranteed to be a power of two.
>>> +                ///
>>> +                /// # Examples
>>> +                ///
>>> +                /// ```
>>> +                /// use kernel::num::PowerOfTwo;
>>> +                ///
>>> +                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
>>> +                /// assert_eq!(v.value(), 16);
>>> +                /// ```
>>> +                #[inline(always)]
>>> +                pub const fn value(self) -> $t {
>>> +                    self.0
>>
>> Let's add:
>>
>>     if !self.0.is_power_of_two() {
>>         core::hint::unreachable_unchecked()
>>     }
>>     self.0
>
> Sure. Is it to enable compiler optimizations by making assumptions about
> the returned value?

Exactly, it will for example turn dividing by it into a right shift and
prevent a branch for the zero check.

---
Cheers,
Benno

[PATCH 2/3] rust: num: add the `last_set_bit` operation

[Alexandre Courbot]

Add an equivalent to the `fls` (Find Last Set bit) C function to Rust
unsigned types.

It is to be first used by the nova-core driver.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
 rust/kernel/num.rs | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/rust/kernel/num.rs b/rust/kernel/num.rs
index 6ecff037893dd25420a6369ea0cd6adbe3460b29..95766201a7cf208989f37ecbc6d54d264385a754 100644
--- a/rust/kernel/num.rs
+++ b/rust/kernel/num.rs
@@ -161,3 +161,41 @@ pub const fn align_up(self, value: $t) -> $t {
 }
 
 power_of_two_impl!(usize, u8, u16, u32, u64, u128);
+
+macro_rules! impl_last_set_bit {
+    ($($t:ty),+) => {
+        $(
+            ::kernel::macros::paste! {
+            /// Last Set Bit: return the 1-based index of the last (i.e. most significant) set bit
+            /// in `v`.
+            ///
+            /// Equivalent to the C `fls` function.
+            ///
+            /// # Examples
+            ///
+            /// ```
+            #[doc = concat!("use kernel::num::last_set_bit_", stringify!($t), ";")]
+            ///
+            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x0), 0);")]
+            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x1), 1);")]
+            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x10), 5);")]
+            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x1f), 5);")]
+            #[doc = concat!(
+                "assert_eq!(last_set_bit_",
+                stringify!($t),
+                "(",
+                stringify!($t),
+                "::MAX), ",
+                stringify!($t), "::BITS);"
+            )]
+            /// ```
+            #[inline(always)]
+            pub const fn [<last_set_bit_ $t>](v: $t) -> u32 {
+                $t::BITS - v.leading_zeros()
+            }
+            }
+        )+
+    };
+}
+
+impl_last_set_bit!(usize, u8, u16, u32, u64, u128);

-- 
2.49.0

Re: [PATCH 2/3] rust: num: add the `last_set_bit` operation

[Benno Lossin]

On Fri Jun 20, 2025 at 3:14 PM CEST, Alexandre Courbot wrote:
> Add an equivalent to the `fls` (Find Last Set bit) C function to Rust
> unsigned types.
>
> It is to be first used by the nova-core driver.
>
> Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>

Reviewed-by: Benno Lossin <lossin@kernel.org>

---
Cheers,
Benno

> ---
>  rust/kernel/num.rs | 38 ++++++++++++++++++++++++++++++++++++++
>  1 file changed, 38 insertions(+)

Re: [PATCH 2/3] rust: num: add the `last_set_bit` operation

[Alice Ryhl]

On Fri, Jun 20, 2025 at 10:14:52PM +0900, Alexandre Courbot wrote:
> Add an equivalent to the `fls` (Find Last Set bit) C function to Rust
> unsigned types.
> 
> It is to be first used by the nova-core driver.
> 
> Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>

Reviewed-by: Alice Ryhl <aliceryhl@google.com>

>  rust/kernel/num.rs | 38 ++++++++++++++++++++++++++++++++++++++
>  1 file changed, 38 insertions(+)
> 
> diff --git a/rust/kernel/num.rs b/rust/kernel/num.rs
> index 6ecff037893dd25420a6369ea0cd6adbe3460b29..95766201a7cf208989f37ecbc6d54d264385a754 100644
> --- a/rust/kernel/num.rs
> +++ b/rust/kernel/num.rs
> @@ -161,3 +161,41 @@ pub const fn align_up(self, value: $t) -> $t {
>  }
>  
>  power_of_two_impl!(usize, u8, u16, u32, u64, u128);
> +
> +macro_rules! impl_last_set_bit {
> +    ($($t:ty),+) => {
> +        $(
> +            ::kernel::macros::paste! {

I think this would read slightly nicer with the paste! invocation on the
outer scope.

$(::kernel::macros::paste! {
    ...
})+

> +            /// Last Set Bit: return the 1-based index of the last (i.e. most significant) set bit
> +            /// in `v`.
> +            ///
> +            /// Equivalent to the C `fls` function.
> +            ///
> +            /// # Examples
> +            ///
> +            /// ```
> +            #[doc = concat!("use kernel::num::last_set_bit_", stringify!($t), ";")]
> +            ///
> +            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x0), 0);")]
> +            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x1), 1);")]
> +            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x10), 5);")]
> +            #[doc = concat!("assert_eq!(last_set_bit_", stringify!($t), "(0x1f), 5);")]
> +            #[doc = concat!(
> +                "assert_eq!(last_set_bit_",
> +                stringify!($t),
> +                "(",
> +                stringify!($t),
> +                "::MAX), ",
> +                stringify!($t), "::BITS);"
> +            )]
> +            /// ```
> +            #[inline(always)]
> +            pub const fn [<last_set_bit_ $t>](v: $t) -> u32 {
> +                $t::BITS - v.leading_zeros()
> +            }
> +            }
> +        )+
> +    };
> +}
> +
> +impl_last_set_bit!(usize, u8, u16, u32, u64, u128);
> 
> -- 
> 2.49.0
> 

[PATCH 3/3] nova-core: use `num` module

[Alexandre Courbot]

Make use of the functions available in the `num` module and remove the
corresponding TODO items.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
 Documentation/gpu/nova/core/todo.rst      | 15 ---------------
 drivers/gpu/nova-core/falcon/hal/ga102.rs |  4 ++--
 drivers/gpu/nova-core/fb.rs               |  6 +++---
 drivers/gpu/nova-core/firmware/fwsec.rs   |  7 ++-----
 drivers/gpu/nova-core/vbios.rs            |  4 ++--
 5 files changed, 9 insertions(+), 27 deletions(-)

diff --git a/Documentation/gpu/nova/core/todo.rst b/Documentation/gpu/nova/core/todo.rst
index 894a1e9c3741a43ad4eb76d24a9486862999874e..01dfa858d11fe377c345b463742c13c37878e334 100644
--- a/Documentation/gpu/nova/core/todo.rst
+++ b/Documentation/gpu/nova/core/todo.rst
@@ -141,21 +141,6 @@ Features desired before this happens:
 | Complexity: Advanced
 | Contact: Alexandre Courbot
 
-Numerical operations [NUMM]
----------------------------
-
-Nova uses integer operations that are not part of the standard library (or not
-implemented in an optimized way for the kernel). These include:
-
-- Aligning up and down to a power of two,
-- The "Find Last Set Bit" (`fls` function of the C part of the kernel)
-  operation.
-
-A `num` core kernel module is being designed to provide these operations.
-
-| Complexity: Intermediate
-| Contact: Alexandre Courbot
-
 Delay / Sleep abstractions [DLAY]
 ---------------------------------
 
diff --git a/drivers/gpu/nova-core/falcon/hal/ga102.rs b/drivers/gpu/nova-core/falcon/hal/ga102.rs
index 664327f75cf4199cca37d22ca18b2b9abac781f8..9158991ec6e30f42fc0c7e49c87e2c04b426189f 100644
--- a/drivers/gpu/nova-core/falcon/hal/ga102.rs
+++ b/drivers/gpu/nova-core/falcon/hal/ga102.rs
@@ -4,6 +4,7 @@
 use core::time::Duration;
 
 use kernel::device;
+use kernel::num::last_set_bit_u32;
 use kernel::prelude::*;
 
 use crate::driver::Bar0;
@@ -69,8 +70,7 @@ fn signature_reg_fuse_version_ga102(
     let reg_fuse_version =
         bar.read32(reg_fuse_base + ((ucode_id - 1) as usize * core::mem::size_of::<u32>()));
 
-    // TODO[NUMM]: replace with `last_set_bit` once it lands.
-    Ok(u32::BITS - reg_fuse_version.leading_zeros())
+    Ok(last_set_bit_u32(reg_fuse_version))
 }
 
 fn program_brom_ga102<E: FalconEngine>(bar: &Bar0, params: &FalconBromParams) -> Result {
diff --git a/drivers/gpu/nova-core/fb.rs b/drivers/gpu/nova-core/fb.rs
index 48003527a2472a4a8b784af0d481a441c8d2426e..ca5e7ef997bc2b2855a1d60e81300fb99fe04cdb 100644
--- a/drivers/gpu/nova-core/fb.rs
+++ b/drivers/gpu/nova-core/fb.rs
@@ -2,6 +2,7 @@
 
 use core::ops::Range;
 
+use kernel::num::PowerOfTwo;
 use kernel::prelude::*;
 use kernel::sizes::*;
 use kernel::types::ARef;
@@ -119,10 +120,9 @@ pub(crate) fn new(chipset: Chipset, bar: &Bar0) -> Result<Self> {
         };
 
         let frts = {
-            const FRTS_DOWN_ALIGN: u64 = SZ_128K as u64;
+            const FRTS_DOWN_ALIGN: PowerOfTwo<u64> = PowerOfTwo::<u64>::new(SZ_128K as u64);
             const FRTS_SIZE: u64 = SZ_1M as u64;
-            // TODO[NUMM]: replace with `align_down` once it lands.
-            let frts_base = (vga_workspace.start & !(FRTS_DOWN_ALIGN - 1)) - FRTS_SIZE;
+            let frts_base = FRTS_DOWN_ALIGN.align_down(vga_workspace.start) - FRTS_SIZE;
 
             frts_base..frts_base + FRTS_SIZE
         };
diff --git a/drivers/gpu/nova-core/firmware/fwsec.rs b/drivers/gpu/nova-core/firmware/fwsec.rs
index 047aab76470ecb0a0486f6917f6fda69b5381391..0edcade5e8b303ee249397736af55c5a6f6fb97f 100644
--- a/drivers/gpu/nova-core/firmware/fwsec.rs
+++ b/drivers/gpu/nova-core/firmware/fwsec.rs
@@ -15,6 +15,7 @@
 use core::ops::Deref;
 
 use kernel::device::{self, Device};
+use kernel::num::PowerOfTwo;
 use kernel::prelude::*;
 use kernel::transmute::FromBytes;
 
@@ -218,11 +219,7 @@ fn dmem_load_params(&self) -> FalconLoadTarget {
         FalconLoadTarget {
             src_start: self.desc.imem_load_size,
             dst_start: self.desc.dmem_phys_base,
-            // TODO[NUMM]: replace with `align_up` once it lands.
-            len: self
-                .desc
-                .dmem_load_size
-                .next_multiple_of(DMEM_LOAD_SIZE_ALIGN),
+            len: PowerOfTwo::<u32>::new(DMEM_LOAD_SIZE_ALIGN).align_up(self.desc.dmem_load_size),
         }
     }
 
diff --git a/drivers/gpu/nova-core/vbios.rs b/drivers/gpu/nova-core/vbios.rs
index a5889eb149a16beabc0ddbdc87666520114c8aec..cac55d1534831775c14f3fed1e939ed89c7eba84 100644
--- a/drivers/gpu/nova-core/vbios.rs
+++ b/drivers/gpu/nova-core/vbios.rs
@@ -8,6 +8,7 @@
 use core::convert::TryFrom;
 use kernel::device;
 use kernel::error::Result;
+use kernel::num::PowerOfTwo;
 use kernel::pci;
 use kernel::prelude::*;
 
@@ -175,8 +176,7 @@ fn next(&mut self) -> Option<Self::Item> {
 
         // Advance to next image (aligned to 512 bytes)
         self.current_offset += image_size;
-        // TODO[NUMM]: replace with `align_up` once it lands.
-        self.current_offset = self.current_offset.next_multiple_of(512);
+        self.current_offset = PowerOfTwo::<usize>::new(512).align_up(self.current_offset);
 
         Some(Ok(full_image))
     }

-- 
2.49.0