Skip to content

Improve detection of loop carries in triton frontend #7200

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

Merged
merged 9 commits into from
Jun 23, 2025

Conversation

Anstow
Copy link
Contributor

@Anstow Anstow commented Jun 17, 2025

Before this change, loop carries weren't correctly detected when @builtin or @core.extern function modified their arguments (which is a particular issue for member functions). This improves the detection of loop carries in for and while loops to handle these cases.

New contributor declaration

  • I am not making a trivial change, such as fixing a typo in a comment.

  • I have written a PR description following these
    rules.

  • I have run pre-commit run --from-ref origin/main --to-ref HEAD.

  • Select one of the following.

    • I have added tests.
      • /test for lit tests
      • /unittest for C++ tests
      • /python/test for end-to-end tests
    • This PR does not need a test because FILL THIS IN.
  • Select one of the following.

    • I have not added any lit tests.
    • The lit tests I have added follow these best practices,
      including the "tests should be minimal" section. (Usually running Python code
      and using the instructions it generates is not minimal.)

@Anstow Anstow requested a review from ptillet as a code owner June 17, 2025 09:09
@Anstow Anstow changed the title Davidw/for loop Improve detection of loop carries in triton frontend Jun 17, 2025
Copy link
Contributor

@peterbell10 peterbell10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! Thanks for fixing this.

Comment on lines +439 to +440
# We must extract the handles before the value is editted in the loop
livehandles = {name: flatten_values_to_ir([v]) for name, v in liveins.items() if _is_triton_value(v)}
Copy link
Contributor Author

@Anstow Anstow Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a little worried that I had to do this. I think it's because liveins is not deep-copied from self.local_defs in enter_sub_region. However I ran into other issues when I tried that. My preference is that figuring out that would be done in a different PR. Unless you have any ideas?

@Mogball Mogball requested review from Mogball June 17, 2025 23:26
@Anstow
Copy link
Contributor Author

Anstow commented Jun 18, 2025

I don't think the gb200 failure is to do with this PR. I hacked the frontend to produce the IR as if it were compiling for b200 and it is identical before and after my change - but I don't have access to a b200 machine so I can't confirm for certain. I'll update the base branch to see if that improves matters

@Mogball Mogball merged commit ba5ac26 into triton-lang:main Jun 23, 2025
9 checks passed
@Mogball
Copy link
Collaborator

Mogball commented Jun 23, 2025

Thanks for this fix

peterbell10 added a commit that referenced this pull request Jun 23, 2025
@Anstow Anstow deleted the davidw/for-loop branch June 23, 2025 16:23
peterbell10 added a commit that referenced this pull request Jun 24, 2025
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.

3 participants