[AMD][Gluon] Expose buffer_load and buffer_store to Gluon #7738
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
peterbell10
reviewed
Aug 1, 2025
| layout = ttgl._unwrap_if_constexpr(layout) | ||
|
|
||
| ret_ty = ttgl.distributed_type(element_type, shape, layout) | ||
| handle = self.builder.create_buffer_load(ret_ty.to_ir(self.builder), ptr, offsets, cache_modifier, mask, other) |
Contributor
There was a problem hiding this comment.
Can you infer the layout from the the offsets, the dtype from the pointer and add defaults for mask and other? You probably also need to add broadcasting.
My general feedback would be to think about this as a language API, not just a wrapper to emit IR.
Comment on lines
358
to
359
| handle = self.builder.create_buffer_store(stored_value, ptr, offsets, cache_modifier, mask) | ||
| return ttgl.tensor(handle, ttgl.void) |
Contributor
There was a problem hiding this comment.
I want to move away from tl.void. I don't think it adds any value.
Suggested change
| handle = self.builder.create_buffer_store(stored_value, ptr, offsets, cache_modifier, mask) | |
| return ttgl.tensor(handle, ttgl.void) | |
| self.builder.create_buffer_store(stored_value, ptr, offsets, cache_modifier, mask) |
|
|
||
|
|
||
| @builtin | ||
| def create_buffer_load(ptr, element_type, offsets, cache, mask, layout, other, _semantic=None): |
Contributor
There was a problem hiding this comment.
The create_ prefix is only used by the builder, not in the language.
Suggested change
| def create_buffer_load(ptr, element_type, offsets, cache, mask, layout, other, _semantic=None): | |
| def buffer_load(ptr, element_type, offsets, cache, mask, layout, other, _semantic=None): |
antiagainst
requested changes
Aug 2, 2025
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
d6fcb54 to
79fdfea
Compare
same style as the implementation in other backend
antiagainst
requested changes
Aug 4, 2025
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
| element_type = ptr.type.scalar.element_ty | ||
|
|
||
| if mask is not None: | ||
| assert mask.shape == shape, "offsets must have the same shape as offsets" |
Collaborator
There was a problem hiding this comment.
"mask must have .." :)
Collaborator
There was a problem hiding this comment.
This is not addressed yet?
antiagainst
approved these changes
Aug 5, 2025
Collaborator
antiagainst
left a comment
There was a problem hiding this comment.
Nice just one final comments from me. Also @peterbell10 to take another look.
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
peterbell10
reviewed
Aug 5, 2025
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
python/triton/experimental/gluon/language/amd/cdna3/__init__.py
Outdated
Show resolved
Hide resolved
antiagainst
requested changes
Aug 5, 2025
antiagainst
approved these changes
Aug 5, 2025
Collaborator
antiagainst
left a comment
There was a problem hiding this comment.
LGTM. I'll land once @peterbell10 is okay.
peterbell10
approved these changes
Aug 5, 2025
ptillet
pushed a commit
that referenced
this pull request
Aug 7, 2025
Expose AMD buffer_load and buffer_store Gluon. Example usage looks like:
```
def buffer_ldst_kernel(x, y):
layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[1, 64], warps_per_cta=[4, 1], order=[1, 0])
offsets = ttgl.arange(0, 64 * 64, layout=layout)
a = ttgl.amd.cdna3.buffer_load(ptr=x, offsets=offsets)
ttgl.amd.cdna3.buffer_store(stored_value=a, ptr=y, offsets=offsets)
```
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Expose AMD buffer_load and buffer_store Gluon. Example usage looks like: