Skip to content

Store transpose attribute in Subgroup2DBlockIO layouts #4549

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

Open
wants to merge 1 commit into
base: alex/use_subgroup_2d_block_encoding_pr
Choose a base branch
from

Conversation

alexbaden
Copy link
Contributor

One of the primary goals of the Subgroup 2D Block IO encoding is to be able to convert to an accurate Linear Layout representation for data produced by Subgroup 2D Block IO instructions. To do this, we need to store transpose information in the layout as an attribute. The attribute tells us to swap the order of the height/width parameters in the layout definition when building the linear layout representation, in effect performing the post-load transpose in registers (this is consistent with the SPV intrinsic documentation, https://github.khronos.org/SPIRV-Registry/extensions/INTEL/SPV_INTEL_2d_block_io. The end result is we get a layout corresponding to the post-transpose loaded data in registers:

./build/cmake.linux-x86_64-cpython-3.10/bin/triton-tensor-layout -l "#ttig.subgroup_2d_block<{warpsPerCTA = [8, 4], instrShape = [32, 16], numBlocks = 1, isTransposed = true, order = [0, 1], kWidth = 2, threadsPerWarp = 16}>" -t "tensor<32x256xf16>" -use-hw-view 

depends on #4500 / #4463

@alexbaden
Copy link
Contributor Author

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