[AMD] Rewrite extract_slice op implementation #7128
Merged
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
antiagainst
requested changes
Jun 12, 2025
third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td
Outdated
Show resolved
Hide resolved
third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td
Outdated
Show resolved
Hide resolved
third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td
Outdated
Show resolved
Hide resolved
third_party/amd/lib/TritonAMDGPUDialectToLLVM/ExtractSliceOpToLLVM.cpp
Outdated
Show resolved
Hide resolved
third_party/amd/lib/TritonAMDGPUDialectToLLVM/ExtractSliceOpToLLVM.cpp
Outdated
Show resolved
Hide resolved
third_party/amd/lib/TritonAMDGPUDialectToLLVM/ExtractSliceOpToLLVM.cpp
Outdated
Show resolved
Hide resolved
680299e to
f961aff
Compare
Contributor
Author
|
@antiagainst I addressed the comments. Thanks for the review! |
antiagainst
approved these changes
Jun 13, 2025
Collaborator
antiagainst
left a comment
There was a problem hiding this comment.
LGTM now. Thanks for tidying up it! Adding @ThomasRaoux to take another look to make sure this also looks good.
Contributor
Hi @ThomasRaoux. Would you have time to have a look at this PR? |
ThomasRaoux
approved these changes
Jun 20, 2025
Collaborator
ThomasRaoux
left a comment
There was a problem hiding this comment.
LGTM, thanks for improving this
tie-pilot-qxw
pushed a commit
to tie-pilot-qxw/triton
that referenced
this pull request
Aug 30, 2025
This PR refactors the extract_slice operation to support two major
improvements:
1) Relaxed Layout Constraints
The operation now allows more flexible source and destination layouts,
aligning better with linear layouts.
2) Support for Arbitrary Tensor Ranks
extract_slice is no longer limited to 2D tensors and can now handle
tensors of any rank.
The "extract_slice" operation enables extracting a slice of a tensor in
registers.
It supports the following arguments:
* source: the base tensor on which to create a view tensor
* offsets: offsets into the base tensor at which to create the view
In distributed layouts, tensors are divided into CTA tiles.
A CTA tile represents the smallest contiguous portion of a tensor that
is distributed across all threads and warps within a workgroup. The
ExtractSlice operation extracts a portion of the tensor that aligns with
CTA tile boundaries.
This op is designed to work on logical tensors directly, avoiding the
need for complex layout reinterpretation or reshaping.
For example, the tt.split operation only supports splitting along the
innermost dimension,
and requires that the resulting innermost dimension provide 2 elements
per thread, distributed across registers.
In contrast, extract_slice op imposes no constraints on the extraction
dimension or the size of dimensions.
---------
Co-authored-by: Ognjen Plavsic <plognjen@amd.com>
Co-authored-by: Lei Zhang <antiagainst@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 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.
This PR refactors the extract_slice operation to support two major improvements:
Relaxed Layout Constraints
The operation now allows more flexible source and destination layouts, aligning better with linear layouts.
Support for Arbitrary Tensor Ranks
extract_slice is no longer limited to 2D tensors and can now handle tensors of any rank.
The "extract_slice" operation enables extracting a slice of a tensor in registers.
It supports the following arguments:
* source: the base tensor on which to create a view tensor
* offsets: offsets into the base tensor at which to create the view
In distributed layouts, tensors are divided into CTA tiles.
A CTA tile represents the smallest contiguous portion of a tensor that is distributed across all threads and warps within a workgroup. The ExtractSlice operation extracts a portion of the tensor that aligns with CTA tile boundaries.
This op is designed to work on logical tensors directly, avoiding the need for complex layout reinterpretation or reshaping.
For example, the tt.split operation only supports splitting along the innermost dimension,
and requires that the resulting innermost dimension provide 2 elements per thread, distributed across registers.
In contrast, extract_slice op imposes no constraints on the extraction dimension or the size of dimensions.