Support for memRef of L1 Interleaved tensors (#1292) #1607
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
odjuricicTT
reviewed
Dec 17, 2024
| std::accumulate(logicalTiledShape.begin(), logicalTiledShape.end(), 1, | ||
| std::multiplies<std::int64_t>()); | ||
| uint64_t numOfGridUnits = | ||
| std::accumulate(grid.getShape().begin(), grid.getShape().end(), 1, |
There was a problem hiding this comment.
Maybe leave an assert here that grid is equal to device max grid? Or at least a comment. I'm not sure that l1 interleaved is supported for non max grid configurations.
There was a problem hiding this comment.
We have to check this, but I suspect that l1 is supported for non max grid configuration.
lib/Dialect/TTNN/IR/TTNNOpsAttrs.cpp
Outdated
Comment on lines
107
to
125
| shardShape.push_back(1); | ||
| shardShape.push_back((numOfTiles + numOfGridUnits - 1) / numOfGridUnits); |
There was a problem hiding this comment.
This is always 2D. Is there a case when it needs more dimensions? Maybe just leave an assert that this is the desired number of dims, so when we introduce 3D this fails in a friendly manner.
There was a problem hiding this comment.
Instead of asserting, we can add n-1 ones at the beginning of merRef so its dimension corresponds to the grid dimension.
odjuricicTT
approved these changes
Dec 17, 2024
fbajraktariTT
force-pushed
the
fix_L1_interleaved_mlir
branch
from
3a99432 to
5158f8b
Compare
fbajraktariTT
requested review from
nobradovictt,
sdjordjevicTT,
svuckovicTT,
mtopalovicTT,
jserbedzijaTT,
jnie-TT and
azecevicTT
as code owners
fbajraktariTT
force-pushed
the
fix_L1_interleaved_mlir
branch
from
5158f8b to
94028e1
Compare
|
Will continue tmrw. |
…ardShapeForInterleaved
fbajraktariTT
force-pushed
the
fix_L1_interleaved_mlir
branch
from
94028e1 to
0e1fc7a
Compare
mtopalovicTT
approved these changes
Dec 18, 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.
This PR fixes the issue of incorrect
memReffor L1 Interleaved layouts.In the previous implementation of
TTNNLayoutAttrwe were modellingmemRefas its memory layout isBlockSharded. By modelling everything asBlockShardedwe hit the error of incorrect memory usage for tensors with the given layout.For example, if we wanted to create a L1 Interleaved
TTNNLayoutoftensor<1x1024xbf16>, we would get amemref<1x4x!tt.tile<32x32, bf16>, #l1_>. This was incorrect since the whole tensor can be tiled into 32 tiles and if grid shape is8x8we could assign only one tile per grid unit. There is no need for 4 tiles to be placed in one grid unit.This PR adds workaround for this issue by using different logic for calculating logical shard shape for L1 Interleaved layouts. This workaround creates a
memref<1x1x!tt.tile<32x32, bf16>, #l1_>fortensor<1x1024xbf16>.Closes issue: #1292