[stdlib] Remove redundant next_power_of_two() from math.mojo and optimize existing function in bit.mojo

[martinvuyk]

Optimization of next_power_of_two

Benchmark results

CPU: i7 7th gen
Measured in miliseconds to do 100 million calculations

bench_next_power_of_two_int_v1,204.92954861666666
bench_next_power_of_two_int_v2,117.86505349666666
bench_next_power_of_two_int_v3,111.60907072111111
bench_next_power_of_two_int_v4,105.20949030055556
bench_next_power_of_two_uint_v1,89.40362252435895
bench_next_power_of_two_uint_v2,82.5889815992857
bench_next_power_of_two_uint_v3,85.60044429065934
bench_next_power_of_two_uint_v4,108.61543930999999

I decided to go for ..._int_v3 and ..._uint_v2 (which use SIMD[bool].select) since they are the most stable in performance between executions, even though the ..._int_v4 and ..._uint_v1 have sometimes given better results.

[JoeLoser]

I'm interested in perf difference (if any) between the two implementations. I think we should just remove it from math entirely rather than having the import around in mojo/stdlib/src/math/__init__.mojo.

[martinvuyk]

@JoeLoser moved the implementation, added a benchmark and changed the whole PR description.

[gryznar]

@JoeLoser moved the implementation, added a benchmark and changed the whole PR description.

Could you please post benchmark results (numbers) as comment / description?

[martinvuyk]

Read the PR description. 2x speedup. 199ms vs 100ms to execute the code, doesn't make much difference to post the exact results for a doubling in perf in a trivial function.

[soraros]

Shall we write it in the same branchless fashion like I did in #4253?

[martinvuyk]

@soraros that is very interesting. Now I'm wondering whether

val = Scalar[DType.index](int_value)
select(val <= 1, 1, 1 << (bitwidthof[Int]() - count_leading_zeros(val - 1)))

is faster or not than using bitwise ops i.e.

(1 << (bitwidthof[Int]() - count_leading_zeros(val - 1))) & -Int(val > 1) | Int(val <= 1)

I'm also wondering whether we should implement this function for UInt as well at some point

[soraros]

I think they ultimately generate the same asm, the select variant is just a bit more readable.

In fact, the select variant is equivalent to using if, though the generated LLVM IR is quite a bit simpler, which is beneficial, since it's an @always_inline function.

[martinvuyk]

Nothing better than over-engineering the most trivial of functions for 2 hours. Anyway, it was fun trying and learning some new bitwise tricks. @soraros it seems like the best (most stable) option is to use SIMD[bool].select, and when the compiler knows that it's a UInt being cast to a DType.index the performance improves noticeably (see the numbers in the PR description)

[martinvuyk]

I personally find g more readable. Also, there's a constructor from scalar to Int, so you could skip the index().

Totally agree, now it looks very pretty :]

[soraros]

I actually find eyeballing the generated IR and assembly really helpful when comparing different implementations. It can greatly enhance the experience of "over-engineering."

[abduld]

lgtm. Thanks a bunch for the contribution and sorry for not searching for the existing implementation. btw. If you want to dump the asm, then you can do

from sys import bitwidthof
from compile import compile_info
from utils._select import _select_register_value as select
from bit import count_leading_zeros


@always_inline
fn next_power_of_two(val: UInt) -> UInt:
    return select(
        val == 0,
        1,
        1 << (bitwidthof[__type_of(val)]() - count_leading_zeros(val - 1)),
    )


def main():
    print(compile_info[next_power_of_two, emission_kind="llvm"]())

    print(compile_info[next_power_of_two, emission_kind="asm"]())

[martinvuyk]

sorry for not searching for the existing implementation.

no worries I also reimplemented the logic and wanted to commit it to the stdlib when I realized bit_ceil was already implemented and in the bit module. I'm not a C++ programmer so I didn't know to look for that function name, that's why I proposed the function name change in #3271

If you want to dump the asm, then you can do

I think I'm going to start doing that more often and not eyeballing benchmarks for small functions so much 😅

[JoeLoser]

!sync