TTNN Op Interface: Op constraints

[rjakovljevicTT]

What Provide a TTNN API to return Op constraints to the TT-MLIR optimizer/legalizer. The goal of this issue is to do an end-to-end PoC and conclude what would be the proper solution. Based on the conclusion we will open a new issue to define the whole work.

PoC Initially, on an example of MNIST NN and all Ops that are parts of MNIST, do the following tasks:

• [ ] Implement Op constraints API in TTNN (we should be in sync with core TTNN team),
  1. Tensor Layout (Row major vs Tile)
  2. Parallelization Constraints (Legal grids, width/height/block shard strategy support, shard orientation)
  3. Other constraints
• [ ] Consume the API in TT-MLIR Optimizer.
• [ ] Conclude on the Op constraints API design in TTNN and define the issue to do the whole work on TTNN side.

The above tasks should be done Iteratively until it works well for both TT-MLIR Optimizer and TTNN.

[ajakovljevicTT]

For this (v0) sprint, we have decided that we are going to support only a subset of operations and their different implementations (colloquially called "flavors"), with the emphasis being on the MNIST model and its operations. With that in mind, we have decided on the following ops and "flavors":

• The Softmax op in its most basic case, without any masking or scaling (ttnn.softmax)
• Only the ReLu unary op constraints will be implemented, as it is the only op used by MNIST
• Only a simple binary op is supported (as constraints are not related to any specific binary op type), which assumes that the two inputs have the same shape, so no internal broadcast is supported.
• In talks with optimizer folks, we have decided to support only two "flavors" of matmul - MatmulMultiCoreReuseMultiCast i MatmulMultiCoreReuseMultiCast1D, without any fused activations or biases.

With this in mind, we have made a set of factories for every op, which construct the corresponding builder objects, which can return a set of valid constraints for optimizer to use. Currently, we have (somewhat arbitrarily) decided that the parameters that define a constraint are:

• DataType of the input/output
• Layout of the input/output (tilized or not)
• StorageType (is it on device/host for now, also has potential support for multiple devices - has not been explored yet)

Other parameters that the op needs, such as shape, sharding, orientation, currently we expect the tilizer to provide. However, this decision is not fixed, and is still susceptible to change.
Most importantly, this has an impact on how we decide "flavors" of an op. Currently, we expect the optimizer to tell the constraint framework which "flavor" it wants and fill out the necessary objects with information for that. However, this logic is better suited to stay on the tt-nn side, which requires information on datatypes. This means that the division of this information between what is provided to the constraint framework and what the framework returns is currently not fully compatible with choosing of operation "flavors" in the constraint framework.

The set of functionalities described above has been implemented in the [feature/mlir-interface] (https://github.com/tenstorrent/tt-metal/tree/feature/mlir-interface) branch, along with the corresponding test, and the overall path of process can be tracked through several PRs to this branch:

• https://github.com/tenstorrent/tt-metal/pull/12458
• https://github.com/tenstorrent/tt-metal/pull/12553
• https://github.com/tenstorrent/tt-metal/pull/12563
• https://github.com/tenstorrent/tt-metal/pull/12581
• https://github.com/tenstorrent/tt-metal/pull/12631
• https://github.com/tenstorrent/tt-metal/pull/12706
• https://github.com/tenstorrent/tt-metal/pull/12700

Still TBD:

• Decide on what parameters will be provided to the constraint framework, and what will the constraint framework "sweep" on.
• Adding additional tests for the compilation paths, so that we currently cover all the validation paths of the supported ops.
• As @mbezuljTT mentioned, it might be beneficial to check why the ops do not support a certain configuration regarding sharding, and if it is possible to implement support for these cases in the ops themselves. This would largely simplify the constraint framework if possible.

As of this point, there is a WIP prototype that plugs the constraint framework into tt-mlir (more details here: https://github.com/tenstorrent/tt-metal/pull/12700).