[TIR] Further robustify floordiv/mod intrin lowering to prevent overflow

[tqchen]

This PR further robustifies floordiv/mod intrin lowering in cases where we can to shift negative values to positive range but also need to carefully prevent overflow in intermediate compiler checks.

[gemini-code-assist]

Summary of Changes

Hello @tqchen, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on improving the stability and correctness of the floordiv and floormod operations within the TIR compiler. It addresses potential overflow issues that can arise when dealing with negative operands, ensuring that the transformations applied during lowering maintain the original semantics of the code. The changes include refined checks and a new utility function to safely shift values into a positive range, along with new test cases to validate the implemented overflow protections.

Highlights

• Floordiv/Mod Intrinsic Lowering: The PR enhances the robustness of floordiv/mod intrinsic lowering by incorporating additional checks to prevent potential overflow issues during the conversion of negative values to a positive range.
• Overflow Prevention: Implements checks to ensure that intermediate calculations, specifically those involving shifts and bounds, do not exceed the maximum value of the data type, thus avoiding non-equivalent transformations due to overflow.
• Code Optimization: Replaces the previous logic with a new TryFindShiftCoefficientForPositiveRange function to determine a suitable shift coefficient, which ensures the numerator remains positive without causing overflow.
• Testing: Adds new test cases to specifically target and validate the overflow checks, ensuring the correctness and stability of the floordiv lowering process under various conditions.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the lowering logic for floordiv and floormod to be more robust against overflows when shifting negative numerators to a positive range. The complex logic is extracted into a new helper function TryFindShiftCoefficientForPositiveRange, which includes several careful checks to prevent intermediate overflows. This improves code clarity and maintainability. Additionally, the test suite is enhanced by generalizing a test helper function and adding new regression tests that specifically target the new overflow checks, ensuring the robustness of the changes.

I have one suggestion to make the optimization more broadly applicable by relaxing an overly conservative overflow check in the new helper function. Overall, this is a great improvement.