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

Introduce triton-ptr-to-memref pass #211

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

Conversation

nhat-nguyen
Copy link
Collaborator

@nhat-nguyen nhat-nguyen commented Jan 2, 2025

This PR introduces the triton-ptr-to-memref pass responsible for converting function signature that uses triton ptr to use memref instead. This is part of the work to allow triton-shared to lower gather / scatter pointer sequences.

Intended lowering pipeline:

  • triton-to-structured (no changes):
    • analyzes structured addptr sequences
      • introduces tts.make_tptr %ptr_arg with offsets and strides
      • introduces tts.load and tts.store
    • leaves unstructured addptr sequences and their corresponding tt.load and tt.store intact
  • fold-unstructured-ptr (PR here):
    • converts all unstructured addptr sequences into sequences that compute pointer offsets
      • introduces tts.make_unstructured_tptr %ptr_arg %offsets
    • removes all tt.addptr
  • structured-to-memref (to be updated in a different PR):
    • currently converts everything to memref including scalar addptr and kernel arguments
    • will change to just convert ops in the tts dialect to memref with the exception of tts.make_unstructured_tptr
  • unstructured-to-memref (to be introduced in a different PR):
    • converts the remaining unstructured tt.load, tt.store, and tts.make_unstructured_tptr into memref
  • triton-ptr-to-memref (this PR):
    • converts kernel arguments with pointer type to memref

@nhat-nguyen nhat-nguyen changed the title Introduce TritonPtrToMemref pass Introduce triton-ptr-to-memref pass Jan 2, 2025
@nhat-nguyen nhat-nguyen marked this pull request as ready for review January 6, 2025 19:21
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.

1 participant