-
Notifications
You must be signed in to change notification settings - Fork 14
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
Gather op implementation [#1015] #1151
Merged
Merged
Conversation
This file contains 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
ddilbazTT
requested review from
jnie-TT,
kmabeeTT,
AleksKnezevic,
pilkicTT,
sdjordjevicTT,
svuckovicTT,
mtopalovicTT,
rpavlovicTT,
jserbedzijaTT,
nobradovictt,
nsmithtt,
mrakitaTT and
tapspatel
as code owners
November 4, 2024 15:54
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
2 times, most recently
from
November 5, 2024 04:31
c05933b
to
b833e79
Compare
rpavlovicTT
requested changes
Nov 5, 2024
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
2 times, most recently
from
November 7, 2024 17:12
fddcfcd
to
7dcf489
Compare
nsmithtt
approved these changes
Nov 7, 2024
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
from
November 8, 2024 19:39
2050dd2
to
e8a402b
Compare
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
3 times, most recently
from
November 15, 2024 20:57
908e39f
to
b589d12
Compare
rpavlovicTT
approved these changes
Nov 18, 2024
rpavlovicTT
requested changes
Nov 18, 2024
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Few additional comments
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
3 times, most recently
from
November 18, 2024 16:55
883de4a
to
9a20bbe
Compare
rpavlovicTT
reviewed
Nov 19, 2024
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
5 times, most recently
from
November 19, 2024 16:00
1d618b7
to
5e41efc
Compare
Gather op implemetation end-to-end. SHLO GatherOp --> TTIR Gather Op --> TTIR Embedding Op (+ if needed, Reshape Op) --> TTNN Embedding Op (+ if needed, Reshape Op) TTNN does not currently support Gather Op, but had Embedding Op support. Given that Torch Embedding is lowered into SHLO Gather Op, and that most Gather Ops in models are simple embeddings, Gather Op is lowered into Embedding Op. If more sophisticated Gather Op implementations are encountered, then Gather could be lowered into slice/ concat/ etc if needed. Gather Op checks certain criteria to ensure lowering into Embedding Op in lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp - Output tensor shape: Multi-dimensional with last dimension as embedding size/ hiddenDim - Slice sizes: Must be [1, hiddenDim], where hiddenDim matches last output dimension - Offset dimensions: Strictly [2] - Collapsed slice dimensions: Strictly [0] - Start indices shape: Must be compatible with output shape - startIndices.size() < output.size() - if startIndices.size() == output.size(), then startIndices[-1] == 1 - Last dimension of start indices can be reduced by reshape op. - - This is due to embedding weights requiring to have smaller size than output shape
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
from
November 19, 2024 16:01
5e41efc
to
a92f222
Compare
Thanks everyone for your feedback! I hope the latest version addresses everybody's concerns. |
rpavlovicTT
approved these changes
Nov 19, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Gather op implemetation end-to-end.
SHLO GatherOp --> TTIR Gather Op --> TTIR Embedding Op (+ if needed, Reshape Op) --> TTNN Embedding Op (+ if needed, Reshape Op)
TTNN does not currently support Gather Op, but had Embedding Op support.
Given that Torch Embedding is lowered into SHLO Gather Op, and that most
Gather Ops in models are simple embeddings, Gather Op is lowered into
Embedding Op. If more sophisticated Gather Op implementations are
encountered, then Gather could be lowered into slice/ concat/ etc if
needed.
Gather Op checks certain criteria to ensure lowering into Embedding Op
in lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
embedding size/ hiddenDim
output dimension
size than output shape