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

[TIR][Transform] Introduce new InjectPermutedLayout pass #16070

Merged

Conversation

Ubospica
Copy link
Contributor

@Ubospica Ubospica commented Nov 4, 2023

This PR introduces the new InjectPermutedLayout pass, partially adapted from the old version.

This pass will permute the visiting indices of the shared memory with xor operations to avoid bank conflicts. For permutation details see p39 of these slides.

The current interface is:

with T.block("..."):
    T.block_attr({"permuted_layout": 1})
    # load or store to shared memory

The load/store operation could be:

  • BufferLoad node
  • BufferStore node
  • ptx_ldmatrix call
  • mma_store call

And the InjectPermutedLayout pass is injected in the lowering process in driver_api.cc. So after this pass in the building stage, the visiting indices will be permuted.

Notably this pass works after the CompactBufferAllocation pass because it relies on the shape of the shared memory in every threadblock.

cc @spectrometerHBH @vinx13 @Hzfengsy

@Ubospica Ubospica marked this pull request as ready for review November 5, 2023 03:51
@Ubospica Ubospica force-pushed the main-dev/2023-11-02-inject-permuted-layout branch 3 times, most recently from 73944d2 to 2cbf342 Compare November 6, 2023 01:21
1104

1106
@Ubospica Ubospica force-pushed the main-dev/2023-11-02-inject-permuted-layout branch from 2cbf342 to 4905883 Compare November 7, 2023 01:10
@Ubospica
Copy link
Contributor Author

Ubospica commented Nov 7, 2023

@tvm-bot rerun

1 similar comment
@Ubospica
Copy link
Contributor Author

Ubospica commented Nov 9, 2023

@tvm-bot rerun

@yzh119 yzh119 merged commit 468bf2d into apache:main Nov 12, 2023
18 checks passed
Ubospica added a commit to Ubospica/tvm-develop that referenced this pull request Nov 13, 2023
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.

None yet

3 participants