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

Gather op implementation [#1015] #1151

Merged
merged 1 commit into from
Nov 19, 2024
Merged

Gather op implementation [#1015] #1151

merged 1 commit into from
Nov 19, 2024

Conversation

ddilbazTT
Copy link
Contributor

@ddilbazTT ddilbazTT commented Nov 4, 2024

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 ddilbazTT force-pushed the ddilbaz/gather branch 2 times, most recently from c05933b to b833e79 Compare November 5, 2024 04:31
@ddilbazTT ddilbazTT force-pushed the ddilbaz/gather branch 2 times, most recently from fddcfcd to 7dcf489 Compare November 7, 2024 17:12
@ddilbazTT ddilbazTT force-pushed the ddilbaz/gather branch 3 times, most recently from 908e39f to b589d12 Compare November 15, 2024 20:57
Copy link
Contributor

@rpavlovicTT rpavlovicTT left a 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 ddilbazTT force-pushed the ddilbaz/gather branch 3 times, most recently from 883de4a to 9a20bbe Compare November 18, 2024 16:55
@ddilbazTT ddilbazTT force-pushed the ddilbaz/gather branch 5 times, most recently from 1d618b7 to 5e41efc Compare November 19, 2024 16:00
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
Copy link
Contributor Author

Thanks everyone for your feedback! I hope the latest version addresses everybody's concerns.

@ddilbazTT ddilbazTT enabled auto-merge (squash) November 19, 2024 17:36
@ddilbazTT ddilbazTT merged commit 2976a24 into main Nov 19, 2024
18 checks passed
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.

7 participants