[Optimizer] Refactor legal gird analysis #1363

odjuricicTT · 2024-11-21T15:12:50Z

PR contains the following:

Rename to LegalLayoutAnalysis
Add generation of both TILE and ROW_MAJOR layouts
- Keep ROW_MAJOR disabled by default and add flag to enable when needed
Explicitly generate all layout related parameters (except for data format).
Adjust overrides to support partial output layout instead of every param (e.g. only override buffer type (dram / l1)
Clean up optimizer mlir tests

Left to do:

Add cpp unit tests for analysis and overrides
Add more partial override tests

github-actions

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp

nobradovictt · 2024-11-21T21:42:01Z

test/ttmlir/Dialect/TTNN/optimizer/partial_output_layout_override.mlir

@@ -0,0 +1,20 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="enable-optimizer=true override-output-layout=add_1=:::row_major:" %s | FileCheck %s


I don't like how this looks, should be in format similar to named parameters in case all options are not specified.

I agree. Are you suggesting that we have multiple options on syntax for overrides or that we switch to always named?

Named might get a bit verbose and we might run out of meaningful separator characters.

E.g:
override-output-layout=add_1=buffer_type=l1:tensor_memory_layout=block_sharded:memory_layout=row_major:data_type=f32

Look how named params work, you can omit them if you specify all params in correct order, if overriding only partially you need to specify param names. I agree verbosity is getting out of hand, maybe we can shorten names a bit. Lets think about it more.

We can shorten param names:
override-output-layout=add_1=bt=l1:tml=block_sharded:ml=row_major:dt=f32 and avoid them if all are specified in order:
override-output-layout=add_1=l1:block_sharded:row_major:f32

As discussed offline, change to support any number of layout params with deducing type of param by value.

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp

nobradovictt · 2024-11-25T09:53:31Z

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp

+
+  // Create element type for the new layout.
+  Type elementType = layout.getScalarElementType();
+  if (override.dataType.has_value()) {


Default is row major?

No, TILE / ROW_MAJOR is always overriden in this path and is controled with the next if: if (override.memoryLayout == Layout::Tile).

This checks dataType which is allowed to not be present and is copied over from the existing tensor if not overriden.

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp

nobradovictt · 2024-11-25T10:54:13Z

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp

+  // Apply partial layout overrides.
+  if (override.has_value()) {
+    analysisResult.erase(
+        std::remove_if(analysisResult.begin(), analysisResult.end(),


I do not like this approach. Can we at the beginning load all overrides and not even add layout not compatible with overrides instead of adding all and then removing them. At the very minimal create a util function instead of this lambda remove_if approach.

I started with that approach and implemented a solution where i first create a search space for each layout param, then apply overrides and then try to itterate over the cros product and generate legal solutions. This was very hard to read because it ended up having a lot of nested if to check for non compatible cases across parameters. So i decided to go back to this which is cleaner an less error prone. I understand we are trading of performance for readability here, but it seems logical atm?

Will extract into a seperate funciton.

svuckovicTT

Dialect-side changes lgtm!

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp

odjuricicTT · 2024-12-10T13:52:35Z

@nsmithtt @tapspatel @vprajapati-tt Can someone take a look at the PR, your approval is required. I think that the only change not reviewed is this file https://github.com/tenstorrent/tt-mlir/pull/1363/files#diff-694f570cfe79f8e3130bde9637283269c8c33048a68ce3b3fd15be810b11d961

Will update codeowners for this in the future.

nsmithtt · 2024-12-11T19:56:58Z

test/unittests/Optimizer/TestOptimizerOverrides.cpp

+  llvm::cl::opt<std::string> OverrideOutputLayoutOption{
+      "override-output-layout"};
+  OutputLayoutOverrideParser parser{OverrideOutputLayoutOption};
+  llvm::StringMap<OutputLayoutOverrideParams> parsedOverride;


@odjuricicTT, I think this might have broken macOS build: https://github.com/tenstorrent/tt-mlir/actions/runs/12269439074/job/34232997171

OptimizerTests: CommandLine Error: Option 'override-output-layout' registered more than once! LLVM ERROR: inconsistency in registered CommandLine options

Perhaps this fixture gets instantiated many times?

Seems so. Thanks for pointing it out. Will take a look asap.

I tried a potential fix here: #1589

But I'm unable to repro locally, it feels like a build config issue actually, I think the usage of parser.parse is correct.

It seems to be some strange behaviour where if two specific tests exist, one of them fails. Once one is deleted, all tests pass. Also unable to repro locally. Will delete one of the tests for now to unblock further work, and take a deeper look later.

Temporary fix: #1598

Issue for future: #1597

nsmithtt · 2024-12-11T19:57:55Z

test/unittests/Optimizer/TestOptimizerOverrides.cpp

+                    "sharded:row_major:f16";
+
+  bool result = parser.parse(OverrideOutputLayoutOption,
+                             "override-output-layout", arg, parsedOverride);


@odjuricicTT, see comment above, maybe this one is the culprit:

[ RUN ] OutputLayoutOverrideTest.ParseMultipleOps OptimizerTests: CommandLine Error: Option 'override-output-layout' registered more than once!

* Rename to LegalLayoutAnalysis * Add generation of both TILE and ROW_MAJOR layouts * Keep ROW_MAJOR disabled by default and add flag to enable when needed * Explicitly generate all layout related parameters (except for data format). * Adjust overrides to support partial output layout instead of every param * Clean up optimizer mlir tests

odjuricicTT added 8 commits November 18, 2024 11:11

Fix legal grid analysis

b092cdb

Add row major flag

168916b

Merge branch 'main' into odjuricic/legal-layout-analysis

21b9628

Rebase and update tests

cc308c5

Rename grid -> layout

e994e19

Cleanup

cb700d7

Fix up mlir tests

26d7b6d

Clean up

a17ac45

odjuricicTT requested review from nobradovictt, sdjordjevicTT, svuckovicTT, mtopalovicTT, rpavlovicTT, jserbedzijaTT and jnie-TT as code owners November 21, 2024 15:12

github-actions bot reviewed Nov 21, 2024

View reviewed changes

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp Outdated Show resolved Hide resolved

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp Show resolved Hide resolved

nobradovictt reviewed Nov 21, 2024

View reviewed changes

odjuricicTT requested a review from vprajapati-tt as a code owner November 22, 2024 09:54

odjuricicTT added 2 commits November 22, 2024 10:56

Clean up / fix up

1474b92

Rename files

56188a9

odjuricicTT force-pushed the odjuricic/legal-layout-analysis branch from 93193c8 to 56188a9 Compare November 22, 2024 11:01

nobradovictt reviewed Nov 25, 2024

View reviewed changes

svuckovicTT approved these changes Nov 25, 2024

View reviewed changes

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp Show resolved Hide resolved

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp Outdated Show resolved Hide resolved

odjuricicTT added 2 commits November 25, 2024 12:34

Fuzzy parsing for layout override and tests

59d3e42

Address comments

45a761c

nobradovictt approved these changes Dec 2, 2024

View reviewed changes

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp Outdated Show resolved Hide resolved

lib/Dialect/TTNN/Analysis/LegalLayoutAnalysis.cpp Outdated Show resolved Hide resolved

Merge branch 'main' into odjuricic/legal-layout-analysis

6791a3f

odjuricicTT requested review from nsmithtt and tapspatel as code owners December 10, 2024 11:25

tapspatel approved these changes Dec 10, 2024

View reviewed changes

Merge branch 'main' into odjuricic/legal-layout-analysis

a4606a5

odjuricicTT enabled auto-merge (squash) December 10, 2024 15:52

odjuricicTT merged commit f2f2e97 into main Dec 10, 2024
20 checks passed

nsmithtt reviewed Dec 11, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Optimizer] Refactor legal gird analysis #1363

[Optimizer] Refactor legal gird analysis #1363

odjuricicTT commented Nov 21, 2024

github-actions bot left a comment

nobradovictt Nov 21, 2024

odjuricicTT Nov 22, 2024

nobradovictt Nov 22, 2024

nobradovictt Nov 25, 2024

odjuricicTT Nov 25, 2024

nobradovictt Nov 25, 2024

odjuricicTT Nov 25, 2024

nobradovictt Nov 25, 2024

odjuricicTT Nov 25, 2024

svuckovicTT left a comment

odjuricicTT commented Dec 10, 2024

nsmithtt Dec 11, 2024

odjuricicTT Dec 13, 2024

nsmithtt Dec 13, 2024

odjuricicTT Dec 16, 2024

nsmithtt Dec 11, 2024

		@@ -0,0 +1,20 @@
		// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="enable-optimizer=true override-output-layout=add_1=:::row_major:" %s \| FileCheck %s

[Optimizer] Refactor legal gird analysis #1363

[Optimizer] Refactor legal gird analysis #1363

Conversation

odjuricicTT commented Nov 21, 2024

github-actions bot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

svuckovicTT left a comment

Choose a reason for hiding this comment

odjuricicTT commented Dec 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment