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

Atomic Load and Store operations for Triton (tl.atomic_store/tl.atomic_load) #5187

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

plotfi
Copy link
Contributor

@plotfi plotfi commented Nov 19, 2024

The following PR adds tl.atomic_load and tl.atomic_store.

This allows for memory semantics that acquire and release in a total ordering for a given scope (cta/gpu/system).

TODO: Will add larger description once I take this PR out of draft mode.

  • 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
    • 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.)

PTX ld and st support scopes { .cta, .cluster, .gpu, .sys } and
sub-opcodes relaxed and acquire.

This patch adds frontend support for them.

(cherry picked from commit ec9fc96)
(cherry picked from commit 7d33470)
…PTX ld/st

This change implements codegen for MemSemantic flags and Scope flags for atomic
load and store ops.

(cherry picked from commit b512a43)
(cherry picked from commit 871e8f7)
@@ -1730,6 +1730,66 @@ def load(pointer, mask=None, other=None, boundary_check=(), padding_option="", c
volatile, _builder)


@builtin
def atomic_load(pointer, mask=None, other=None, sem="", scope="", boundary_check=(), padding_option="",
cache_modifier="", eviction_policy="", volatile=False, _builder=None):
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Comment needs updating for sem/scope.

@builtin
def atomic_store(pointer, value, mask=None, sem="", scope="", boundary_check=(), cache_modifier="", eviction_policy="",
_builder=None):
"""
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Comment needs updating for sem/scope.

@Jokeren
Copy link
Contributor

Jokeren commented Nov 19, 2024

@peterbell10 is familiar with this

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