-
Notifications
You must be signed in to change notification settings - Fork 13
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Replacing LayoutAttr with TensorConfigAttr (#1231)
* Moving layout to TTNN * Addressing comments and merging with main. * Fixing build * Align attribute name with IR name
- Loading branch information
1 parent
14d87b7
commit 5c1f2a9
Showing
68 changed files
with
1,576 additions
and
536 deletions.
There are no files selected for viewing
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| // SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC | ||
| // | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| #ifndef TTMLIR_DIALECT_TT_UTILS_OPERANDCONSTRAINTS_H | ||
| #define TTMLIR_DIALECT_TT_UTILS_OPERANDCONSTRAINTS_H | ||
|
|
||
| #include "ttmlir/Dialect/TT/IR/TT.h" | ||
| #include "ttmlir/Dialect/TTIR/IR/TTIROps.h" | ||
|
|
||
| namespace mlir::tt { | ||
|
|
||
| inline OperandConstraint | ||
| memorySpaceAsOperandConstraint(MemorySpace memorySpace) { | ||
| switch (memorySpace) { | ||
| case MemorySpace::System: | ||
| case MemorySpace::SystemMMIO: | ||
| return OperandConstraint::System; | ||
| case MemorySpace::DeviceDRAM: | ||
| return OperandConstraint::DRAM; | ||
| case MemorySpace::DeviceL1: | ||
| return OperandConstraint::L1; | ||
| } | ||
| } | ||
|
|
||
| inline OperandConstraint | ||
| memoryLayoutAsOperandConstraint(TensorMemoryLayout memoryLayout) { | ||
| switch (memoryLayout) { | ||
| case TensorMemoryLayout::None: | ||
| return OperandConstraint::None; | ||
| case TensorMemoryLayout::Interleaved: | ||
| return OperandConstraint::Interleaved; | ||
| case TensorMemoryLayout::SingleBank: | ||
| return OperandConstraint::SingleBank; | ||
| case TensorMemoryLayout::HeightSharded: | ||
| return OperandConstraint::HeightSharded; | ||
| case TensorMemoryLayout::WidthSharded: | ||
| return OperandConstraint::WidthSharded; | ||
| case TensorMemoryLayout::BlockSharded: | ||
| return OperandConstraint::BlockSharded; | ||
| } | ||
| } | ||
|
|
||
| inline MemorySpace getLegalMemorySpace(OperandConstraint operandConstraint, | ||
| MemorySpace defaultMemorySpace) { | ||
| if (bitEnumContainsAny(operandConstraint, | ||
| memorySpaceAsOperandConstraint(defaultMemorySpace))) { | ||
| return defaultMemorySpace; | ||
| } | ||
| if (bitEnumContainsAny(operandConstraint, OperandConstraint::DRAM)) { | ||
| return MemorySpace::DeviceDRAM; | ||
| } | ||
| if (bitEnumContainsAny(operandConstraint, OperandConstraint::L1)) { | ||
| return MemorySpace::DeviceL1; | ||
| } | ||
| return MemorySpace::System; | ||
| } | ||
|
|
||
| inline TensorMemoryLayout | ||
| getLegalTensorMemoryLayout(OperandConstraint operandConstraint, | ||
| MemorySpace targetMemorySpace, | ||
| TensorMemoryLayout defaultDeviceMemLayout) { | ||
| if (defaultDeviceMemLayout == TensorMemoryLayout::None) { | ||
| return TensorMemoryLayout::None; | ||
| } | ||
|
|
||
| if (isSystemMemorySpace(targetMemorySpace)) { | ||
| return TensorMemoryLayout::None; | ||
| } | ||
|
|
||
| assert(isDeviceMemorySpace(targetMemorySpace)); | ||
| if (bitEnumContainsAny(operandConstraint, memoryLayoutAsOperandConstraint( | ||
| defaultDeviceMemLayout))) { | ||
| return defaultDeviceMemLayout; | ||
| } | ||
|
|
||
| std::map<OperandConstraint, TensorMemoryLayout> validLayoutsMap = { | ||
| {OperandConstraint::Interleaved, TensorMemoryLayout::Interleaved}, | ||
| {OperandConstraint::SingleBank, TensorMemoryLayout::SingleBank}, | ||
| {OperandConstraint::HeightSharded, TensorMemoryLayout::HeightSharded}, | ||
| {OperandConstraint::WidthSharded, TensorMemoryLayout::WidthSharded}, | ||
| {OperandConstraint::BlockSharded, TensorMemoryLayout::BlockSharded}}; | ||
|
|
||
| for (const auto &[constraintLayout, memLayout] : validLayoutsMap) { | ||
| if (bitEnumContainsAny(operandConstraint, constraintLayout)) { | ||
| return memLayout; | ||
| } | ||
| } | ||
|
|
||
| return TensorMemoryLayout::None; | ||
| } | ||
|
|
||
| } // namespace mlir::tt | ||
|
|
||
| #endif // TTMLIR_DIALECT_TT_UTILS_OPERANDCONSTRAINTS_H |
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
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
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
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
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
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
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
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
Oops, something went wrong.