Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[stdlib] Optimize normalize_index for unsigned types #3957

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

yinonburgansky
Copy link
Contributor

@yinonburgansky yinonburgansky commented Jan 18, 2025

Use the Indexer trait in normalize_index to optimize for UInt, UInt8, UInt16, UInt32, and UInt64 types.

Optimize normalize_index for UInt Indexer type.

Signed-off-by: Yinon Burgansky <[email protected]>
@yinonburgansky
Copy link
Contributor Author

I didn't use it in inline_array since the current implementation doesn't perform bound checks at all and since normalize_index does bound checks with debug_assert[assert_mode="safe", cpu_only=True] it will affect performance.
If this is something we want, I can migrate to it in where needed in the stdlib in another PR.

@yinonburgansky
Copy link
Contributor Author

yinonburgansky commented Jan 18, 2025

Maybe to make it more generic, a change to the signature taking the length directly instead of the sized container, might be needed.

@yinonburgansky
Copy link
Contributor Author

With regards to UInt vs Int:

  • Currently there is an assumption that all container lengths are not greater than Int.MAX this allows ignoring overflow having more performant compiled code.
  • For this reason I think even If we use UInt for length, when comparing against negative integers, it is a bit more performant to assume that the length is still not grater than Int.MAX.
  • If the user code indexes just with UInt, and all the container methods are specialized for UInt, then indexes up to UInt.MAX can probably be used safely.
  • Generally using UInt allows the compiler to make more optimizations since it can assume the number is positive.
  • Using debug_assert[assert_mode="safe", cpu_only=True] probably means we want to optimize the debug_assert too, since indexing is a common operation.
  • Illustration of the compilation differences in C: https://godbolt.org/z/bM4Menc4h

@yinonburgansky
Copy link
Contributor Author

As a side note:
Currently there is an implicit casting between UInt and Int:

fn main():
    # binary operators use the method of left side type
    print(Int(-1) < UInt(1<<63)) # False 
    print(UInt(1<<63) > Int(-1)) # False
    print(UInt(1) > Int(-1)) # False
    print(Int(-1) < UInt(1)) # True

@yinonburgansky
Copy link
Contributor Author

The second commit addresses all of the issues above.

In the future, if needed this pattern can be extended to support SIMD sized containers without upcasting and/or normalizing multiple indexes at once using SIMD.

@yinonburgansky
Copy link
Contributor Author

yinonburgansky commented Jan 23, 2025

I've found out that by doing the bounds check after the index normalization we can avoid a comparison
and support negative indexing up to UInt.MAX in a performant way.
The compiler is smart enough to optimize the index normalization into branchless index += (index >> 63) & len and the bounds check for Int and UInt are the same.
so the only cost of using negative indexing now is just a few ALU operations.

This is also True for SIMD so if needed we can add support UInt8.MAX sized containers and use negative Int8 indexes without upcasting or doing more expensive bounds check, the only cost is the normalization which can be done branchless.

See https://godbolt.org/z/nMhz6dWhh

@yinonburgansky yinonburgansky changed the title [stdlib] Use Indexer in normalize_index [stdlib] Optimize normalize_index for unsigned types Jan 23, 2025
@skongum02 skongum02 deleted the branch modular:main January 29, 2025 18:59
@skongum02 skongum02 closed this Jan 29, 2025
@skongum02 skongum02 reopened this Jan 29, 2025
@skongum02 skongum02 changed the base branch from nightly to main January 29, 2025 20:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants