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

Support i1 datatype with an experimental flag. #18713

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

lialan
Copy link
Contributor

@lialan lialan commented Oct 7, 2024

Enable packed i1 datatype storage

This commit introduces support for packed storage of the i1 (bit) datatype. When subbyte type packing is enabled via the --iree-experimental-packed-i1-storage option, vectors of i1 elements will be stored in a compact packed representation.

For example, a vector<6xi1> will occupy a single byte of memory with the 6 bit elements packed together and 2 padding bits. A vector<3x3xi1> will take up 2 bytes, with the 9 bit elements packed across the bytes and 7 padding bits.

Limitations:

  • To ensure correct behavior, the tiling configuration aligns the innermost dimension data loads with byte boundaries. This is necessitated by the current lack of emulation for unaligned subbyte vector loading/storing.
  • Unaligned subbyte emulation support can be added in the future, though it may incur some performance overhead.

This change requires corresponding updates in the frontend to utilize the packed i1 storage format.

Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alan and I had an offline sync, and he is revisiting the codegen side changes. I'm not an expert of host side changes, so we need some inputs from Ben.

@lialan lialan force-pushed the lialan/i1 branch 3 times, most recently from a6f19b2 to 36561a9 Compare October 15, 2024 16:11
@lialan lialan marked this pull request as ready for review October 15, 2024 16:19
@lialan lialan force-pushed the lialan/i1 branch 3 times, most recently from 6b327fe to 07de9c4 Compare October 15, 2024 16:58
@hanhanW
Copy link
Contributor

hanhanW commented Nov 6, 2024

I'll revisit the PR because now we have the support for cross-byte i1 loads. @lialan you're still working on the store part, right? Also, could you add an e2e i1 load test to the PR? I think we can iterate on that for the load support.

@lialan
Copy link
Contributor Author

lialan commented Nov 6, 2024

I'll revisit the PR because now we have the support for cross-byte i1 loads. @lialan you're still working on the store part, right? Also, could you add an e2e i1 load test to the PR? I think we can iterate on that for the load support.

@hanhanW Looks like the stores are not our highest priority at the moment. We want to have i1 support for loads primarily

@lialan lialan force-pushed the lialan/i1 branch 2 times, most recently from ba8c833 to 872d19e Compare November 6, 2024 15:08
@hanhanW
Copy link
Contributor

hanhanW commented Nov 6, 2024

@lialan I'm not sure if the PR is ready or not because you added few commits. Can you mark it draft and turn it open when it's ready for review?

If you need some inputs and it's not ready for review yet, feel free to tag me or ping me on discord.

@lialan lialan marked this pull request as draft November 6, 2024 18:11
@lialan
Copy link
Contributor Author

lialan commented Nov 6, 2024

@lialan I'm not sure if the PR is ready or not because you added few commits. Can you mark it draft and turn it open when it's ready for review?

If you need some inputs and it's not ready for review yet, feel free to tag me or ping me on discord.

Apologies, not sure what automation or my mis-click marked this as open. I turned it back to draft. There are some issues with e2e tests and I am fixing it.

@lialan
Copy link
Contributor Author

lialan commented Nov 7, 2024

Turns out, when converting to flow dialect, some operations such as tensor.insert_slice will be turned into vector stores to memref. For that we need emulated vector store support.

@lialan lialan force-pushed the lialan/i1 branch 2 times, most recently from b31fc83 to df005bf Compare November 8, 2024 22:41
Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I took a quick look at TypePropagationPass.cpp and I think we are not able to remove it because we can't handle i3/i5/etc at this moment.

The changes look okay to me, and I think there is a bug in the PR. Please take a look.

compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp Outdated Show resolved Hide resolved
compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp Outdated Show resolved Hide resolved
@lialan lialan force-pushed the lialan/i1 branch 3 times, most recently from 3020ae6 to df13e54 Compare November 9, 2024 03:38
@lialan lialan requested a review from hanhanW November 9, 2024 04:12
@lialan lialan marked this pull request as ready for review November 9, 2024 04:12
Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me, just few comments about tests

@lialan lialan force-pushed the lialan/i1 branch 4 times, most recently from c0ee0a7 to 04ca8c7 Compare November 12, 2024 16:26
@hanhanW hanhanW changed the title Support i1 datatype Support i1 datatype with an experimental flag. Nov 12, 2024
Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, please update the checks in encode_host_tensors_pack.mlir.

Copy link
Contributor

@MaheshRavishankar MaheshRavishankar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think would be good to have @benvanik review the changes to ElementPackingUtils.cpp . I am unblocking it, but please wait for Ben to review.

@lialan lialan requested a review from benvanik November 13, 2024 15:06
@lialan
Copy link
Contributor Author

lialan commented Nov 14, 2024

@benvanik Ben can you review it for us?

* Turn on by `--iree-enable-i1` option.
* also added e2e tests.

Signed-off-by: Alan Li <[email protected]>
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.

Plumb i1 datatype through the compilation pipeline
4 participants