Introduction

The following document provides an overview of the TT-MLIR project, with a focus on the technical specifications of an MLIR-based compiler stack. So what exactly is an MLIR-based compiler stack? MLIR (Multi Level Intermediate Representation) is a subproject coming out of the LLVM Project. It seeks to introduce extensibility and sustainable code design to a very modular compiler framework. This essentially means to take a much larger more involved compiler (like LLVM) and split it into sub-compilers that each produce their own Intermediate Representation (IR) of what you've fed the compiler.

Disclaimer: This is intended to be a working document, if you find something incorrect or incomplete please feel free to create a PR.

Motivations

The idea of having a multi-level IR might not seem so far fetched, in fact it resembles some of our current software stacks. The idea of going from a High Level TVM Graph → Lowered PyBUDA Graph → Netlist, with each layer having their own level of optimizations is quite a familiar concept. However, there are problems with the reusability and integration of optimizations for the current software compiler stack. Currently, users are almost forced to choose between a top-down optimization or bottom-up optimization, with both requiring "expert-level" expertise to optimize for desired performance. Developing 2 entirely different projects is taxing, and it's hard to translate the benefits of BUDA over to metal (or the other way around). One of the primary goals of tt-mlir is to enable a consistent programming model between software stacks, concepts for improving optimizations in the compiler stack should 1-1 carry over to hand-written TTNN.

The benefits grow even further when one can understand all the possible entry points that multiple IRs present. Existing MLIR based projects like OpenXLA and torch-mlir can natively output MLIR in a dialect that can be transcribed to the TTIR dialect as well!

What is MLIR and why use it?

MLIR is a compiler infrastructure that is designed to be modular and extensible. The main benefits the tt-mlir project hopes to gain by using MLIR include:

• Industry Standard Compiler Framework
  • Lots of boilerplate algorithms, data structures, and useful software that is common to compiler development
• Ecosystem
  • Hook into existing front-end MLIR projects
• Testing framework
  • A battle-tested test infrastructure that will enable us to write fine grained tests and rely less on end-to-end testing
  • Common IR Serialization Format that's easy to test, debug, and edit

Additional documentation to highlight the benefits of MLIR can be found here:

• MLIR Whitepaper
• MLIR Documentation

MLIR: Overview

MLIR is at it's root an interpreter that can parse "readable" text in some .mlir format. The unique properties lie in the modularity of the parsing itself. MLIR is built upon a collection of Dialects, each of these Dialects define a collection of Operations, Types, and Attributes. These dialects follow their own syntax, and they can encode any amount of information. The benefit is that MLIR provides bindings and hooks such that a user can directly translate these IRs into usable artifacts for that layer of complexity. An example of this would be the relatively high level TOSA Dialect, which is used to represent computation over tensors, and then lowering that to a more hardware specific dialect that closely models the programming model of the hardware or underlying backend. It is the dialect system itself which powers the multi-level functionality of MLIR, with different dialects a user can essentially "lower" through their software stack by just transforming between the different dialects for their layers. Dialects can exist in a broad range from purely mathematical dialects, to a LinAlg Dialect, or a Tensorflow Dialect defined for ML Graphs. Each dialect encodes its own information and their operations can use the Types/Attributes of other dialects as parameters. Multiple dialects are possible in one module, and encouraged to highlight optimizations of different dialects. In our usecase for the TT Stack, MLIR acts a "mid-level" compiler which makes the task of joining together various entry points and backends much simpler.

MLIR Primitives

So what does MLIR look like, how does it work and get parsed? The hierarchy of an MLIR module is as shown:

#permutation = array<i64: 0, 2, 1>

module {
  func.func @forward(%input: tensor<32x64x128xf32>) -> tensor<32x128x64xf32> {
    %output = ttir.empty() : tensor<32x128x64xf32>
    %result = "ttir.permute"(%input, %output) <{permutation = #permutation}> : (tensor<32x64x128xf32>, tensor<32x128x64xf32>) -> tensor<32x128x64xf32>
    return %result : tensor<32x128x64xf32>
  }
}

• Attributes (defined using #)

  • The syntax of actually creating an attribute is modular, and custom assembly instructions for different attributes can be applied.

• Operations

  • These operations are accessed with the . method, so you'll see some examples like func.func or ttir.empty. Each operation also provides it's own assembly instructions but often strictly defines the type of result

  • Quotes are added around ttir.permute since it's part of a custom dialect.

  • Operations typically have operands (arguments) and results which are highlighted with %, these results and operands help to show the relationship between operations

• Types

  • Types are shown as dataformats throughout this compiled mlir module, where tensor and array are some examples.

  • They help to demonstrate the transformation of information and its representation as it's processed across this module.

MLIR Workflow

The overall MLIR workflow doesn't necessarily involve writing .mlir files or even modifying them. The Intermediate Representations are truly just representations, we can parse them to demonstrate what the graph looks like at that current stage of optimization, or run a pass through them to optimize certain functions. The overall framework is designed with the following architecture in mind:

1. Graph Information exists

2. Graph Information is transformed (through any which method) into a high-level MLIR representation

3. Passes are run on the high-level implementation to lower into TTIR, a common IR that can be lowered into multiple backends

4. Depending on the usecase more passes are run to lower to whatever backend the user would like (ex: TTNN Backend)

What are Passes?

Transformations in MLIR are represented as passes that occur during the parsing of some information. These passes can be executed when parsing or generating MLIR modules. These transformations can have a myriad of purposes, and are completely user defined as to how they modify the module. Some examples of passes can be for lowering purposes as mentioned before, where a dialect is parsed and then each operation is transformed to a lowered dialect following some set of user defined rules. Passes are also used for optimizations and backend code transformation in the context of this project. They're a powerful tool and provide most of the functionality to transform between layers of dialects, and they provide a simple platform for modifications of an MLIR module.

Why not make our own?

Now that I've described the functionality of the MLIR framework, it seems like making an in house multi level Intermediate Representation system would be pretty similar, so why are we going through the effort of implementing this framework?

One of the biggest reason can be attributed to the active developer community surrounding the project, being a part of the LLVM Project means that there is solid developer support, and the framework is designed to be a tool for many different paradigms of compute. This scalability and strong mission statement lend to the strengths of MLIR being a solid platform to use as a middle layer in our compiler stack. Furthermore, as a functional benefit of being part of a larger open source project, MLIR has a whole library of tests and infrastructure that we can leverage for solid code health while starting a new project.

Automation

It's not only about developer support, another key benefit of MLIR is that it's built with autogeneration in mind. Through TableGen a lot of the boilerplate of creating this multi-level IR become abstracted away to truly focus on implementation and execution. This automation is built on top of a pre-existing robust framework with a lot of implementations and support from other large players in the ML scene. By integrating with these automation pipelines, we allow for external developers to have a much simpler entry-point into our software stack!

TT-MLIR: Bringing MLIR to the TT Stack

Now that we have defined this pretty cool project, let's look at the implementation details of bringing MLIR (and related optimizations) into the TT Stack. Since it acts as a mid-level compiler we can start by defining the "bottom" and "top" layers of the compiler. BUDA already has a well defined set of frontend optimizations to some TVM defined graph and is knowledgeable of the hardware that these models want to run on. We want to interrupt the BUDA stack to only give us the frontend compiled graph before any hardware specific lowering is to occur. What this will produce is information that is agnostic to different backends and their execution on TT hardware, but this is still valid information to optimize at different levels for later compilation. The "bottom" of our graph is now defined as the backend that will produce the machine-specific code to be executed. While MLIR could allow for any level of complexity downwards for the bottom, we will define a very aggressive TTNN backend for the MVP. Desired Optimization List:

• Forge-FE (frontend)

  • Graph Optimizations, Constant Folding, Operation Fusion

• TT-MLIR (mid-level)

  • Data Storage, Memory Configuration, Grid Configuration

• TTNN (backend)

  • Kernel Configuration*, Network Optimization

*Subject to Change / Be Moved to TT-MLIR

TT-MLIR Dialects

Now that we have defined the series of optimizations that we would like to see implemented in TT-MLIR, we can begin to help define the dialects that would help to support these different levels of optimizations. For more detail on each of these dialects, please refer to the GitHub Wiki and TableGen descriptors. I think that Nick does a great job of documenting the key functionality.

TT Dialect

The TT Dialect is only for common Types and Attributes used throughout the many levels of the mid level compiler.

TTIR Dialect

The TTIR Dialect is defined as the common dialect for TT-MLIR, as such it doesn't define anything hardware/backend specific. It lists out general actions that would take place on TT hardware such as dispatch, layout, and kernel operations.

Generic Operation

This is one of two operations that's crucial to understand the intended optimization characteristics of the TTIR Dialect. The generic operation dictates the actions that would be taken to dispatch some instruction to TT hardware such that it executes some instruction. Parametrically, the operation consumes inputs, outputs, maps to read the tensors, and access-types to the memory. These parameters highlight the optimizations that can be performed at this level to change the location of the memory, transpose using variant access maps, or even the grid upon which the computation takes place. The operation also contains a block in which the exact behaviour for that operation to occur is stored.

Layout Operation

The layout operation is key in describing the storage of memory throughout the execution graph. Layout determines the sharding spec, location of the memory, data types, and tile sizes of some tensor. While generic describes the dispatch for some data-wise transformation to take place, the data itself is laid out across the chip through the layout operation.

Both of these operations describe the key functionality of the TTIR dialect and the optimization space that it provides.

Built-in MLIR Dialects

The functionality of TT-MLIR Dialects also depends / is inspired by the functionality of Built-in MLIR Dialects like Affine and LinAlg. Below are summaries of some of the key members of these Dialects

Affine Dialect

[Reference] Affine maps help to describe transformations on coordinate systems, while this may not really make sense, imagine trying to index a rank 2 tensor. By getting t[x, y] I can access the element in the Xth row and Yth column, but if I wanted to transpose the tensor I might have to re-layout the entire tensor such that the data would be accessible using t[x, y] to get the element in the Yth row and Xth column. This transpose can also be represented using an Affine Map to transform (x, y) -> (y, x) and this would let the tensor data remain in place while the access method is modified. This extends even further to more complex transformations such that stride lengths or unique indexing methods can be implemented without complicated manipulation.

Tensor Dialect

[Reference] The tensor dialect defines the functionality and Type of the fundamental Tensor. This dialect contains members that would represent manipulation and representation of tensors as multi-dimensional data with shapes and datatypes. Not much else is different about this dialect, the reference covers key topics if implementation details are needed.

Func Dialect

[Reference]

TOSA Dialect

[Reference]

SCF Dialect

[Reference]

EmitC Dialect

[Reference]

tt-explorer - Performance Optimization Tool

A unique project related to TT-MLIR is the integration of Performance Optimization Tools such that users are easily able to visualize and readily tune their models without needing an expert level understanding of the tech stack. 'tt-explorer' is built with Google AI's Model Explorer as a base for the visualization tool, and a custom adapter to parse TT-MLIR projects. This would allow users to readily tune their models, and optimize for the TTIR layer (ex: they can change certain memory to be laid out in L1 instead of DRAM, or change the grid layout of an operation to be larger than what was previously assigned). After compilation with these overrides, the runtime information can then be fed directly into a Tracy Performance Analysis for the user to visualize the impacts of their tuning, seeing which operations were least performant and continuing in a gamified design loop for iterative performance tuning!

Getting Started

This page walks you through the steps required to set up tt-mlir.

NOTE: If you have a build issue, you can file a bug here.

Prerequisites

Hardware Setup

Use this guide to set up your hardware - Hardware Setup.

System Dependencies

You can use tt-mlir with Ubuntu or Mac OS, however the runtime does not work on Mac OS. tt-mlir project has the following system dependencies:

• Ubuntu 22.04 OS or Mac OS
• Clang >= 14 & <= 18
• Ninja
• CMake 3.24 or higher
• Python 3.11
• python3.11-venv

Ubuntu

Install Clang, Ninja, CMake, and python3.11-venv:

sudo apt install git clang cmake ninja-build pip python3.11-venv

You should now have the required dependencies installed.

NOTE: If you intend to build with runtime enabled (-DTTMLIR_ENABLE_RUNTIME=ON), you also need to install tt-metal dependencies which can be found here.

Full developer dependencies as packaged in our docker image:

apt-get update
apt-get install -y \
    software-properties-common \
    build-essential \
    python3-pip \
    git \
    libhwloc-dev \
    pandoc \
    libtbb-dev \
    libcapstone-dev \
    pkg-config \
    linux-tools-generic \
    ninja-build \
    wget \
    libgtest-dev \
    cmake \
    ccache \
    doxygen \
    graphviz \
    libyaml-cpp-dev \
    libboost-all-dev \
    curl \
    jq \
    sudo \
    gh \
    lcov \
    zstd \
    unzip

# Install Python 3.11
add-apt-repository ppa:deadsnakes/ppa && \
    apt-get update && \
    apt-get install -y python3.11 python3.11-dev python3.11-venv python3.11-distutils

# Setup / install metal dependencies
wget https://raw.githubusercontent.com/tenstorrent/tt-metal/${TT_METAL_DEPENDENCIES_COMMIT}/{install_dependencies.sh,tt_metal/sfpi-info.sh,tt_metal/sfpi-version}
chmod u+x sfpi-info.sh
bash install_dependencies.sh --docker

Mac OS

On MacOS we need to install the latest version of cmake, and ninja which can be done using Homebrew with (Docs for installing Homebrew: https://brew.sh).

brew install cmake ninja

Clone the tt-mlir Repo

1. Clone the tt-mlir repo:

git clone https://github.com/tenstorrent/tt-mlir.git

2. Navigate into the tt-mlir folder.

Environment Setup

There are two ways to set up the environment, either using a docker image or building the environment manually. The docker image is recommended since it is easier to set up and use.

Using a Docker Image

Please see Docker Notes for details on how to set up and use the docker image.

Once you have the docker image running and you are logged into the container, you should be ready to build.

Setting up the Environment Manually

This section explains how to manually build the environment so you can use tt-mlir. You only need to build this once, it builds llvm, flatbuffers, and a Python virtual environment. You can specify the LLVM build type by using -DLLVM_BUILD_TYPE=*. The default is MinSizeRel, and available options are listed here.

1. Navigate into the tt-mlir folder.

2. The environment gets installed into a toolchain directory, which is by default set to /opt/ttmlir-toolchain, but can be overrideen by setting (and persisting in your environment) the environment variable TTMLIR_TOOLCHAIN_DIR. You need to manually create the toolchain directory as follows:

export TTMLIR_TOOLCHAIN_DIR=/opt/ttmlir-toolchain/
sudo mkdir -p "${TTMLIR_TOOLCHAIN_DIR}"
sudo chown -R "${USER}" "${TTMLIR_TOOLCHAIN_DIR}"

3. Please ensure that you do not already have an environment (venv) activated before running the following commands:

cmake -B env/build env
cmake --build env/build
source env/activate

NOTE: The last command takes time to run, so give it time to complete.

Building the tt-mlir Project

In this step, you build the tt-mlir project:

source env/activate
cmake -G Ninja -B build
cmake --build build

You have now configured tt-mlir.

You can add different flags to your build. Here are some options to consider:

• To enable the ttnn/metal runtime add -DTTMLIR_ENABLE_RUNTIME=ON. Clang 17 is the minimum required version when enabling the runtime.
• To enable the ttnn/metal perf runtime add -DTT_RUNTIME_ENABLE_PERF_TRACE=ON.
• To accelerate the builds with ccache use -DCMAKE_CXX_COMPILER_LAUNCHER=ccache.
• To workaround OOM issues it can be useful to decrease the number of parallel jobs with -DCMAKE_BUILD_PARALLEL_LEVEL=4.
• If Python bindings aren't required for your project, you can accelerate builds further with the command -DTTMLIR_ENABLE_BINDINGS_PYTHON=OFF.
• To enable tt-explorer add the -DTT_RUNTIME_ENABLE_PERF_TRACE=ON, -DTTMLIR_ENABLE_RUNTIME=ON, and -DTT_RUNTIME_DEBUG=ON.
• To enable optimizer pass that uses the op model library, add -DTTMLIR_ENABLE_OPMODEL=ON.
• The TTNN build is automatically integrated / handled by the tt-mlir cmake build system. For debugging and further information regarding the TTNN backend build step, please refer to TTNN Documentation.
• The runtime build depends on the TT_METAL_RUNTIME_ROOT variable, which is also set in env/activate script. For more information, please refer to TT-NN and TT-Metailium installation documentation.

Test the Build

Use this step to check your build. Do the following:

source env/activate
cmake --build build -- check-ttmlir

Lint

Set up lint so you can spot errors and stylistic issues before runtime:

source env/activate
cmake --build build -- clang-tidy

Note for developers: You can run:

source env/activate
cmake --build build -- clang-tidy-ci

This reproduces the Lint (clang-tidy) CI job. It runs clang-tidy only on committed files that have been modified relative to the origin/main branch.

Pre-Commit

Pre-Commit applies a git hook to the local repository such that linting is checked and applied on every git commit action. Install from the root of the repository using:

source env/activate
pre-commit install

If you have already committed before installing the pre-commit hooks, you can run on all files to "catch up":

pre-commit run --all-files

For more information visit pre-commit

Docs

Build the documentation by doing the following:

1. Make sure you have mdbook, doxygen, sphinx, and sphinx-markdown-builder installed.

2. Build the docs:

source env/activate
cmake --build build -- docs
mdbook serve build/docs

NOTE: mdbook serve will by default create a local server at http://localhost:3000.

For more information about building the docs please read the full guide on building the docs.

Common Build Errors

TTMLIRPythonCAPI target requires changing an RPATH

CMake Error at /opt/ttmlir-toolchain/lib/cmake/llvm/AddLLVM.cmake:594 (add_library):
  The install of the TTMLIRPythonCAPI target requires changing an RPATH from
  the build tree, but this is not supported with the Ninja generator unless
  on an ELF-based or XCOFF-based platform.  The
  CMAKE_BUILD_WITH_INSTALL_RPATH variable may be set to avoid this relinking
  step.

If you get the above error, it means you tried to build with an old version of cmake or ninja and there is a stale file. To fix this, rm -rf your build directory, install a newer version of cmake/ninja, and then rebuild. If you installed ninja via sudo apt install ninja-build, it might still be not up-to-date (v1.10.0). You may use ninja in the python virtual environment, or install it via pip3 install -U ninja, either way the version 1.11.1.git.kitware.jobserver-1 should work.

clang++ is not a full path and was not found in the PATH

CMake Error at CMakeLists.txt:2 (project):
  The CMAKE_CXX_COMPILER:
    clang++
  is not a full path and was not found in the PATH.
  Tell CMake where to find the compiler by setting either the environment
  variable "CXX" or the CMake cache entry CMAKE_CXX_COMPILER to the full path
  to the compiler, or to the compiler name if it is in the PATH.
CMake Error at CMakeLists.txt:2 (project):
  The CMAKE_C_COMPILER:
    clang
  is not a full path and was not found in the PATH.
  Tell CMake where to find the compiler by setting either the environment
  variable "CC" or the CMake cache entry CMAKE_C_COMPILER to the full path to
  the compiler, or to the compiler name if it is in the PATH.

If you get the following error, it means you need to install clang which you can do with sudo apt install clang on Ubuntu.

tt-metal Update Failures

Failed to unstash changes in: '/path/to/tt-metal/src/tt-metal'
You will have to resolve the conflicts manually

This error occurs during CMake's ExternalProject update of tt-metal. The build system tries to apply changes using Git's stash mechanism, but fails due to conflicts. This can happen even if you haven't manually modified any files, as the build process itself may leave behind artifacts or partial changes from previous builds.

To resolve, run the following command:

rm -rf third_party/tt-metal

Then retry your build command. If the error persists, you may need to do the following:

1. Remove the build directory: rm -rf build

2. Run CMake commands again.

3. Run the above.

Common Runtime Errors

Debugging Python on Mac OS

When debugging python on macOS via lldb you may see an error like:

(lldb) r
error: process exited with status -1 (attach failed (Not allowed to attach to process.  Look in the console messages (Console.app), near the debugserver entries, when the attach failed.  The subsystem that denied t
he attach permission will likely have logged an informative message about why it was denied.))

For preinstalled macOS binaries you must manually codesign with debug entitlements.

Create file debuggee-entitlement.xml:

<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
        <key>com.apple.security.cs.disable-library-validation</key>
        <true/>
        <key>com.apple.security.get-task-allow</key>
        <true/>
</dict>
</plist>

Sign the binary:

sudo codesign -f -s - --entitlements debuggee-entitlement.xml /opt/ttmlir-toolchain/venv/bin/python

Working with Docker Images

Components:

• Dockerfile
• Workflow for building Docker image
• Project build using Docker image

Overview

We use docker images to prepare the project environment, install dependencies, tooling and prebuild toolchain. Project builds four docker images:

• Base image tt-mlir-base-ubuntu-22-04 Dockerfile.base
• CI image tt-mlir-ci-ubuntu-22-04 Dockerfile.ci
• Base IRD image tt-mlir-base-ird-ubuntu-22-04Dockerfile.ird
• IRD image tt-mlir-ird-ubuntu-22-04 Dockerfile.ird

Base image starts with a supported base image (Ubuntu 22.04) and installs dependencies for project build. From there, we build the CI image that contains the prebuild toolchain and is used in CI to shorten the build time. The IRD image contains dev tools such as GDB, vim and ssh which are used in IRD environments.

During the CI Docker build, the project is built and tests are run to ensure that everything is set up correctly. If any dependencies are missing, the Docker build will fail.

Using the Docker Image

Here is a typical command to run the latest developer (ird) docker image:

sudo docker run -it -d --rm \
  --name my-docker \
  --cap-add ALL \
  --device /dev/tenstorrent/0:/dev/tenstorrent/0 \
  -v /dev/hugepages:/dev/hugepages \
  -v /dev/hugepages-1G:/dev/hugepages-1G \
  ghcr.io/tenstorrent/tt-mlir/tt-mlir-ird-ubuntu-22-04:latest bash

Special attention should be paid to flags:

• --device /dev/tenstorrent/0:/dev/tenstorrent/0: this is required to map the hardware device into the container. For machines with multiple devices, this flag can be specified multiple times or adjusted with the appropriate device number.
• -v /dev/hugepages:/dev/hugepages / -v /dev/hugepages-1G:/dev/hugepages-1G: this is required to map the hugepages volume into the container. For more information on hugepages, please refer to the Getting Started Guide.

The base or CI image can also be used in the same way, but the IRD image is recommended for development.

Using the Docker Image via IRD (Internal Developers Only)

Internally we use a tool called IRD. As part of your reserve command, you can specify the docker image to use:

ird reserve \
  --docker-image ghcr.io/tenstorrent/tt-mlir/tt-mlir-ird-ubuntu-22-04:latest

See ird reserve --help for more information on the reserve command. Typical ird usage might look like:

# list machine availability
ird list-machines

# reserve a machine
ird reserve \
  --volumes /localdev/$USER:/localdev/$USER \
  --docker-image ghcr.io/tenstorrent/tt-mlir/tt-mlir-ird-ubuntu-22-04:latest \
  --timeout 720 \
  wormhole_b0 \
  --machine [MACHINE_NAME]

# list your currently reserved machines
ird list

# connect to the first reserved machine
ird connect-to 1

# release the first reserved machine
ird release 1

Building the Docker Image using GitHub Actions

The GitHub Actions workflow Build and Publish Docker Image builds the Docker images and uploads them to GitHub Packages at https://github.com/orgs/tenstorrent/packages?repo_name=tt-mlir. We use the git SHA we build from as the tag.

Building the Docker Image Locally

To test the changes and build the image locally, use the following command:

docker build -f .github/Dockerfile.base -t ghcr.io/tenstorrent/tt-mlir/tt-mlir-base-ubuntu-22-04:latest .
docker build -f .github/Dockerfile.ci -t ghcr.io/tenstorrent/tt-mlir/tt-mlir-ci-ubuntu-22-04:latest .
docker build -f .github/Dockerfile.ird --build-arg FROM_IMAGE=base -t ghcr.io/tenstorrent/tt-mlir/tt-mlir-ird-base-ubuntu-22-04:latest .
docker build -f .github/Dockerfile.ird --build-arg FROM_IMAGE=ci -t ghcr.io/tenstorrent/tt-mlir/tt-mlir-ird-ubuntu-22-04:latest .

Using the Image in GitHub Actions Jobs

The GitHub Actions workflow Build in Docker uses a Docker container for building:

    container:
      image: ghcr.io/${{ github.repository }}/tt-mlir-ci-ubuntu-22-04:latest
      options: --user root

Running Virtualized Ubuntu VM on macOS

In some cases, like running a software simulated device, it can be beneficial to run the stack on a local macOS machine. This document covers the necessary setup and configuration steps to get a performant Ubuntu VM setup on Apple Silicon.

Prerequisite Steps

1. UTM is the VM application we'll be using in this guide, so the first step is to download and install UTM.
2. Ubuntu 22.04 ARM image download.

• Direct link: 64-bit ARM (ARMv8/AArch64) server install image

UTM Setup

1. Launch UTM and click the + button to start a new VM.
2. Choose Virtualize (emulation works, but is unusably slow).
3. Under Preconfigured choose Linux.
4. Check box Use Apple Virtualization and select the ubuntu iso we just downloaded for Boot ISO Image.

• Optionally check Enable Rosetta which can enable running ELF's compiled for x86 if you're interested. It's not required and additional steps are required for it to work.

5. This step depends on your machine's capabilities, but it's recommended to give 16GB of memory and to use the defualt CPU Cores setting. Note this can be changed after initial setup if you want to go back and tweak settings.
6. It's recommended to at least 128GB of storage, with LLVM installation and full SW stack we quickly reach 80 gigs of storage.
7. Optionally choose a shared host/VM directory.
8. Optionally name your new VM ubuntu 22.04 arm64

VM Setup

1. Boot your newly created VM!
2. Run through the Ubuntu setup as you see fit, be sure that openssh is enabled which simplifies logging into your VM, but the rest of the defaults are sufficient.
3. If you plan on using your VM via ssh you can retrieve the ip address ip a and looking at the inet row under enp0s1. Should look something like inet 192.168.64.3. Another tip is to add this to the host's ~/.ssh/config.
4. Install your normal developer tools as you see fit.

Software Stack Installation

The majority of the software install flow is the same, with the exception of a few caveats called out here.

1. Installing metal deps needs the additional flags below:

git clone git@github.com:tenstorrent/tt-metal.git
cd tt-metal
sudo bash install_dependencies.sh --docker --no-distributed

• --docker: Despite not being in a docker, this is the flag that turns off configuring hugepages which is not required for VM.
• --no-distributed: Currently the metal distributed feature requires a package version of openmpi that only supports x86.

2. Install tt-mlir system dependencies as outlined by this step.
3. The environment needs to be built manually as outlined here.
4. We can then build tt-mlir per usual.
5. If planning to run tests on software sim, let's build the ttrt tool.

• The following all works per usual:
  • Querying a system desc.
  • Running a flatbuffer

Testing

To run tests:

source env/activate
cmake --build build -- check-ttmlir

Lit testing

llvm-lit tool is used for MLIR testing. With it you can:

# Query which tests are available
llvm-lit -sv ./build/test --show-tests

# Run an individual test:
llvm-lit -sv ./build/test/ttmlir/Dialect/TTIR/test_allocate.mlir

# Run a sub-suite:
llvm-lit -sv ./build/test/ttmlir/Dialect/TTIR

See the full llvm-lit documentation for more information.

EmitC testing

NOTE: This is a developer's guide on how to test EmitC as a feature. For usage of EmitC, please refer to ttnn-standalone docs.

Prerequisites

• Built ttmlir

• Built ttrt

• Activated virtual environment:

  source env/activate
  

• Saved system descriptor file:

  ttrt query --save-artifacts
  

Table of Contents

1. Generate all EmitC tests and run them
2. Generate a single EmitC test and run it
3. Generate EmitC tests with Builder

Generate all EmitC tests and run them

1. Generate flatbuffers and .cpp files for EmitC tests

   If you don't have SYSTEM_DESC_PATH environment variable exported, you can run:

   SYSTEM_DESC_PATH=/path/to/system_desc.ttsys llvm-lit -sv test/ttmlir/EmitC/TTNN
   

   Or if you have SYSTEM_DESC_PATH exported, you can omit it:

   llvm-lit -sv test/ttmlir/EmitC/TTNN
   

2. Compile generated .cpp files to shared objects

   tools/ttnn-standalone/ci_compile_dylib.py
   

3. Run flatbuffers + shared objects and compare results

   ttrt run --emitc build/test/ttmlir/EmitC/TTNN
   

Generate EmitC tests with Builder

Builder offers support for building EmitPy modules from ttir or stablehlo ops. Refer to Builder documentation.

Generate a single EmitC test and run it

1. Generate flatbuffers and .cpp files for EmitC test

   SYSTEM_DESC_PATH=/path/to/system_desc.ttsys llvm-lit -sv test/ttmlir/EmitC/TTNN/eltwise_binary/add.mlir
   

2. Compile generated .cpp files to shared objects

   Assuming default build directory path:

   tools/ttnn-standalone/ci_compile_dylib.py --file build/test/ttmlir/EmitC/TTNN/eltwise_binary/add.mlir.cpp
   

3. Run the flatbuffer + shared object and compare results

   ttrt emitc build/test/ttmlir/EmitC/TTNN/eltwise_binary/add.mlir.so --flatbuffer build/test/ttmlir/EmitC/TTNN/eltwise_binary/add.mlir.ttnn
   

Tools

The ttmlir project currenly exposes the following tools:

• ttmlir-opt: The ttmlir optimizer driver. This tool is used to run the ttmlir compiler passes on a .mlir source files and is central to developing and testing the compiler.
• ttmlir-translate: The ttmlir translation tool. This tool can convert from IR to external representation (and inverse). For example, IR in EmitC dialect can be converted into C++ code.
• ttrt: This tool is intended to be a swiss army knife for working with flatbuffers generated by the compiler. Its primary role is to inspect and run flatbuffer files.
• ttir-builder: This tool is for creating ttir operations. It provides support for those ops to be compiled into modules or directly to flatbuffer files.
• tt-explorer: Visualizer tool for ttmlir-powered compiler results. Visualizes from emitted .mlir files to display compiled model, attributes, performance results, and provide a platform for human-driven overrides to gamify model tuning.
• ttnn-standalone: This tool is used to run C++ TTNN code outside of the compiler environment.

ttmlir-opt

The ttmlir optimizer driver. This tool is used to run the ttmlir compiler passes on a .mlir source files and is central to developing and testing the compiler.

Simple Test

./build/bin/ttmlir-opt --ttir-to-ttnn-backend-pipeline test/ttmlir/Dialect/TTNN/simple_multiply.mlir
# Or
./build/bin/ttmlir-opt --ttir-to-ttmetal-pipeline test/ttmlir/Dialect/TTNN/simple_multiply.mlir

ttmlir-translate

The ttmlir-translate translation utility. Unlike ttmlir-opt tool which is used to run passes within the MLIR world, ttmlir-translate allows us to ingest something (e.g. code) into MLIR world, and also produce something (e.g. executable binary, or even code again) from MLIR.

Generate C++ code from MLIR

# First, let's run `ttmlir-opt` to convert to proper dialect
./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/Dialect/TTNN/eltwise/binary/multiply/simple_multiply.mlir -o c.mlir

# Now run `ttmlir-translate` to produce C++ code
./build/bin/ttmlir-translate --mlir-to-cpp c.mlir

Bonus: These two commands can be piped, to avoid writing a mlir file to disk, like so:

./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/Dialect/TTNN/eltwise/binary/multiply/simple_multiply.mlir | ./build/bin/ttmlir-translate -mlir-to-cpp

Generate flatbuffer file from MLIR

# First run `ttmlir-opt` to convert to proper dialect
./build/bin/ttmlir-opt --ttir-to-ttnn-backend-pipeline test/ttmlir/Dialect/TTNN/eltwise/binary/multiply/simple_multiply.mlir -o ttnn.mlir

# Now run `ttmlir-translate` to produce flatbuffer file
./build/bin/ttmlir-translate --ttnn-to-flatbuffer ttnn.mlir -o out.ttnn

ttrt

This tool is intended to be a swiss army knife for working with flatbuffers generated by the compiler. Its primary role is to inspect and run flatbuffer files. It enables the running of flatbuffer files without a front-end runtime.

Building

1. Build ttmlir
2. Build ttrt:

source env/activate
cmake --build build
ttrt --help

Building runtime mode

Add the following flags when building the compiler

-DTTMLIR_ENABLE_RUNTIME=ON

Building perf mode

Add the following flags when building the compiler

-DTTMLIR_ENABLE_RUNTIME=ON
-DTT_RUNTIME_ENABLE_PERF_TRACE=ON

LOGGER Levels

ttrt support logging at different logger levels. You will need to set env var TTRT_LOGGER_LEVEL in command line or a python script. By default, it's set to INFO.

TTRT_LOGGER_LEVEL=INFO
TTRT_LOGGER_LEVEL=CRITICAL
TTRT_LOGGER_LEVEL=ERROR
TTRT_LOGGER_LEVEL=WARNING
TTRT_LOGGER_LEVEL=DEBUG

tt-metal logging

ttrt runtime uses tt-metal for op execution and device interfacing. For more detailed logs, which can help in troubleshooting build or runtime issues, set env var TT_METAL_LOGGER_LEVEL. By default, it is set to FATAL.

export TT_METAL_LOGGER_LEVEL=DEBUG

Installing ttrt as python whls

Every time ttrt is built, it creates a whls file in build/tools/ttrt/build. Ex filename: ttrt-0.0.235-cp311-cp311-linux_x86_64.whl. You can take this whls file and install it in any docker container and in any venv outside of ttmlir. After which, you can use all the following functionality as the same.

1. Download whls
2. Create a python venv

python -m venv ttrt_env
source ttrt_env/bin/activate

3. Install whls (replace with your version of the whls)

pip install build/tools/ttrt/build/ttrt-0.0.235-cp311-cp311-linux_x86_64.whl

Generating a flatbuffer

tt-mlir exposes a few ways to generate flatbuffers.

Generate a flatbuffer file from ttir-builder

ttir-builder is a tool for creating TTIR ops, converting them into MLIR modules, running passes to lower modules into backends, and translating to flatbuffers. See documentation for further instructions.

Generate a flatbuffer file from compiler

The compiler supports a pass to load a system descriptor to compile against. You can feed this pass into ttmlir-opt.

1. Build ttmlir
2. Generate ttsys file from the system you want to compile for using ttrt. This will create a system_desc.ttsys file under ttrt-artifacts folder.

ttrt query --save-artifacts

3. Use ttmlir-opt tool in compiler to feed system descriptor. See the ttmlir-opt documentation for more information on how to generate .mlir files.

./build/bin/ttmlir-opt --ttcore-register-device="system-desc-path=/path/to/system_desc.ttsys" --ttir-to-ttnn-backend-pipeline test/ttmlir/Dialect/TTNN/simple_subtract.mlir -o ttnn.mlir
or (pipe path directly into ttir-to-ttnn-backend-pipeline)
./build/bin/ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=/path/to/system_desc.ttsys" test/ttmlir/Dialect/TTNN/simple_subtract_to_add.mlir -o ttnn.mlir

4. Use ttmlir-translate tool in compiler to generate the flatbuffer executable. See the ttmlir-translate documentation for more information on how to generate flatbuffer files.

./build/bin/ttmlir-translate --ttnn-to-flatbuffer ttnn.mlir -o out.ttnn

5. Run your test cases using ttrt

ttrt run /path/to/out.ttnn

Generate flatbuffer files using llvm-lit

There are already existing .mlir test cases under test/ttmlir/Silicon. You can use llvm-lit tool to generate the corresponding ttnn and ttm files.

1. Build ttmlir
2. Generate ttsys file from the system you want to compile for using ttrt. This will create a system_desc.ttsys file under ttrt-artifacts folder.

ttrt query --save-artifacts

3. Export this file in your environment using export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys. When llvm-lit is run, it will query this variable and generate the ttnn and ttm files using this system. Optionally, you can also provide this manually when running llvm-lit.
4. Generate your test cases. This will generate all your ttnn and ttm files under build/test/ttmlir/Silicon. ttnn files have a .ttnn file extension and ttmetal files have a .ttm extension.

cmake --build build -- check-ttmlir

5. (Optional) If you have a single .mlir file (or a directory of custom .mlir files) that you created using the compiler, and you want to generate the corresponding ttnn and ttm files for it, you can run llvm-lit standalone to the path of your .mlir file or directory of .mlir files to generate the flatbuffer executables. You will have to make sure you add in the correct llvm-lit configs into your .mlir file. See section on adding llvm-lit config options inside a .mlir file to create flatbuffer binaries for more info. You must also make sure your .mlir test is found within test/ttmlir/Silicon folder (and point lit to the build folder)!

llvm-lit -v ./build/test/ttmlir/Silicon
or
SYSTEM_DESC_PATH=/path/to/system_desc.ttsys llvm-lit -v ./build/test/ttmlir/Silicon

6. Run your test cases using ttrt

ttrt run /path/to/test.ttnn
ttrt run /path/to/dir/of/flatbuffers

Adding llvm-lit config options inside a .mlir file to create flatbuffer binaries

Inside of your .mlir file, you can add certain config options that llvm-lit will use when running against that test case. For the purpose of generating flatbuffer executables, you can add --ttcore-register-device="system-desc-path=%system_desc_path%" which will tell llvm-lit to parse the system desc found from the environment flag set by export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys. You can also paste a custom path to a system desc file as well.

// RUN: ttmlir-opt --ttcore-register-device="system-desc-path=%system_desc_path%" --ttnn-layout --convert-ttir-to-ttnn %s  > %t.mlir
// RUN: FileCheck %s --input-file=%t.mlir
// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn

Adding new mlir test cases

You can copy your .mlir test file (with the appropriate llvm-lit config options for generating flatbuffer binaries) into test/ttmlir/Silicon. Then, follow generating flatbuffer files using llvm-lit to generate the executables to run!

Versioning

ttrt and flatbuffers have strict versioning check. When running a flatbuffer against ttrt, you have to make sure the flatbuffer was generated using the same version as ttrt (or vice versa). Major and Minor versions are manually set using github tags when releases are made. Patch versioning is the number of commits from the last major/minor tag.

vmajor.minor.patch

The flag --ignore-version can be used to bypass versioning checks. Use at your own risk; it can cause unpredictable errors.

Application APIs

ttrt --help
ttrt read
ttrt run
ttrt query
ttrt perf
ttrt check
ttrt emitpy

Command line usage

There are different ways you can use the APIs under ttrt. The first is via the command line as follows. All artifacts are saved under ttrt-artifacts folder under TT_MLIR_HOME environment variable. By default, all logging is printed to the terminal. You can specify a log file to dump output to.

read

Read sections of a binary file

ttrt read --help
ttrt read --section version out.ttnn
ttrt read --section system_desc out.ttnn
ttrt read --section mlir out.ttnn
ttrt read --section inputs out.ttnn
ttrt read --section outputs out.ttnn
ttrt read --section op_stats out.ttnn
ttrt read --section mesh_shape out.ttnn
ttrt read --section all out.ttnn --clean-artifacts
ttrt read --section all out.ttnn --save-artifacts
ttrt read --section all /dir/of/flatbuffers
ttrt read system_desc.ttsys
ttrt read --section system_desc system_desc.ttsys
ttrt read system_desc.ttsys --log-file ttrt.log
ttrt read out.ttnn --save-artifacts --artifact-dir /path/to/some/dir
ttrt read out.ttnn --result-file result.json

run

Run a binary file or a directory of binary files Note: It's required to be on a system with silicon and to have a runtime enabled build -DTTMLIR_ENABLE_RUNTIME=ON.

ttrt run --help
ttrt run out.ttnn
ttrt run out.ttnn --seed 0
ttrt run out.ttnn --init arange
ttrt run out.ttnn --identity
ttrt run out.ttnn --identity --rtol 1 --atol 1
ttrt run out.ttnn --clean-artifacts
ttrt run out.ttnn --save-artifacts
ttrt run out.ttnn --loops 10
ttrt run --program-index all out.ttnn
ttrt run --program-index 0 out.ttnn
ttrt run /dir/of/flatbuffers
ttrt run /dir/of/flatbuffers --loops 10
ttrt run /dir/of/flatbuffers --log-file ttrt.log
ttrt run out.ttnn --save-artifacts --artifact-dir /path/to/some/dir
ttrt run out.ttnn --load-kernels-from-disk
ttrt run out.ttnn --result-file result.json
ttrt run out.ttnn --disable-golden
ttrt run out.ttnn --save-golden-tensors
ttrt run out.ttnn --print-input-output-tensors
ttrt run out.ttnn --debugger
ttrt run out.ttnn --memory --save-artifacts
ttrt run out.ttnn --memory --check-memory-leak

For info on running EmitC tests, see EmitC testing.

Run results

The run api saves a run_results.json file that records information about the run including any errors that were thrown and location of other saved run data.

{
[
  {
    "file_path": "ttnn/test_tan[f32-shape0]_ttnn.mlir.ttnn",
    "result": "pass",
    "exception": "",
    "log_file": "ttrt.log",
    "artifacts": "/home/$USER/tt-mlir/ttrt-artifacts",
    "program_index": "all",
    "program_results": {
      "program_index_0": {
        "loop_0": {
          "total_duration_ns": 3269341588,
          "total_ttnn_api_duration_ns": null,
          "total_device_kernel_duration_ns": null
        }
      }
    }
  }
]

Golden checks

Golden checks are used to verify runtime op accuracy. They are run by default during the golden callback unless flag --disable-golden is used. If flag --save-artifacts is used, a golden results report will be saved under the artifacts directory.

{
    "loc(\"/home/$USER/tt-mlir/test/python/golden/test_ttir_ops.py:74:id(0)\")": {
        "expected_pcc": 0.99,
        "actual_pcc": 0.0015917614829425491,
        "atol": 1e-08,
        "rtol": 1e-05,
        "allclose": false,
        "max": 8529.765625,
        "mean_absolute_error": 6.644593238830566,
        "root_mean_square_error": 100.30211639404297,
        "cosine_similarity": 0.0016297339461743832
    }
}

Memory

Memory callback functions are run when flag --memory is used. A memory report will be written under the artifacts directory that contains information on op memory usage.

{
    "0": {
        "loc": "loc(\"/home/$USER/tt-mlir/test/python/golden/test_ttir_ops.py:74:id(0)\")",
        "debug_str": "%0 = \"ttnn.tan\"(%arg0) : (tensor<128x128xf32, #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<4x4x!ttcore.tile<32x32, f32>, #ttnn.buffer_type<dram>>, <interleaved>>>) -> tensor<128x128xf32, #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<4x4x!ttcore.tile<32x32, f32>, #ttnn.buffer_type<dram>>, <interleaved>>> loc(\"/home/$USER/tt-mlir/test/python/golden/test_ttir_ops.py:74:id(0)\")",
        "dram": {
            "num_banks": 12,
            "total_bytes_per_bank": 1071181792,
            "total_bytes_allocated_per_bank": 16384,
            "total_bytes_free_per_bank": 1071167456,
            "largest_contiguous_bytes_free_per_bank": 1071165408,
            "block_table": [
                {
                    "allocated": "yes",
                    "nextID": "1",
                    "prevID": "-1",
                    "size": "8192",
                    "address": "0",
                    "blockID": "0"
                },
                {
                    "allocated": "yes",
                    "nextID": "3",
                    "prevID": "0",
                    "size": "8192",
                    "address": "8192",
                    "blockID": "1"
                },
                {
                    "allocated": "no",
                    "nextID": "-1",
                    "prevID": "1",
                    "size": "1071165408",
                    "address": "16384",
                    "blockID": "3"
                }
            ]
        },
        "l1": {
            "num_banks": 64,
            "total_bytes_per_bank": 1369120,
            "total_bytes_allocated_per_bank": 0,
            "total_bytes_free_per_bank": 1369120,
            "largest_contiguous_bytes_free_per_bank": 1369120,
            "block_table": [
                {
                    "allocated": "no",
                    "nextID": "-1",
                    "prevID": "-1",
                    "size": "1369120",
                    "address": "0",
                    "blockID": "0"
                }
            ]
        },
        "l1_small": {
            "num_banks": 64,
            "total_bytes_per_bank": 32768,
            "total_bytes_allocated_per_bank": 0,
            "total_bytes_free_per_bank": 32768,
            "largest_contiguous_bytes_free_per_bank": 32768,
            "block_table": [
                {
                    "allocated": "no",
                    "nextID": "-1",
                    "prevID": "-1",
                    "size": "32768",
                    "address": "0",
                    "blockID": "0"
                }
            ]
        },
        "trace": {
            "num_banks": 12,
            "total_bytes_per_bank": 0,
            "total_bytes_allocated_per_bank": 0,
            "total_bytes_free_per_bank": 0,
            "largest_contiguous_bytes_free_per_bank": 0,
            "block_table": [
                {
                    "allocated": "no",
                    "nextID": "-1",
                    "prevID": "-1",
                    "size": "0",
                    "address": "0",
                    "blockID": "0"
                }
            ]
        }
    }
}

Debugger

Enabling the --debugger flag sets a pbd trace to run after each op during the callback hook.

query

Query the system to obtain the system desc file (optionally store it to disk) Note: It's required to be on a system with silicon and to have a runtime enabled build -DTTMLIR_ENABLE_RUNTIME=ON.

ttrt query --help
ttrt query
ttrt query --quiet
ttrt query --save-artifacts
ttrt query --clean-artifacts
ttrt query --save-artifacts --log-file ttrt.log
ttrt query --save-artifacts --artifact-dir /path/to/some/dir
ttrt query --result-file result.json

perf

Run performance mode of a binary file or a directory of binary files Note: It's required to be on a system with silicon and to have a runtime enabled build -DTTMLIR_ENABLE_RUNTIME=ON. Also need perf enabled build -DTT_RUNTIME_ENABLE_PERF_TRACE=ON. Note: You can collect host only related performance data via --host-only flag. By default, host and device side performance data are both collected. If the saving artifacts flag is provided, perf mode will dump the following files in the artifacts directory

ops_perf_results.csv : compiled op performance results

OP CODE,OP TYPE,GLOBAL CALL COUNT,DEVICE ID,ATTRIBUTES,MATH FIDELITY,CORE COUNT,PARALLELIZATION STRATEGY,HOST START TS,HOST END TS,HOST DURATION [ns],DEVICE FW START CYCLE,DEVICE FW END CYCLE,OP TO OP LATENCY [ns],OP TO OP LATENCY BR/NRISC START [ns],DEVICE FW DURATION [ns],DEVICE KERNEL DURATION [ns],DEVICE KERNEL DURATION DM START [ns],DEVICE KERNEL DURATION PER CORE MIN [ns],DEVICE KERNEL DURATION PER CORE MAX [ns],DEVICE KERNEL DURATION PER CORE AVG [ns],DEVICE KERNEL FIRST TO LAST START [ns],DEVICE BRISC KERNEL DURATION [ns],DEVICE NCRISC KERNEL DURATION [ns],DEVICE TRISC0 KERNEL DURATION [ns],DEVICE TRISC1 KERNEL DURATION [ns],DEVICE TRISC2 KERNEL DURATION [ns],DEVICE ERISC KERNEL DURATION [ns],DEVICE COMPUTE CB WAIT FRONT [ns],DEVICE COMPUTE CB RESERVE BACK [ns],DISPATCH TOTAL CQ CMD OP TIME [ns],DISPATCH GO SEND WAIT TIME [ns],INPUT_0_W,INPUT_0_Z,INPUT_0_Y,INPUT_0_X,INPUT_0_LAYOUT,INPUT_0_DATATYPE,INPUT_0_MEMORY,OUTPUT_0_W,OUTPUT_0_Z,OUTPUT_0_Y,OUTPUT_0_X,OUTPUT_0_LAYOUT,OUTPUT_0_DATATYPE,OUTPUT_0_MEMORY,METAL TRACE ID,METAL TRACE REPLAY SESSION ID,COMPUTE KERNEL SOURCE,COMPUTE KERNEL HASH,DATA MOVEMENT KERNEL SOURCE,DATA MOVEMENT KERNEL HASH,BRISC MAX KERNEL SIZE [B],NCRISC MAX KERNEL SIZE [B],TRISC 0 MAX KERNEL SIZE [B],TRISC 1 MAX KERNEL SIZE [B],TRISC 2 MAX KERNEL SIZE [B],ERISC MAX KERNEL SIZE [B],PM IDEAL [ns],PM COMPUTE [ns],PM BANDWIDTH [ns],PM REQ I BW,PM REQ O BW,PM FPU UTIL (%),NOC UTIL (%),DRAM BW UTIL (%),NPE CONG IMPACT (%),LOC,CONST_EVAL_OP,PROGRAM_METADATA
UnaryDeviceOperation,tt_dnn_device,1024,0,{'bfp8_pack_precise': 'false'; 'fp32_dest_acc_en': 'true'; 'op_chain': '{UnaryWithParam(op_type=UnaryOpType::TAN;param={})}'; 'output_dtype': 'DataType::FLOAT32'; 'output_memory_config': 'MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED;buffer_type=BufferType::DRAM;shard_spec=std::nullopt;nd_shard_spec=std::nullopt;created_with_nd_shard_spec=0)'; 'preserve_fp32_precision': 'true'},HiFi4,16,,4556959654,4557518500,558846,9815181939513,9815181946491,0,0,6978,6314,6126,4982,6216,5652,335,6087,1375,1656,4957,465,,,,,,1,1,128,128,TILE,FLOAT32,DEV_1_DRAM_INTERLEAVED,1,1,128,128,TILE,FLOAT32,DEV_1_DRAM_INTERLEAVED,,,['ttnn/cpp/ttnn/operations/eltwise/unary/device/kernels/compute//eltwise_sfpu.cpp'],['eltwise_sfpu/3265258334475852953/'],['ttnn/cpp/ttnn/operations/eltwise/unary/device/kernels/dataflow/reader_unary_interleaved_start_id.cpp'; 'ttnn/cpp/ttnn/operations/eltwise/unary/device/kernels/dataflow/writer_unary_interleaved_start_id.cpp'],['reader_unary_interleaved_start_id/1146610629329498539/'; 'writer_unary_interleaved_start_id/1727642094059197364/'],708,736,1344,1568,1380,0,1,1,1,[],[],0.016,,,,"loc(""/home/$USER/tt-mlir/test/python/golden/test_ttir_ops.py:74:id(0)"")",false,"{'loop_number': 0, 'program_index': 0, 'disable_eth_dispatch': False, 'enable_program_cache': False, 'dump_device_rate': 1000}"

profile_log_device.csv : dump of all device side profiled results
tracy_ops_data.csv : op data results dumped in a readable format
tracy_ops_times.csv : op time results dumped in a readable format
tracy_profile_log_host.tracy : tracy profiled results file, this file can be fed into the tracy GUI

check

Check a binary file or a directory of binary files against a system desc (by default, uses the host machine) Note: It's required to be on a system with silicon and to have a runtime enabled build -DTTMLIR_ENABLE_RUNTIME=ON.

ttrt check --help
ttrt check out.ttnn
ttrt check out.ttnn --system-desc /path/to/system_desc.ttsys
ttrt check out.ttnn --clean-artifacts
ttrt check out.ttnn --save-artifacts
ttrt check out.ttnn --log-file ttrt.log
ttrt check /dir/of/flatbuffers --system-desc /dir/of/system_desc
ttrt check --save-artifacts --artifact-dir /path/to/some/dir out.ttnn
ttrt check out.ttnn --result-file result.json

emitpy

Run a python file or a directory of python files. Optionally provide a binary file or directory of binary files for output tensor comparison. Note: It's required to be on a system with silicon and to have a runtime enabled build -DTTMLIR_ENABLE_RUNTIME=ON.

ttrt emitpy --help
ttrt emitpy out.py
ttrt emitpy out.py --clean-artifacts
ttrt emitpy out.py --save-artifacts
ttrt emitpy out.py --loops 10
ttrt emitpy --program-index all out.py
ttrt emitpy --program-index 0 out.py
ttrt emitpy /dir/of/emitpy_modules
ttrt emitpy /dir/of/emitpy_modules --loops 10
ttrt emitpy /dir/of/emitpy_modules --log-file ttrt.log
ttrt emitpy /dir/of/emitpy_modules --flatbuffer /path/to/flatbuffer
ttrt emitpy out.py --save-artifacts --artifact-dir /path/to/some/dir
ttrt emitpy out.py --result-file result.json
ttrt emitpy out.py --print-input-output-tensors
ttrt emitpy out.py --memory --save-artifacts

For info on generating EmitPy tests through ttmlir-opt and ttmlir-translate, see EmitPy. For info on generating EmitPy tests through ttir-builder, see ttir-builder.

emitpy results

The emitpy api saves a emitpy_results.json file that records information about the run including any errors that were thrown and location of other saved data.

[
  {
    "file_path": "ttir-builder-artifacts/emitpy/test_binary_ops[add-emitpy-f32-128x128]_ttnn.mlir.py",
    "result": "pass",
    "exception": "",
    "log_file": "ttrt.log",
    "artifacts": "/home/$USER/tt-mlir/ttrt-artifacts",
    "program_index": "all"
  }
]

emitc

Run a .so file or a directory of .so files. Optionally provide a binary file or directory of binary files for output tensor comparison. Note: It's required to be on a system with silicon and to have a runtime enabled build -DTTMLIR_ENABLE_RUNTIME=ON.

ttrt emitc --help
ttrt emitc out.py
ttrt emitc out.py --clean-artifacts
ttrt emitc out.py --save-artifacts
ttrt emitc out.py --loops 10
ttrt emitc --program-index all out.py
ttrt emitc --program-index 0 out.py
ttrt emitc /dir/of/emitc_modules
ttrt emitc /dir/of/emitc_modules --loops 10
ttrt emitc /dir/of/emitc_modules --log-file ttrt.log
ttrt emitc /dir/of/emitc_modules --flatbuffer /path/to/flatbuffer
ttrt emitc out.py --save-artifacts --artifact-dir /path/to/some/dir
ttrt emitc out.py --result-file result.json
ttrt emitc out.py --print-input-output-tensors
ttrt emitc out.py --memory --save-artifacts

For info on generating EmitC tests through ttnn-standalone, see EmitC testing documentation. For info on generating EmitC tests through ttir-builder, see ttir-builder documentation.

emitc results

The emitc api saves a emitc_results.json file that records information about the run including any errors that were thrown and location of other saved data.

[
  {
    "file_path": "ttir-builder-artifacts/emitc/test_reciprocal[emitc-f32-128x128]_ttnn.mlir.so",
    "result": "pass",
    "exception": "",
    "log_file": "ttrt.log",
    "artifacts": "/home/$USER/tt-mlir/ttrt-artifacts",
    "program_index": "all"
  }
]

gdb

You can relaunch ttrt inside of gdb which can be useful for debugging C++ runtime components.

ttrt --gdb run ...
ttrt --gdb perf ...

Using as a python package

The other way to use the APIs under ttrt is importing it as a library. This allows the user to use it in custom scripts.

Import ttrt as a python package

from ttrt.common.api import API

Setup API and register all features

API.initialize_apis()

Setup arguments

You can specify certain arguments to pass to each API, or use the default arguments provided

Args

This can be a dictionary of values to set inside your API instance. These are the same options as found via the command line. You can get the total list of support arguments via the ttrt --help command. Any argument not provided will be set to the default.

custom_args = {}
custom_args["--clean-artifacts"] = True
query_instance = API.Query(args=custom_args)

Logging

You can specify a specific logging module you want to set inside your API instance. The rationale behind this is to support different instances of different APIs, all being able to be logged to a different file. You can also customize the level of detail your log file contains.

from ttrt.common.util import Logger
import os

os.environ["LOGGER_LEVEL"] = "DEBUG"
log_file_name = "some_file_name.log"
custom_logger = Logger(log_file_name)
read_instance = API.Read(logger=custom_logger)

Artifacts

You can specify a specific artifacts directory to store all the generate metadata during the execution of any API run. This allows you to specify different artifact directories if you wish for different instances of APIs.

from ttrt.common.util import Artifacts

log_file_name = "some_file_name.log"
artifacts_folder_path = "/opt/folder"
custom_logger = Logger(log_file_name)
custom_artifacts = Artifacts(logger=custom_logger, artifacts_folder_path=artifacts_folder_path)
run_instance = API.Run(artifacts=custom_artifacts)

Execute API

Once all the arguments are setup, you can run your API instance with all your provided arguments. Note, APIs are stateless. Thus, subsequent calls to the same API instance will not preserve previous call artifacts. You can generate a new artifacts directory for subsequent runs if you wish to call the APIs multiple times, for example.

result_code, results = query_instance()
result_code, results = read_instance()
result_code, results = run_instance()

Putting it all together

You can do interesting stuff when combining all the above features into your python script

from ttrt.common.api import API
from ttrt.common.util import Logger
from ttrt.common.util import Artifacts

API.initialize_apis()

custom_args = {}
custom_args["--clean-artifacts"] = True
custom_args["--save-artifacts"] = True
custom_args["--loops"] = 10
custom_args["--init"] = "randn"
custom_args["binary"] = "/path/to/subtract.ttnn"

log_file_name = "some_file_name.log"
custom_logger = Logger(log_file_name)

artifacts_folder_path = "/opt/folder"
custom_artifacts = Artifacts(logger=custom_logger, artifacts_folder_path=artifacts_folder_path)

run_instance = API.Run(args=custom_args, logger=custom_logger, artifacts=custom_artifacts)
result_code, results = run_instance()

Runtime integration

The full set of ttrt.runtime exposed APIs and types can be found in runtime/python/runtime/runtime.cpp, however only the ones intended to be used for runtime customization through callback hooks are outlined here.

Callback hooks

MLIR Runtime exposes a feature to register python callback functions. Any two python fuctions can be provided - the first function will be executed before every op in MLIR Runtime, the second after every op. The following steps describe how to extend your application to register python functions. Callback functions are already implemented by default for pbd debugger implementation and gathering memory and golden check data as outlined in the run API section.

1. Pybind DebugHooks C++ class, specifically tt::runtime::debug::Hooks::get. See runtime/python/runtime/runtime.cpp for an example of how ttrt pybinds it.

tt::runtime::debug::Hooks
tt::runtime::debug::Hooks::get

2. Register callback functions in your python script. The following is registering the two callback functions written in tools/ttrt/common/callback.py. The Debug Hooks get function has been pybinded to ttrt.runtime.DebugHooks.get

import ttrt.runtime

callback_env = ttrt.runtime.DebugHooks.get(pre_op_callback_runtime_config, post_op_callback_runtime_config)

3. The callback function has a particular function signature, which looks like the following

def pre_op_callback_runtime_config(binary, program_context, op_context):

binary: reference to the binary you are currently running, ttrt.binary Binary object program_context: reference to the program currently running, ttrt.runtime ProgramContext object op_context: reference to the op that is currently running, ttrt.runtime OpContext object

4. Each of these parameters has certain runtime APIs exposed which can only be called within the callback functions since they rely on the op_context variable that is only available from runtime during callbacks.

import ttrt.runtime

loc = ttrt.runtime.get_op_loc_info(op_context) : get the location of the op as a string which is used as the key when indexing the golden tensors stored in the flatbuffer
op_debug_str = ttrt.runtime.get_op_debug_str(op_context) : get the op debug str (contains op metadata inculding op type, attributes, input tensor shapes and dtypes, memref with layout and buffer type, and loc)
op_golden_tensor = ttrt.runtime.get_debug_info_golden(binary, loc) : get the golden tensor from the binary as a ttrt.binary GoldenTensor object
op_output_tensor = ttrt.runtime.get_op_output_tensor(op_context, program_context) : get the currently running output tensor from device as a ttrt.runtime Tensor object, if this is called in a preOp function or the op doesn't output a tensor, an empty tensor will be returned.

Note: ttrt is not needed to implement this callback feature. It aims to provide an example of how this callback feature can be implemented for golden application.

FAQ

Flatbuffer version does not match ttrt version!

ttrt and flatbuffer have strict versioning that is checked during ttrt execution. You will have to generate a flatbuffer using the same version of ttrt (or vice versa). This mean you might have to build on the same branch on which the flatbuffer was generated or regenerate the flatbuffer using your current build.

System desc does not match flatbuffer!

Flatbuffers are compiled using a specific system desc (or default values if no system desc is provided). During runtime, the flatbuffer system desc is checked against the current system to ensure the system being run on supports the flatbuffer that was compiled. If you get this error, you will have to regenerate the flatbuffer using the system you want to run on. See generate a flatbuffer file from compiler section on how to do this.

I just want to test and push my commit! What do I do!

Follow these steps (on n150, n300, and llmbox)

1. Build ttmlir (sample instructions - subject to change)

source env/activate
cmake -G Ninja -B build -DCMAKE_BUILD_TYPE=Release -DCMAKE_C_COMPILER=clang-17 -DCMAKE_CXX_COMPILER=clang++-17 -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DTTMLIR_ENABLE_RUNTIME=ON -DTT_RUNTIME_ENABLE_PERF_TRACE=ON
cmake --build build

2. Query system

ttrt query --save-artifacts

3. Export system desc file

export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys (path dumped in previous command)

4. Generate test cases

cmake --build build -- check-ttmlir

5. Run test cases

ttrt run build/test/ttmlir/Silicon

6. (Optional) Run perf test cases

ttrt perf build/test/ttmlir/Silicon

TTRT yields an ambiguous segmentation fault!

The ttrt toolchain has specific behaviors and requirements that can lead to build and runtime issues, particularly when dealing with version mismatches or out-of-sync dependencies.

Version Mismatch Due to Local Commits

The ttrt toolchain verifies whether the current system configuration matches the model’s compilation environment. This verification involves tracking the number of commits since the last synchronization. When local commits are made in your branch, it may trigger a version mismatch between the compiled model and the current environment. This mismatch may not be handled properly by the runtime (rt), leading to potential issues.

To resolve issues stemming from these synchronization problems, follow this workflow:

1. Incremental build

# make some changes
# commit
cmake --build build
# note you need to generate system_desc and flatbuffer again once you do this

This incremental build should be sufficient. If it does not resolve the error, please file an issue and proceed with the following steps for now.

2. Clear the existing build and dependencies:

rm -rf build third_party/tt-metal

This ensures that all previous build artifacts and dependencies are removed, preventing conflicts or stale files from affecting the new build.

3. Rebuild from scratch: After clearing the build directories, rebuild the project from the ground up. This ensures that the build process incorporates all the necessary components without any remnants of previous builds. Build Instructions

4. Switch build configurations: If switching from a Debug to a Release build (or vice versa), ensure that you clean the build environment before transitioning. This avoids inconsistencies between build configurations and potential issues with optimization levels or debugging symbols.

5. Re-acquire the IRD: By relinquishing and re-acquiring the IRD, you ensure that the correct toolchain is used for the new build. This step ensures synchronization between the model and the toolchain.

6. Enable Debug Logging for tt-metal: To gain more insight into potential issues, enable debugging by setting the TT_METAL_LOGGER_LEVEL to DEBUG. This will provide detailed logs, which can help in troubleshooting build or runtime issues.

export TT_METAL_LOGGER_LEVEL=DEBUG

ttir-builder

ttir-builder is a tool for creating TTIR operations. It provides support for MLIR modules to be generated from user-constructed ops, lowered into TTNN or TTMetal backends, and finally translated into executable flatbuffers. Or you can do all three at once!

Building

1. Build tt-mlir
2. Build ttrt
3. Generate ttsys file from the system you want to compile for using ttrt. This will create a ttrt-artifacts folder containing a system_desc.ttsys file.

ttrt query --save-artifacts

4. Export this file in your environment using export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys. builder.base.builder_utils uses the system_desc.ttsys file as it runs a pass over an MLIR module to the TTNN or TTMetal backend.

Getting started

TTIRBuilder is a builder class providing the API for creating TTIR ops. The python package builder contains everything needed to create ops through a TTIRBuilder object. builder.base.builder_utils contains the APIs for wrapping op-creating-functions into MLIR modules and flatbuffers files.

from builder.ttir.ttir_builder import TTIRBuilder
from builder.base.builder_utils import compile_ttir_to_flatbuffer

Creating a TTIR module

build_ttir_module defines an MLIR module specified as a python function. It wraps fn in a MLIR FuncOp then wraps that in an MLIR module, and finally ties arguments of that FuncOp to test function inputs. It will instantiate and pass a TTIRBuilder object as the last argument of fn. Each op returns an OpView type which is a type of Operand that can be passed into another builder op as an input.

def build_ttir_module(
    fn: Callable,
    inputs_shapes: List[Shape],
    inputs_types: Optional[List[Union[torch.dtype, TypeInfo]]] = None,
    mesh_name: str = "mesh",
    mesh_dict: OrderedDict[str, int] = OrderedDict([("x", 1), ("y", 1)]),
    module_dump: bool = False,
    base: Optional[str] = None,
    output_root: str = ".",
) -> Tuple[Module, TTIRBuilder]:

Example

from builder.base.builder import Operand
from builder.ttir.ttir_builder import TTIRBuilder
from builder.base.builder_utils import build_ttir_module

shapes = [(32, 32), (32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, in2: Operand, builder: TTIRBuilder):
    add_0 = builder.add(in0, in1)
    multiply_1 = builder.multiply(in1, add_0)
    return builder.multiply(multiply_1, in2)

module, builder = build_ttir_module(model, shapes)

Returns

An MLIR module containing an MLIR op graph defined by fn and the TTIRBuilder object used to create it

module {
  func.func @model(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<32x32xf32>) -> tensor<32x32xf32> {
    %0 = ttir.empty() : tensor<32x32xf32>
    %1 = "ttir.add"(%arg0, %arg1, %0) : (tensor<32x32xf32>, tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32>
    %2 = ttir.empty() : tensor<32x32xf32>
    %3 = "ttir.multiply"(%arg1, %1, %2) : (tensor<32x32xf32>, tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32>
    %4 = ttir.empty() : tensor<32x32xf32>
    %5 = "ttir.multiply"(%3, %arg2, %4) : (tensor<32x32xf32>, tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32>
    return %5 : tensor<32x32xf32>
  }
}

Running a pipeline

run_ttir_pipeline runs a pass on the TTIR module to lower it into a backend, using pipeline_fn. You can pass pipeline_fn in as one of the following: ttir_to_ttnn_backend_pipeline, ttir_to_ttmetal_backend_pipeline (both found in ttmlir.passes), or a custom pipeline built with create_custom_pipeline_fn. The default if none is provided is the TTNN pipeline.

def run_ttir_pipeline(
    module,
    pipeline_fn: Callable,
    pipeline_options: List[str] = [],
    dump_to_file: bool = True,
    output_file_name: str = "test.mlir",
    system_desc_path: Optional[str] = None,
    mesh_dict: OrderedDict[str, int] = None,
    argument_types_string: Optional[str] = None,
)

TTNN example

Let's expand on our previous example

from ttmlir.passes import ttir_to_ttnn_backend_pipeline
from builder.base.builder import Operand
from builder.ttir.ttir_builder import TTIRBuilder
from builder.base.builder_utils import build_ttir_module, run_ttir_pipeline

shapes = [(32, 32), (32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, in2: Operand, builder: TTIRBuilder):
    add_0 = builder.add(in0, in1)
    multiply_1 = builder.multiply(in1, add_0)
    return builder.multiply(multiply_1, in2)

module, builder = build_ttir_module(model, shapes)
ttnn_module = run_ttir_pipeline(module, ttir_to_ttnn_backend_pipeline)

Returns

An MLIR module lowered into TTNN

#dram = #ttnn.buffer_type<dram>
#system_desc = #ttcore.system_desc<[{role = host, target_triple = "x86_64-pc-linux"}], [{arch = <wormhole_b0>, grid = 8x8, coord_translation_offsets = 18x18, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32, l1_unreserved_base = 97248, erisc_l1_unreserved_base = 69632, dram_unreserved_base = 32, dram_unreserved_end = 1073158336, physical_helper_cores = {dram = [ 0x0,  0x1,  0x2,  0x3,  0x4,  0x5,  0x6,  0x7,  0x8,  0x9,  0x10,  0x11] eth_inactive = [ 16x18,  16x19,  16x20,  16x21,  16x22,  16x23,  16x24,  16x25,  17x19,  17x20,  17x22,  17x23,  17x24]}, supported_data_types = [<f32>, <f16>, <bf16>, <bfp_f8>, <bfp_bf8>, <bfp_f4>, <bfp_bf4>, <bfp_f2>, <bfp_bf2>, <u32>, <u16>, <u8>, <si32>], supported_tile_sizes = [ 4x16,  16x16,  32x16,  4x32,  16x32,  32x32], num_cbs = 32, num_compute_threads = 1, num_datamovement_threads = 2}], [0], [3 : i32], [ 0x0x0x0]>
#ttnn_layout = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!ttcore.tile<32x32, f32>, #dram>, <interleaved>>
module {
  ttcore.device_module {
    builtin.module attributes {ttcore.system_desc = #system_desc} {
      ttcore.device @default_device = <workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d0, d1)>, l1Map = (d0, d1, d2)[s0] -> (0, d0, d1, d2 + s0), dramMap = (d0, d1, d2)[s0, s1, s2, s3, s4, s5] -> (0, 0, (((d0 * s1) * (s2 * s3) + d1 * (s2 * s3) + d2) floordiv s4) mod 12, ((d0 * s1) * (s2 * s3) + d1 * (s2 * s3) + d2) floordiv (s4 * 12) + ((d0 * s1) * (s2 * s3) + d1 * (s2 * s3) + d2) mod s4 + s5), meshShape = , chipIds = [0]>
      func.func @model(%arg0: tensor<32x32xf32, #ttnn_layout>, %arg1: tensor<32x32xf32, #ttnn_layout>, %arg2: tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout> {
        %0 = "ttnn.abs"(%arg0) : (tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout>
        "ttnn.deallocate"(%arg0) <{force = false}> : (tensor<32x32xf32, #ttnn_layout>) -> ()
        %1 = "ttnn.multiply"(%arg1, %0) : (tensor<32x32xf32, #ttnn_layout>, tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout>
        "ttnn.deallocate"(%0) <{force = false}> : (tensor<32x32xf32, #ttnn_layout>) -> ()
        "ttnn.deallocate"(%arg1) <{force = false}> : (tensor<32x32xf32, #ttnn_layout>) -> ()
        %2 = "ttnn.multiply"(%1, %arg2) : (tensor<32x32xf32, #ttnn_layout>, tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout>
        "ttnn.deallocate"(%1) <{force = false}> : (tensor<32x32xf32, #ttnn_layout>) -> ()
        "ttnn.deallocate"(%arg2) <{force = false}> : (tensor<32x32xf32, #ttnn_layout>) -> ()
        return %2 : tensor<32x32xf32, #ttnn_layout>
      }
    }
  }
}

TTMetal example

Let's use the same code for TTMetal that was used in the TTNN example but change the pipeline_fn to ttir_to_ttmetal_backend_pipeline. Only one or the other can be run on a module since run_ttir_pipeline modifies the module in place. Note that while all TTIR ops supported by builder can be lowered to TTNN, not all can be lowered to TTMetal yet. Adding documentation to specify what ops can be lowered to TTMetal is in the works.

from ttmlir.passes import ttir_to_ttmetal_backend_pipeline
ttmetal_module = run_ttir_pipeline(module, ttir_to_ttmetal_backend_pipeline)

Returns

An MLIR module lowered into TTMetal

#l1 = #ttcore.memory_space<l1>
#system_desc = #ttcore.system_desc<[{role = host, target_triple = "x86_64-pc-linux-gnu"}], [{arch = <wormhole_b0>, grid = 8x8, coord_translation_offsets = 18x18, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32, l1_unreserved_base = 1024, erisc_l1_unreserved_base = 1024, dram_unreserved_base = 1024, dram_unreserved_end = 1073741824, physical_helper_cores = {dram = [ 8x0,  9x0,  10x0,  8x1,  9x1,  10x1,  8x2,  9x2,  10x2,  8x3,  9x3,  10x3]}, supported_data_types = [<f32>, <f16>, <bf16>, <bfp_f8>, <bfp_bf8>, <bfp_f4>, <bfp_bf4>, <bfp_f2>, <bfp_bf2>, <u32>, <u16>, <u8>, <si32>], supported_tile_sizes = [ 4x16,  16x16,  32x16,  4x32,  16x32,  32x32], num_cbs = 32, num_compute_threads = 1, num_datamovement_threads = 2}], [0], [3 : i32], [ 0x0x0x0]>
module {
  ttcore.device_module {
    builtin.module attributes {ttcore.system_desc = #system_desc} {
      ttcore.device @default_device = <workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d0, d1)>, l1Map = (d0, d1, d2)[s0] -> (0, d0, d1, d2 + s0), dramMap = (d0, d1, d2)[s0, s1, s2, s3, s4, s5] -> (0, 0, (((d0 * s1) * (s2 * s3) + d1 * (s2 * s3) + d2) floordiv s4) mod 12, ((d0 * s1) * (s2 * s3) + d1 * (s2 * s3) + d2) floordiv (s4 * 12) + ((d0 * s1) * (s2 * s3) + d1 * (s2 * s3) + d2) mod s4 + s5), meshShape = , chipIds = [0]>
      func.func @model(%arg0: memref<32x32xf32>, %arg1: memref<32x32xf32>, %arg2: memref<32x32xf32>) -> memref<32x32xf32> {
        %0 = "ttmetal.create_buffer"() <{address = 9216 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        %1 = "ttmetal.create_buffer"() <{address = 1024 : i64}> : () -> memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>
        "ttmetal.enqueue_write_buffer"(%arg0, %1) : (memref<32x32xf32>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        "ttmetal.enqueue_program"(%1, %0, %1, %0) <{cb_ports = array<i64: 0, 1>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel0, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, noc0>, #ttmetal.compute_config<@compute_kernel1, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 2, 2>}> : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%1) : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        %2 = "ttmetal.create_buffer"() <{address = 1024 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        %3 = "ttmetal.create_buffer"() <{address = 5120 : i64}> : () -> memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>
        "ttmetal.enqueue_write_buffer"(%arg1, %3) : (memref<32x32xf32>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        "ttmetal.enqueue_program"(%3, %2, %3, %2) <{cb_ports = array<i64: 0, 1>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel2, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, noc0>, #ttmetal.compute_config<@compute_kernel3, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 2, 2>}> : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%3) : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        %4 = "ttmetal.create_buffer"() <{address = 13312 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        "ttmetal.enqueue_program"(%0, %2, %4, %0, %2, %4) <{cb_ports = array<i64: 0, 1, 2>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel4, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, noc0>, #ttmetal.noc_config<@datamovement_kernel5, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, noc1>, #ttmetal.compute_config<@compute_kernel6, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 3, 3>}> : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%0) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%2) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        %5 = "ttmetal.create_buffer"() <{address = 1024 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        %6 = "ttmetal.create_buffer"() <{address = 5120 : i64}> : () -> memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>
        "ttmetal.enqueue_write_buffer"(%arg1, %6) : (memref<32x32xf32>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        "ttmetal.enqueue_program"(%6, %5, %6, %5) <{cb_ports = array<i64: 0, 1>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel7, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, noc0>, #ttmetal.compute_config<@compute_kernel8, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 2, 2>}> : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%6) : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        %7 = "ttmetal.create_buffer"() <{address = 17408 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        "ttmetal.enqueue_program"(%5, %4, %7, %5, %4, %7) <{cb_ports = array<i64: 0, 1, 2>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel9, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, noc0>, #ttmetal.noc_config<@datamovement_kernel10, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, noc1>, #ttmetal.compute_config<@compute_kernel11, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 3, 3>}> : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%5) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%4) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        %8 = "ttmetal.create_buffer"() <{address = 9216 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        %9 = "ttmetal.create_buffer"() <{address = 1024 : i64}> : () -> memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>
        "ttmetal.enqueue_write_buffer"(%arg2, %9) : (memref<32x32xf32>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        "ttmetal.enqueue_program"(%9, %8, %9, %8) <{cb_ports = array<i64: 0, 1>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel12, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, noc0>, #ttmetal.compute_config<@compute_kernel13, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 2, 2>}> : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%9) : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        %10 = "ttmetal.create_buffer"() <{address = 5120 : i64}> : () -> memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>
        "ttmetal.enqueue_program"(%7, %8, %10, %7, %8, %10) <{cb_ports = array<i64: 0, 1, 2>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel14, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, noc0>, #ttmetal.noc_config<@datamovement_kernel15, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, noc1>, #ttmetal.compute_config<@compute_kernel16, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>, <cb_port[2]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 3, 3>}> : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%8) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%7) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        %alloc = memref.alloc() : memref<32x32xf32>
        %11 = "ttmetal.create_buffer"() <{address = 1024 : i64}> : () -> memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>
        "ttmetal.enqueue_program"(%10, %11, %10, %11) <{cb_ports = array<i64: 0, 1>, kernelConfigs = [#ttmetal.noc_config<@datamovement_kernel17, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, noc0>, #ttmetal.compute_config<@compute_kernel18, #ttmetal.core_range<0x0, 1x1>, #ttmetal.kernel_args< ct_args = [<cb_port[0]>, <cb_port[1]>]>, hifi4, false, false, [default]>], operandSegmentSizes = array<i32: 2, 2>}> : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>, memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        "ttmetal.deallocate_buffer"(%10) : (memref<1x1x1x1x!ttcore.tile<32x32, f32>, #ttcore.shard<4096x4096>, #l1>) -> ()
        "ttmetal.enqueue_read_buffer"(%11, %alloc) : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>, memref<32x32xf32>) -> ()
        "ttmetal.finish"() : () -> ()
        "ttmetal.deallocate_buffer"(%11) : (memref<1x1x32x32xf32, #ttcore.shard<128x4>, #l1>) -> ()
        return %alloc : memref<32x32xf32>
      }
      func.func private @datamovement_kernel0() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel1() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %2 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "tilize_init"(%1, %0, %2) : (!emitc.opaque<"::tt::CB">, i32, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "experimental::tilize_block"(%1, %2, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, i32, i32) -> ()
        emitc.call_opaque "cb_push_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel2() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel3() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %2 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "tilize_init"(%1, %0, %2) : (!emitc.opaque<"::tt::CB">, i32, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "experimental::tilize_block"(%1, %2, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, i32, i32) -> ()
        emitc.call_opaque "cb_push_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel4() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel5() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel6() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 0 : index}> : () -> !emitc.size_t
        %1 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        emitc.call_opaque "tile_regs_acquire"() : () -> ()
        %2 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %3 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        %4 = emitc.literal "get_compile_time_arg_val(2)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%3, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "binary_op_init_common"(%2, %3, %4) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "add_tiles_init"(%2, %3) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "add_tiles"(%2, %3, %0, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, !emitc.size_t, !emitc.size_t, !emitc.size_t) -> ()
        emitc.call_opaque "tile_regs_commit"() : () -> ()
        emitc.call_opaque "tile_regs_wait"() : () -> ()
        emitc.call_opaque "pack_tile"(%0, %4, %0) {template_args = [true]} : (!emitc.size_t, !emitc.opaque<"::tt::CB">, !emitc.size_t) -> ()
        emitc.call_opaque "tile_regs_release"() : () -> ()
        emitc.call_opaque "cb_push_back"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%3, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel7() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel8() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %2 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "tilize_init"(%1, %0, %2) : (!emitc.opaque<"::tt::CB">, i32, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "experimental::tilize_block"(%1, %2, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, i32, i32) -> ()
        emitc.call_opaque "cb_push_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel9() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel10() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel11() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 0 : index}> : () -> !emitc.size_t
        %1 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        emitc.call_opaque "tile_regs_acquire"() : () -> ()
        %2 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %3 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        %4 = emitc.literal "get_compile_time_arg_val(2)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%3, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "binary_op_init_common"(%2, %3, %4) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "mul_tiles_init"(%2, %3) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "mul_tiles"(%2, %3, %0, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, !emitc.size_t, !emitc.size_t, !emitc.size_t) -> ()
        emitc.call_opaque "tile_regs_commit"() : () -> ()
        emitc.call_opaque "tile_regs_wait"() : () -> ()
        emitc.call_opaque "pack_tile"(%0, %4, %0) {template_args = [true]} : (!emitc.size_t, !emitc.opaque<"::tt::CB">, !emitc.size_t) -> ()
        emitc.call_opaque "tile_regs_release"() : () -> ()
        emitc.call_opaque "cb_push_back"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%3, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel12() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel13() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %2 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "tilize_init"(%1, %0, %2) : (!emitc.opaque<"::tt::CB">, i32, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "experimental::tilize_block"(%1, %2, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, i32, i32) -> ()
        emitc.call_opaque "cb_push_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel14() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel15() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel16() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>, <arg_type = cb_port, operand_index = 2>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 0 : index}> : () -> !emitc.size_t
        %1 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        emitc.call_opaque "tile_regs_acquire"() : () -> ()
        %2 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %3 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        %4 = emitc.literal "get_compile_time_arg_val(2)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%3, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "binary_op_init_common"(%2, %3, %4) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "mul_tiles_init"(%2, %3) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "mul_tiles"(%2, %3, %0, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, !emitc.size_t, !emitc.size_t, !emitc.size_t) -> ()
        emitc.call_opaque "tile_regs_commit"() : () -> ()
        emitc.call_opaque "tile_regs_wait"() : () -> ()
        emitc.call_opaque "pack_tile"(%0, %4, %0) {template_args = [true]} : (!emitc.size_t, !emitc.opaque<"::tt::CB">, !emitc.size_t) -> ()
        emitc.call_opaque "tile_regs_release"() : () -> ()
        emitc.call_opaque "cb_push_back"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%3, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%4, %1) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @datamovement_kernel17() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<noc>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_push_back"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
      func.func private @compute_kernel18() attributes {ttkernel.arg_spec = #ttkernel.arg_spec< ct_args = [<arg_type = cb_port, operand_index = 0>, <arg_type = cb_port, operand_index = 1>]>, ttkernel.thread = #ttkernel.thread<compute>} {
        %0 = "emitc.constant"() <{value = 1 : i32}> : () -> i32
        %1 = emitc.literal "get_compile_time_arg_val(0)" : !emitc.opaque<"::tt::CB">
        %2 = emitc.literal "get_compile_time_arg_val(1)" : !emitc.opaque<"::tt::CB">
        emitc.call_opaque "cb_reserve_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "untilize_init"(%1) : (!emitc.opaque<"::tt::CB">) -> ()
        emitc.call_opaque "experimental::untilize_block"(%1, %2, %0, %0) : (!emitc.opaque<"::tt::CB">, !emitc.opaque<"::tt::CB">, i32, i32) -> ()
        emitc.call_opaque "cb_push_back"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_wait_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%1, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        emitc.call_opaque "cb_pop_front"(%2, %0) : (!emitc.opaque<"::tt::CB">, i32) -> ()
        return
      }
    }
  }
}

Compiling into flatbuffer

compile_ttir_to_flatbuffer compiles a TTIRBuilder function fn straight to flatbuffer. This decorator is mainly a wrapper around the following functions, with each next function called on the output of the last: build_ttir_module, run_ttir_pipeline, and ttnn_to_flatbuffer_file, ttmetal_to_flatbuffer_file, ttir_to_emitpy_pipeline, or ttir_to_ttnn_emitc_pipeline as dictated by the target parameter.

def compile_ttir_to_flatbuffer(
    fn: Callable,
    inputs_shapes: List[Shape],
    inputs_types: Optional[List[Union[torch.dtype, TypeInfo]]] = None,
    system_desc_path: Optional[str] = None,
    test_base: str = "test",
    output_root: str = ".",
    target: Literal["ttnn", "ttmetal", "emitc", "emitpy"] = "ttnn",
    mesh_name: str = "mesh",
    mesh_dict: OrderedDict[str, int] = OrderedDict([("x", 1), ("y", 1)]),
    module_dump: bool = True,
    argument_types_string: Optional[str] = None,
    custom_pipeline: Optional[Union[Callable, str]] = None,
    pipeline_options: List[str] = [],
    print_ir: Union[bool, str] = False,
) -> str:

The executable flatbuffer is written to a file, compile_ttir_to_flatbuffer returns the file address of that flatbuffer.

TTNN example

Let's use our previous model function.

from builder.base.builder import Operand
from builder.ttir.ttir_builder import TTIRBuilder
from builder.base.builder_utils import compile_ttir_to_flatbuffer

shapes = [(32, 32), (32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, in2: Operand, builder: TTIRBuilder):
    add_0 = builder.add(in0, in1)
    multiply_1 = builder.multiply(in1, add_0)
    return builder.multiply(multiply_1, in2)

compile_ttir_to_flatbuffer(
    model,
    shapes,
    target="ttnn",
)

TTMetal example

Let's once again use the same code for TTMetal that was used in the TTNN example but change the target to "ttmetal". Just as with run_ttir_pipeline, only one or the other can be run on a module since compile_ttir_to_flatbuffer modifies the module in place.

compile_ttir_to_flatbuffer(
    model,
    shapes,
    target="ttmetal",
)

Integrating with other tt-mlir tools

Alternatives for file creation

1. The ttmlir-opt tool runs a compiler pass on an .mlir file.
2. The ttmlir-translate can generate a flatbuffer from an .mlir file.
3. llvm-lit can also be used to generate a flatbuffer from an existing .mlir file.

Running models

ttrt

ttrt is intended to be a swiss army knife for working with flatbuffers.

tt-explorer

tt-explorer is a visualizer tool for ttmlir-powered compiler results.

ttnn-standalone

ttnn-standalone is a post-compile tuning/debugging tool.

llvm-lit

llvm-lit can also be used for MLIR testing.

Golden mode

Golden dataclass

TTIRBuilder provides support to code golden tensors into flatbuffers which will be used for comparison with TT device output in ttrt runtime. Golden is the dataclass used to store information about a golden tensor. Each TTIR op should have a matching PyTorch op (or golden function built from PyTorch ops) which should perform exactly the same operation, generating the same outputs given the same inputs. You can use TTIRBuilder helper functions to store input, intermediate, and output tensors within the flatbuffer. Input and output goldens are mapped with keys "input_" and "output_" followed by a tensor index: input_0. Intermediate output tensors are mapped to the location of the respective op creation.

GoldenCheckLevel Enum

TTIRBuilder stores an instance of the class GoldenCheckLevel(Enum) that dictates golden handling. It defaults to GoldenCheckLevel.OP_LEVEL. The exception is that TTIRBuilder CCL ops force the golden level to be set to GRAPH_LEVEL.

DISABLED : do not store goldens
OP_LEVEL : check every single op level goldens
GRAPH_LEVEL : check graph level goldens only

Check and set GoldenCheckLevel with TTIRBuilder APIs.

from builder.base.builder import Operand, GoldenCheckLevel
from builder.ttir.ttir_builder import TTIRBuilder

def model(in0: Operand, in1: Operand, in2: Operand, builder: TTIRBuilder):
    builder.golden_check_level = GoldenCheckLevel.GRAPH_LEVEL
    add_0 = builder.add(in0, in1)
    multiply_1 = builder.multiply(in1, add_0)
    return builder.multiply(multiply_1, in2)

Getting golden data

Unless otherwise specified in the GoldenCheckLevel, all input and output tensors will generate and store a golden in TTIRBuilder as a Golden type.

The TTIRBuilder API get_golden_map(self) is used to export golden data for flatbuffer construction. It returns a dictionary of golden tensor names and GoldenTensor objects.

To get info from a GoldenTensor object, use the attributes supported by ttmlir.passes: name, shape, strides, dtype, data.

from ttmlir.passes import GoldenTensor
from builder.ttir.ttir_builder import TTIRBuilder

shapes = [(32, 32), (32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, in2: Operand, builder: TTIRBuilder):
    add_0 = builder.add(in0, in1)
    builder.print_goldens()
    print(builder.get_golden_map())
    return add0

Golden tensor:
tensor([[ 4.0450e+00,  1.4274e+00,  5.9156e-01,  ..., -5.9834e-01,
         -1.1830e-01,  1.2837e-01],
        [ 2.3788e+00,  2.9242e-03, -5.2838e-02,  ...,  1.8294e+00,
          5.0348e+00,  9.7179e-01],
        [ 1.5168e-02,  1.0577e-01, -3.0682e-01,  ...,  6.7212e-01,
          9.4523e-02,  5.3765e+00],
        ...,
        [ 1.4241e-01,  1.1838e+00, -1.0601e+00,  ...,  4.9099e-01,
          4.2267e+00,  4.0610e-01],
        [ 5.6630e-01, -1.3068e-01, -1.7771e-01,  ...,  2.3862e+00,
          3.9376e-01,  7.3140e-01],
        [ 4.2420e+00,  1.7006e-01, -3.4861e-01,  ...,  1.1471e-01,
          1.6189e+00, -6.9106e-01]])
{'input_0': <ttmlir._mlir_libs._ttmlir.passes.GoldenTensor object at 0x7f77c70fa0d0>, 'output_0': <ttmlir._mlir_libs._ttmlir.passes.GoldenTensor object at 0x7f77c6fc9590>}

Setting golden data

Use TTIRBuilder API set_graph_input_output to set your own input and output golden tensors using PyTorch tensors. Keep in mind that this also sets graph inputs and outputs. There are some functions for which setting custom input tensors is required to pass PCC accuracy checks: ttir.tan, ttir.log, ttir.log1p. See example implementation and explanation in test/python/golden/test_ttir_ops.py.

set_graph_input_output(
        self,
        inputs: List[torch.Tensor],
        outputs: Optional[List[torch.Tensor]] = None,
        override: bool = False,
    )

import torch

input_0 = torch.ones((32, 32))
output_0 = torch.zeros((32, 32))
builder.set_graph_input_output([input_0], [output_0], override=True)

Running flatbuffer with golden data in ttrt

Running flatbuffers in ttrt requires building and setting up the environment. Run these commands before creating MLIR modules or flatbuffers so the system description in the flatbuffers match your device.

cmake --build build
ttrt query --save-artifacts
export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys

Set environment variable TTRT_LOGGER_LEVEL to DEBUG so ttrt logs golden comparison results and prints graph level golden tensors.

export TTRT_LOGGER_LEVEL=DEBUG

Finally run ttrt. Our example flatbuffer file (since we didn't specify otherwise) defaulted to file path ./builder-artifacts/ttir-builder/test_ttnn/test_ttnn.mlir.ttnn. --log-file ttrt.log and --save-golden-tensors are both optional flags. They ensure that all golden data produced by the ttrt run gets written to files.

ttrt run builder-artifacts/ttir-builder/test_ttnn/test_ttnn.mlir.ttnn --log-file ttrt.log --save-golden-tensors

Golden callbacks

The ttrt documentation contains a section on the callback function feature. Callback functions run between each op execution during runtime and contain op level golden analysis. They are also customizable and provide the flexibility for you to get creative with your golden usage.

builder.apis

Adding a new op to ttir-builder

ttir-builder is designed to only create ops supported in TTIR. At the moment, most but not all ops are supported, and new ops are still occasionally added to TTIR. Creating ttir-builder support for an op entails writing a function in tools/builder/ttir/ttir_builder.py that will create the op and its golden counterpart.

TTIR op factories

All ops are created when their relevant information is run through the _op_proxy function which provides a general interface for proxy-ing and creating ops.

def _op_proxy(
    self,
    op_ttir_function: Callable,
    inputs: List[Operand],
    unit_attrs: List[str] = None,
    organize_ttir_args: Optional[Callable] = None,
    organize_golden_args: Optional[Callable] = None,
    output_shape: Optional[Shape] = None,
    output_type: Optional[Type] = None,
    output_create_fn: Optional[Callable] = None,
    golden_kwargs: dict = {},
    ttir_kwargs: dict = {},
    skip_golden: bool = False,
)

Start by finding the TTIR op you wish to replicate in include/ttmlir/Dialect/TTIR/IR/TTIROps.td or the TTIR dialect documentation.

All op attributes should be included as arguments in your function and passed into a proxy function as keyword arguments using ttir_kwargs.

All input operands should be passed into a proxy function using the argument inputs. Output operands are considered inputs and can optionally be passed into inputs if their shape or datatype is relevant to the op's result operand. organize_ttir_args dictates what information gets passed into autogenerated file build/python_packages/ttmlir/dialects/_ttir_ops_gen.py and can be used if operand arguments require special handling.

Golden functions

Golden functions provide the reference implementation for TTIR operations using PyTorch. They are centralized in tools/builder/base/builder_golden.py and must be mapped to their corresponding TTIR operations. The _op_proxy function automatically retrieves the appropriate golden function based on the TTIR operation class. The skip_golden argument omits golden tensor creation and addition to the golden map. Since goldens are relied upon to set output_shape and output_type, setting skip_golden=True requires passing in output_shape and output_type to _op_proxy.

Writing a golden function

Before writing a golden function, you need to know exactly what the TTIR op does to its input data because you will have to replicate that exactly using PyTorch operations. This information is usually covered in TTIR documentation, but if not, you may have to do some detective work and trial and error. Get creative with keyword argument handling, using similar Pytorch operations, and maybe multiple operations. Google is your friend. If you have to figure out how to do something Pytorch doesn't, odds are someone online has encountered the same situation.

Golden functions should be implemented in builder_golden.py and follow this pattern:

1. Simple operations: If PyTorch has an identical function, you can directly use it in the mappings
2. Complex operations: Define a custom golden function that implements the behavior using PyTorch operations

Adding golden function mappings

All golden functions must be registered in the GOLDEN_MAPPINGS dictionary in builder_golden.py:

# In builder_golden.py
def cbrt_golden(input: torch.Tensor) -> torch.Tensor:
    """Golden function for cube root operation."""
    golden_sign = torch.sign(input)
    golden_cbrt = torch.pow(torch.abs(input), 1 / 3)
    return golden_sign * golden_cbrt

# Add to GOLDEN_MAPPINGS dictionary
GOLDEN_MAPPINGS: Dict[type, Callable] = {
    # ... other mappings ...
    ttir.CbrtOp: cbrt_golden,
    # ... more mappings ...
}

Using golden functions in ops.py

In your operation implementation in ops.py, simply pass the TTIR operation class to _op_proxy. The golden function is automatically retrieved internally:

# In ops.py
def cbrt(self, in0: Operand, unit_attrs: Optional[List[str]] = None) -> OpView:
    return self._op_proxy(
        ttir.CbrtOp,  # Golden function automatically retrieved from GOLDEN_MAPPINGS
        [in0],
        unit_attrs=unit_attrs,
    )

Adding Silicon tests

Silicon tests are created in the test/python/golden directory.

pytest test/python/golden/test_ttir_ops.py

Be sure to file an issue for failing tests and add a pytest mark for any failing or unsupported tests. The pytest marks instruct CI to ignore tests.

pytest.mark.skip("Issue number") : skip flatbuffer creation for this test
pytest.mark.skip_config(config, ... reason=None): skip test if all of the specified targets/backends per config are present

The skip_config mark here is a little nuanced. By passing in a list of strings representing targets and/or systems (e.g. ["ttmetal", "p150"]) this mark will intelligently skip tests with that configuration. The example given will skip tests lowered to ttmetal iff we are runing on a p150 (i.e. blackhole). This functionality will be expanded to include other axes of test configuration, but target and system are sufficient for our needs at the moment.

For tests exclusive to n300 or llmbox, use the following pytest marks or add them to their respective test files.

pytestmark = pytest.mark.n300
pytestmark = pytest.mark.llmbox

Running Silicon tests

Follow these steps. The directory test/python/golden contains tests for modules, individual ops, and various machines.

1. Build ttmlir
source env/activate
cmake -G Ninja -B build -DCMAKE_BUILD_TYPE=Release -DCMAKE_C_COMPILER=clang-17 -DCMAKE_CXX_COMPILER=clang++-17 -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DTTMLIR_ENABLE_RUNTIME=ON -DTT_RUNTIME_ENABLE_PERF_TRACE=ON
cmake --build build

2. Query system
ttrt query --save-artifacts

3. Export system desc file
export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys (path dumped in previous command)

4. Generate test cases
pytest test/python/golden/test_ttir_ops.py

5. Run test cases
ttrt run builder-artifacts

Sphinx documentation

Docstrings

Sphinx generates documentation for builder ops from the docstrings in TTIRBuilder functions. This is the structure to follow when writing your docstring

"""
Creates ``ttir.add``.

*Elementwise addition operation.*

Performs elementwise addition between two tensors.
For each pair of corresponding elements, adds the element in the second
tensor to the element in the first tensor.

Mathematical definition: add(x, y) = x + y

.. code-block:: mlir

    // Add corresponding elements
    %result = ttir.add(%lhs, %rhs, %output) : tensor<3xf32>, tensor<3xf32>, tensor<3xf32> -> tensor<3xf32>
    // Input tensors:
    // lhs: [3.5, 0.0, -1.2]
    // rhs: [1.5, 2.0, -3.2]
    // Output tensor:
    // [5.0, 2.0, -4.4]

Parameters
----------
in0 : Operand
    First input tensor
in1 : Operand
    Second input tensor
unit_attrs : *Optional[List[str]]*, optional
    Optional list of unit attributes

Returns
-------
*OpView*
    A tensor containing the elementwise sum of the inputs
"""

Autogen skip

All functions in TTIRBuilder are included in documentation by default. If your op is failing any of the tests, it can't yet be added to the documentation. Custom golden functions also must be excluded. Tag those functions with autodoc_skip.

@autodoc_skip
def bitwise_not(
    self, in0: Operand, unit_attrs: Optional[List[str]] = None
) -> OpView:

stablehlo-builder

stablehlo-builder is a tool for creating stableHLO operations. It provides support for MLIR modules to be generated from user-constructed ops.

Getting started

StableHLOBuilder is a builder class providing the API for creating stableHLO ops. The python package builder contains everything needed to create ops through a StableHLOBuilder object. builder.base.builder_utils contains the APIs for wrapping op-creating-functions into MLIR modules and flatbuffers files.

from builder.stablehlo.stablehlo_builder import StableHLOBuilder
from builder.base.builder_utils import build_stablehlo_module, compile_stablehlo_to_flatbuffer

Creating a StableHLO module

build_stablehlo_module defines an MLIR module specified as a python function. It wraps fn in a MLIR FuncOp then wraps that in an MLIR module, and finally ties arguments of that FuncOp to test function inputs. It will instantiate and pass a StableHLOBuilder object as the last argument of fn. Each op returns an OpView type which is a type of Operand that can be passed into another builder op as an input.

def build_stablehlo_module(
    fn: Callable,
    inputs_shapes: List[Shape],
    inputs_types: Optional[List[Union[torch.dtype, TypeInfo]]] = None,
    mesh_name: str = "mesh",
    mesh_dict: OrderedDict[str, int] = OrderedDict([("x", 1), ("y", 1)]),
    module_dump: bool = False,
    base: Optional[str] = None,
    output_root: str = ".",
) -> Tuple[Module, StableHLOBuilder]:

Example

from builder.base.builder import Operand
from builder.stablehlo.stablehlo_builder import StableHLOBuilder
from builder.base.builder_utils import build_stablehlo_module

shapes = [(32, 32), (32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, in2: Operand, builder: StableHLOBuilder):
    return builder.add(in0, in1)

module, builder = build_stablehlo_module(model, shapes)

Returns

An MLIR module containing an MLIR op graph defined by fn and the StableHLOBuilder object used to create it

module {
  func.func @model(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<32x32xf32>) -> tensor<32x32xf32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<32x32xf32>
    return %0 : tensor<32x32xf32>
  }
}

Creating a StableHLO module with Shardy annotations

StableHLOBuilder allows you to attach shardy annotations to the generated mlir graph.

Example

from builder.base.builder import Operand
from builder.stablehlo.stablehlo_builder import StableHLOBuilder
from builder.base.builder_utils import build_stablehlo_module

shapes = [(32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, shlo_builder: StableHLOBuilder):
    tensor_sharding_attr = shlo_builder.tensor_sharding_attr(
        mesh_name="mesh",
        dimension_shardings=[
            shlo_builder.dimension_sharding_attr(
                axes=[shlo_builder.axis_ref_attr(name="x")],
                is_closed=True,
            ),
            shlo_builder.dimension_sharding_attr(
                axes=[shlo_builder.axis_ref_attr(name="y")],
                is_closed=False,
            )
        ]
    )

    shlo_builder.sharding_constraint(in0, tensor_sharding_attr=tensor_sharding_attr)
    return shlo_builder.add(in0, in1)

module, shlo_builder = build_stablehlo_module(model, shapes, mesh_name="mesh", mesh_dict=OrderedDict([("x", 1), ("y", 8)]))

Returns

An MLIR module containing shardy annotations.

module {
  sdy.mesh @mesh = <["x"=1, "y"=8]>
  func.func @model(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>) -> tensor<32x32xf32> {
    %0 = sdy.sharding_constraint %arg0 <@mesh, [{"x"}, {"y", ?}]> : tensor<32x32xf32>
    %1 = stablehlo.add %arg0, %arg1 : tensor<32x32xf32>
    return %1 : tensor<32x32xf32>
  }
}

Compiling into flatbuffer

compile_stablehlo_to_flatbuffer compiles a StableHLOBuilder function fn straight to flatbuffer. This decorator is mainly a wrapper around the following functions, with each next function called on the output of the last: build_stablehlo_module, _run_ttir_pipeline, and ttnn_to_flatbuffer_file, ttmetal_to_flatbuffer_file, or ttir_to_ttnn_emitc_pipeline as dictated by the target parameter.

def compile_stablehlo_to_flatbuffer(
    fn: Callable,
    inputs_shapes: List[Shape],
    inputs_types: Optional[List[Union[torch.dtype, TypeInfo]]] = None,
    system_desc_path: Optional[str] = None,
    test_base: str = "test",
    output_root: str = ".",
    target: Literal["ttnn", "ttmetal", "emitc"] = "ttnn",
    mesh_name: str = "mesh",
    mesh_dict: OrderedDict[str, int] = OrderedDict([("x", 1), ("y", 1)]),
    module_dump: bool = True,
    argument_types_string: Optional[str] = None,
    custom_pipeline: Optional[Union[Callable, str]] = None,
    ttir_pipeline_options: List[str] = [],
    shlo_pipeline_options: List[str] = [],
    shlo_to_ttir_pipeline_options: List[str] = [],
    print_ir: Union[bool, str] = False,
) -> str:

The executable flatbuffer is written to a file, compile_stablehlo_to_flatbuffer returns the file address of that flatbuffer.

TTNN example

Let's use our previous model function.

from builder.base.builder import Operand
from builder.stablehlo.stablehlo_builder import StableHLOBuilder
from builder.base.builder_utils import compile_stablehlo_to_flatbuffer

shapes = [(32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, shlo_builder: StableHLOBuilder):
    tensor_sharding_attr = shlo_builder.tensor_sharding_attr(
        mesh_name="mesh",
        dimension_shardings=[
            shlo_builder.dimension_sharding_attr(
                axes=[shlo_builder.axis_ref_attr(name="x")],
                is_closed=True,
            ),
            shlo_builder.dimension_sharding_attr(
                axes=[shlo_builder.axis_ref_attr(name="y")],
                is_closed=False,
            )
        ]
    )

    shlo_builder.sharding_constraint(in0, tensor_sharding_attr=tensor_sharding_attr)
    return shlo_builder.add(in0, in1)

compile_stablehlo_to_flatbuffer(
    model,
    shapes,
    mesh_name="mesh",
    mesh_dict=OrderedDict([("x", 1), ("y", 8)]),
    target="ttnn",
)

TTMetal example

Let's once again use the same code for TTMetal that was used in the TTNN example but change the target to "ttmetal". Just as with _run_ttir_pipeline, only one or the other can be run on a module since compile_stablehlo_to_flatbuffer modifies the module in place.

compile_stablehlo_to_flatbuffer(
    model,
    shapes,
    mesh_name="mesh",
    mesh_dict=OrderedDict([("x", 1), ("y", 8)]),
    target="ttmetal",
)

Integrating with other tt-mlir tools

Alternatives for file creation

1. The ttmlir-opt tool runs a compiler pass on an .mlir file.
2. The ttmlir-translate can generate a flatbuffer from an .mlir file.
3. llvm-lit can also be used to generate a flatbuffer from an existing .mlir file.

Running models

ttrt

ttrt is intended to be a swiss army knife for working with flatbuffers.

tt-explorer

tt-explorer is a visualizer tool for ttmlir-powered compiler results.

ttnn-standalone

ttnn-standalone is a post-compile tuning/debugging tool.

llvm-lit

llvm-lit can also be used for MLIR testing.

Golden mode

Golden dataclass

StableHLOBuilder provides support to code golden tensors into flatbuffers which will be used for comparison with TT device output in ttrt runtime. Golden is the dataclass used to store information about a golden tensor. Each StableHLOBuilder op should have a matching PyTorch op (or golden function built from PyTorch ops) which should perform exactly the same operation, generating the same outputs given the same inputs. You can use StableHLOBuilder helper functions to store input, intermediate, and output tensors within the flatbuffer. Input and output goldens are mapped with keys "input_" and "output_" followed by a tensor index: input_0. Intermediate output tensors are mapped to the location of the respective op creation.

GoldenCheckLevel Enum

StableHLOBuilder stores an instance of the class GoldenCheckLevel(Enum) that dictates golden handling. It defaults to GoldenCheckLevel.OP_LEVEL.

DISABLED : do not store goldens
OP_LEVEL : check every single op level goldens
GRAPH_LEVEL : check graph level goldens only

Check and set GoldenCheckLevel with StableHLOBuilder APIs.

from builder.base.builder import Operand, GoldenCheckLevel
from builder.stablehlo.stablehlo_builder import StableHLOBuilder

def model(in0: Operand, in1: Operand, in2: Operand, builder: StableHLOBuilder):
    builder.golden_check_level = GoldenCheckLevel.GRAPH_LEVEL
    add_0 = builder.add(in0, in1)
    multiply_1 = builder.multiply(in1, add_0)
    return builder.multiply(multiply_1, in2)

Getting golden data

Unless otherwise specified in the GoldenCheckLevel, all input and output tensors will generate and store a golden in StableHLOBuilder as a Golden type.

The StableHLOBuilder API get_golden_map(self) is used to export golden data for flatbuffer construction. It returns a dictionary of golden tensor names and GoldenTensor objects.

To get info from a GoldenTensor object, use the attributes supported by ttmlir.passes: name, shape, strides, dtype, data.

from ttmlir.passes import GoldenTensor
from builder.stablehlo.stablehlo_builder import StableHLOBuilder

shapes = [(32, 32), (32, 32), (32, 32)]

def model(in0: Operand, in1: Operand, in2: Operand, builder: StableHLOBuilder):
    add_0 = builder.add(in0, in1)
    builder.print_goldens()
    print(builder.get_golden_map())
    return add0

Golden tensor:
tensor([[ 4.0450e+00,  1.4274e+00,  5.9156e-01,  ..., -5.9834e-01,
         -1.1830e-01,  1.2837e-01],
        [ 2.3788e+00,  2.9242e-03, -5.2838e-02,  ...,  1.8294e+00,
          5.0348e+00,  9.7179e-01],
        [ 1.5168e-02,  1.0577e-01, -3.0682e-01,  ...,  6.7212e-01,
          9.4523e-02,  5.3765e+00],
        ...,
        [ 1.4241e-01,  1.1838e+00, -1.0601e+00,  ...,  4.9099e-01,
          4.2267e+00,  4.0610e-01],
        [ 5.6630e-01, -1.3068e-01, -1.7771e-01,  ...,  2.3862e+00,
          3.9376e-01,  7.3140e-01],
        [ 4.2420e+00,  1.7006e-01, -3.4861e-01,  ...,  1.1471e-01,
          1.6189e+00, -6.9106e-01]])
{'input_0': <ttmlir._mlir_libs._ttmlir.passes.GoldenTensor object at 0x7f77c70fa0d0>, 'output_0': <ttmlir._mlir_libs._ttmlir.passes.GoldenTensor object at 0x7f77c6fc9590>}

Setting golden data

Use StableHLOBuilder API set_graph_input_output to set your own input and output golden tensors using PyTorch tensors. Keep in mind that this also sets graph inputs and outputs.

set_graph_input_output(
        self,
        inputs: List[torch.Tensor],
        outputs: Optional[List[torch.Tensor]] = None,
        override: bool = False,
    )

import torch

input_0 = torch.ones((32, 32))
output_0 = torch.zeros((32, 32))
builder.set_graph_input_output([input_0], [output_0], override=True)

Running flatbuffer with golden data in ttrt

Running flatbuffers in ttrt requires building and setting up the environment. Run these commands before creating MLIR modules or flatbuffers so the system description in the flatbuffers match your device.

cmake --build build
ttrt query --save-artifacts
export SYSTEM_DESC_PATH=/path/to/system_desc.ttsys

Set environment variable TTRT_LOGGER_LEVEL to DEBUG so ttrt logs golden comparison results and prints graph level golden tensors.

export TTRT_LOGGER_LEVEL=DEBUG

Finally run ttrt. Our example flatbuffer file (since we didn't specify otherwise) defaulted to file path ./builder-artifacts/stablehlo-builder/test_ttnn/test_ttnn.mlir.ttnn. --log-file ttrt.log and --save-golden-tensors are both optional flags. They ensure that all golden data produced by the ttrt run gets written to files.

ttrt run builder-artifacts/stablehlo-builder/test_ttnn/test_ttnn.mlir.ttnn --log-file ttrt.log --save-golden-tensors

Golden callbacks

The ttrt documentation contains a section on the callback function feature. Callback functions run between each op execution during runtime and contain op level golden analysis. They are also customizable and provide the flexibility for you to get creative with your golden usage.

tt-explorer

Welcome to the tt-explorer wiki! The Wiki will serve as a source for documentation, examples, and general knowledge related to the TT-MLIR visualization project. The sidebar will provide navigation to relevant pages. If this is your first time hearing about the project, take a look at Project Architecture for an in-depth introduction to the tool and motivations behind it. 🙂

Overview

Visualizer tool for ttmlir-powered compiler results. Visualizes from emitted .mlir files to display compiled model, attributes, performance results, and provides a platform for human-driven overrides to gameify model tuning.

Quick Start

tt-explorer comes packaged as a tool in the tt-mlir repo. If you haven't done so yet, please refer to "Setting up the environment manually" section from the Getting Started Guide to build the environment manually.

Here is a summary of the steps needed:

1. Clone tt-mlir and build the environment
2. Run source env/activate to be in tt-mlir virtualenv for the following steps
3. Ensure tt-mlir is built with atleast these flags:
   • -DTT_RUNTIME_ENABLE_PERF_TRACE=ON
   • -DTTMLIR_ENABLE_RUNTIME=ON
   • -DTT_RUNTIME_DEBUG=ON
   • -DTTMLIR_ENABLE_STABLEHLO=ON
4. Build explorer target in tt-mlir using cmake --build build -- explorer
5. Run tt-explorer in terminal to start tt-explorer instance. (Refer to CLI section in API for specifics)
   • Note: tt-explorer requires Pandas in addition to the tt-mlir System Dependencies.
6. Ensure server has started in tt-explorer shell instance (check for message below)

   Starting Model Explorer server at:
   http://localhost:8080
   

Building tt-explorer

To build tt-explorer you need first to clone and configure the environment for tt-mlir. Please refer to the Getting Started Guide.

After building and activating the virtualenv, build tt-mlir and ensure the following flags are present, as they are needed for executing models in tt-explorer and without them it won't build.

Flags required:

• -DTT_RUNTIME_ENABLE_PERF_TRACE=ON
• -DTTMLIR_ENABLE_RUNTIME=ON
• -DTT_RUNTIME_DEBUG=ON
• -DTTMLIR_ENABLE_STABLEHLO=ON

Then build the explorer target by running the following command:

cmake --build build -- explorer

After it finishes building, start the explorer server by running the following command:

tt-explorer

The server should then start and show a message similar to this:

Starting Model Explorer server at:
http://localhost:8080

Running tt-explorer CI Tests Locally

Note: CI tests are ran like described below. Here we provide the steps needed to reproduce it and debug failing CI tests locally.

tt-explorer relies on tests that are present in the tests/ directory as well as tests dynamically created through llvm-lit. Below are the steps to replicate the testing procedure seen in CI:

1. Make sure you're in the tt-mlir directory
2. You need to build the explorer target with cmake --build build -- explorer
3. Run and save the system descriptor ttrt query --save-artifacts
4. Save the system variable export SYSTEM_DESC_PATH=$(pwd)/ttrt-artifacts/system_desc.ttsys
5. Run and generate ttnn + MLIR tests: cmake --build build -- check-ttmlir
6. Save the relevant test directories:
   • export TT_EXPLORER_GENERATED_MLIR_TEST_DIRS=$(pwd)/build/test/python/golden/ttnn,$(pwd)/build/test/ttmlir/Silicon/TTNN/n150/perf
   • export TT_EXPLORER_GENERATED_TTNN_TEST_DIRS=$(pwd)/build/test/python/golden/ttnn
7. Run the pytest for tt-explorer with pytest tools/explorer/test/run_tests.py

or in a concise shell script:

# Ensure you are present in the tt-mlir directory
source env/activate

# Build Tests
cmake --build build -- explorer
ttrt query --save-artifacts
export SYSTEM_DESC_PATH=$(pwd)/ttrt-artifacts/system_desc.ttsys
cmake --build build -- check-ttmlir

# Load Tests
export TT_EXPLORER_GENERATED_MLIR_TEST_DIRS=$(pwd)/build/test/python/golden/ttnn,$(pwd)/build/test/ttmlir/Silicon/TTNN/n150/perf
export TT_EXPLORER_GENERATED_TTNN_TEST_DIRS=$(pwd)/build/test/python/golden/ttnn

# Run Tests
pytest tools/explorer/test/run_tests.py

tt-explorer UI

For general reference of the UI, refer to the model-explorer wiki. This section will highlight specific UI elements added to the Tenstorrent fork of model-explorer.

Model Execution

In the top right of the screen an additional button has been added to the top bar, it sends the model to the server for execution and updates the visualization once it has been executed. Once the model has executed, overlays are also created. These overlays provide information on how the execution went.

Performance Overlay

The performance overlay is generated on every execution, it highlights the time it took to execute each node on the graph. This is visualized with a gradient from Yellow -> Red, with Yellow being the lowest time amongst all nodes on the graph, and Red being highest.

Accuracy Overlay

The accuracy overlay is only generated when executing from a compatible flatbuffer (.ttnn file extension with Debug Info). The overlay consists of either Green or Red node overlays. Green if the node passed a "golden" test, Red if not.

The value for the overlay is the actual Pearson Correlation Coefficient (PCC) value with the "golden" tensor subtracted by the expected PCC value. If the number is < 0 we know it doesn't match the expected PCC, otherwise it is an accurate comparison.

Advanced Settings

This menu will open a window with some advanced settings for Model execution.

Opt. Policy

This dropdown provides a list of Optimization Policies which will be used when the model is executed. These policies are applied when lowering from a ttir module to an executable ttnn module.

Generate C++ Code

This toggle will run the EmitC pass in the tt-mlir compiler to generate TTNN C++ Code and make it available to you after running a model. Default value for this toggle is Off.

"Play" Button

This button invokes the execute function which will compile and execute the model. The button will then be "loading" until execution is finished. Once execution is finished a performance trace should be overlayed on the graph and it should reload.

"Code" Button

If the Generate C++ Code option is enabled, this button will become available to view and download the C++ code in a window within explorer.

"Logs" Button

This button will open a window to view the shell logs while execution is running. If any errors occur they will be displayed here.

Overridden Fields

Certain Nodes on the graph will have attributes that are presented as editable fields. These are fields which have overrides available. This value can be changed and then sent to be recompiled, invalid configurations will result in errors.

tt-explorer CLI

This section provides a details about the usage of tt-explorer.

Input Models

Currently tt-explorer supports 3 types of models that can be executed/visualized.

CLI

The CLI for tt-explorer provides a simple suite of options to start the UI:

tt-explorer -p <port> -u <url> -q

Options

-p, --port

Port that model-explorer server will be exposed to. Default is 8080.

-u, --url

Host URL Address for server. Default is "localhost".

-q, --no-browser

Create server without opening a browser tab.

-x, --no-model-execution

Disable execution of models from the UI.

Example usage:

tt-explorer -p 8000 -u 0.0.0.0 -q

This command will start the tt-explorer server on port 8000, accessible at the address 0.0.0.0, and without opening a browser tab.

tt-explorer - API

TT-Adapter

The following is a reference for the REST API provided by TT-Adapter.

First, a short info-dump on how an extensible API can be built on top of Model Explorer.

Building an API using Model Explorer

The /apipost/v1/send_command endpoint provides an extensible platform with which commands are sent to be executed directly by the adapter specified. This becomes the main endpoint through which communication is facilitated between the server and client, the commands respond with an "adapter response".

Sending Commands

The body of the command must be JSON, and conform to the following interface (described below as a Typescript interface). Specific commands may narrow the field types or extend this interface providing extra information. But all interfaces should be based on this.

interface ExtensionCommand {
	cmdId: string;
	extensionId: string;
	modelPath: string;
	settings: Record<string, any>;
	deleteAfterConversion: boolean;
}

More often than not, functions do not need all of these fields, but they must all be present to properly process the command sent into the handling function on the server.

On the server side, the signature that all function that handle commands have to follow is:

class TTAdapter(Adapter):
  # ...
  def my_adapter_fn(self, model_path: str, settings: dict):
    # Parse model_path and settings objects as they are fed from send_command endpoint.
    pass

This function is invoked and called from a new instance every time. This is important to understand for the idea of persisting information on the server. As all requests to the server are stateless, the onus is often on the end-user to keep track of important information such as the path of a model they've uploaded, or the paths of important artifacts that the server has produced. tt-explorer aims to make this as easy as possible, but this may not always be possible due to the very nature of how the server works.

Information can be processed in this function as defined by the user, and often settings becomes a versatile endpoint to provide more information and context for the execution of some function. As an example, refer to ModelRunner:initialize, this function doesn't use any of the parameter, as such they are not processed at all, and the function only executes a static initialization process regardless of the parameters passed into the command.

Example request

Below is an example of the JSON request sent from the UI to the server:

{
	// tt_adapter to invoke functions from TT-Adapter
	"extensionId": "tt_adapter",
	// Name of function to be run, "convert" is built into all adapters to convert some model to graph
	"cmdId": "convert",
	// Path to model on server to be fed into function
	"modelPath": "/tmp/tmp80eg73we/mnist_sharding.mlir",
	// Object holding custom settings to be fed into function
	"settings": {
		"const_element_count_limit": 16,
		"edge_label_font_size": 7.5,
		"artificial_layer_node_count_threshold": 1000,
		"keep_layers_with_a_single_child": false,
		"show_welcome_card": false,
		"disallow_vertical_edge_labels": false,
		"show_op_node_out_of_layer_edges_without_selecting": false,
		"highlight_layer_node_inputs_outputs": false,
		"hide_empty_node_data_entries": false
	},
	// `true` if file at `modelPath` is to be deleted after function run
	"deleteAfterConversion": true
}

Adapter Response

Model Explorer was not made to allow for such an extensible framework to be tacked onto it. As such, the adapter response is processed in a very particular way before it is sent back to the user.

In particular, refer to model_explorer.utils.convert_adapter_response which is run on the output of every function.

This means that for compatibility reasons (i.e. to not stray too much from the upstream implementation that we are based off of) responses sent from the server must be in JSON format only and wrap the data on a graph property.

Below is the base typescript interface that the UI expects for the json response. Commands can define custom data inside the graph property.

/** A response received from the extension. */
interface ExtensionResponse<
	G extends Array<unknown> = Graph[],
	E extends unknown = string
> {
	graphs: G;
	error?: E;
}

For custom adapter responses. This limits the transfer of raw bytes data through different MIME Types, and requires the tt_adapter.utils.to_adapter_format which turns any dict object into a model explorer adapter compatible response. While this framework works well for graphs, it makes an "extensible" API difficult to implement.

Current API Reference

convert

Standard built-in conversion function, converts TTIR Module into Model Explorer Graph. Also provides settings as a platform for overrides to be applied to the graph.

Request

// As this is the base request everything is based off,
// this interface only narrows down the command to be "convert".
interface AdapterConvertCommand extends ExtensionCommand {
	cmdId: 'convert';
}

Response

// As this is the base response everything is based off,
// it is exactly the same as `ExtensionResponse`.
type AdapterConvertResponse = ExtensionResponse;

{
	"graphs": [{
		// Model Explorer Graph JSON Object
	}]
}

initialize

Called from TTAdapter.__init__, used to Load SystemDesc into environment.

Request

interface InitializeCommand extends ExtensionCommand {
	cmdId: 'initialize';
}

Response

type AdapterInitializeResponse = ExtensionResponse<[{
	system_desc_path: string
}]>;

{
	"graphs": [{
		"system_desc_path": "<path to system_desc.ttsys>"
	}]
}

execute

Called from TTAdapter.execute, executes a model.

Request

interface AdapterExecuteCommand extends ExtensionCommand {
	cmdId: 'execute';
}

Response

// When the request is successful, we don't expect any response back.
// Thus, an empty array is returned for `graphs`.
type AdapterExecuteResponse = ExtensionResponse<[]>;

{
	"graphs": []
}

status-check

Called from TTExplorer.status_check, it is used for checking the execution status of a model and update the UI accordingly.

Request

interface AdapterStatusCheckCommand extends ExtensionCommand {
	cmdId: 'status_check';
}

Response

type AdapterStatusCheckResponse = ExtensionResponse<[{
	isDone: boolean,
	progress: number,
	total?: number,
	timeElapsed?: number,
	currentStatus?: string,
	error?: string,
	stdout?: string,
	log_file?: string
}]>;

{
	"graphs": [{
		"isDone": false,
		"progress": 20,
		"total": 100,
		"timeElapsed": 234,
		"stdout": "Executing model...\nPath: /path/to/model",
		"log_file": "/path/to/log/on/the/server"
	}]
}

Editable attributes

To enable an attribute to be edited, a response coming from the server should contain the editable field on the attribute.

The typescript interface is as follows:

interface Graph {
	nodes: GraphNode[];
	// ...
}

interface GraphNode {
	attrs?: Attribute[];
	// ...
}

type EditableAttributeTypes = EditableIntAttribute | EditableValueListAttribute | EditableGridAttribute; // Attribute types are defined below...

interface Attribute {
	key: string;
	value: string;
	editable?: EditableAttributeTypes; // <- the editable attribute information
	// ...
}

EditableIntAttribute

This editable attribute represents a list of integer values. It expects the attribute value to be formatted as a string, starting with [ and ending with ], with all values separated by ,. Like the example below:

[1, 2, 3]

The typescript interface for the editable attribute is this:

interface EditableIntAttribute {
  input_type: 'int_list';
  min_value?: number = 0;
  max_value?: number = 100;
  step?: number = 1;
}

Both min_value and max_value define the accepted range of values, and step define the number to increment or decrement per step.

The default range of values is between 0 and 100, inclusive, and the default step is 1. Thus by default, the value will increment or decrement by 1 each time to a minimum of 0 and a maximum of 100.

Here is an example of what this attribute look like:

{
	"graphs": [{
		"nodes": [
			{
				"attrs": [
					{
						"key": "shape",
						"value": "[8, 8]",
						"editable": {
							"input_type": "int_list",
							"min_value": 8,
							"max_value": 64,
							"step": 8
						}
					}
				]
			}
		]
	}]
}

EditableValueListAttribute

This editable attribute define a fixed list of string values to display.

The typescript interface for the editable attribute is this:

interface EditableValueListAttribute {
	input_type: 'value_list';
	options: string[];
}

The options property provides the list of options to be displayed. The current value will be added to this list and any duplicates will be removed.

Here is an example of what this attribute look like:

{
	"graphs": [{
		"nodes": [
			{
				"attrs": [
					{
						"key": "chip_arch",
						"value": "wormhole",
						"editable": {
							"input_type": "value_list",
							"options": [
								"wormhole",
								"grayskull"
							]
						}
					}
				]
			}
		]
	}]
}

EditableGridAttribute

The grid attribute is similar to to the integer list, with the main difference that you can specify a separator for the place the list will be split, and it doesn't need to be enclosed in bracket ([ and ]). The data for a grid attribute looks like this:

4x4x2

The typescript interface for the editable attribute is this:

interface EditableGridAttribute {
  input_type: 'grid';
  separator?: string = 'x';
  min_value?: number = 0;
  max_value?: number = 100;
  step?: number = 1;
}

Both min_value and max_value define the accepted range of values, and step define the number to increment or decrement per step.

The default range of values is between 0 and 100, inclusive, and the default step is 1. Thus by default, the value will increment or decrement by 1 each time to a minimum of 0 and a maximum of 100.

The separator attribute defines the character used to split the string, it defaults to "x".

Here is an example of what this attribute look like:

{
	"graphs": [{
		"nodes": [
			{
				"attrs": [
					{
						"key": "grid",
						"value": "4x4",
						"editable": {
							"input_type": "grid",
							"min_value": 4,
							"max_value": 64,
							"step": 4,
							"separator": "x"
						}
					}
				]
			}
		]
	}]
}

Attribute display type

To change how the attribute is displayed from plain text to something else, we do extend the attribute interface (presented above) with the display_type optional field.

type AttributeDisplayType = 'memory';

interface Attribute {
	key: string;
	value: string;
	display_type?: AttributeDisplayType; // <- Optional, add a different display type.
	// ...
}

If the display_type attribute is present, and it matches one of the available values, then the attribute will display differently than the others.

In the example below, the two attributes have different display types, one shows the regular, plain text display; and the other shows the memory display type, which renders it as a progress bar.

memory

Setting the display type to memory will make the attribute try to render as a progress bar.

The UI will then check the value property in the attribute for the following conditions:

• Is a double precision floating point number
• Is not NaN
• Is grater than or equal to 0
• Is less than or equal to 1

If all of the conditions are true, then the value will be rendered as a progress bar.

tt-explorer Roadmap

Milestone 1 (v0.1)

Main Goal - Visualize & Execute

This will highlight half of the essential work that this tool should be able to do in both visualizing a model and executing it using the current TT-Forge stack. The frontend transformation of a model → TTIR will be done outside of the scope of tt-explorer at the moment. For this milestone tt-explorer will be able to spin up a host-side and a client-side instance. The tool will be able to ingest TTIR modules to produce a visual result, and be able to execute this module. Ambitiously, the performance traces should be collected back into tt-explorer to be displayed.

Tasks:

• Load TTIR Modules and Visualize TTIR-Ops in Model Explorer
• Create Extensible Notebook UX allowing for visualization and scripting capabilities
• Add functionality to Model Explorer to load from re-compiled TTIR Modules (might be from JSON)
• Add functionality to TT-MLIR to execute from Python Bindings
• Create REST API skeleton in TT-Adapter
• From REST API Call, Invoke python bindings to execute TTIR module using TT-Adapter
• (If possible) Parse Perf Trace Artifact and visualize performance in Model-Explorer (as Node Data)

Milestone 2 (v0.2)

Main Goal - Model Editor

The primary function of tt-explorer is to visualize and edit the model according to what the user defines as overrides the automatically generated compiler results. This milestone highlights that functionality in tt-explorer, focusing around providing UI, TT-MLIR, and tt-explorer features that enable the user to edit and tune a model “in-loop” with the TT-Forge compiler.

Tasks:

• Flesh out and test locations ID such that operations can be tracked through the compiler stack.
• Use Loc IDs to bind TTIR Ops with Tracy Perf Trace Artifact, and send to Model-Explorer to visualize.
• Implement Overrides Functionality into TT-MLIR, tracking based on Loc IDs.
• Overhaul UI to enable editing node attributes, use these updated fields to send information back to tt-explorer via REST API (in the form of an Overrides JSON)
• Parse Overrides JSON and apply Overrides over a REST API Call, visualize re-compiled graph now.
• Provide REST API endpoint to provide “legal” options attached to Graph JSON.

Milestone 3 (v0.3+)

Main Goal - Matured Tool and Extensibility

The focus of this milestone is to transition tt-explorer from a prototype tool into a mature visualization and editing tool for “Human-In-Loop” compilation. The tool is now planned to made extensible for other dialects and entry points forecast into TT-MLIR (Jax, StableHLO, etc…) and development of the visualization components of the tool provide feedback to upstream repos like model-explorer. Here the focus is on providing extensible interfaces for new UI elements (in supporting multi-chip and beyond), REST API, and Overrides.

Tasks:

• Begin adding new dialects like .ttm, .ttnn to Model Explorer so that complied results can be inspected and analyzed to optimize at different steps of the compiler.
• Add Accuracy/Performance Overlays as Node Data into the Model Explorer graph to visualize execution results
• Enable interaction with ttnn-visualizer and other TT Visualizer tools to provide a more detailed view of execution results.
• Start introducing InterOp with builtin adapters in model-explorer to support visualizing models from FE.
• Use split panes to display graph transformations occurring through compiler, leveraging multiple dialects.
• To be defined later, depending on the growth of the MLIR Project

tt-alchemist

tt-alchemist is a code generation tool that converts MLIR models to executable C++ or Python solutions for Tenstorrent AI accelerators.

Table of Contents

• Support Matrix
• Usage
  • Using via CLI
  • Usage via lib

Support Matrix

Note: The tool is currently in development and is subject to frequent changes. Please refer to this document for most up-to-date information. Support matrix is provided below.

The following table summarizes the current support for code generation modes in tt-alchemist:

Modes:

• standalone: Generates a self-contained solution with all necessary dependencies copied into the output directory. Useful for deployment and sharing.
• local: Generates code that uses libraries from the source tree, minimizing duplication and disk usage. Useful for development and debugging.

Note: Python codegen currently supports a small subset of operations compared to C++. Full support is being actively worked on and is coming soon.

Usage

The tool is compiled into a C++ library, with a thin CLI wrapper written in Python. This means that it can be distributed both as a C++ library, and as a CLI tool via Python wheel mechanism.

Using via CLI

To use via CLI, it is suggested to build the tool from source. Alternatively, look for tt-alchemist artifacts within CI runs.

# Assuming the user had already built the tt-mlir compiler and turned on the python virtual env

# Build the tt-alchemist lib, package into Python wheel, and install to active env
cmake --build build -- tt-alchemist

For all available CLI options and usage instructions, run:

tt-alchemist --help

All APIs today accept a .mlir file that describe a model in TTIR dialect. Example usage:

# Generate a whole standalone C++ solution and run
tt-alchemist generate-cpp tools/tt-alchemist/test/models/mnist.mlir -o mnist_cpp --standalone
cd mnist_cpp
./run

# Similar to above, but use "local" libs from source dir - this saves on memory by not copying the whole dev package to the output dir
tt-alchemist generate-cpp tools/tt-alchemist/test/models/mnist.mlir -o mnist_cpp --local
cd mnist_cpp
./run

# Similarly for python
tt-alchemist generate-python tools/tt-alchemist/test/models/mnist.mlir -o mnist_python --local
cd mnist_python
./run

# Following APIs are intended to be used for debugging purposes

# Convert a mlir file to C++ code and print to console
tt-alchemist model-to-cpp tools/tt-alchemist/test/models/mnist.mlir

# Same, but for python (current support limited to few ops)
tt-alchemist model-to-python tools/tt-alchemist/test/models/mnist.mlir

Usage via lib

To use within another project (e.g. a frontend like tt-xla), build the library from source:

# Assuming the user had already built the tt-mlir compiler and turned on the python virtual env

# Build the tt-alchemist lib
cmake --build build -- tt-alchemist-lib

Then, you may call any of the APIs listed here.

ttnn-standalone

ttnn-standalone is a post-compile tuning/debugging tool.

Forge and third party ML models (PyTorch, Jax, ONNX, ...) can be compiled to a set of TTNN library op calls in C++. This generated code can then be used outside of the compiler environment. ttnn-standalone tool offers all the scaffolding needed to run the C++ code on device (build & run scripts).

Usage

# 1. Convert a model from TTIR dialect to EmitC dialect using ttmlir-opt
# 2. Translate the resulting EmitC dialect to C++ code using ttmlir-translate
# 3. Pipe the generated C++ code to a .cpp file
ttmlir-opt \
  --ttir-to-emitc-pipeline \
  test/ttmlir/EmitC/TTNN/sanity_add.mlir | \
ttmlir-translate \
  --mlir-to-cpp > \
  tools/ttnn-standalone/ttnn-standalone.cpp

# 1. Change dir to `tools/ttnn-standalone`
# 2. Use `run` script to compile and run the compiled binary
cd tools/ttnn-standalone
./run

Note: if you receive this error

-bash: ./run: Permission denied

running chmod +x run will set the execute permission on the script.

EmitPy

EmitPy is part of the tt-mlir compiler project. Its primary function is to translate MLIR IR from various dialects into human-readable, executable Python source code.

By representing Python laguage constructs within a dedicated EmitPy dialect, the project provides a structured pathway for lowering high-level computational graphs (e.g., from machine learning frameworks) into a familiar and flexible Python language, enabling rapid prototyping, debugging, and integration with Tenstorrent's TTNN open source library.

Current implementation enables support for MNIST and ResNet models.

Prerequisites

• Built ttmlir

• Activated virtual environment:

  source env/activate
  

Usage

ttmlir-opt

# 1. Convert a model from TTIR dialect to EmitPy dialect using ttmlir-opt
# 2. Translate the resulting EmitPy dialect to Python code using ttmlir-translate
# 3. Pipe the generated Python code to a .py file
ttmlir-opt --ttir-to-emitpy-pipeline test/ttmlir/Dialect/EmitPy/ttir_to_emitpy_pipeline_sanity.mlir | \
ttmlir-translate --mlir-to-python > example.py

builder

Builder offers support for building EmitPy modules. ttrt offers support for running EmitPy modules.

Optimizer

Optimizer is the main component responsible for performance. It is a collection of passes with the two most important purposes being optimizing tensor memory layouts and selecting optimal operation configurations.

Prerequisites

To use the optimizer:

• A physical Tenstorrent device must be present on the machine
• Build of tt-mlir must be with OpModel support enabled:

cmake -G Ninja -B build -DTTMLIR_ENABLE_OPMODEL=ON

Basic Usage

Optimizer is disabled by default. To enable it, use the enable-optimizer option:

ttmlir-opt --ttir-to-ttnn-backend-pipeline="enable-optimizer=true" input.mlir

Optimizer Options

The optimizer provides additional configuration options:

• enable-optimizer (default: false)

  • Enables the optimizer pass
  • Must be set to true to use any other optimizer options

• memory-layout-analysis-enabled (default: true)

  • Enables memory layout optimization
  • Shards tensors to maximize usage of fast L1 memory instead of DRAM

• max-legal-layouts (default: 64)

  • Maximum number of different layouts to generate for each operation during analysis
  • Higher values may provide better results but increase compile time

Example

# Enable optimizer with default settings
ttmlir-opt --ttir-to-ttnn-backend-pipeline="enable-optimizer=true memory-layout-analysis-enabled=true max-legal-layouts=8" input.mlir

PyKernel Guide

PyKernel is a Python interface for developing custom TTNN operations for Tenstorrent's AI accelerators. This guide explains how to use the PyKernel interface to implement your own TTNN operations.

Introduction to PyKernel

PyKernel provides a Python-based framework to define hardware-specific kernels that can be used with the TTNN framework. It allows developers to implement custom operations by defining compute kernels, reader/writer kernels, and control logic in a high-level Python interface.

The PyKernel framework consists of:

• PyKernelOp: Base class that manages kernel selection, compilation, and execution
• AST module: Decorators and utilities for defining kernels
• Types module: Type definitions for PyKernel operations

PyKernel Architecture

Foundationally, PyKernel is a compiler built on top of 3 core components, described below.

Python ast Frontend

The frontend of PyKernel is made to parse Python code and is enabled through using the ast (Abstract Syntax Tree) parser builtin to Python. By walking through the AST produced by this module, a MLIR module is created with the ttkernel dialect (among others such as arith, memref, scf). This MLIR module is then piped into the next step of the PyKernel compiler. For more information about the type of kernel code that can be parsed by the Frontend, refer to the ttkernel Op spec.

Direct To Metal (D2M) Kernel Code Generation

Another component of the tt-mlir project that PyKernel is built on is the D2M compiler infrastructure. This infrastructure enables dynamic generation of kernels to performantly execute ML models and is leveraged by providing the custom MLIR module created by the PyKernel frontend. The compilation flow runs a series of transformations on the MLIR module and lowers to the emitc dialect to translate the module into C++ code. This C++ code is the artifact that is consumed by the runtime to execute on Tenstorrent Hardware.

TTNN Generic Op

TTNN consists of Python bindings to precompiled kernels and operator factories that maintain API parity with PyTorch. The Generic Op extends this by operating directly on TTNN tensors and primitives but does not define its own factory or kernels. Instead, these must be supplied to the Generic Op to enable execution. PyKernel leverages this flexibility by injecting dynamically compiled C++ kernels into the Generic Op, allowing them to interface with TTNN data as if they were native “custom” ops. This mechanism serves as the integration layer that connects the compiler to TTNN.

Prerequisites

Before using PyKernel, ensure your environment is set up with:

• TT-MLIR built and installed
• Python 3.11 or newer
• Required Python packages
• TTMLIR_ENABLE_RUNTIME and TTMLIR_ENABLE_PYKERNEL flags set during build

Creating a Custom PyKernel Operation

To create a custom PyKernel operation, you need to:

1. Create a class that inherits from PyKernelOp
2. Implement the define_core_ranges method to specify the grid of cores for the operation
3. Define kernels using the @compute_thread(), @reader_thread(), or @writer_thread() decorators
4. Implement the invoke method to create and connect kernels
5. Define necessary circular buffers
6. Create a program descriptor that combines kernels and circular buffers

Basic Structure

from pykernel.kernel_ast import *
from pykernel.op import PyKernelOp
from pykernel.kernel_types import *

import ttnn
import torch

class MyCustomOp(PyKernelOp):
    # Define Core Grid
    def define_core_ranges(self, tensors, options):
        # Your logic to determine the core ranges
        core_1 = ttnn.CoreCoord(0, 0)
        core_2 = ttnn.CoreCoord(1, 1)
        return ttnn.CoreRangeSet([ttnn.CoreRange(core_1, core_2)])

    # Define compute kernel with appropriate decorator
    @compute_thread()
    def my_compute_kernel(cb_in: CircularBuffer, cb_out: CircularBuffer,
                         per_core_block_cnt: CompileTimeValue,
                         per_core_block_dim: CompileTimeValue):
        # Kernel processing code here
        return

    # Define reader kernel
    @reader_thread()
    def reader_kernel(cb_in: CircularBuffer, cb_out: CircularBuffer,
                     src_addr, num_tiles, start_id,
                     src_is_dram: CompileTimeValue):
        # Reader kernel code here
        return

    # Define writer kernel
    @writer_thread()
    def writer_kernel(cb_in: CircularBuffer, cb_out: CircularBuffer,
                     dst_addr, num_tiles, start_id,
                     dst_is_dram: CompileTimeValue):
        # Writer kernel code here
        return

    # The invoke method is the main entry point for kernel execution
    def invoke(self, in_tensor, out_tensor, **options):
        # Create circular buffers for input and output tensors
        cb_in = self.create_cb(in_tensor, 0)
        cb_out = self.create_cb(out_tensor, 1)

        # Prepare parameters for kernels
        start_id = 0
        is_dram = in_tensor.memory_config().buffer_type == ttnn.BufferType.DRAM
        num_tiles = options["num_tiles"]

        # Create kernels with appropriate parameters
        kernels = [
            self.create_kernel(
                MyCustomOp.my_compute_kernel,
                cb_in, cb_out,
                per_core_block_cnt=num_tiles,
                per_core_block_dim=1
            ),
            self.create_kernel(
                MyCustomOp.writer_kernel,
                cb_in, cb_out,
                out_tensor.buffer_address(),
                num_tiles, start_id,
                dst_is_dram=is_dram
            ),
            self.create_kernel(
                MyCustomOp.reader_kernel,
                cb_in, cb_out,
                in_tensor.buffer_address(),
                num_tiles, start_id,
                src_is_dram=is_dram
            )
        ]

        # Create and return the program descriptor
        return self.create_program(kernels, [cb_in, cb_out])

Kernel Types

PyKernel supports different types of kernels:

1. Compute Kernels: Process data on the compute units (e.g., SFPU - Scalar Floating-Point Unit)
2. Reader Kernels: Transfer data from memory to circular buffers
3. Writer Kernels: Transfer data from circular buffers to memory

Each kernel type has a specific decorator:

• @compute_thread() - For compute kernels that run on Tensix cores
• @reader_thread() - For reader kernels that transfer data from memory to circular buffers
• @writer_thread() - For writer kernels that transfer data from circular buffers to memory

These decorators handle the compilation of Python code into hardware-specific kernels. You can also use the older style decorators if needed:

• @ttkernel_tensix_compile() - Equivalent to @compute_thread()
• @ttkernel_noc_compile() - For both reader and writer kernels

Runtime Arguments

In PyKernel, you can pass runtime arguments to your kernels to control their behavior on a per-core basis. There are two types of runtime arguments:

1. Single-Core Arguments (Common Runtime Arguments): These are scalar values (integers) that are broadcast to all cores in the grid. They are passed as common_runtime_args to the create_kernel method.

2. Multi-Core Arguments (Runtime Arguments): These are lists of lists of integers, where each inner list corresponds to a core in the grid. This allows you to provide different values for each core. They are passed as runtime_args to the create_kernel method.

Single-Core Arguments

Single-core arguments are useful when all cores need the same value for a particular parameter. For example, num_tiles_per_core in the VecAdd example is a single-core argument because each core processes the same number of tiles.

Multi-Core Arguments

Multi-core arguments are necessary when each core requires a unique value. A common use case is distributing work across cores, where each core needs a different start_id to process its portion of the data. In the VecAdd example, start_id_multicore is a multi-core argument.

Default Core Range Behavior

If you do not override the define_core_ranges method in your PyKernelOp class, it will default to a single core at (0, 0). This is suitable for single-core operations like the EltwiseSFPU demo, where the entire operation runs on a single core.

Circular Buffers

Circular buffers are used to transfer data between kernels and memory. In the PyKernel framework, there are two aspects of circular buffers:

1. CircularBuffer class: Used in kernel definitions to represent a circular buffer
2. CB Descriptors: Used at runtime to configure the actual hardware circular buffers

CircularBuffer Class

The CircularBuffer class is defined in pykernel.types and is used in kernel definitions:

class CircularBuffer:
    def __init__(self, cb_id, tensor_shape=(8, 128, 128), dtype="Float32"):
        self.cb_id = cb_id
        self.tensor_shape = tensor_shape
        self.tile_shape = 32  # default to 32x32 tile shape
        self.tilized_shape = self.get_tilized_memref_shape()
        self.dtype = dtype

Creating Circular Buffers in the Invoke Method

In your custom operation's invoke method, you can create circular buffers using the create_cb helper method from the PyKernelOp base class:

def invoke(self, in_tensor, out_tensor, **options):
    cb_in = self.create_cb(in_tensor, 0)  # buffer_index=0
    cb_out = self.create_cb(out_tensor, 1)  # buffer_index=1

    # Use cb_in and cb_out in kernel creation
    # ...

    return self.create_program(kernels, [cb_in, cb_out])

The create_cb method handles the creation of the necessary format descriptors and buffer descriptors based on the tensor properties:

Kernel Decorator Options

The kernel decorators (@compute_thread, @reader_thread, and @writer_thread) accept two optional boolean arguments:

• verbose: When set to True, the PyKernel compiler will print the generated MLIR and the Python AST (Abstract Syntax Tree) during compilation. This is useful for debugging.
• optimize: When set to True, the PyKernel compiler will run an optimization pipeline on the generated MLIR before converting it to C++. This can improve the performance of your kernel.

Example: Vector Add Operation

The VecAdd operation adds two tensors element-wise. Let's examine a complete implementation based on the demo in test/pykernel/demo/vecadd_multicore_demo.py:

1. Define the Operation Class

from pykernel.kernel_ast import *
from pykernel.op import PyKernelOp
from pykernel.kernel_types import *

import ttnn
import torch

class VecAddMulticorePyKernelOp(PyKernelOp):
    # Kernel implementations will go here

2. Define Core Ranges

The define_core_ranges method specifies the grid of cores that the operation will run on.

def define_core_ranges(self, tensors, options):
    core_0 = ttnn.CoreCoord(0, 0)
    if self.max_core_ranges is None:
        core_1 = ttnn.CoreCoord(1, 1)
    else:
        core_1 = self.max_core_ranges
    return ttnn.CoreRangeSet([ttnn.CoreRange(core_0, core_1)])

3. Define the Compute Kernel

@compute_thread()
def add_multicore(
    cb_in0: CircularBuffer,
    cb_in1: CircularBuffer,
    cb_out: CircularBuffer,
    num_tiles,
    start_tile_id,
):
    binary_op_init_common(cb_in0, cb_in1, cb_out)
    add_tiles_init(cb_in0, cb_in1)

    end_tile_id = start_tile_id + num_tiles
    dst_reg = 0

    for i in range(start_tile_id, end_tile_id, 1):
        cb_wait_front(cb_in0, 1)
        cb_wait_front(cb_in1, 1)
        tile_regs_acquire()
        add_tiles(cb_in0, cb_in1, 0, 0, dst_reg)
        tile_regs_commit()

        cb_reserve_back(cb_out, 1)
        tile_regs_wait()
        pack_tile(dst_reg, cb_out, 0)
        tile_regs_release()

        cb_push_back(cb_out, 1)
        cb_pop_front(cb_in0, 1)
        cb_pop_front(cb_in1, 1)
        tile_regs_release()
    return

4. Define Writer Kernel

@writer_thread()
def writer_multicore(
    cb_out: CircularBuffer,
    dst_addr,
    num_tiles,
    start_id,
    dst_is_dram: CompileTimeValue,
):
    onetile = 1
    tile_bytes = get_tile_size(cb_out)
    dataformat = get_dataformat(cb_out)

    s0 = get_interleaved_addr_gen_fast(
        dst_is_dram, dst_addr, tile_bytes, dataformat
    )

    end_id = start_id + num_tiles
    for i in range(start_id, end_id, onetile):
        cb_wait_front(cb_out, onetile)
        l1_read_addr = get_read_ptr(cb_out)
        noc_async_write_tile(i, s0, l1_read_addr)
        noc_async_write_barrier()
        cb_pop_front(cb_out, onetile)
    return

5. Define Reader Kernel

@reader_thread()
def reader_binary_interleaved(
    cb_in0: CircularBuffer,
    cb_in1: CircularBuffer,
    src_addr0,
    src_addr1,
    num_tiles,
    start_id,
    src0_is_dram: CompileTimeValue,
    src1_is_dram: CompileTimeValue,
):
    onetile = 1
    tile_bytes0 = get_tile_size(cb_in0)
    dataformat0 = get_dataformat(cb_in0)

    s0 = get_interleaved_addr_gen_fast(
        src0_is_dram, src_addr0, tile_bytes0, dataformat0
    )

    tile_bytes1 = get_tile_size(cb_in1)
    dataformat1 = get_dataformat(cb_in1)

    s1 = get_interleaved_addr_gen_fast(
        src1_is_dram, src_addr1, tile_bytes1, dataformat1
    )

    end_id = start_id + num_tiles
    for i in range(start_id, end_id, onetile):
        cb_reserve_back(cb_in0, onetile)
        cb_reserve_back(cb_in1, onetile)

        src0_write_addr = get_write_ptr(cb_in0)
        src1_write_addr = get_write_ptr(cb_in1)

        noc_async_read_tile(i, s0, src0_write_addr)
        noc_async_read_tile(i, s1, src1_write_addr)

        noc_async_read_barrier()
        cb_push_back(cb_in0, onetile)
        cb_push_back(cb_in1, onetile)
    return

6. Implement the Invoke Method

The invoke method is the critical part that connects the kernels together and creates the program descriptor:

def invoke(self, a_tensor, b_tensor, out_tensor):
    # Create circular buffers
    cb_in0 = self.create_cb(a_tensor, 0)
    cb_in1 = self.create_cb(b_tensor, 1)
    cb_out = self.create_cb(out_tensor, 2)

    # Set up parameters
    is_a_dram = a_tensor.memory_config().buffer_type == ttnn.BufferType.DRAM
    is_b_dram = b_tensor.memory_config().buffer_type == ttnn.BufferType.DRAM
    is_out_dram = out_tensor.memory_config().buffer_type == ttnn.BufferType.DRAM

    num_tiles = ceil(max(map(lambda t: t.volume(), [a_tensor, b_tensor, out_tensor])) / 1024)
    num_cores = self.get_core_ranges().num_cores()
    num_tiles_per_core = int(num_tiles / num_cores)

    # Define the multicore runtime arguments
    start_id = 0
    start_id_multicore = []
    bb = self.get_core_ranges().bounding_box()
    for i in range(bb.start.x, bb.end.x + 1):
        start_id_multicore.append([])
        for j in range(bb.start.y, bb.end.y + 1):
            start_id_multicore[-1].append([start_id])
            start_id += 1

    # Create kernels with appropriate parameters
    kernels = [
        self.create_kernel(
            VecAddMulticorePyKernelOp.add_multicore,
            cb_in0,
            cb_in1,
            cb_out,
            num_tiles_per_core,
            start_id_multicore,
        ),
        self.create_kernel(
            VecAddMulticorePyKernelOp.writer_multicore,
            cb_out,
            out_tensor.buffer_address(),
            num_tiles_per_core,
            start_id_multicore,
            dst_is_dram=is_out_dram,
        ),
        self.create_kernel(
            VecAddMulticorePyKernelOp.reader_binary_interleaved,
            cb_in0,
            cb_in1,
            a_tensor.buffer_address(),
            b_tensor.buffer_address(),
            num_tiles_per_core,
            start_id_multicore,
            src0_is_dram=is_a_dram,
            src1_is_dram=is_b_dram,
        ),
    ]

    # Create and return the program descriptor
    return self.create_program(kernels, [cb_in0, cb_in1, cb_out])

Running the VecAdd Demo

The VecAdd demo demonstrates adding two tensors element-wise. This can be run using the pykernel-demo target:

source env/activate
# Ensure the TTMLIR_ENABLE_RUNTIME and TTMLIR_ENABLE_PYKERNEL flags are set during build
cmake --build build -- pykernel-demo

Demo Breakdown

Let's examine how to use the PyKernel operation in practice:

# Open a device
device = ttnn.open_device(device_id=0)

# Define tensor shapes and data
num_tiles = 4
shape = [1, num_tiles, 32, 32]
data = torch.rand(shape).to(torch.bfloat16)
data2 = torch.rand(shape).to(torch.bfloat16)


# Configure memory
dram_memory_config = ttnn.DRAM_MEMORY_CONFIG

# Create input tensors
a_tensor = ttnn.from_torch(
    data,
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=device,
    memory_config=dram_memory_config,
)

b_tensor = ttnn.from_torch(
    data2,
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=device,
    memory_config=dram_memory_config,
)


# Create output tensor
output_tensor = ttnn.allocate_tensor_on_device(
    ttnn.Shape(shape),
    ttnn.bfloat16,
    ttnn.TILE_LAYOUT,
    device,
    dram_memory_config,
)

# Create the custom operation
vecadd_op = VecAddMulticorePyKernelOp()

# Execute the operation with the tensors and options
output = vecadd_op(a_tensor, b_tensor, output_tensor)

# Compare with the built-in add operation
golden = ttnn.add(a_tensor, b_tensor)

# Convert to torch tensors for comparison
torch_golden = ttnn.to_torch(golden)
torch_output = ttnn.to_torch(output)

# Verify results
matching = torch.allclose(torch_golden, torch_output)
print(f"Tensors are matching: {matching}")
assert matching

This demo shows the complete workflow:

1. Opens a device
2. Creates input and output tensors with appropriate memory configuration
3. Instantiates the VecAddMulticorePyKernelOp class
4. Executes the operation by calling the op with tensors
5. Compares the result with the built-in TTNN implementation

Comparison with Native TTNN Operations

PyKernel operations integrate seamlessly with native TTNN operations. As shown in the demo, you can compare your custom PyKernel operation with built-in TTNN operations:

# Execute your custom PyKernel operation
output = vecadd_op(a_tensor, b_tensor, output_tensor)

# Execute the equivalent built-in TTNN operation
golden = ttnn.add(a_tensor, b_tensor)

# Convert both to torch tensors for comparison
torch_golden = ttnn.to_torch(golden)
torch_output = ttnn.to_torch(output)

# Verify the results match
matching = torch.allclose(torch_golden, torch_output)
print(f"Tensors are matching: {matching}")
assert matching

This approach allows you to:

1. Validate your custom operation against known implementations
2. Benchmark performance differences between custom and built-in operations
3. Extend the TTNN framework with operations not available in the standard library

Building and Testing

To build and test PyKernel, you need to enable both the runtime and PyKernel components:

source env/activate

# Configure with PyKernel enabled
cmake -G Ninja -B build \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_C_COMPILER=clang-17 \
    -DCMAKE_CXX_COMPILER=clang++-17 \
    -DTTMLIR_ENABLE_RUNTIME=ON \
    -DTTMLIR_ENABLE_PYKERNEL=ON

# Build the project
cmake --build build

# Run the PyKernel demo
cmake --build build -- pykernel-demo

The TTMLIR_ENABLE_RUNTIME and TTMLIR_ENABLE_PYKERNEL flags are essential for PyKernel functionality. Without these flags, the PyKernel components will not be built.

Best Practices

When developing with PyKernel, follow these best practices:

1. Separate concerns: Keep compute, reader, and writer kernels separate for better maintainability and reusability

2. Use appropriate decorators: Apply the correct decorator for each kernel type:

   • @compute_thread() for compute kernels
   • @reader_thread() for reader kernels
   • @writer_thread() for writer kernels

3. Implement the invoke method properly: The invoke method is critical as it connects all components:

   • Create circular buffers with appropriate parameters
   • Set up kernel parameters correctly
   • Create kernels with the right arguments
   • Return a program descriptor that includes all kernels and circular buffers

4. Handle memory configurations: Be aware of memory types (DRAM vs L1) when creating kernels

5. Reuse kernels: Create reusable kernels for common operations to avoid code duplication

6. Leverage caching: PyKernelOp automatically caches compiled kernels for performance

7. Test thoroughly: Always compare results with reference implementations or built-in TTNN operations

8. Document parameters: Clearly document the expected parameters for your PyKernel operation

Summary

PyKernel provides a flexible and powerful way to implement custom operations for Tenstorrent hardware. By following the pattern outlined in this guide, you can create your own operations that integrate seamlessly with the TTNN framework.

Key components of the PyKernel framework:

1. PyKernelOp base class: Handles kernel management, compilation, and caching
2. Kernel decorators: @compute_thread(), @reader_thread(), and @writer_thread()
3. CircularBuffer class: Represents circular buffers in kernel definitions
4. invoke method: The critical implementation that connects kernels and creates the program

The workflow for creating a custom PyKernel operation is:

1. Create a class that inherits from PyKernelOp
2. Define compute, reader, and writer kernels with appropriate decorators
3. Implement the invoke method to create circular buffers and connect kernels
4. Use the operation by instantiating your class and calling it with tensors and options

With PyKernel, you can extend the TTNN framework with custom operations that leverage the full power of Tenstorrent hardware while maintaining a clean, high-level Python interface.

Creating Bug Repros for TTNN Using TT-MLIR Codegen

While developing in tt-mlir, it's not uncommon to encounter bugs originating in the TTNN library. To isolate and report such bugs, a practical approach is to use the C++ codegen feature (EmitC) to generate a minimal repro. This guide walks you through how to create such repros and integrate them into the tt-metal repository, where TTNN is developed.

Step-by-Step Guide

Note: If you run into issues while following these steps, check the Known Issues section at the end of this guide for common problems and solutions.

1. Generate C++ Code from TT-MLIR

Use the ttnn-standalone tool to run the compiler and emit C++ code.

📖 See ttnn-standalone for instructions on how to generate C++ code from your MLIR input using EmitC.

2. Scope Down the Repro

Once you've generated the C++ code:

• Use the ttnn-standalone tool to run and debug it in isolation.
• Reduce the repro to the minimal example that still triggers the bug.
• Confirm the issue still reproduces reliably.

3. Clone the TT-Metal Repository

Clone the tt-metal repo:

git clone git@github.com:tenstorrent/tt-metal.git
cd tt-metal

4. Add the Repro to the GTest Infrastructure

Place your .cpp file in:

tests/ttnn/unit_tests/gtests/emitc/

and add it to the cmake file:

tests/ttnn/unit_tests/gtests/CMakeLists.txt

like so:

set(EMITC_UNIT_TESTS_SRC
    ${CMAKE_CURRENT_SOURCE_DIR}/emitc/test_sanity.cpp
    ${CMAKE_CURRENT_SOURCE_DIR}/emitc/your_test_name.cpp  # <<<===
)

Use the existing file test_sanity.cpp in that directory as a reference.

5. Modify the Repro for GTest

There are some modifications that need to be made in order to fit the GTest infra:

• Convert the main() function to a TEST(...) macro:

TEST(EmitC, YourTestName) {
    // Your original main function body here
}

• Remove any return statements from the TEST(...) function body.
• Replace #include "ttnn-precompiled.hpp" with #include "emitc.hpp"

6. Build the TTNN EmitC Tests

First, activate the python virtual env, and set some env variables:

source python_env/bin/activate
export TT_METAL_RUNTIME_ROOT=$(pwd)
export PYTHONPATH=$(pwd)

Then, build the tests:

./build_metal.sh --build-ttnn-tests

Note: some unrelated gtests might fail here, we can ignore them.

7. Run the EmitC Unit Tests

To run all EmitC tests:

./build/test/ttnn/unit_tests_ttnn_emitc

To run a specific test:

./build/test/ttnn/unit_tests_ttnn_emitc --gtest_filter=EmitC.YourTestName

8. Share the Repro

• Create a branch with your changes.
• Open a GitHub issue or comment on an existing one.
• Link to your branch and include the instructions for running the repro

./build_metal.sh --build-ttnn-tests
./build/test/ttnn/unit_tests_ttnn_emitc
./build/test/ttnn/unit_tests_ttnn_emitc --gtest_filter=EmitC.YourTestName

Known Issues

• Missing sfpi compiler or other dependencies If you encounter errors about a missing sfpi compiler or other system-level dependencies, refer to the tt-metal installation guide for instructions on installing the required packages.

• TTNN test compilation failures If the build fails when compiling TTNN tests, inspect the specific tests that caused the failure. If the failures are unrelated to EmitC tests, they can typically be ignored — this is a known issue.

Python Bindings

This page aims to clarify, document, and de-mystify the tt-mlir python bindings. It will do so by first highlighting the mechanism with which these bindings are generated and exposed to users. It will then document the nuances of nanobind, and the different parts of these bindings that must be written in by hand. Finally, it will go through a hands-on example of how to add your own functionality to the tt-mlir python bindings.

nanobind

Nanobind is the successor of the ubiquitous pybind project. In almost the same syntactical form, it provides a framework to define InterOp between C++ and Python. For more information about nanobind specifically, I'd recommend reading through the documentation. MLIR (and by extension: tt-mlir) leverages nanobind to create bindings for the C++ framework of Dialects, Ops, Types, Attributes, and Passes to be used in Python.

MLIR in Python

This section highlights the machinery and configuration with which MLIR can be exposed to Python, while still maintaining functional interop with the C++ code. For more context and information feel free to read the MLIR Python Documentation.

C-API

While the documentation provides a very lack-lustre explanation as to why the C-API exists, I am here to provide my take on the existence and purpose of the MLIR CAPI.

RTTI

MLIR, being a part of the llvm-project, follows their "custom" RTTI. For this reason, the entire C++ portion of the project isn't built with RTTI to enable to custom functionality. nanobind, however, requires RTTI to perform a lot of the casting and transformation required to interop with Python. This conflict leads to the natural desire for an alternative.

C doesn't have RTTI, it's a stable language without the extra convenience and machinery presented in C++. If a C-API were present, the python bindings can link against the C-API, relying on externally defined NanobindAdaptors to do the type conversions using nanobind mechanisms instead of relying on the C++/LLVM RTTI for the Python bindings.

C++ ABI

The C++ Application Boundary Interface (ABI) proves to be a challenging barrier to accessing functionality from C++. Without a defined stable ABI, it becomes difficult to deal with some of the complexity required to package and InterOp with Python. Specifically, dealing with templates, inheritance, and RTTI can prove quite the challenge.

To simplify this process, C provides a relatively stable ABI. The C-API also acts as a wrapper around the complex C++ functions, providing a simple "trampoline" for Python to link against.

nanobind x C-API Functionality

In the previous section, I mentioned NanobindAdaptors. This file helps to define some of the key design decisions made when linking the Python bindings against the C-API instead of the underlying C++ API. Functionally, the Python bindings act as a "wrapper" around the CAPI, exposing the functionality through python.

include/mlir-c/Bindings/Python/Interop.h

This file is key to defining the InterOp between the C-API and Python w.r.t. maintaining and accessing information in a pointer. It exposes an AI that interfaces immediate data pointers with python capsules. PyCapsules are essentially thin wrappers around data pointers in Python. The critically contain data (void*), destructor method, and a name.

Within the Interop, the assumption is that the data's ownership and lifetime is managed by some bound object that was created in C++. This file merely provides the API with which the underlying data pointer is passed around as either a PyCapsule or the raw pointer, and this file provides the type conversion utilities to convert between Python and C from an underlying object.

include/mlir/CAPI/Wrap.h

This header defines the API to InterOp between C-API objects and their C++ equivalent. By calling wrap() on a C++ MLIR object to have the underlying data create a C-API object on the same memory, and unwrap() does it the other way around.

They key caveat with this wrapping/unwrapping is the ownership over the lifetime of the data itself. The constructors for almost all of the primitives have already been defined in C++. As such the syntax for creating a new C-API object is more the syntax of creating an object in C++ and wrapping it into a CAPI object. The lifetime of the pointer is therefore maintained by the CAPI object as it gets passed around in return objects.

include/mlir/Bindings/Python/NanobindAdaptors.h

As the CAPI object gets bounced around in memory, the ownership and lifetime of the data must eventually reach python to be controlled by the user. The implementation details are not relevant to this component as to how the data reaches python. This component provides the utility to create copies of the underlying data and send them through nanobind, effectively framing itself as the InterOp component between CAPI objects and their nanobind equivalents.

Through the carefully created contract between these components of the MLIR project, the IR primitives are exposed to Python, created in C++, and bounced off of the C-API. While I may have gleaned over the other supporting mechanisms in this explanation, explore the parent directories for these three files for a more detailed look into the semantics of ownership and such.

Defining the C-API.

For primitives to be defined for use in Python, they must first be implemented in C++. This is outside of the scope of the Python specific code, please refer to the rest of tt-mlir documentation for references on this. Once the C++ functionality is defined, the C-API must be constructed on top of this to serve as the "InterOp" layer.

get & Constructing C-API Objects

Since most constructors for IR primitives are created in C++, the goal is to construct objects in C++, but have the ownership exposed to Python. We do this through the creation of a Get function. The get function will essentially intake primitive C-types, and invoke the ::get operator in C++ to construct the object. A simple code example for the ttkernel.TileType is shown below:

include/ttmlir-c/TTTypes.h

// We export the function outside of the scope of "C" such that it can be defined later using C++ methods.

MLIR_CAPI_EXPORTED MlirType ttmlirTTTileTypeGet(MlirContext get, unsigned height, unsigned width, uin32_t dataType);

lib/CAPI/TTTypes.cpp

MlirType ttmlirTTTileTypeGet(MlirContext ctx, unsigned height, unsigned width, uint32_t dataType) {
    // We return the **wrapped** created C++ object, transferring Ownership to the C-API
    return wrap(
        TileType::get(
            unwrap(ctx), // Now we unwrap the MlirContext object to cast it to a mlir::MLIRContext object (w/o affecting ownership)
            llvm::SmallVector<std::int64_t>{height, width}, // We construct the list here since a list isn't natively defined in the C-API,
            static_cast<ttcore::DataType>(dataType) // Here we cast the int value to get the Enum value from `ttcore::DataType`
        ) // Invoking the builtin get operator to create and get the pointer for some object
    );
}

The key details to note are the reliance on C++ methods in the get definition like intiializer lists. By leveraging the InterOp the get method will return a pointer which can easily be represented in the C-API and owned as such, while masking the complexities of the C++ underneath from nanobind. Definitions such as these must either be written by hand (as shown above), or they can automatically be generated for certain IR primitives. We will learn more about that below.

Generating Bindings

This section will outline the mechanism with which bindings are generated, and the intricacies of this step.

Declaring Python Bindings

The first step to kicking off binding generation is to declare that they should exist for some dialect. MLIR provides a CMake module (AddMLIRPython) which exposes the following utility functions which can be declared to state what Python bindings are generated. For more information about the specific arguments and expected structure of these CMake functions refer to the AddMLIRPython module and python/CMakeLists.txt.

declare_mlir_python_sources

Overview
This function provides an interface to directly copy .py source files into the final built python module.

Key Arguments

• ADD_TO_PARENT defines the Parent name to which this source will be added to, inheriting the location.

Usecases

• We use it to declare generic "Parents" which contain the generated/declared python files from many of the submodules within the dialects.
• We use it to directly copy over key test infrastructure like ttir_builder as purely python programmed modules.

declare_mlir_dialect_python_bindings

Overview
This function is the key to invoking the mechanism to generate python bindings from Tablegen Definitions.

Key Arguments

• TD_FILE Relative to ROOT_DIR, where the Tablegen Definition file to build bindings off of is located. Note: This currently just forwards the TD files from include/ttmlir/Dialect.
• SOURCES Raw python files associated with bindings. Note: These files will essentially forward the generated modules forward.
• GEN_ENUM_BINDINGS_TD_FILE if GEN_ENUM_BINDINGS is ON, this will build enum bindings from the defined Tablegen file.
• DIALECT_NAME What name the dialects should be generated under.

Usecases

• We use this CMake function to define and generate the bindings for the ttkernel, ttir, tt, and ttnn dialects.

declare_mlir_python_extension

Overview
This is the CMake function used to link C++ Source Files + declared nanobinds into the generated python module.

Key Arguments

• EMBED_CAPI_LINKS_LIBS This is to declare the libraries used to link against the CAPI in the bindings. Learn more in the CAPI section below.
• PRIVATE_LINK_LIBS Declares other libraries that are linked against the Python bindings.

Usecases

• We use this function to build and link all of our custom nanobinds and hand-written Type/Attr bindings into the ttmlir module.

add_mlir_python_common_capi_library

Overview
This function adds a shared library embedding all of the core CAPI libs needed to link against extensions.

add_mlir_python_modules

Overview
This is the final packaging function of the python bindings, linking all of the sources together and packaging it into a built module.

Building MLIR Primitives from Tablegen

The declare_mlir_dialect_python_bindings leverages a mechanism of the mlir-tblgen to build the python bindings for some defined dialect. What are the intricacies of this functionality?

mlir-tblgen

This tool parses .td Tablegen files to automatically generate C++ code to implement that functionality in MLIR. We leverage the Tablegen language to define our dialects in tt-mlir, and this tool is exactly what gets invoked to build and generate the code to functionally use this dialect in our codebase.

Trivial Constructors

To deal with automatically generating the functionality around an Operation, a certain amount of generality is needed to deem the problem trivial enough to generate. All of the IR primitives are thankfully able to be constructed from .td to their relevant C++ implementations. However, as shown in the TileType example, the conversion from simple C primitives (+ pre-defined MLIR C-API types) to C++ get functions isn't trivial. For this reason, we can start to analyze the IR primitives and deem which ones are trivial for C-API generation, and which must be implemented by hand.

• enum
  • The enum type can be considered very generic. With the underlying data storage type being integral values, and an optional String representation in MLIR. By iterating over all of the user defined enum values, a very trivial constructor can be made to automatically generate enums.
• operation
  • Operations are a unique case where the constructor isn't often generic enough; however, the OperationState exists as a strictly defined struct which contains all of the relevant construction details and implementation requirements for an operation. For this reason, while it is not trivial, it is generic enough that the OperationState can be relied on to form a mechanism which automatically generates C-API builders.
• Types/Attributes
  • Types and Attributes unfortunately receive the short end of the stick. Their constructors are wildly generic, and there is no baseline for what is required in the construction of a Type/Attr. For this reason, at the current moment these primitives aren't supported for automatic generation in mlir-tblgen, and must be defined by hand.

Writing Bindings

With the understanding that not all bindings can be automatically generated for us, we can head into the intricacies of defining your own bindings for Types/Attrs.

LLVM-Style Pythonic "Type Information" + Casting

An important caveat to introduce before entering the domain of writing our own bindings is the understanding of how MLIR approaches the problem of downcasting w.r.t. IR primitives. Considering the C-API doesn't have an inheritance structure, Python is required to uphold the inheritance structure and hold the type information such that casting is possible between primitives and their specific implementation (ex: going from MlirAttribute -> TTNNLayoutAttr).

This mechanism can be exposed to Python in multiple different ways, where MLIR supports a specific implementation of an mlir_attribute_class and mlir_type_class which intake 2 additional C-API functions. To initialize a class using this structure the following functions are required:

• myAttributeGet: to construct the Type/Attr
• myAttributeGetTypeID: provides a unique static TypeID for myAttribute
• isAMyAttribute: boolean to see if higher level type is of the same type.

This will then provide an interface where in python a type can be cast by calling the constructor method of some downcasted type:

# Example to downcast using MLIR provided methods.
my_attribute = mlir.MyAttribute(attr: _mlir.ir.MlirAttribute)

Choosing a direct C++ structure instead of C-API

Those who are familiar with the tt-mlir python bindings may be aware that our code structure looks drastically different from this, why is that? The answer lies in the redundancy and lack of extensive use of the nanobind mechanisms around tt-mlir Python bindings.

As mentioned in the C-API section, the C-API is required to form the contract between C++ -> Python, to reduce the collisions with RTTI and the unstable ABI from C++. That being said, it's not unsupported to still directly access C++ members from nanobind and skip the C-API Builder functions, instead just opting to create in C++ directly and then wrap that result. This is the approach taken "consciously" in the tt-mlir python bindings.

What are the consequences of this design decision? The advantages?

Direct MLIR Casting Support

Instead of relying on Python for casting, and defining C-API functions to support this functionality; this approach allows us to directly use mlir::isa, mlir::cast, etc... in it's place.

For example, we support tt_attribute_class and tt_type_class, which leverage isa and dyn_cast to downcast to Types and Attrs by wrapping the Python types and operating on the underlying C++ types.

This also brings about some potential collisions with RTTI from nanobind. None are present in the bindings (as far as I know), but the bindings are exposed to this problem moving forward.

Simpler Initialization Structures

Instead of having to invoke a C-API function to define the get method in nanobind we can directly invoke the wrap(CppType::get(...)) functionality that the C-API ends up calling. The primary difference is the native support for complex data structures like vector and map through nanobind. Take for example the following initialization for an attribute:

// C-API Definition for myAttributeGet
MlirAttribute myAttributeGet(MlirContext ctx, int* array, size_t arraySize) {
    return wrap(MyAttribute::get(ctx, std::vector<int>{array, array + arraySize}));
}

// nanobind direct invocation
tt_attribute_class(m, "MyAttribute")
    .def_static("get", [](MlirContext ctx, std::vector<int> arr) {
        return wrap(MyAttribure::get(ctx, arr));
    })

// nanobind invocation through C-API
mlir_attribute_class(m, "MyAttribute", myAttributeGetTypeId, isAMyAttribute)
    .def_static("get", [](MlirContext ctx, std::vector<int> arr) {
        return myAttributeGet(ctx, arr.data(), arr.size());
    })
// Note: While this may seem like a trivial change, the cost for retaining the function signature in C begins to grow very quickly. Especially when considering maps and more complex data structures.

Again, this does come with some nuances w.r.t. the ABI, but for our simple usecase of the bindings it can be considered acceptable...

Wait... why are we still defining the CAPI Builders Then?

This leads to an underlying question: What's the point of still defining the CAPI functions if we actually never end up using them? The answer is that we would ideally still maintain the infrastructure to backtrack our changes if we end up making more extensive use of the Python bindings and come across nasty ABI/RTTI issues, or MLIR upstreams significant changes to the Python bindings where we would have to leverage their architecture. With regards to the latter, I have asked some of the contributors and received "iffy" responses, with the general answer being that major changes are not planned for the MLIR Python bindings infrastructure.

That being said, for the low low cost of a few redundant functions being defined, we have a clear backup route in case the Python bindings blow up in our faces. I do think this argument is built on significant personal opinion, in the future we may change the strategy for the bindings. For now, it makes the structure of our python code cleaner, while having a clear route forward if something breaks.

Each MLIR project I've used as a reference approaches the problems differently. AFAIK the bindings are generally defined however the end user desires to invoke them :)

General Structure

Considering that mlir-tblgen will handle the generation of the underlying C++ code, we only need to define the C Builders and the nanobinds for each of the Types/Attrs we would like to add.

This often comprises of the following contributions:

• Declaring the C-API Header Function(s) in include/ttmlir-c
• Defining the C-API Function(s) in lib/CAPI
• Writing out the nanobind for that Type/Attr in python/.

Example: Defining ttkernel Python Bindings

In this section, we will go through a worked example on the different steps required to expose functionality for the TTKernel dialect.

1. We will continue while assuming that the TTKernel dialect has been defined using Tablegen and already has a valid target that compiles the C++ functionality. We will also assume that the current CMake build targets and functionality that uphold the rest of the ttmlir dialects already exists.
2. Declare and register the TTKernel dialect in the C-API by calling the MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(TTKernel, ttkernel); macro in include/ttmlir-c/Dialects.h:

// File: include/ttmlir-c/Dialects.h

#include "mlir-c/IR.h"

#ifdef __cplusplus
extern "C" {
#endif

MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(TTKernel, ttkernel);

#ifdef __cplusplus
}
#endif

3. Declare CAPI Builder for all of the Types (namely only CBType needs to be implemented) in include/ttmlir-c/TTKernelTypes.h

// File: include/ttmlir-c/TTKernelTypes.h

#include "ttmlir-c/Dialects.h"

#ifdef __cplusplus
extern "C" {
#endif

MLIR_CAPI_EXPORTED MlirType ttmlirTTKernelCBTypeGet(
    MlirContext ctx, uint64_t port, uint64_t address,
    MlirType memrefType);

#ifdef __cplusplus
}
#endif

4. Declare the CAPI builder target in lib/CAPI/CMakeLists.txt by adding TTKernelTypes.cpp as a source to TTMLIRCAPI.
5. Define the Dialect by formalling applying the generated Dialect type into the CAPI_DIALECT_REGISTRATION macro.

// File: lib/CAPI/Dialects.cpp

#include "ttmlir-c/Dialects.h"

#include "mlir/CAPI/Registration.h"
#include "ttmlir/Dialect/TTKernel/IR/TTKernel.h"

MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(
    TTKernel, ttkernel, mlir::tt::ttkernel::TTKernelDialect)

6. Define the CAPI get method for CBType

// File: lib/CAPI/TTKernelTypes.cpp

#include "ttmlir-c/TTKernelTypes.h"
#include "mlir/CAPI/IR.h"
#include "mlir/CAPI/Support.h"

#include "ttmlir/Dialect/TTKernel/IR/TTKernelOpsTypes.h"

using namespace mlir::tt::ttkernel;

MlirType ttmlirTTKernelCBTypeGet(MlirContext ctx, MlirType memrefType) {
  return wrap(CBType::get(unwrap(ctx), mlir::cast<mlir::MemRefType>(unwrap(memrefType))));
}

7. Define the nanobind build target in python/CMakeLists.txt by adding ttkernel as a dialect, and providing TTkernelModule.cpp as a source for TTMLIRPythonExtensions.Main.

# Define ttkernel dialect
declare_mlir_dialect_python_bindings(
  ADD_TO_PARENT TTMLIRPythonSources.Dialects
  ROOT_DIR "${TTMLIR_PYTHON_ROOT_DIR}"
  TD_FILE dialects/TTKernelBinding.td
  SOURCES dialects/ttkernel.py
  DIALECT_NAME ttkernel
)

8. Create python/dialects/TTKernelBindings.td to forward the tablegen for TTKernel to the CMake dialect target:

include "ttmlir/Dialect/TTKernel/IR/TTKernelOps.td"

9. Create nanobind module for TTKernel Dialect in python/TTMLIRModule.cpp

// Representation of the Delta you have to add to TTMLIRModule.cpp in the correct locations
NB_MODULE(_ttmlir, m) {
  m.doc() = "ttmlir main python extension";

  m.def(
      "register_dialect",
      [](MlirContext context, bool load) {
        MlirDialectHandle ttkernel_handle mlirGetDialectHandle__ttkernel__();
        mlirDialectHandleRegisterDialect(ttkernel_handle, context);
        if (load) {
          mlirDialectHandleLoadDialect(ttkernel_handle, context);
        }
      },
      py::arg("context"), py::arg("load") = true);

  auto ttkernel_ir = m.def_submodule("ttkernel_ir", "TTKernel IR Bindings");
  mlir::ttmlir::python::populateTTKernelModule(ttkernel_ir);
}

10. Define populateTTKernelModule in python/TTKernelModule.cpp

// File: python/TTKernelModule.cpp
#include <vector>

#include "ttmlir/Bindings/Python/TTMLIRModule.h"

#include "mlir/CAPI/IR.h"
#include "ttmlir-c/TTKernelTypes.h"

#include "ttmlir/Dialect/TTKernel/IR/TTKernelOpsTypes.h"

namespace mlir::ttmlir::python {
void populateTTKernelModule(py::module &m) {
  tt_type_class<tt::ttkernel::CBType>(m, "CBType")
      .def_static("get",
                  [](MlirContext ctx, uint64_t port, uint64_t address,
                     MlirType memrefType) {
                    return ttmlirTTKernelCBTypeGet(ctx, port, address,
                                                   memrefType);
                    // Note that for more complex constructors / out of ease this could also be defined using the wrap(CBType::get) style constructor.
                  })
      .def_prop_ro("shape", [](tt::ttkernel::CBType &cb) {
            cb.getShape().vec();
        })
      .def_prop_ro("memref", &tt::ttkernel::CBType::getMemref);
}
} // namespace mlir::ttmlir::python

11. Finally, expose the built python bindings using a "trampoline" python file in python/dialects/ttkernel.py

from ._ttkernel_ops_gen import *
from .._mlir_libs._ttmlir import register_dialect, ttkernel_ir as ir

# Import nanobind defined targets into ttkernel.ir, and the rest of the generated Ops into the top-level ttkernel python module.

Concluding The Example

While there are quite a few steps for adding a whole new dialect, often times more than not you will only need a subset of these steps to add a new Type/Attr to some existing dialect. Even less to modify the signature of some existing Type/Attr in the bindings.

Using the Python Bindings

This section will cover the basics of using the Python bindings. I think the folks at MLIR have produced documentation that can help you get up to speed quickly. This section will go over some of the nuances of using the python bindings that ttmlir has defined explicitly.

Interfacing with Generated Op Classes

The unfortunate reality is that documentation for autogenerated Ops isn't present. Fortunately, argument names are preserved and the function structure can be invoked by leveraging the help function in python. Iteratively running through the functions you want to implement can be helpful.

MLIRModuleLogger

Almost all of the python bindings behave exactly as expected coming from the ttmlir python bindings. A weird addition I think would provide some more context on nanobind and managed memory would be the MLIRModuleLogger.

This class is defined in C++ to attach to an existing MLIRContext, adding hooks to save the module to a std::vector<std::string, std::string>. Binding this forward through nanobind requires some delicacy about the state of this MLIRModuleLogger object. It needs to modify memory managed by C++, but it attaches to a context that exists in Python. This state management is done through nanobind owning and managing a thinly wrapped pointer to the C++ object by setting the return_value policy.

Using the Python bindings when traversing frequently through memory outside of the IR primitives requires some delicacy to ensure data is preserved and the code functions as intended.

Flatbuffers

Flatbuffers are the binary serialization format used by TTMLIR and they currently come in a few flavors (designated by the file extension):

• .ttsys: A system description file that is the mechanism for supplying target information to the compiler. These can be collected on a target machine and downloaded to a development machine to enable cross-compilation.
• .ttnn: A compiled binary file intended to be loaded and executed by the TTNN backend runtime.
• .ttb: A compiled binary file intended to be loaded and executed by the TTMetal backend runtime (Unsupported).

LLVM Dependency Update

The TT-MLIR compiler has the following LLVM-related dependencies:

Dependency Hierarchy

• LLVM - Core LLVM infrastructure
• StableHLO - High-level operations dialect
  • Depends on LLVM as an external dependency
• Shardy - Sharding and partitioning dialect
  • Depends on StableHLO as an external dependency

Dependency Management

These projects are actively developed upstream, and our compiler code needs to occasionally update these dependencies to:

• Incorporate new features and optimizations
• Apply security patches and bug fixes
• Maintain compatibility with the broader LLVM ecosystem

Updating Cadence

We should update our compiler LLVM dependencies at least once every three months. The schedule for the next updates should be around:

• November 2025
• February 2026
• May 2026

Updating Dependencies

1. Identify compatible versions across the dependency chain:

• Select the latest Shardy version from the main branch and record it as SHARDY_COMMIT
• On this commit of Shardy, we can obtain the used version of StableHLO:
  • STABLEHLO_COMMIT = "..."
• On this commit of StableHLO, we can obtain the used version of LLVM:
  • LLVM_COMMIT = "..."

2. Prepare the development environment using the base tt-mlir IRD Ubuntu image. This clean Docker image lacks prebuilt dependencies, enabling fresh builds for dependency updates and troubleshooting:

ghcr.io/tenstorrent/tt-mlir/tt-mlir-base-ird-ubuntu-22-04:latest

3. Synchronize with the latest tt-mlir main branch:

git pull

4. Create uplift branch in the following format:

git checkout -b [alias]/[year]_[month]_llvm_dependency_update
# For example: sdjordjevic/2025_august_llvm_dependency_update

5. Update dependency versions in the CMakeLists.txt configuration:

• LLVM_PROJECT_VERSION with LLVM_COMMIT obtained from step 1
• STABLEHLO_VERSION with STABLEHLO_COMMIT obtained from step 1
• SHARDY_VERSION with SHARDY_COMMIT obtained from step 1

6. Build the local environment following this section of doc:

cmake -B env/build env
cmake --build env/build

7. Build the tt-mlir compiler with runtime, optimizer and StableHLO enabled following this section of doc:

source env/activate
cmake -G Ninja -B build -DTTMLIR_ENABLE_RUNTIME=ON -DTTMLIR_ENABLE_OPMODEL=ON -DTTMLIR_ENABLE_STABLEHLO=ON -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
cmake --build build

8. Resolving Shardy patch compatibility issues:

• The most challenging aspect of dependency uplifts involves maintaining the Shardy patch. Since Shardy only supports Bazel builds, we maintain a custom patch to ensure CMake compatibility. Due to active upstream development, this patch typically requires updates during each dependency uplift. The most common error is something like this

CMake Error at /opt/ttmlir-toolchain/lib/cmake/llvm/AddLLVM.cmake:568 (add_library):
Cannot find source file:

  constant_merger.cc

Tried extensions .c .C .c++ .cc .cpp .cxx .cu .mpp .m .M .mm .ixx .cppm
.ccm .cxxm .c++m .h .hh .h++ .hm .hpp .hxx .in .txx .f .F .for .f77 .f90
.f95 .f03 .hip .ispc
Call Stack (most recent call first):
/opt/ttmlir-toolchain/lib/cmake/mlir/AddMLIR.cmake:386 (llvm_add_library)
/opt/ttmlir-toolchain/src/shardy/shardy/dialect/sdy/transforms/export/CMakeLists.txt:35 (add_mlir_library)

• This indicates that library files have been renamed, moved, or deleted, requiring patch regeneration. To resolve this, navigate to the Shardy environment directory (e.g., /opt/ttmlir-toolchain/src/shardy/shardy/dialect/sdy/transforms/export/CMakeLists.txt) and update the library definition to reflect current .cc files. For example:

add_mlir_library(SdyTransformsExportPasses
close_shardings.cc
constant_or_scalar_merger.cc // Renamed from constant_merger.cc
drop_sharding_rules.cc
explicit_reshards_util.cc // Added as a new file, previously didn't exist
export_pipeline.cc
insert_explicit_reshards.cc
remove_propagation_debug_info.cc // Added as a new file, previously didn't exist
remove_sharding_groups.cc
reshard_to_collectives.cc
sharding_constraint_to_reshard.cc
sink_data_flow_edges.cc
update_non_divisible_input_output_shardings.cc // Removed temp_explicit_reshards_for_optimizations.cc as it doesn't exist anymore

• Another common error involves deleted libraries. If libraries were removed in the previous step, update the SHARDY_LIBS configuration for the TTMLIRCompilerStatic target:

CMake Error at /opt/ttmlir-toolchain/lib/cmake/llvm/AddLLVM.cmake:605 (add_dependencies):
The dependency target "SdyRoundtripImportShardyAttrs" of target
"obj.TTMLIRCompilerStatic" does not exist.
Call Stack (most recent call first):
/opt/ttmlir-toolchain/lib/cmake/mlir/AddMLIR.cmake:386 (llvm_add_library)
lib/CMakeLists.txt:71 (add_mlir_library)

• Missing symbol errors represent another category of common issues, typically manifesting as:

ld.lld: error: undefined symbol: mlir::sdy::log::LogMessageFatal::LogMessageFatal(mlir::sdy::log::LogMessageData)
>>> referenced by dialect.cc:265 (/opt/ttmlir-toolchain/src/shardy/shardy/dialect/sdy/ir/dialect.cc:265)
>>>               dialect.cc.o:(mlir::sdy::MeshAttr::getAxisSize(llvm::StringRef) const) in archive lib/libSdyDialect.a

ld.lld: error: undefined symbol: mlir::sdy::log::LogMessage::stream()
>>> referenced by dialect.cc:265 (/opt/ttmlir-toolchain/src/shardy/shardy/dialect/sdy/ir/dialect.cc:265)
>>>               dialect.cc.o:(mlir::sdy::MeshAttr::getAxisSize(llvm::StringRef) const) in archive lib/libSdyDialect.a

ld.lld: error: undefined symbol: mlir::sdy::log::LogMessageFatal::~LogMessageFatal()
>>> referenced by dialect.cc:265 (/opt/ttmlir-toolchain/src/shardy/shardy/dialect/sdy/ir/dialect.cc:265)
>>>               dialect.cc.o:(mlir::sdy::MeshAttr::getAxisSize(llvm::StringRef) const) in archive lib/libSdyDialect.a
clang++-17: error: linker command failed with exit code 1 (use -v to see invocation)

• This error pattern is similar to the file renaming issue but occurs when new files are added to Bazel build targets without corresponding updates to our CMake equivalents. To resolve this, examine the Shardy environment directory (e.g., /opt/ttmlir-toolchain/src/shardy/shardy/common/CMakeLists.txt) and ensure all .cc files are included in the library definition. For example:

add_mlir_library(SdyCommonFileUtils
file_utils.cc
logging.cc // Added new file, previously didn't exist
save_module_op.cc

9. After fixing the shardy patch errors locally in your shardy environment folder, we need to create a patch diff with all your changes in shardy directory. Follow these steps to produce the updated shardy patch:

# First we need to undo the commit from already applied patch so we can get the diffs
git reset --soft HEAD~1
# Produce the full diff patch including the staged changes from the previous patch (--cached)
git diff --cached > shardy.patch
# Copy generated patch to env folder to update the patch
cp shardy.patch ../../../tt-mlir/env/patches/

• Since patches are applied during environment builds, verification requires rebuilding both the environment and project:

  cmake -B env/build env
  cmake --build env/build
  source env/activate
  cmake -G Ninja -B build -DTTMLIR_ENABLE_RUNTIME=ON -DTTMLIR_ENABLE_OPMODEL=ON -DTTMLIR_ENABLE_STABLEHLO=ON -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
  cmake --build build

10. Address MLIR and LLVM API compatibility issues:

• LLVM updates frequently introduce breaking changes to API signatures and interfaces
• Build failures typically manifest as incorrect function signatures or deprecated API usage
• Review compilation errors and update code to match current LLVM APIs

11. Create and submit the update pull request:

• Commit all changes with descriptive messages
• Push the branch to the remote repository
• Create a PR with the title format: "[Year] [Month] LLVM Dependency Update"

12. Verify CI pipeline success:

• Ensure all CI checks pass
• Confirm tt-mlir test suite completion without failures

13. Validate frontend compatibility:

• Obtain the commit ID from your latest uplift branch commit:
• Execute On PR actions for each frontend repository:
  • forge-fe
  • tt-xla
• Run CI workflows from the main branch using the tt-mlir commit from step 1

CI

Our CI infrastructure is currently hosted in the cloud. Cloud machines are used and linked as GitHub runners.

Overview

CI automatically triggers on:

• Pull requests - validates code changes before merging
• Pushes to main - typically when PRs are merged
• Nightly runs - comprehensive testing with all components
• Uplift PRs - special PRs that update tt-metal to the latest version

The CI system automatically collects analytics data from each workflow run, including test results and code coverage. It also publishes the latest documentation to GitHub.

Builds

CI performs several types of builds:

Release Builds

• speedy - optimized for performance and speed
• tracy - includes runtime tracing and debug capabilities with performance measurements

Development Builds

• Debug build - includes unit tests and code coverage collection
• MacOS build - ensures cross-platform compatibility
• Wheels - Python package distributions
• Clang-tidy - static code analysis

Release builds include the runtime needed to execute on TT hardware, making them suitable for integration testing. The debug build runs unit tests and generates code coverage reports that are published to Codecov with detailed results linked in PR comments.

Release Build Components

Release builds do more than just compile tt-mlir - they also prepare tests, build tools, and create wheels. Components are configured in .github/settings/build.json:

{ "image": "tracy", "script": "explorer.sh" }

• image: Specifies which release build to use (speedy or tracy)
• script: Build script located in .github/build_scripts/
• if (optional): Links to optional components - only builds when that component is enabled

Before running build scripts, the workflow will activate the default TT-MLIR Python venv and set a number of useful environment variables:

• WORK_DIR - set to repo root
• BUILD_DIR - set to build artifacts
• INSTALL_DIR - set to install artifacts
• BUILD_NAME - name of the build image

Uploading Artifacts from Build Scripts

Build scripts can upload their output files as artifacts for later use in testing. This is especially important for optional components that may not always run.

How it works:

• Build scripts write artifact information to a JSON file specified by the $UPLOAD_LIST environment variable
• Each artifact entry contains a name (identifier) and path (file location)
• The CI system automatically uploads these artifacts after the build completes

Example:

echo "{\"name\":\"ttrt-whl-$BUILD_NAME\",\"path\":\"$WORK_DIR/build/tools/ttrt/build/ttrt*.whl\"}," >> $UPLOAD_LIST

This example uploads Python wheel files with a descriptive name that includes the build type.

Please note >> as it appends to existing list.

Optional Components

As the codebase grows, CI can become slow and bloated. To keep development efficient, some components are made optional and only run when needed:

When optional components run:

• ✅ Nightly builds (full testing)
• ✅ Uplift PRs (tt-metal version updates)
• ✅ PRs that modify the component's code
• ❌ Regular PRs (unless component files changed)

Good candidates for optional components:

• Mature, stable features that rarely break
• Legacy code that's still supported but not actively developed
• Rarely-used functionality where breakage isn't critical
• Time-intensive wheel builds (when functionality is tested elsewhere)

Configuring Optional Components

Define components in .github/settings/optional-components.yml:

component_name:
  - path/to/files/*.py
  - specific/file.py
  - another/directory/**

The component name can then be referenced in build and test configurations in "if" field:

Example:

emitc:
  - test/ttmlir/EmitC/**
  - tools/ttnn-standalone/ci_compile_dylib.py
  - include/ttmlir/Conversion/TTNNToEmitC/**
  - lib/Conversion/TTNNToEmitC/**

Build example:

{ "image": "speedy", "script": "emitc.sh", "if": "emitc" }

Test example:

{ "runs-on": "n150", "image": "speedy", "script": "emitc.sh", "if": "emitc" }

This makes the EmitC component optional - it only builds/tests when EmitC-related files are modified in a PR.

Testing

Testing is performed inside the call-test.yml workflow as run-tests jobs. It uses a matrix strategy, which means that multiple jobs are created and executed on multiple machines using the same job task.

Tests

The tests are defined by a JSON file tests.json inside the .github/settings directory. Each row in the JSON array represents a test that will execute on a specific machine using a specified (release) build image. Example:

  { "runs-on": "n150",   "image": "tracy",  "script": "pykernel.sh" },
  { "runs-on": "n300",   "image": "speedy", "script": "ttrt.sh", "args": ["run", "Silicon", "--non-zero"] },

runs-on

Specifies the machine on which the test suite will be executed. Currently supported runners are:

• n150 - Wormhole 1 chip card
• n300 - Wormhole 2 chip card
• llmbox - Loudbox machine with 4 N300 cards
• tg - Galaxy box
• p150 - Blackhole 1 chip card

It is expected that the list will expand soon as machines with blackhole chip family are added to the runner pool.

image

Specifies which release build image to use. It can be:

• speedy
• tracy

Please take a look at the Builds section for a more detailed description of the builds.

script

Test type. It is the name of the BASH script that executes the test. Scripts are located in the .github/test_scripts directory, and it is possible to create new test types simply by adding scripts to the directory.

args (optional)

This field represents the arguments for the script. This can be omitted, a string, or a JSON array.

reqs (optional)

Specifies additional requirements for test execution. These arguments are passed as the REQUIREMENTS environment variable to the test script.

if (optional)

Specifies the name of optional component. The test will be executed only if optional component is enabled.

Using JSON arrays

The runs-on and image fields can be passed as JSON arrays. With arrays, one can define a test to execute on multiple machines and images. Examples:

{ "runs-on": ["n150","n300"],
    "image": ["speedy","tracy"],
        "script": "unit" }

Adding New Test

Usually, it is enough to add a single line to the test matrix and your tests will become part of the tt-mlir CI. Here is a checklist of what you should decide before adding it:

• On which TT hardware should your tests run? Put the specific hardware in the "runs-on" field.
• Do your tests run with ttrt or pytest or any other standard type that other tests also use? Put this decision in the "script" field.
• Refer to the test script you've put in the type for interpretation of args and reqs parameters.

Each line in the matrix MUST be unique! There is no point in running the same test with the same build image on the same type of hardware.

Consider

Here are a few things to consider:

• Design your ttrt test so it is generated with a -- check_ttmlir CMake target. These will be generated at compile time and will be available for test jobs.
• For pytest, use pytest test discovery to run all tests in subdirectories. In most cases, there is no need for two sets of tests.
• If you want to have separate test reports, do not add additional XML file paths and steps to upload these. Use ${TTRT_REPORT_PATH} (for ttrt JSON files) or ${TEST_REPORT_PATH} (for JUnit XML) because it will be automatically picked up and sent to analytics.
• If separate reports are required, treat them as different tests. Add an additional line to the test matrix. You can use a construct from this example: cp run_results.json ${TTRT_REPORT_PATH%_*}_ttir_${TTRT_REPORT_PATH##*_}

Adding New Test Script

To design a new test script, add new files to .github/test_scripts

Make sure you set the execution flag for the new test script file (chmod +x <file>)!

Before running test scripts, the workflow will activate the default TT-MLIR Python venv and set a number of useful environment variables:

• WORK_DIR - set to repo root
• BUILD_DIR - set to build artifacts
• INSTALL_DIR - set to install artifacts
• LD_LIBRARY_PATH - set to install artifacts lib and toolchain lib directories
• SYSTEM_DESC_PATH - set to system_desc.ttsys system descriptor file generated by ttrt
• TT_METAL_RUNTIME_ROOT - set to tt-metal install directory
• RUNS_ON - machine label script runs on
• IMAGE_NAME - name of the image script runs on
• RUN_ID - id of workflow run, see below.

Also, a soft link is created inside the build directory to the install directory.

Please make sure you implement cleanup logic inside your script and leave the script with the repo in the same state as before execution!

A good practice is to put some comments on how the script interprets arguments (and requirements if applicable).

For example builder.sh:

# arg $1: path to pytest test files
# arg $2: pytest marker expression to select tests to run
# arg $3: "run-ttrt" or predefined additional flags for pytest and ttrt

runttrt=""
TTRT_ARGS=""
PYTEST_ARGS=""

[[ "$RUNS_ON" != "n150" ]] && PYTEST_ARGS="$PYTEST_ARGS --require-exact-mesh"
[[ "$RUNS_ON" == "p150" ]] && TTRT_ARGS="$TTRT_ARGS --disable-eth-dispatch"

for flag in $3; do
    [[ "$flag" == "run-ttrt" ]] && runttrt=1
    [[ "$flag" == "require-opmodel" ]] && PYTEST_ARGS="$PYTEST_ARGS --require-opmodel"
done

pytest "$1" -m "$2" $PYTEST_ARGS -v --junit-xml=$TEST_REPORT_PATH
if [[ "$runttrt" == "1" ]]; then
    ttrt run $TTRT_ARGS ttir-builder-artifacts/
    cp run_results.json ${TTRT_REPORT_PATH%_*}_ttir_${TTRT_REPORT_PATH##*_} || true
    ttrt run $TTRT_ARGS stablehlo-builder-artifacts/
    cp run_results.json ${TTRT_REPORT_PATH%_*}_stablehlo_${TTRT_REPORT_PATH##*_} || true
fi

This script has several types of flags that can be stated concurrently. Arguments are parsed as run-ttrt and other possible flags for pytest or ttrt. This test uses TTRT_REPORT_PATH, but due to the fact that it has two ttrt runs, it inserts its type inside the filename.

The second example is pytest.sh script:

if [ -n "$REQUIREMENTS" ]; then
    eval "pip install $REQUIREMENTS"
fi
export TT_EXPLORER_GENERATED_MLIR_TEST_DIRS=$BUILD_DIR/test/ttmlir/Silicon/TTNN/n150/perf,$BUILD_DIR/test/python/golden/ttnn
export TT_EXPLORER_GENERATED_TTNN_TEST_DIRS=$BUILD_DIR/test/python/golden/ttnn
pytest -ssv "$@" --junit-xml=$TEST_REPORT_PATH

This script uses $REQUIREMENTS to specify additional wheels to be installed. Note how it uses the eval command to expand bash variables where suitable. It also defines some additional environment variables using the provided ones.

Downloading Artifacts

If you need to download artifacts (e.g., wheels) from a workflow run, you can use the following command:

gh run download $RUN_ID --repo tenstorrent/tt-mlir --name <artifact_name>

This command downloads the specified artifact to the current directory. You can specify a different download location using the --dir <directory> option.

To download multiple artifacts matching a pattern, use the --pattern option instead of --name:

gh run download $RUN_ID --repo tenstorrent/tt-mlir --pattern "tt_*.whl"

Note: When using --pattern, artifacts are downloaded into separate subdirectories, even when --dir is specified.

CI Run (under the hood)

Test runs are prepared in the prepare-run job when the input test matrix is transformed into a job test matrix that will be used for test runs. All jobs are grouped based on runs-on and image fields and then split (and balanced) into several runs based on test durations and target total time. This is done to make efficient use of resources because there are many tests that last for just several seconds while preparation can take ~4 minutes. So, tests are run in batches in the Run Test step with clear separation and summary. Lists of tests are displayed at the beginning, and one can search for test <number> (number ranges from 1 to total number of tests) when needing to see test flow and test results for a particular test. Also, it is possible during development to comment out tests in the JSON file of Test Matrix using the # character in a development branch and make test runs much faster, but please do not forget to remove comments when a PR is created or finalized. Test durations are collected after each push to main, and these are automatically used on each subsequent PR, Push, and other runs.

Additional Reading

This section contains pointers to reading material that may be useful for understanding the project.

MLIR

• https://llvm.org/docs/tutorial/MyFirstLanguageFrontend/index.html
• https://mlir.llvm.org/docs/Tutorials/Toy/
• https://www.jeremykun.com/2023/08/10/mlir-getting-started/
• https://arxiv.org/pdf/2002.11054
• https://ieeexplore.ieee.org/abstract/document/9370308

Dialects

• affine dialect
  • Affine map is a really powerful primitive that can be used to describe most data movement patterns.
  • It can also be used to describe memory layouts.
• linalg dialect
• tosa dialect
• tosa spec
• memref dialect
• torch-mlir
• onnx-mlir
• triton-mlir

Tablegen

• docs
• Passes in TableGen
• Ops in TableGen (we use this)

LLVM Testing Framework Tools

• Lit
• Filecheck

Jax
Flatbuffer
Openxla Website
openxla
StableHLO

Contributor Covenant Code of Conduct

Our Pledge

We as members, contributors, and leaders pledge to make participation in our community a harassment-free experience for everyone, regardless of age, body size, visible or invisible disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, religion, or sexual identity and orientation.

We pledge to act and interact in ways that contribute to an open, welcoming, diverse, inclusive, and healthy community.

Our Standards

Examples of behavior that contributes to a positive environment for our community include:

• Demonstrating empathy and kindness toward other people
• Being respectful of differing opinions, viewpoints, and experiences
• Giving and gracefully accepting constructive feedback
• Accepting responsibility and apologizing to those affected by our mistakes, and learning from the experience
• Focusing on what is best not just for us as individuals, but for the overall community

Examples of unacceptable behavior include:

• The use of sexualized language or imagery, and sexual attention or advances of any kind
• Trolling, insulting or derogatory comments, and personal or political attacks
• Public or private harassment
• Publishing others' private information, such as a physical or email address, without their explicit permission
• Other conduct which could reasonably be considered inappropriate in a professional setting

Enforcement Responsibilities

Community leaders are responsible for clarifying and enforcing our standards of acceptable behavior and will take appropriate and fair corrective action in response to any behavior that they deem inappropriate, threatening, offensive, or harmful.

Community leaders have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, and will communicate reasons for moderation decisions when appropriate.

Scope

This Code of Conduct applies within all community spaces, and also applies when an individual is officially representing the community in public spaces. Examples of representing our community include using an official e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event.

Enforcement

Instances of abusive, harassing, or otherwise unacceptable behavior may be reported to the community leaders responsible for enforcement at nsmith@tenstorrent.com or staylor@tenstorrent.com. All complaints will be reviewed and investigated promptly and fairly.

All community leaders are obligated to respect the privacy and security of the reporter of any incident.

Enforcement Guidelines

Community leaders will follow these Community Impact Guidelines in determining the consequences for any action they deem in violation of this Code of Conduct:

1. Correction

Community Impact: Use of inappropriate language or other behavior deemed unprofessional or unwelcome in the community.

Consequence: A private, written warning from community leaders, providing clarity around the nature of the violation and an explanation of why the behavior was inappropriate. A public apology may be requested.

2. Warning

Community Impact: A violation through a single incident or series of actions.

Consequence: A warning with consequences for continued behavior. No interaction with the people involved, including unsolicited interaction with those enforcing the Code of Conduct, for a specified period of time. This includes avoiding interactions in community spaces as well as external channels like social media. Violating these terms may lead to a temporary or permanent ban.

3. Temporary Ban

Community Impact: A serious violation of community standards, including sustained inappropriate behavior.

Consequence: A temporary ban from any sort of interaction or public communication with the community for a specified period of time. No public or private interaction with the people involved, including unsolicited interaction with those enforcing the Code of Conduct, is allowed during this period. Violating these terms may lead to a permanent ban.

4. Permanent Ban

Community Impact: Demonstrating a pattern of violation of community standards, including sustained inappropriate behavior, harassment of an individual, or aggression toward or disparagement of classes of individuals.

Consequence: A permanent ban from any sort of public interaction within the community.

Attribution

This Code of Conduct is adapted from the Contributor Covenant, version 2.0, available at https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.

Community Impact Guidelines were inspired by Mozilla's code of conduct enforcement ladder.

For answers to common questions about this code of conduct, see the FAQ at https://www.contributor-covenant.org/faq. Translations are available at https://www.contributor-covenant.org/translations.

Project Structure

• env: Contains the environment setup for building project dependencies, such as LLVM and Flatbuffers
• include/ttmlir: Public headers for the TTMLIR library
  • Dialect: MLIR dialect interfaces and definitions, dialects typically follow a common directory tree structure:
    • IR: MLIR operation/type/attribute interfaces and definitions
    • Passes.[h|td]: MLIR pass interfaces and definitions
    • Transforms: Common MLIR transformations, typically invoked by passes
  • Target: Flatbuffer schema definitions. This defines the binary interface between the compiler and the runtime
• lib: TTMLIR library implementation
  • CAPI: C API for interfacing with the TTMLIR library, note this is needed for implementing the python bindings. Read more about it here: https://mlir.llvm.org/docs/Bindings/Python/#use-the-c-api
  • Dialect: MLIR dialect implementations
• runtime: Device runtime implementation
  • include/tt/runtime: Public headers for the runtime interface
  • lib: Runtime implementation
  • tools/python: Python bindings for the runtime, currently this is where ttrt is implemented
• test: Test suite
• tools/ttmlir-opt: TTMLIR optimizer driver

Namespaces

• mlir: On the compiler side, we use the MLIR namespace for all MLIR types and operations and subnamespace for our dialects.
  • mlir::tt: Everything ttmlir related is underneath this namespace. Since we need to subnamespace under mlir, just mlir::tt seemed better than mlir::ttmlir which feels redundant.
    • mlir::tt::ttir: The TTIR dialect namespace
    • mlir::tt::ttnn: The TTNN dialect namespace
    • mlir::tt::ttmetal: The TTMetal dialect namespace
    • mlir::tt::ttkernel: The TTKernel dialect namespace
• tt::runtime: On the runtime side, we use the tt::runtime namespace for all runtime types and operations.
  • tt::runtime::ttnn: The TTNN runtime namespace
  • tt::runtime::ttmetal: The TTMetal runtime namespace (not implemented)

Dialects Overview

Here is a brief overview of the dialects in the project, please refer to the individual dialect documentation for more details.:

• ttcore: Common types such as, ttcore.tile, ttcore.metal_layout, ttcore.grid, etc. and enums such as, data formats, memory spaces, iterator types etc.
• ttir: A high level dialect that models the tensor compute graph on tenstorrent devices. Accepts tosa and linalg input.
  • ttir.generic: Generically describe compute work.
  • ttir.to_layout: Convert between different tensor memory layouts and transfer between different memory spaces.
  • tensor.pad: Pad a tensor with a value (ie. convs)
  • ttir.yield: return result memref of computation in dispatch region body, lowers to ttkernel.yield
  • ttir.kernel: lowers to some backend kernel
• ttnn: A TTNN dialect that models ttnn API.
• ttkernel: Tenstorrent kernel library operations.
  • ttkernel.noc_async_read
  • ttkernel.noc_async_write
  • ttkernel.cb_push_back
  • ttkernel.[matmul|add|multiply]: Computations on tiles in source register space, store the result in dest register space.
  • ttkernel.sfpu_*: Computations on tiles in dest register space using sfpu coprocessor.
• ttmetal: Operations that dispatch work from host to device.
  • ttmetal.enqueue_program: Dispatch a grid of compute work.

Guidelines

This page contains a collection of guidelines to help maintain consistency and quality across our project. Please refer to the following documents for detailed instructions on coding practices, as well as specific dialect guidelines.

• Coding guidelines
• TTNN Dialect guidelines

TT-MLIR Coding Guidelines

This document outlines the coding standards used in the tt-mlir project. These guidelines are designed to enhance the readability and maintainability of our shared codebase. While these guidelines are not strict rules for every situation, they are essential for maintaining consistency across the repository.

Our long-term aim is to have the entire codebase adhere to these conventions.

Since our compiler is built on the LLVM MLIR framework, we strive to align closely with the LLVM coding style guidelines outlined here: LLVM Coding Standards.

Naming

Clear and descriptive names are crucial for code readability and preventing bugs. It’s important to choose names that accurately reflect the semantics and purpose of the underlying entities, within reason. Avoid abbreviations unless they are widely recognized. Once you settle on a name, ensure consistent capitalization throughout the codebase to avoid confusion.

The general naming rule is to use camel case for most names (for example, WorkaroundPass, isRankedTensor())

• Type Names
  • Applies to classes, structs, enums, and typedefs.
  • Should be nouns that describe the entity's purpose.
  • Use upper camel case (for example, TTNNOptimizerOptions, DecompositionPass).
• Variable Names
  • Should be nouns, as they represent state.
  • Use lower camel case (for example, inputLayout).
• Function Names
  • Represent actions and should be verb phrases
  • Use lower camel case (for example, createTTNNOptimizer(), emitTTNNAsCpp()).

Includes

We prefer #includes to be listed in this order:

1. Main Module Header
2. Local/Private Headers
3. LLVM project/subproject headers (clang/..., lldb/..., llvm/..., etc)
4. System #includes

Each category should:

• Be sorted lexicographically by the full path.
• Be separated by a single blank line for clarity.

Only the standard lib header includes should use <> whereas all the others should use quotes "". Additionally, all project headers must use absolute paths (rooted at ttmlir) to prevent preprocessor and namespacing issues. For example, the following is preferred:

#include "ttmlir/module/something.h"

over:

#include "something.h"

Using TTIRToTTNN.cpp as an example, this is what includes would look like for us:

#include "ttmlir/Conversion/TTIRToTTNN/TTIRToTTNN.h"  # main header

#include "ttmlir/Dialect/TTCore/IR/TTCoreOpsTypes.h"  # these are local/private headers
#include "ttmlir/Dialect/TTNN/Utils/Utils.h"

#include "mlir/Dialect/MemRef/IR/MemRef.h"  # llvm project/subproj headers
#include "llvm/Support/LogicalResult.h"

#include <cstdio>  # system includes
#include <algorithm>

Comments

Write comments as full sentences, starting with a capital letter and ending with a period. Comments should explain why the code exists, not just what it does. Use comments to clarify logic, assumptions, or any non-obvious aspects of the code.

Example of a comment:

// Initialize the buffer to store incoming data from the network.

In general, C++ style comments (//) should be used. Use C-style comments (/**/) only for when documenting the significance of constants used as actual parameters in a call:

object.callFunction(/*arg0=*/nullptr);

Every function, class, or non-trivial piece of logic should have a comment. Avoid redundant comments for self-explanatory code, but never leave complex code unexplained. Example of redundant comment:

// Increment the counter by 1.  // Redundant, avoid.
counter++;

Ensure comments are accurate and reflect the current state of the code. Outdated or misleading comments can be worse than no comments at all.

All TODO comments should be marked with an alias as follows:

// TODO (your-alias): Refactor this loop for clarity. Issue: https://github.com/tenstorrent/tt-mlir/issues/XYZ

A TT-MLIR issue should be created and linked inline to track the TODO.

Code Denesting (Inversion)

Strive to minimize unnecessary indentation without compromising code clarity. One effective way to achieve this is by using early exits and the continue keyword in long loops.

Consider following example:

void doSomething(Operation *op)
{
    if (op->getNumOperands() > 0
        && isDpsOp(op)
        && doSomethingDifferent(op))
    {
        // ... some long code ...
    }
}

It is strongly recommended to format the code as follows:

void doSomething(Operation *op)
{
    // ...
    // We need to do something with the op that has more than 0 operands
    if (op->getNumOperands() <= 0 ) return;

    // We need something to do with the DPS op
    if (!isDpsOp(op)) return;

    // Just for example purposes
    if (!doSomethingDifferent(op)) return;

    // .. some long code ...
}

This reduces loop nesting, makes the reasoning behind the conditions clearer, and signals to the reader that there is no subsequent else to worry about, reducing cognitive load. This can significantly improve code readability and comprehension.

Function Declaration and Definition Order

To improve code readability and maintainability, we should adopt a consistent approach for organizing function declarations and definitions within a file. The goal is to make it easier for readers to follow the logical flow of function dependencies.

Follow a bottom-up call order:

• Arrange functions so that lower-level helper functions are defined first, followed by higher-level functions that call them.
• This allows each function to be defined after its dependencies, making it clear which functions rely on which.
• For example, if function A calls A1 and A2, then the preferred order is:

void A1();
void A2();
void A(){
  A1();
  A2();
}

Group related functions together:

• If functions are only relevant to a specific “parent” function (for example, A1 and A2 are only called by A), place them directly before the “parent” function.
• If a function (like A2) is also called by other functions (for example, B), place it where it fits the overall bottom-up order.

Avoid mixed ordering:

• Mixing top-down and bottom-up call orders within the same file can make the code hard to read and maintain.

Example of a preferred order:

void A1() {
  /*...*/
}
void A2() {
  /*...*/
}
void B() {
  A2(); // A2 is defined before B, so dependencies are clear.
}
void A() {
  A1();
  A2();
  B();
}

Helper Functions

These coding guidelines address visibility and linkage of simple helper functions to ensure clarity, prevent linking errors, and improve maintainability:

• If a helper function needs to be defined in a .cpp file, it should be declared static or wrapped inside an anonymous namespace.

• If a helper function needs to be defined in a header file (for example, for templated or performance-critical code), it should be marked as inline.

[!NOTE] A significant concern with declaring functions as non-public (for example, static functions or functions in unnamed namespaces) is that they cannot be unit tested in isolation. This limitation hinders our ability to write focused, granular tests that verify the correctness of individual components and it also reduces test coverage.

Using Namespaces

Namespaces are an important part of C++ programming, providing a way to organize code and avoid naming conflicts. Choose namespace names that reflect the purpose or functionality of the code contained within.

Follow these guidelines when defining namespaces:

• Use lower-case letters for short, single-word names or those with a clear acronym (for example, ttnn, mlir).
• Use nested namespaces to group logically related code, avoiding too deep or unnecessarily complex hierarchy

Follow these guidelines when using namespaces:

• Do not use a using-directive to make all names from a namespace available because it pollutes the namespace.

// Forbidden -- This pollutes the namespace.
using namespace std;

• Avoid placing code in the global namespace to reduce the potential for name conflicts and ambiguity. Always use specific namespaces. If necessary to use something from the global namespace (such as std), use an explicit std:: prefix rather than importing everything using using namespace std;.
• Do not use namespace aliases at namespace scope in header files except in explicitly marked internal-only namespaces, because anything imported into a namespace in a header file becomes part of the public API exported by that file.
• Try to avoid mixing concepts from different namespaces in a single function or class. If a function belongs to one namespace but calls classes from others, ensure the relationships are clear.
• Wrap classes/structs declared in .cpp files inside of an anonymous namespace to avoid violating ODR. See LLVM docs for more detailed information.

Using Alternative Tokens (and, or, xor, etc.)

Although they are standard, we should avoid their use. They are very rarely used in practice and the C++ community widely uses the standard operators (&&, ||, !, etc.), as they are more familiar and easily recognizable to most C++ developers. Their usage can make the code harder to read and maintain, especially for developers who are not familiar with these alternatives. We should stick to the standard operators (&&, ||, !, etc.) for clarity, consistency, and compatibility with other C++ developers and tools.

Type Aliasing

When declaring type aliases in C++ prefer using over typedef. using provides better readability, especially for complex types, and supports alias templates. Here is an example:

// Preferred
using Callback = void(*)(int, double);

// Avoid
typedef void (*Callback)(int, double);

Choose alias names that clarify their role in the code. Avoid overly generic names that might obscure the type’s purpose. Do not create a type alias unless it significantly improves clarity or simplifies complex types.

Using auto to Deduce Type

Use auto only when it enhances code readability or maintainability. Avoid defaulting to “always use auto.” Instead, apply it thoughtfully in the following scenarios:

• When the type is immediately clear from the initializer, such as in cast(...).
• When the type is obvious from the context, making the code cleaner and more concise.
• When the type is already abstracted, such as with container typedefs like std::vector::iterator.

In all other cases, prefer explicit type declarations to maintain clarity and ensure the code remains easy to understand.

Python Coding Guidelines

Python Version and Source Code Formatting

The current minimum version of Python required is 3.11 or higher. Python code in the tt-mlir repository should only use language features available in this version of Python.

The Python code within the tt-mlir repository should adhere to the formatting guidelines outlined in PEP 8.

For consistency and to limit churn, code should be automatically formatted with the black utility, which is PEP 8 compliant. Use its default rules. For example, avoid specifying --line-length even though it does not default to 80. The default rules can change between major versions of black. In order to avoid unnecessary churn in the formatting rules, we currently use black version 23.x.

When contributing a patch unrelated to formatting, you should format only the Python code that the patch modifies. When contributing a patch specifically for reformatting Python files, use black, which currently only supports formatting entire files.

Here is a quick example, but see the black documentation for details:

$ black test.py                    # format entire file

TTNN Dialect Contribution Guidelines

This document provides clear and consistent guidelines for contributing to the TTNN dialect, including operations, attributes, types, and other components. Following these ensures a streamlined development process, faster code reviews, and higher-quality code with fewer bugs.

General Principle: Model TTNN Library Closely

The TTNN dialect should closely reflect the TTNN library wherever practical, serving as the core guiding principle when contributing to the dialect. Whenever there's a need to deviate from this principle, it should be discussed with stakeholders.

Ops and Operands

Signature Selection

Ops in TTNN may have multiple signatures available - it's important to choose the right one when creating its model in the TTNN dialect. Going through an example, these are the available signatures for the ttnn::transpose op:

struct ExecuteTranspose {
    static ttnn::Tensor invoke(
        uint8_t queue_id,
        const ttnn::Tensor& input_tensor,
        const int64_t& dim1,
        const int64_t& dim2,
        const std::optional<MemoryConfig>& memory_config_arg,
        const std::optional<float>& pad_value = 0.0f);

    static ttnn::Tensor invoke(
        const ttnn::Tensor& input_tensor,
        const int64_t& dim1,
        const int64_t& dim2,
        const std::optional<MemoryConfig>& memory_config,
        const std::optional<float>& pad_value = 0.0f);

    static ttnn::Tensor invoke(
        const ttnn::Tensor& input_tensor,
        const int64_t& dim1,
        const int64_t& dim2,
        const std::optional<float>& pad_value = 0.0f);
};

The first and second signature differ only in the queue_id parameter - we don't model queues today, so the second signature has priority here. The second and third signature differ in memory_config parameter - the second signature is preferred as it is more robust: the parameter is optional so it can remain unused if it isn't needed.

Only one signature should be chosen. If the need would arise for more than one signature, it would be a precedent, and should be discussed with stakeholders.

Operand ordering

Operands in the TTNN dialect ops should match the ordering of the signature of the op being modelled. For the chosen signature of the ttnn::transpose op, the operands should look like this:

let arguments = (ins AnyRankedTensor:$input,
                     SI64Attr:$dim0,
                     SI64Attr:$dim1,
                     OptionalAttr<TTNN_MemoryConfigAttr>:$memory_config,
                     OptionalAttr<FloatAttr>:$pad_value);

Mixing types and attributes within the ordering is not an issue, this is valid:

let arguments = (ins TTNN_ShapeAttr:$shape,
                     OptionalAttr<TT_DataTypeAttr>:$dtype,
                     OptionalAttr<TTNN_LayoutAttr>:$layout,
                     Optional<TT_Device>:$device,
                     OptionalAttr<TTNN_MemoryConfigAttr>:$memory_config);

Following this guideline provides consistency with the TTNN lib.

Optional operands

If an operand is optional in the TTNN lib, it should be modelled as optional in the dialect.

Default-valued operands

If an operand has a default value in the TTNN lib, it should have a default value in the dialect.

ttnn::permute as an example:

static ttnn::Tensor invoke(
    const ttnn::Tensor& input_tensor,
    ttsl::Span<const int64_t> dims,
    const std::optional<MemoryConfig>& memory_config,
    const std::optional<float>& pad_value = 0.0f);

let arguments = (ins AnyRankedTensor:$input,
                     DenseI64ArrayAttr:$permutation,
                     OptionalAttr<TTNN_MemoryConfigAttr>:$memory_config,
                     DefaultValuedOptionalAttr<F32Attr, "0.0f">:$pad_value);

Numerical operands

Numerical operands should match in signedness and bit width. If an operand is a signed integer of width of 32 bits, SI32Attr should be used to model it.

Pointers and references

Pointers and references should be ignored. We do not want to model this level of detail at this point in time.

There were very few issues with these previously, and they were caused by inconsistencies in TTNN lib APIs.

Attrs vs Types

General guideline is that if a value is known at compile time, it should probably be an Attr. Example: dims in transpose op, pooling windows in a conv, etc. If the value is unknown at compile time (e.g. tensor) it should be a Type.

There's another consideration to account for: does the value need its own SSA? Remember, Attrs need something to latch onto, like an op or a Type, but Types need to be constructed, i.e. have their own SSA, in order to exist. Let's look at ttnn::Shape for example - in TTNN lib, these need to be constructed, so it naturally follows that they should have their own SSA value within the IR, implying that they should be implemented as Types. However, there are several downsides to this:

• More IR is produced
• Diminished readability as they're not attached to the object whose shape they're describing
• Not as easy to construct in code
• Runtime would need to keep track of all the Shape objects (it currently maps all SSAs, which are currently only tensors and devices)

One upside for implementing ttnn::Shape as a Type is that it would enable optimizing out multiple constructor calls for the same Shape.

It is agreed that we should prefer using Attrs in these scenarios. However, this guideline is not set in stone - stakeholders should be notified if anyone believes there's a need to implement an object as a Type.

Destination-passing style (DPS)

If the op in TTNN lib has the destination tensor, is should be modelled as DPS op.

An example signature, where the last operand is a destination tensor:

static Tensor invoke(
    const Tensor& input_tensor,
    float exponent,
    const std::optional<MemoryConfig>& memory_config = std::nullopt,
    const std::optional<Tensor>& optional_output_tensor = std::nullopt);

Variadic operands

Variadic<> type constraint should only be used for operands that are variadic in nature, e.g. a vector of tensors, like in ttnn::concat:

static ttnn::Tensor invoke(
    const std::vector<ttnn::Tensor>& input_tensors,
    int dim,
    const std::optional<MemoryConfig>& memory_config = std::nullopt,
    const std::optional<ttnn::Tensor>& optional_output_tensor = std::nullopt,
    unsigned int groups = 1);

Operand naming

Operands should be named as they are in the TTNN lib. However, this guideline is not strict, and some reasonable deviations are acceptable.

Operand namespaces

Some operands are defined in a namespace nested within the TTNN namespace, i.e. ttnn::ccl::Topology, and some are in other but related namespaces, i.e. tt::tt_metal::MemoryConfig. While it would be ideal to model these completely accurately, it doesn’t provide value and we should pretend they’re all in the ttnn:: namespace for the sake of simplicity.

Adding an Op

This guide will walk you through the process of adding a new Op end to end in tt-mlir, in this case we will be adding a matmul operation. Note that the matmul op was added as part of the same changeset as this guide, it could be useful to reference the diff alongside this guide to see the changes in full.

This guide will cover the following steps:

• Adding an Op
  • 1. Define the Op in the TTIR frontend dialect
  • 2. Define the Op in the TTNN backend dialect
    • TTNNOps.td
    • TTNNOps.cpp
    • Adding constraint/runtime APIs
  • 3. Convert / Implement the Op in the TTNN passes
  • 4. Add a compiler unit test for the Op
    • test/ttmlir/Dialect/TTNN/matmul/simple_matmul.mlir
  • 5. Define flatbuffer schema for the Op
    • include/ttmlir/Target/TTNN/CMakeLists.txt
    • include/ttmlir/Target/TTNN/operations/matmul.fbs
    • include/ttmlir/Target/TTNN/program.fbs
  • 6. Serialize the Op in the flatbuffer format
  • 7. Add runtime support for the Op
    • runtime/lib/ttnn/operations/matmul/matmul.cpp
    • runtime/lib/ttnn/operations/CMakeLists.txt
    • runtime/lib/ttnn/program_executor.cpp
  • 8. Add a silicon unit test for the Op
    • test/ttmlir/Silicon/TTNN/matmul/simple_matmul.mlir
  • 9. Add an EmitC test for the Op
    • test/ttmlir/EmitC/TTNN/matmul/matmul.mlir

1. Define the Op in the TTIR frontend dialect

We will start by defining the Op in the TTIR dialect. The TTIR Ops are defined in a tablegen file located at include/ttmlir/Dialect/TTIR/IR/TTIROps.td.

Tablegen is a domain-specific language for defining ops/types/attributes in MLIR and LLVM, these definitions constitute the dialect's Operation Definition Specification (ODS).

Here is an example of defining matmul in the TTIR dialect:

def TTIR_MatmulOp : TTIR_NamedOp<"matmul"> {
    let summary = "Matrix multiplication operation.";
    let description = [{
      The `matmul` operation computes the matrix multiplication of two tensors.

      This operation performs matrix multiplication between tensors `a` and `b`. It supports optional
      transposition of either input tensor before multiplication. For 2D tensors, this computes the standard
      matrix product. For tensors with more dimensions, it applies batched matrix multiplication.

      Example:
      ```mlir
      // Basic matrix multiplication of 2D tensors
      %a = ... : tensor<3x4xf32>  // Matrix A with shape [3,4]
      %b = ... : tensor<4x5xf32>  // Matrix B with shape [4,5]
      %output = ttir.empty() : tensor<3x5xf32>  // Output matrix shape
      %result = ttir.matmul(%a, %b, %output) :
          tensor<3x4xf32>, tensor<4x5xf32>, tensor<3x5xf32> -> tensor<3x5xf32>

      // Batched matrix multiplication with transposition
      %a = ... : tensor<2x3x4xf32>  // Batch of 2 matrices with shape [3,4]
      %b = ... : tensor<2x5x4xf32>  // Batch of 2 matrices with shape [5,4]
      %output = ttir.empty() : tensor<2x3x5xf32>  // Output shape
      %result = ttir.matmul(%a, %b, %output) {
          transpose_a = false,  // Don't transpose A
          transpose_b = true    // Transpose B before multiplication
      } : tensor<2x3x4xf32>, tensor<2x5x4xf32>, tensor<2x3x5xf32> -> tensor<2x3x5xf32>
      ```

      Inputs:
      - `a` (Tensor): The first input tensor.
      - `b` (Tensor): The second input tensor.

      Attributes:
      - `transpose_a` (Boolean, default=false): Whether to transpose tensor `a` before multiplication.
      - `transpose_b` (Boolean, default=false): Whether to transpose tensor `b` before multiplication.

      Outputs:
      - `result` (Tensor): The result of the matrix multiplication.

      Note: The inner dimensions of the input tensors must be compatible for matrix multiplication.
      If `a` has shape [..., m, k] and `b` has shape [..., k, n], then the result will have shape [..., m, n].
      If `transpose_a` is true, then `a` is treated as having shape [..., k, m].
      If `transpose_b` is true, then `b` is treated as having shape [..., n, k].
    }];

    let arguments = (ins AnyRankedTensor:$a,
                         AnyRankedTensor:$b,
                         AnyRankedTensor:$output,
                         DefaultValuedAttr<BoolAttr, "false">:$transpose_a,
                         DefaultValuedAttr<BoolAttr, "false">:$transpose_b);

    let results = (outs AnyRankedTensor:$result);

    let hasVerifier = 1;

    let hasCanonicalizer = 1;
}

There are many things to break down here, starting from the top:

• def in tablegen is used to define a concrete type, this will have a 1-1 mapping to a C++ generated class, and for this particular case the build will end up generating file build/include/ttmlir/Dialect/TTIR/IR/TTIROps.h.inc.
• It inherits from class TTIR_DPSOp, classes in tablegen don't define a concrete type, but rather an interface that augment or constrain inherited defs. TTIR_DPSOp is a class that defines the common attributes for all TTIR Ops that implement Destination Passing Style (DPS) semantics. DPS just means that the result tensor is passed as an argument to the operation which will be critical for modeling buffer allocation / lifetimes. Note the 3rd argument AnyRankedTensor:$output.
• Next we have a list of arguments. These arguments consist of a mixture of Types (i.e. AnyRankedTensor) and Attributes. Read more about Types & Attributes here.
  • AnyRankedTensor is part of a tablegen standard library which type aliases to MLIR's builtin Tensor type, with the added constraint that the tensor has a static rank. As much as possible we want to use the builtin types and infrastructure provided by MLIR.
• Next we have a list of results in this case just 1, which aliases the output tensor. One drawback of DPS is that the result tensor and the output tensor will appear to have different SSA names in the IR, but they really alias the same object. This can make writing some passes more cumbersome.
• Next we have extraClassDeclaration, which enables us to inject member functions, written directly in C++, into the generated class. We are doing this for this particular case in order to satisfy the DPS interface which requires an implementation for getting the mutated output tensor.
• Finally, we have hasVerifier = 1, this tells MLIR that we have a verifier function that will be called to validate the operation. This is a good practice to ensure that the IR is well formed.

We can now try building and opening the TTIROps.h.inc file to see the generated C++ code. We will actually get a linker error because we have hasVerifier = 1 which automatically declared a verifier function, but we need to go implement.

Let's head over to lib/Dialect/TTIR/IR/TTIROps.cpp and implement the verifier.

// MatmulOp verification
::mlir::LogicalResult mlir::tt::ttir::MatmulOp::verify() {
  ::mlir::RankedTensorType inputAType = getA().getType();
  ::mlir::RankedTensorType inputBType = getB().getType();
  ::mlir::RankedTensorType outputType = getOutput().getType();

  llvm::ArrayRef<int64_t> outputShape = outputType.getShape();
  llvm::SmallVector<int64_t> inputAShape(inputAType.getShape());
  llvm::SmallVector<int64_t> inputBShape(inputBType.getShape());

  // Verify that the input A is at least 1D tensor.
  if (inputAType.getRank() < 1) {
    return emitOpError("Input A must be at least a 1D tensor");
  }

  // Verify that the input B is at least 1D tensor.
  if (inputBType.getRank() < 1) {
    return emitOpError("Input B must be at least a 1D tensor");
  }

  // If input A is a vector (1D tensor), 1 is prepended to its dimensions for
  // the purpose of the matrix multiplication. After the matrix
  // multiplication, the prepended dimension is removed. Otherwise, check if
  // the LHS needs to be transposed.
  if (inputAType.getRank() == 1) {
    inputAShape.insert(inputAShape.begin(), 1);
  } else if (getTransposeA()) {
    std::swap(inputAShape[inputAShape.size() - 1],
              inputAShape[inputAShape.size() - 2]);
  }

  // If input B is a vector (1D tensor), a 1 is appended to its dimensions for
  // the purpose of the matrix-vector product and removed afterwards.
  // Otherwise, check if the RHS needs to be transposed.
  if (inputBType.getRank() == 1) {
    inputBShape.push_back(1);
  } else if (getTransposeB()) {
    std::swap(inputBShape[inputBShape.size() - 1],
              inputBShape[inputBShape.size() - 2]);
  }

  // Verify that the input A and input B has matching inner dimensions.
  if (inputAShape[inputAShape.size() - 1] !=
      inputBShape[inputBShape.size() - 2]) {
    return emitOpError("Input A[-1](")
           << inputAShape[inputAShape.size() - 1] << ") and B[-2]("
           << inputBShape[inputBShape.size() - 2]
           << ") must have matching inner dimensions";
  }

  llvm::SmallVector<int64_t> expectedOutputShape;
  // Verify that the batch dimensions are broadcast compatible and construct
  // the expected output shape. If either of input A or input B is at most 2D
  // tensors, the batch dimensions are trivially broadcast compatible.
  if (inputAShape.size() > 2 || inputBShape.size() > 2) {
    llvm::SmallVector<int64_t> inputABatchDims(inputAShape.begin(),
                                               inputAShape.end() - 2);
    llvm::SmallVector<int64_t> inputBBatchDims(inputBShape.begin(),
                                               inputBShape.end() - 2);

    // Verify that the batch dimensions of input A and B are broadcast
    // compatible.
    llvm::SmallVector<int64_t, 4> broadcastedShape;
    if (!mlir::OpTrait::util::getBroadcastedShape(
            inputABatchDims, inputBBatchDims, broadcastedShape)) {

      return emitOpError("Batch dimensions of input A(" +
                         ttmlir::utils::join(inputABatchDims, ",") +
                         ") and B(" +
                         ttmlir::utils::join(inputBBatchDims, ",") +
                         ") are not broadcast compatible");
    }

    // Insert the broadcasted batch dimensions in the expected output shape.
    expectedOutputShape = std::move(broadcastedShape);
  }

  // Insert the input A and B inner dimensions in expected output shape
  // Consider the case where input A and B are vectors. In that case,
  // the dimension 1 is ommited from the output shape.
  if (inputAType.getRank() > 1) {
    expectedOutputShape.push_back(inputAShape[inputAShape.size() - 2]);
  }

  if (inputBType.getRank() > 1) {
    expectedOutputShape.push_back(inputBShape[inputBShape.size() - 1]);
  }

  // Check the case of a vector-vector product. At this moment we don't
  // support scalars in IR, hence check that the output is at least 1D tensor
  // of size 1.
  if (expectedOutputShape.size() == 0) {
    if (outputType.getRank() < 1) {
      return emitOpError("Scalar output is not supported, output must be at "
                         "least a 1D tensor");
    }

    if (outputType.getRank() > 1 || outputType.getShape()[0] != 1) {
      return emitOpError("Scalar output must be a 1D tensor of size 1");
    }

    return success();
  }

  // Verify that the output shape is correct.
  if (outputShape.size() != expectedOutputShape.size()) {
    return emitOpError("Output shape rank(")
           << outputShape.size()
           << ") must match the expected output shape rank("
           << expectedOutputShape.size() << ")";
  }

  // Verify each dim of the output shape.
  for (auto [index, outputDim, expectedDim] : llvm::zip(
           llvm::seq(outputShape.size()), outputShape, expectedOutputShape)) {
    if (outputDim != expectedDim) {
      return emitOpError("Output shape dimension[")
             << index << "](" << outputDim
             << ") doesn't match the expected output shape dimension[" << index
             << "](" << expectedDim << ")";
    }
  }

  return success();
}

2. Define the Op in the TTNN backend dialect

Next we will define the Op in the TTNN dialect. TTNN Ops are defined in the same way, but in their respective set of dialect files. Refer to the previous section for details, the process is the same.

TTNNOps.td

def TTNN_MatmulOp : TTNN_Op<"matmul"> {
    let arguments = (ins AnyRankedTensor:$a,
                         AnyRankedTensor:$b,
                         DefaultValuedAttr<BoolAttr, "false">:$transpose_a,
                         DefaultValuedAttr<BoolAttr, "false">:$transpose_b,
                         OptionalAttr<AnyAttrOf<[
                            TTNN_MatmulMultiCoreReuseProgramConfigAttr,
                            TTNN_MatmulMultiCoreReuseMultiCastProgramConfigAttr,
                            TTNN_MatmulMultiCoreReuseMultiCast1DProgramConfigAttr,
                            TTNN_MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfigAttr
                         ]>>:$matmul_program_config);

    let results = (outs AnyRankedTensor:$result);

    let hasVerifier = 1;
}

TTNNOps.cpp

// MatmulOp verification
::mlir::LogicalResult mlir::tt::ttnn::MatmulOp::verify() {
  ::mlir::RankedTensorType inputAType = getA().getType();
  ::mlir::RankedTensorType inputBType = getB().getType();
  ::mlir::RankedTensorType outputType = getResult().getType();

  llvm::ArrayRef<int64_t> outputShape = outputType.getShape();
  llvm::SmallVector<int64_t> inputAShape(inputAType.getShape());
  llvm::SmallVector<int64_t> inputBShape(inputBType.getShape());

  // Verify that the input A is at least 1D tensor.
  if (inputAType.getRank() < 1) {
    return emitOpError("Input A must be at least a 1D tensor");
  }

  // Verify that the input B is at least 1D tensor.
  if (inputBType.getRank() < 1) {
    return emitOpError("Input B must be at least a 1D tensor");
  }

  // If input A is a vector (1D tensor), 1 is prepended to its dimensions for
  // the purpose of the matrix multiplication. After the matrix multiplication,
  // the prepended dimension is removed. Otherwise, check if the LHS needs to be
  // transposed.
  if (inputAType.getRank() == 1) {
    inputAShape.insert(inputAShape.begin(), 1);
  } else if (getTransposeA()) {
    std::swap(inputAShape[inputAShape.size() - 1],
              inputAShape[inputAShape.size() - 2]);
  }

  // If input B is a vector (1D tensor), a 1 is appended to its dimensions for
  // the purpose of the matrix-vector product and removed afterwards. Otherwise,
  // check if the RHS needs to be transposed.
  if (inputBType.getRank() == 1) {
    inputBShape.push_back(1);
  } else if (getTransposeB()) {
    std::swap(inputBShape[inputBShape.size() - 1],
              inputBShape[inputBShape.size() - 2]);
  }

  // Verify that the input A and input B has matching inner dimensions.
  if (inputAShape[inputAShape.size() - 1] !=
      inputBShape[inputBShape.size() - 2]) {
    return emitOpError("Input A[-1](")
           << inputAShape[inputAShape.size() - 1] << ") and B[-2]("
           << inputBShape[inputBShape.size() - 2]
           << ") must have matching inner dimensions";
  }

  llvm::SmallVector<int64_t> expectedOutputShape;
  // Verify that the batch dimensions are broadcast compatible and construct the
  // expected output shape. If either of input A or input B is at most 2D
  // tensors, the batch dimensions are trivially broadcast compatible.
  if (inputAShape.size() > 2 || inputBShape.size() > 2) {
    llvm::SmallVector<int64_t> inputABatchDims(inputAShape.begin(),
                                               inputAShape.end() - 2);
    llvm::SmallVector<int64_t> inputBBatchDims(inputBShape.begin(),
                                               inputBShape.end() - 2);

    // Verify that the batch dimensions of input A and B are broadcast
    // compatible.
    llvm::SmallVector<int64_t, 4> broadcastedShape;
    if (!OpTrait::util::getBroadcastedShape(inputABatchDims, inputBBatchDims,
                                            broadcastedShape)) {

      return emitOpError("Batch dimensions of input A(" +
                         ttmlir::utils::join(inputABatchDims, ",") +
                         ") and B(" +
                         ttmlir::utils::join(inputBBatchDims, ",") +
                         ") are not broadcast compatible");
    }

    // Insert the broadcasted batch dimensions in the expected output shape.
    expectedOutputShape = std::move(broadcastedShape);
  }

  // Insert the input A and B inner dimensions in expected output shape
  // Consider the case where input A and B are vectors. In that case,
  // the dimension 1 is ommited from the output shape.
  if (inputAType.getRank() > 1) {
    expectedOutputShape.push_back(inputAShape[inputAShape.size() - 2]);
  }

  if (inputBType.getRank() > 1) {
    expectedOutputShape.push_back(inputBShape[inputBShape.size() - 1]);
  }

  // Check the case of a vector-vector product. At this moment we don't support
  // scalars in IR, hence check that the output is at least 1D tensor of size 1.
  if (expectedOutputShape.size() == 0) {
    if (outputType.getRank() < 1) {
      return emitOpError("Scalar output is not supported, output must be at "
                         "least a 1D tensor");
    }

    if (outputType.getRank() > 1 || outputType.getShape()[0] != 1) {
      return emitOpError("Scalar output must be a 1D tensor of size 1");
    }

    return success();
  }

  // Verify that the output shape is correct.
  if (outputShape.size() != expectedOutputShape.size()) {
    return emitOpError("Output shape rank(")
           << outputShape.size()
           << ") must match the expected output shape rank("
           << expectedOutputShape.size() << ")";
  }

  // Verify each dim of the output shape.
  for (auto [index, outputDim, expectedDim] : llvm::zip(
           llvm::seq(outputShape.size()), outputShape, expectedOutputShape)) {
    if (outputDim != expectedDim) {
      return emitOpError("Output shape dimension[")
             << index << "](" << outputDim
             << ") doesn't match the expected output shape dimension[" << index
             << "](" << expectedDim << ")";
    }
  }

  return success();
}

For more details on adding ops to the TTNN dialect, refer to TTNN Dialect Contribution Guidelines.

Adding constraint/runtime APIs

We need to implement two APIs when adding a TTNN Op, namely getOpConstraints and getOpRuntime. More details about this can be found here.

3. Convert / Implement the Op in the TTNN passes

TTIR to TTNN

Next we will implement the conversion from the TTIR matmul Op to the TTNN matmul Op. This is a trivial conversion, as the Ops are identical in their semantics, so the changeset isn't going to be very instructive, but will at least point to the files involved. The conversion is implemented in the ConvertTTIRToTTNNPass pass in file lib/Conversion/TTIRToTTNN/TTIRToTTNNPass.cpp.

Zooming into class ConvertTTIRToTTNNPass we can see we implement the pass interface via member function void runOnOperation() final. This function will be called for every operation matching the type specified in the pass tablegen file. A quick look at include/ttmlir/Conversion/Passes.td we can see:

def ConvertTTIRToTTNN: Pass<"convert-ttir-to-ttnn", "::mlir::ModuleOp"> {

This means that runOnOperation will be called for every ModuleOp in the graph, usually there is only one ModuleOp which serves as the root of the graph.

Inside runOnOperation is usually where we define a rewrite pattern set that can match much more complicated patterns (nested inside of the ModuleOp's regions) than just a single operation. In runOperation method you will see the call to method populateTTIRToTTNNPatterns(...) that actually generates rewrite patterns. Method populateTTIRToTTNNPatterns(...) is defined in lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp.

  patterns
      .add<TensorEmptyConversionPattern,
           NamedFullConversionPattern<ttir::ZerosOp, ttnn::ZerosOp>,
           NamedFullConversionPattern<ttir::OnesOp, ttnn::OnesOp>,
           FullOpConversionPattern,
           ToLayoutOpConversionPattern,
           QuantizationOpConversionPattern<ttir::QuantizeUnrolledOp, ttnn::QuantizeOp>,
           QuantizationOpConversionPattern<ttir::DequantizeUnrolledOp, ttnn::DequantizeOp>,
           RequantizeOpConversionPattern,
           ElementwiseBinaryOpConversionPattern<ttir::AddOp, ttnn::AddOp>,
           ElementwiseBinaryOpConversionPattern<ttir::LogicalRightShiftOp, ttnn::LogicalRightShiftOp>,
           ElementwiseBinaryOpConversionPattern<ttir::SubtractOp, ttnn::SubtractOp>,
           ElementwiseBinaryOpConversionPattern<ttir::MultiplyOp, ttnn::MultiplyOp>,
           ElementwiseBinaryOpConversionPattern<ttir::DivOp, ttnn::DivideOp>,
           ElementwiseBinaryOpConversionPattern<ttir::EqualOp, ttnn::EqualOp>,
           ElementwiseBinaryOpConversionPattern<ttir::NotEqualOp, ttnn::NotEqualOp>,
           ElementwiseBinaryOpConversionPattern<ttir::GreaterEqualOp, ttnn::GreaterEqualOp>,
           ElementwiseBinaryOpConversionPattern<ttir::GreaterThanOp, ttnn::GreaterThanOp>,
           ElementwiseBinaryOpConversionPattern<ttir::LessEqualOp, ttnn::LessEqualOp>,
           ElementwiseBinaryOpConversionPattern<ttir::LessThanOp, ttnn::LessThanOp>,
           ElementwiseBinaryOpConversionPattern<ttir::LogicalAndOp, ttnn::LogicalAndOp>,
           ElementwiseBinaryOpConversionPattern<ttir::LogicalOrOp, ttnn::LogicalOrOp>,
           ElementwiseBinaryOpConversionPattern<ttir::LogicalXorOp, ttnn::LogicalXorOp>,
           ElementwiseOpConversionPattern<ttir::BitwiseAndOp, ttnn::BitwiseAndOp>,
           ElementwiseOpConversionPattern<ttir::LogicalLeftShiftOp, ttnn::LogicalLeftShiftOp>,
           ElementwiseOpConversionPattern<ttir::BitwiseOrOp, ttnn::BitwiseOrOp>,
           ElementwiseOpConversionPattern<ttir::BitwiseXorOp, ttnn::BitwiseXorOp>,
           ElementwiseOpConversionPattern<ttir::MaximumOp, ttnn::MaximumOp>,
           ElementwiseOpConversionPattern<ttir::MinimumOp, ttnn::MinimumOp>,
           ElementwiseOpConversionPattern<ttir::RemainderOp, ttnn::RemainderOp>,
           ElementwiseOpConversionPattern<ttir::Atan2Op, ttnn::Atan2Op>,
           ElementwiseOpConversionPattern<ttir::AbsOp, ttnn::AbsOp>,
           ElementwiseOpConversionPattern<ttir::CbrtOp, ttnn::CbrtOp>,
           ElementwiseOpConversionPattern<ttir::FloorOp, ttnn::FloorOp>,
           ElementwiseOpConversionPattern<ttir::IsFiniteOp, ttnn::IsFiniteOp>,
           ElementwiseOpConversionPattern<ttir::LogicalNotOp, ttnn::LogicalNotOp>,
           ElementwiseOpConversionPattern<ttir::BitwiseNotOp, ttnn::BitwiseNotOp>,
           ElementwiseOpConversionPattern<ttir::NegOp, ttnn::NegOp>,
           ElementwiseOpConversionPattern<ttir::ReluOp, ttnn::ReluOp>,
           ElementwiseOpConversionPattern<ttir::Relu6Op, ttnn::Relu6Op>,
           ElementwiseOpConversionPattern<ttir::GeluOp, ttnn::GeluOp>,
           ElementwiseOpConversionPattern<ttir::SqrtOp, ttnn::SqrtOp>,
           ElementwiseOpConversionPattern<ttir::RsqrtOp, ttnn::RsqrtOp>,
           ElementwiseOpConversionPattern<ttir::SignOp, ttnn::SignOp>,
           ElementwiseOpConversionPattern<ttir::SigmoidOp, ttnn::SigmoidOp>,
           ElementwiseOpConversionPattern<ttir::HardsigmoidOp, ttnn::HardsigmoidOp>,
           ElementwiseOpConversionPattern<ttir::SiluOp, ttnn::SiluOp>,
           ElementwiseOpConversionPattern<ttir::Log1pOp, ttnn::Log1pOp>,
           ElementwiseOpConversionPattern<ttir::ReciprocalOp, ttnn::ReciprocalOp>,
           ElementwiseOpConversionPattern<ttir::ExpOp, ttnn::ExpOp>,
           ElementwiseOpConversionPattern<ttir::ErfOp, ttnn::ErfOp>,
           ElementwiseOpConversionPattern<ttir::ErfcOp, ttnn::ErfcOp>,
           ElementwiseOpConversionPattern<ttir::LogOp, ttnn::LogOp>,
           ElementwiseOpConversionPattern<ttir::CeilOp, ttnn::CeilOp>,
           ElementwiseOpConversionPattern<ttir::SinOp, ttnn::SinOp>,
           ElementwiseOpConversionPattern<ttir::CosOp, ttnn::CosOp>,
           ElementwiseOpConversionPattern<ttir::Expm1Op, ttnn::Expm1Op>,
           ElementwiseOpConversionPattern<ttir::WhereOp, ttnn::WhereOp>,
           ElementwiseOpConversionPattern<ttir::TanOp, ttnn::TanOp>,
           ElementwiseOpConversionPattern<ttir::TanhOp, ttnn::TanhOp>,
           ElementwiseOpConversionPattern<ttir::AtanOp, ttnn::AtanOp>,
           Pooling2dOpConversionPattern<ttir::MaxPool2dOp, ttnn::MaxPool2dOp>,
           Pooling2dOpConversionPattern<ttir::AvgPool2dOp, ttnn::AvgPool2dOp>,
           GlobalAvgPool2dOpConversionPattern,
           ReductionOpConversionPattern<ttir::SumOp, ttnn::SumOp>,
           ReductionOpConversionPattern<ttir::MeanOp, ttnn::MeanOp>,
           ReductionOpConversionPattern<ttir::MaxOp, ttnn::MaxOp>,
           ReductionOpConversionPattern<ttir::MinOp, ttnn::MinOp>,
           ReductionProdOpConversionPattern,
           ReductionArgMaxOpConversionPattern,
           ElementwiseUnaryWithFloatParameterOpConversionPattern<ttir::LeakyReluOp, ttnn::LeakyReluOp>,
           BroadcastOpConversionPattern,
           PadOpConversionPattern,
           PowOpConversionPattern,
           EmbeddingOpConversionPattern,
           EmbeddingBackwardOpConversionPattern,
           RepeatOpConversionPattern,
           CumSumOpConversionPattern,
           RepeatInterleaveOpConversionPattern,
           SoftmaxOpConversionPattern,
           SortOpConversionPattern,
           TypecastOpConversionPattern,
           ClampOpConversionPattern<ttir::ClampScalarOp, ttnn::ClampScalarOp>,
           ClampOpConversionPattern<ttir::ClampTensorOp, ttnn::ClampTensorOp>,
           ConcatOpConversionPattern,
           ReshapeOpConversionPattern,
           SliceOpConversionPattern<ttir::SliceStaticOp, ttnn::SliceStaticOp>,
           SliceOpConversionPattern<ttir::SliceDynamicOp, ttnn::SliceDynamicOp>,
           SqueezeOpConversionPattern,
           UnsqueezeOpConversionPattern,
           ConstantOpConversionPattern,
           LinearOpConversionPattern,
           BatchNormInferenceOpConversionPattern,
           BatchNormTrainingOpConversionPattern,
           RMSNormOpConversionPattern,
           MatmulOpConversionPattern,
           Conv2dOpConversionPattern,
           ConvTranspose2dOpConversionPattern,
           MeshShardOpConversionPattern,
           AllReduceOpConversionPattern,
           AllGatherOpConversionPattern,
           ReduceScatterOpConversionPattern,
           CollectivePermuteOpConversionPattern,
           ArangeOpConversionPattern,
           RandOpConversionPattern,
           UpdateCacheOpConversionPattern,
           FillCacheOpConversionPattern,
           ScatterInDimOpConversionPattern,
           PermuteOpConversionPattern,
           UpsampleOpConversionPattern,
           AllToAllOpConversionPattern,
           CollectiveBroadcastOpConversionPattern,
           ConcatenateHeadsOpConversionPattern,
           ScaledDotProductAttentionOpConversionPattern,
           ScaledDotProductAttentionDecodeOpConversionPattern,
           SplitQueryKeyValueAndSplitHeadsOpConversionPattern
           >(typeConverter, ctx);

More information on rewrite patterns and their capabilities can be found in the MLIR documentation here and here.

For matmul, we defined a new conversion pattern that's generic to all binary ops with arguments named a and b:

namespace {
class MatmulOpConversionPattern : public OpConversionPattern<ttir::MatmulOp> {
public:
  using OpConversionPattern<ttir::MatmulOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(ttir::MatmulOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<ttnn::MatmulOp>(
        op, this->getTypeConverter()->convertType(op.getType()), adaptor.getA(),
        adaptor.getB(), adaptor.getTransposeA(), adaptor.getTransposeB(),
        nullptr);
    return success();
  }
};
} // namespace

Invoked as part of the rewrite set:

MatmulOpConversionPattern

TTNN to EmitC

Similarly, we also need to add a pattern to convert from TTNN dialect to EmitC dialect.

Method to populate rewrite patterns can be found in lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp:

void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
                                 mlir::RewritePatternSet &patterns,
                                 TypeConverter &typeConverter) {
  // Device ops
  //
  patterns.add<TTDeviceOpConversionPattern>(typeConverter, ctx);
  patterns.add<GetDeviceOpConversionPattern>(typeConverter, ctx);

  // Memory ops
  //
  // clang-format off
  patterns.add<ToLayoutOpConversionPattern,
               ToMemoryConfigOpConversionPattern,
               ToDTypeOpConversionPattern,
               TypecastOpConversionPattern,
               ToDeviceOpConversionPattern,
               FromDeviceOpConversionPattern,
               DeallocateOpConversionPattern>(typeConverter, ctx);
  // clang-format on

  // Tensor ops
  //
  // clang-format off
  patterns.add<EmptyOpConversionPattern,
               NamedFullOpConversionPattern<mlir::tt::ttnn::ZerosOp>,
               NamedFullOpConversionPattern<mlir::tt::ttnn::OnesOp>,
               FullOpConversionPattern,
               DefaultOpConversionPattern<mlir::tt::ttnn::ArangeOp>,
               DefaultOpConversionPattern<mlir::tt::ttnn::ConstantOp>,
               RandOpConversionPattern>(typeConverter, ctx);
  // clang-format on

  // Eltwise unary ops
  //
  patterns
      .add<EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::AbsOp>,
           EltwiseUnaryCompositeOpConversionPattern<mlir::tt::ttnn::CbrtOp>,
           ClampOpConversionPattern<::mlir::tt::ttnn::ClampScalarOp>,
           ClampOpConversionPattern<mlir::tt::ttnn::ClampTensorOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::FloorOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::IsFiniteOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::LogicalNotOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::BitwiseNotOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::NegOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::ReluOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::RsqrtOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::Relu6Op>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::HardsigmoidOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::SiluOp>,
           ElementwiseUnaryWithFloatParameterOpConversionPattern<
               mlir::tt::ttnn::LeakyReluOp>,
           EltwiseUnaryWithFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::GeluOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::SqrtOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::SignOp>,
           EltwiseUnaryWithVectorAndFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::SigmoidOp>,
           EltwiseUnaryCompositeWithFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::Log1pOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::ReciprocalOp>,
           EltwiseUnaryWithFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::ExpOp>,
           EltwiseUnaryWithFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::ErfOp>,
           EltwiseUnaryWithFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::ErfcOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::CeilOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::SinOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::CosOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::Expm1Op>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::TanOp>,
           EltwiseUnaryWithOutputAndApproxModeOpConversionPattern<
               mlir::tt::ttnn::TanhOp>,
           EltwiseUnaryOpConversionPattern<mlir::tt::ttnn::AtanOp>,
           EltwiseUnaryWithFastAndApproximateModeOpConversionPattern<
               mlir::tt::ttnn::LogOp>>(typeConverter, ctx);

  // Eltwise binary ops
  //
  patterns.add<
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::AddOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::LogicalRightShiftOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::SubtractOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::MultiplyOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::LogicalAndOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::LogicalOrOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::LogicalXorOp>,
      EltwiseBinaryCompositeOpConversionPattern<mlir::tt::ttnn::BitwiseAndOp>,
      EltwiseBinaryCompositeOpConversionPattern<mlir::tt::ttnn::BitwiseOrOp>,
      EltwiseBinaryCompositeOpConversionPattern<mlir::tt::ttnn::BitwiseXorOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::EqualOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::NotEqualOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::GreaterEqualOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::GreaterThanOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::LessEqualOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::LessThanOp>,
      EltwiseBinaryNGCompositeOpConversionPattern<mlir::tt::ttnn::MaximumOp>,
      EltwiseBinaryNGCompositeOpConversionPattern<mlir::tt::ttnn::MinimumOp>,
      EltwiseBinaryOpConversionPattern<mlir::tt::ttnn::DivideOp>,
      EltwiseBinaryCompositeOpConversionPattern<
          mlir::tt::ttnn::LogicalLeftShiftOp>,
      EltwiseBinaryCompositeOpConversionPattern<mlir::tt::ttnn::RemainderOp>,
      EltwiseBinaryNGCompositeOpConversionPattern<mlir::tt::ttnn::PowTensorOp>,
      EltwiseBinaryCompositeOpConversionPattern<mlir::tt::ttnn::Atan2Op>,
      PowScalarOpConversionPattern>(typeConverter, ctx);

  // Eltwise ternary ops
  //
  patterns.add<EltwiseTernaryOpConversionPattern<mlir::tt::ttnn::WhereOp>>(
      typeConverter, ctx);

  // Tensor manipulation ops
  //
  patterns.add<TransposeOpConversionPattern, ConcatOpConversionPattern,
               ReshapeOpConversionPattern, RepeatOpConversionPattern,
               RepeatInterleaveOpConversionPattern,
               SliceStaticOpConversionPattern, SliceDynamicOpConversionPattern,
               SortOpConversionPattern, PermuteOpConversionPattern,
               DefaultOpConversionPattern<mlir::tt::ttnn::PadOp>>(typeConverter,
                                                                  ctx);

  // Quantization ops.
  //
  patterns.add<QuantizationOpConversionPattern<mlir::tt::ttnn::QuantizeOp>,
               QuantizationOpConversionPattern<mlir::tt::ttnn::DequantizeOp>,
               RequantizeOpConversionPattern>(typeConverter, ctx);

  // Matmul ops
  //
  patterns.add<LinearOpConversionPattern, MatmulOpConversionPattern>(
      typeConverter, ctx);

  // Reduction ops
  //
  patterns.add<ReductionOpConversionPattern<mlir::tt::ttnn::SumOp>,
               ReductionOpConversionPattern<mlir::tt::ttnn::MeanOp>,
               ReductionOpConversionPattern<mlir::tt::ttnn::MaxOp>,
               ReductionOpConversionPattern<mlir::tt::ttnn::MinOp>,
               ProdOpConversionPattern, ArgMaxOpConversionPattern>(
      typeConverter, ctx);

  // Pooling ops
  //
  patterns.add<AvgPool2dOpConversionPattern>(typeConverter, ctx);
  patterns.add<MaxPool2dOpConversionPattern>(typeConverter, ctx);
  patterns.add<GlobalAvgPool2dOpConversionPattern>(typeConverter, ctx);
  patterns.add<UpsampleOpConversionPattern>(typeConverter, ctx);

  // Convolution ops
  //
  patterns.add<PrepareConv2dWeightsOpConversionPattern>(typeConverter, ctx);
  patterns.add<PrepareConv2dBiasOpConversionPattern>(typeConverter, ctx);
  patterns.add<Conv2dOpConversionPattern>(typeConverter, ctx);
  patterns.add<ConvTranspose2dOpConversionPattern>(typeConverter, ctx);

  // Other ops
  //
  patterns.add<
      SoftmaxOpConversionPattern, EmbeddingOpConversionPattern,
      DefaultOpConversionPattern<mlir::tt::ttnn::EmbeddingBackwardOp>,
      MorehCumSumOpConversionPattern, BatchNormInferenceOpConversionPattern,
      BatchNormTrainingOpConversionPattern, RMSNormOpConversionPattern>(
      typeConverter, ctx);

  // CCL ops
  //
  patterns.add<AllGatherOpConversionPattern>(typeConverter, ctx);
  patterns.add<ReduceScatterOpConversionPattern>(typeConverter, ctx);
  patterns.add<ScatterOpConversionPattern>(typeConverter, ctx);
  patterns.add<CollectivePermuteOpConversionPattern>(typeConverter, ctx);
  patterns.add<MeshShardOpConversionPattern>(typeConverter, ctx);
  patterns.add<PointToPointOpConversionPattern>(typeConverter, ctx);

  // KV Cache ops
  //
  patterns.add<UpdateCacheOpConversionPattern>(typeConverter, ctx);
  patterns.add<DefaultOpConversionPattern<mlir::tt::ttnn::FillCacheOp>>(
      typeConverter, ctx);

  // Tensor serialization ops
  //
  patterns.add<DumpTensorOpConversionPattern>(typeConverter, ctx);
  patterns.add<LoadTensorOpConversionPattern>(typeConverter, ctx);

  // Trace ops
  //
  patterns.add<WriteTensorOpConversionPattern>(typeConverter, ctx);
  patterns.add<BeginTraceCaptureOpConversionPattern>(typeConverter, ctx);
  patterns.add<EndTraceCaptureOpConversionPattern>(typeConverter, ctx);
  patterns.add<CaptureOrExecuteTraceOpConversionPattern>(typeConverter, ctx);
  patterns.add<ExecuteTraceOpConversionPattern>(typeConverter, ctx);

  // Arith ops
  //
  patterns.add<ArithConstantOpConversionPattern>(typeConverter, ctx);

  // Tuple ops
  //
  patterns.add<GetTupleElementOpConversionPattern>(typeConverter, ctx);
  patterns.add<TupleOpConversionPattern>(typeConverter, ctx);

  // LoadCached op
  //
  patterns.add<LoadCachedOpConversionPattern>(typeConverter, ctx);

  // Module op
  //
  patterns.add<ModuleOpConversionPattern>(typeConverter, ctx);

  // FuncOp
  //
  patterns.add<FuncOpConversionPattern>(typeConverter, ctx);

  // Transformers ops
  //
  patterns.add<ConcatenateHeadsOpConversionPattern>(typeConverter, ctx);
  patterns.add<SplitQueryKeyValueAndSplitHeadsOpConversionPattern>(
      typeConverter, ctx);
  patterns.add<RotaryEmbeddingLlamaOpConversionPattern>(typeConverter, ctx);
  patterns.add<NLPConcatHeadsDecodeOpConversionPattern>(typeConverter, ctx);
  patterns.add<ScaledDotProductAttentionDecodeOpConversionPattern>(
      typeConverter, ctx);
  patterns.add<ScaledDotProductAttentionOpConversionPattern>(typeConverter,
                                                             ctx);
  patterns.add<NLPCreateQKVHeadsDecodeOpConversionPattern>(typeConverter, ctx);
}

Writing conversion patterns to EmitC is a little tricky at first. In general case, we will be converting an op that has operands (SSAs) and attributes (e.g. data type) as arguments. We want to flatten these arguments at call site.

We'll use EmitC's CallOpaqueOp as the target op. Let's take a look at our matmul IR within TTNN dialect:

"ttnn.matmul"(%2, %4, %5) : (tensor<64x128xbf16, #ttnn_layout4>, tensor<128x96xbf16, #ttnn_layout6>, tensor<64x96xbf16, #ttnn_layout7>) -> tensor<64x96xbf16, #ttnn_layout7>

Now let's look at matmul's call signature in TTNN lib:

    static Tensor invoke(
        const Tensor& input_tensor_a,
        const Tensor& input_tensor_b,
        const bool transpose_a = false,
        const bool transpose_b = false,
        const std::optional<const MemoryConfig>& memory_config = std::nullopt,
        const std::optional<const DataType> dtype = std::nullopt,
        const std::optional<const MatmulProgramConfig>& program_config = std::nullopt,
        const std::optional<const std::string>& activation = std::nullopt,
        const std::optional<const DeviceComputeKernelConfig> compute_kernel_config = std::nullopt,
        const std::optional<const CoreGrid> core_grid = std::nullopt,
        const std::optional<const tt::tt_metal::Tile>& output_tile = std::nullopt,
        std::optional<Tensor> optional_output_tensor = std::nullopt,
        const std::optional<const DeviceGlobalCircularBuffer>& global_cb = std::nullopt);

If we look closely, we'll notice that the IR has way less arguments than can be seen in the actual signature of the op - as we're lowering to EmitC, which gets translated into actual C++ code, we need to correct for this (ideally the op would be perfectly modelled with all the arguments, but that is not the case today).

We do this by filling in the gaps. EmitC's CallOpaqueOp takes in an array of attributes, and an array of operands, which need to be combined. The combining is done by extending the array of attributes with "pointers" into operands, like so:

    llvm::SmallVector<mlir::Attribute> args{
        emitter.emit(srcOp.getA()),
        emitter.emit(srcOp.getB()),
        emitter.emit(srcOp.getTransposeA()),
        emitter.emit(srcOp.getTransposeB()),
        emitter.emit(std::nullopt) | emitter.getMemoryConfig(srcOp.getResult()),
    };

Pointers are denoted with IndexTypes, wrapped into IntegerAttrs. Attributes are converted into EmitC's OpaqueAttr which can, for practical purposes, be treated as strings: a BoolAttr carrying "false" as value needs to be converted into an OpaqueAttr whose value is a string "false", which is what the convertBoolAttr function does.

This is our final converted EmitC CallOpaqueOp:

emitc.call_opaque "ttnn::matmul"(%3, %6, %9) {args = [0 : index, 1 : index, #emitc.opaque<"false">, #emitc.opaque<"false">, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, 2 : index]} : (!emitc.opaque<"ttnn::Tensor">, !emitc.opaque<"ttnn::Tensor">, !emitc.opaque<"ttnn::Tensor">) -> !emitc.opaque<"ttnn::Tensor">

which, when translated to C++ code, looks like:

ttnn::matmul(v6, v9, false, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, v12);

Full conversion pattern for matmul op:

namespace {
class MatmulOpConversionPattern
    : public TTNNToEmitCBaseOpConversionPattern<mlir::tt::ttnn::MatmulOp> {

public:
  using TTNNToEmitCBaseOpConversionPattern<
      mlir::tt::ttnn::MatmulOp>::TTNNToEmitCBaseOpConversionPattern;

  LogicalResult
  matchAndRewrite(mlir::tt::ttnn::MatmulOp srcOp,
                  mlir::tt::ttnn::MatmulOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    ttnn_to_emitc::EmitCTTNNEmitter<mlir::tt::ttnn::MatmulOp> emitter(
        srcOp, adaptor, rewriter);

    llvm::SmallVector<mlir::Attribute> args{
        emitter.emit(srcOp.getA()),
        emitter.emit(srcOp.getB()),
        emitter.emit(srcOp.getTransposeA()),
        emitter.emit(srcOp.getTransposeB()),
        emitter.emit(std::nullopt) | emitter.getMemoryConfig(srcOp.getResult()),
    };

    emitter.replaceOp(*this, args);

    return success();
  }
};
} // namespace

4. Add a compiler unit test for the Op

So far we have defined the Op in the TTIR and TTNN dialects, implemented verifiers, and have conversion passes. Now we need to add a unit test to ensure that the pass is working correctly. The compiler unit tests are located in test/ttmlir/Dialect area. In this case we'll add a test under the TTNN subdirectory since we are testing the ConvertTTIRToTTNNPass.

test/ttmlir/Dialect/TTNN/matmul/simple_matmul.mlir

// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline -o %t %s
// RUN: FileCheck %s --input-file=%t
module {
  func.func @forward(%arg0: tensor<64x128xbf16>, %arg1: tensor<128x96xbf16>) -> tensor<64x96xbf16> {
    %0 = ttir.empty() : tensor<64x96xbf16>
    // CHECK: "ttnn.matmul"
    %1 = "ttir.matmul"(%arg0, %arg1, %0) : (tensor<64x128xbf16>, tensor<128x96xbf16>, tensor<64x96xbf16>) -> tensor<64x96xbf16>
    return %1 : tensor<64x96xbf16>
  }
}

Unit tests in MLIR are typically written using a tool called FileCheck, please refer to the llvm FileCheck documentation for a tutorial and more information about the RUN and CHECK directives.

A few things to point out specifically regarding tt-mlir dialects:

• ttcore.system_desc: This is a 1-1 mapping to the SystemDesc flatbuffer schema that is used to describe the system configuration. This is a required attribute tagged on the top level module for all tt-mlir dialects.
• Pass --ttnn-layout is a prerequisite before running convert-ttir-to-ttnn. This pass is responsible for converting the input tensors to device memory space and tile layout before lowering to TTNN.
• This test is asserting that ttir.matmul converts to ttnn.matmul.

To run the test, you can use the following command:

cmake --build build -- check-ttmlir

You can also manually run ttmlir-opt on the test file to see the resulting output:

./build/bin/ttmlir-opt --ttcore-register-device="system-desc-path=<PATH_TO_SYSTEM_DESC>" --ttir-to-ttnn-backend-pipeline test/ttmlir/Dialect/TTNN/matmul/simple_matmul.mlir

5. Define flatbuffer schema for the Op

Next we will define the flatbuffer schema for the Op. The schema must capture all tensor inputs, outputs, and attributes of the Op, i.e. everything the runtime needs to execute the Op.

The schema can be placed in an existing .fbs file located in the include/ttmlir/Target/TTNN/operations directory.

If no suitable .fbs file exists for the operation category, feel free to create new .fbs files as needed. After creating a new .fbs file, remember to add a corresponding cmake target in the include/ttmlir/Target/TTNN/CMakeLists.txt file.

include/ttmlir/Target/TTNN/CMakeLists.txt

  operations/matmul.fbs

In our case, we can add our schema to include/ttmlir/Target/TTNN/operations/matmul.fbs directly, without needing to create a new file.

include/ttmlir/Target/TTNN/operations/matmul.fbs

table MatmulOp {
  a: tt.target.ttnn.TensorRef;
  b: tt.target.ttnn.TensorRef;
  out: tt.target.ttnn.TensorRef;
  transpose_a: bool;
  transpose_b: bool;
  matmul_program_config: tt.target.ttnn.MatmulProgramConfig;
}

Type TensorRef, flatbuffer tables with suffix Ref are used to represent live values during the runtime, decoupled from the underlying Desc suffixes which carry the type and attribute information for the object.

After creating the schema for our new operation type, we need to register it in the OpType union within program.fbs. This file serves as the main entry point for all program information, where the OpType union collects and defines all supported operation types and their corresponding schemas.

include/ttmlir/Target/TTNN/program.fbs

  MatmulOp,

If a new .fbs file was created, don't forget to include the new file in include/ttmlir/Target/TTNN/program.fbs.

include "ttmlir/Target/TTNN/operations/matmul.fbs";

More information about writing flatbuffer schemas can be found in the flatbuffers documentation

6. Serialize the Op in the flatbuffer format

In the previous section we defined the flatbuffer schema for the matmul Op, now let's put our new schema definition to use. The schema is used as input to a program called flatc which generates C++ code (or any language for that matter) for serializing and deserializing the schema. This generated code can be found in build/include/ttmlir/Target/TTNN/program_generated.h.

Let's head over to lib/Target/TTNN/TTNNToFlatbuffer.cpp to define a createOp overloaded function that does the conversion from MLIR to flatbuffer:

::flatbuffers::Offset<::tt::target::ttnn::MatmulOp>
createOp(FlatbufferObjectCache &cache, MatmulOp op) {
  auto a = cache.at<::tt::target::ttnn::TensorRef>(
      getOperandThroughDPSOps(op.getA()));
  auto b = cache.at<::tt::target::ttnn::TensorRef>(
      getOperandThroughDPSOps(op.getB()));
  auto output = cache.getOrCreate(op.getResult(), tensorValueToFlatbuffer);

  using MatmulConfigType = ::tt::target::ttnn::MatmulProgramConfig;
  MatmulConfigType matmulProgramConfigType = MatmulConfigType::NONE;
  ::flatbuffers::Offset<void> matmulProgramConfigDesc;
  if (auto matmulProgramConfig = op.getMatmulProgramConfigAttr()) {
    if (auto config =
            mlir::dyn_cast<ttnn::MatmulMultiCoreReuseProgramConfigAttr>(
                matmulProgramConfig)) {
      matmulProgramConfigType =
          MatmulConfigType::MatmulMultiCoreReuseProgramConfig;
      matmulProgramConfigDesc = toFlatbuffer(cache, config).Union();
    } else if (auto config = mlir::dyn_cast<
                   ttnn::MatmulMultiCoreReuseMultiCastProgramConfigAttr>(
                   matmulProgramConfig)) {
      matmulProgramConfigType =
          MatmulConfigType::MatmulMultiCoreReuseMultiCastProgramConfig;
      matmulProgramConfigDesc = toFlatbuffer(cache, config).Union();
    } else if (auto config = mlir::dyn_cast<
                   ttnn::MatmulMultiCoreReuseMultiCast1DProgramConfigAttr>(
                   matmulProgramConfig)) {
      matmulProgramConfigType =
          MatmulConfigType::MatmulMultiCoreReuseMultiCast1DProgramConfig;
      matmulProgramConfigDesc = toFlatbuffer(cache, config).Union();
    } else if (
        auto config = mlir::dyn_cast<
            ttnn::MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfigAttr>(
            matmulProgramConfig)) {
      matmulProgramConfigType = MatmulConfigType::
          MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig;
      matmulProgramConfigDesc = toFlatbuffer(cache, config).Union();
    }
  }

  return ::tt::target::ttnn::CreateMatmulOp(
      *cache.fbb, a, b, output, op.getTransposeA(), op.getTransposeB(),
      matmulProgramConfigType, matmulProgramConfigDesc);
}

Lots of things are happening here, let's break it down:

• FlatbufferObjectCache: This is a helper class that is used to cache objects in the flatbuffer that are created during the serialization process. This is necessary for managing value lifetimes and identifiers, at the same time it is an optimization to avoid having multiple copies of the same object. For example, a TensorRef with multiple uses could naively be recreated, one for each use, but with the cache we can ensure that the object is only created once and all uses point to the same flatbuffer offset. The cache is passed around to all serialization functions and should be used whenever creating a new object.
• getOperandThroughDPSOps: In section 1. we discussed DPS semantics and the drawback of having the result alias the output tensor. This is one of those cases where we need to use a helper function to trace through the output operands to find the original SSA name in order to associate it with the original TensorRef.
• CreateMatmulOp: The autogenerated function from the flatbuffer schema that actually serializes the data into the flatbuffer format.

We can finally generate a binary with our new Op! We can use the following command:

./build/bin/ttmlir-opt --ttcore-register-device="system-desc-path=<PATH_TO_SYSTEM_DESC>" --ttir-to-ttnn-backend-pipeline test/ttmlir/Dialect/TTNN/matmul/simple_matmul.mlir | ./build/bin/ttmlir-translate --ttnn-to-flatbuffer -o out.ttnn

And we can inspect the with ttrt:

ttrt read out.ttnn

Note: If the above ttrt command yields a segfault, a clean build of your workspace may be required: Build Instructions

7. Add runtime support for the Op

Next, we want to add runtime support for the Op by parsing the flatbuffer and invoking the TTNN API.

runtime/lib/ttnn/operations/matmul/matmul.cpp

void run(const ::tt::target::ttnn::MatmulOp *op, ProgramContext &context) {
  ProgramTensorPool &tensorPool = context.getTensorPool();
  const ::ttnn::Tensor &lhs = tensorPool.getTTNNTensorAndValidate(op->a());
  const ::ttnn::Tensor &rhs = tensorPool.getTTNNTensorAndValidate(op->b());

  auto outputMemoryConfig =
      ::tt::runtime::ttnn::utils::createMemoryConfigIfNeeded(
          ::tt::runtime::ttnn::utils::getTensorRefMemoryConfig(op->out()));
  LOG_ASSERT(::tt::runtime::ttnn::utils::inSystemMemory(op->out()) ||
                 outputMemoryConfig,
             "Memory config must exist for device tensors");

  ::ttnn::DataType outputDataType = utils::getDataType(op->out());

  std::optional<::ttnn::operations::matmul::MatmulProgramConfig>
      matmulProgramConfig = utils::createMatmulProgramConfigIfNeeded(op);

  ::ttnn::Tensor output = ::ttnn::matmul(
      lhs, rhs, op->transpose_a(), op->transpose_b(), outputMemoryConfig,
      outputDataType, matmulProgramConfig,
      /*activation=*/std::nullopt, /*compute_kernel_config=*/std::nullopt,
      /*core_grid=*/std::nullopt, /*output_tile=*/std::nullopt,
      /* optional_output_tensor=*/std::nullopt);

  tensorPool.insertTTNNTensorAndValidate(op->out(), output);
}

A couple things to note from above:

• Most runtime op functions will follow a similar pattern, they will take in some additional datastructures for managing the program context.
  • Program context tracks the state of the current program. It stores intermediate tensors and devices.
• tensorPool.at(op->in0()->global_id()): global_id is a unique identifier for the tensor that was generated and managed by the FlatbufferObjectCache. This is how it's intended to be used by the runtime.
• Some operations may belong to a larger set of operations. For example, any eltwise unary operations can be added in runtime/lib/ttnn/operations/eltwise/unary.cpp directly without needing to create a new file.

If a new file is created for the op, we need to add a new source to runtime/lib/ttnn/operations/CMakeLists.txt and a new case to runtime/lib/ttnn/program_executor.cpp.

To update runtime/lib/ttnn/operations/CMakeLists.txt, include the path to the source file in TTNN_OPS_SRCS:

runtime/lib/ttnn/operations/CMakeLists.txt

  ${CMAKE_CURRENT_SOURCE_DIR}/matmul/matmul.cpp

To update runtime/lib/ttnn/program_executor.cpp, add a new case to the runOperation method of ProgramExecutor:

runtime/lib/ttnn/program_executor.cpp

  case ::tt::target::ttnn::OpType::MatmulOp: {
    return operations::matmul::run(op->type_as_MatmulOp(), getContext());
  }

We can test our changes with ttrt (don't forget to rebuild ttrt):

ttrt run out.ttnn

8. Add a silicon unit test for the Op

After adding runtime support, we're ready to test our Op on silicon. All silicon tests are located under test/ttmlir/Silicon. The process is similar to adding a compiler unit test.

In our specific case, we create a unit test here:

test/ttmlir/Silicon/TTNN/matmul/simple_matmul.mlir

// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" -o %t.mlir %s
// RUN: FileCheck %s --input-file=%t.mlir
// RUN: ttmlir-translate --ttnn-to-flatbuffer -o %t.ttnn %t.mlir
module {
  func.func @forward(%arg0: tensor<64x128xbf16>, %arg1: tensor<128x96xbf16>) -> tensor<64x96xbf16> {
    %0 = ttir.empty() : tensor<64x96xbf16>
    // CHECK: "ttnn.matmul"
    %1 = "ttir.matmul"(%arg0, %arg1, %0) : (tensor<64x128xbf16>, tensor<128x96xbf16>, tensor<64x96xbf16>) -> tensor<64x96xbf16>
    return %1 : tensor<64x96xbf16>
  }

  func.func @matmul_transpose_lhs(%arg0: tensor<64x128xbf16>, %arg1: tensor<64x128xbf16>) -> tensor<128x128xbf16> {
    %0 = ttir.empty() : tensor<128x128xbf16>
    // CHECK: "ttnn.matmul"
    %1 = "ttir.matmul"(%arg0, %arg1, %0) <{transpose_a = true}>: (tensor<64x128xbf16>, tensor<64x128xbf16>, tensor<128x128xbf16>) -> tensor<128x128xbf16>
    return %1 : tensor<128x128xbf16>
  }

  func.func @matmul_transpose_rhs(%arg0: tensor<64x128xbf16>, %arg1: tensor<64x128xbf16>) -> tensor<64x64xbf16> {
    %0 = ttir.empty() : tensor<64x64xbf16>
    // CHECK: "ttnn.matmul"
    %1 = "ttir.matmul"(%arg0, %arg1, %0) <{transpose_b = true}>: (tensor<64x128xbf16>, tensor<64x128xbf16>, tensor<64x64xbf16>) -> tensor<64x64xbf16>
    return %1 : tensor<64x64xbf16>
  }
}

Couple things to point out about this process:

• Tests placed under test/ttmlir/Dialect will only test the compiler's capability of compiling the module. If you want the module to run on silicon in CI, the test must be placed under test/ttmlir/Silicon.
• Notice the differences between the compilation headers of test/ttmlir/Silicon/TTNN/simple_matmul.mlir and test/ttmlir/Dialect/TTNN/matmul/simple_matmul.mlir
  • --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%": The system-desc-path option specifies the location of the system descriptor required for compiling the module. This is crucial for silicon tests, as modules compiled with different system descriptors may vary in silicon compatibility. Ensuring the system descriptor accurately reflects the target hardware is essential for running the module correctly.
  • // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn: This runs ttmlir-translate that serializes the output mlir module to a flatbuffer binary. We added the logic for this serialization in the Serialize the Op in the flatbuffer format section.

9. Add an EmitC test for the Op

Op should be tested in the EmitC (C++ codegen) path as well.

TTNN EmitC tests live in the test/ttmlir/EmitC/TTNN path. In our case, the test is in test/ttmlir/EmitC/TTNN/matmul/matmul.mlir.

test/ttmlir/EmitC/TTNN/matmul/matmul.mlir

// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" -o %t.mlir %s
// RUN: ttmlir-translate --ttnn-to-flatbuffer -o %basename_t.ttnn %t.mlir
// RUN: ttmlir-opt --ttnn-backend-to-emitc-pipeline -o %t2.mlir %t.mlir
// RUN: ttmlir-translate --mlir-to-cpp -o %basename_t.cpp %t2.mlir

func.func @matmul(%arg0: tensor<64x128xbf16>, %arg1: tensor<128x96xbf16>) -> tensor<64x96xbf16> {
  %0 = ttir.empty() : tensor<64x96xbf16>
  %1 = "ttir.matmul"(%arg0, %arg1, %0) : (tensor<64x128xbf16>, tensor<128x96xbf16>, tensor<64x96xbf16>) -> tensor<64x96xbf16>
  return %1 : tensor<64x96xbf16>
}

The first two RUN lines create a flatbuffer. The third and forth convert to EmitC dialect, translate to C++, then output the result to matmul.mlir.cpp file.

Additionally, the op's header file operations/matmul/matmul.hpp should be added to the list of includes in tools/ttnn-standalone/ttnn-precompiled.hpp:

#include "operations/ccl/ccl_host_types.hpp"
#include "operations/conv/conv2d/conv2d.hpp"
#include "operations/conv/conv2d/prepare_conv2d_weights.hpp"
#include "operations/conv/conv_transpose2d/conv_transpose2d.hpp"
#include "operations/core/core.hpp"
#include "operations/creation.hpp"
#include "operations/data_movement/concat/concat.hpp"
#include "operations/data_movement/permute/permute.hpp"
#include "operations/data_movement/repeat/repeat.hpp"
#include "operations/data_movement/repeat_interleave/repeat_interleave.hpp"
#include "operations/data_movement/scatter/scatter.hpp"
#include "operations/data_movement/slice/slice.hpp"
#include "operations/data_movement/sort/sort.hpp"
#include "operations/data_movement/transpose/transpose.hpp"
#include "operations/eltwise/binary/binary.hpp"
#include "operations/eltwise/binary/binary_composite.hpp"
#include "operations/eltwise/quantization/quantization.hpp"
#include "operations/eltwise/unary/unary_composite.hpp"
#include "operations/embedding/embedding.hpp"
#include "operations/embedding_backward/embedding_backward.hpp"
#include "operations/experimental/transformer/nlp_concat_heads/nlp_concat_heads.hpp"
#include "operations/kv_cache/kv_cache.hpp"
#include "operations/matmul/matmul.hpp"
#include "operations/moreh/moreh_cumsum/moreh_cumsum.hpp"
#include "operations/normalization/batch_norm/batch_norm.hpp"
#include "operations/normalization/rmsnorm/rmsnorm.hpp"
#include "operations/normalization/softmax/softmax.hpp"
#include "operations/pool/generic/generic_pools.hpp"
#include "operations/pool/global_avg_pool/global_avg_pool.hpp"
#include "operations/pool/upsample/upsample.hpp"
#include "operations/rand/rand.hpp"
#include "operations/reduction/argmax/argmax.hpp"
#include "operations/reduction/generic/generic_reductions.hpp"
#include "operations/reduction/prod/prod.hpp"
#include "operations/trace.hpp"
#include "operations/transformer/concatenate_heads/concatenate_heads.hpp"
#include "operations/transformer/sdpa/sdpa.hpp"
#include "operations/transformer/sdpa_decode/sdpa_decode.hpp"
#include "operations/transformer/split_query_key_value_and_split_heads/split_query_key_value_and_split_heads.hpp"
#include "tt-metalium/bfloat16.hpp"
#include "ttnn/common/queue_id.hpp"
#include "ttnn/core.hpp"
#include "ttnn/device.hpp"
#include "ttnn/operations/copy/typecast/typecast.hpp"
#include "ttnn/operations/experimental/transformer/nlp_concat_heads_decode/nlp_concat_heads_decode.hpp"
#include "ttnn/operations/experimental/transformer/nlp_create_qkv_heads_decode/nlp_create_qkv_heads_decode.hpp"
#include "ttnn/operations/experimental/transformer/rotary_embedding_llama/rotary_embedding_llama.hpp"
#include "ttnn/tensor/serialization.hpp"
#include "ttnn/tensor/tensor.hpp"
#include "ttnn/tensor/types.hpp"
#include "ttnn/types.hpp"
#include "workarounds.hpp"

Adding OpConstraints and OpRuntime APIs to TTNN Operations

Overview

The TTNN Op Model Interface provides two key APIs for analyzing and optimizing operations:

• getOpConstraints: Returns constraint information including memory requirements, layout compatibility, and operation feasibility
• getOpRuntime: Returns performance metrics including execution time estimates

These APIs enable the compiler to make informed decisions about operation placement, memory allocation, and performance optimization.

This guide walks you through best practices for implementing these APIs. It will cover the following steps:

1. Architecture
2. Implementation Steps
   • Step 1: Implement Operation-Specific Methods
   • Step 2: Add Core Model Implementation
   • Step 3: Add Unit Tests
   • Step 4: Add Integration Tests
3. Key Considerations
4. Example: Complete Implementation

Architecture

The implementation follows a layered architecture:

TTNNOpModelInterface.cpp (Operation-specific implementations)
    ↓
TTNNOpModel.h/.cpp (Core model implementations and helpers)
    ↓
Metal Backend (Runtime execution and constraint validation)

Important note: getOpConstraints and getOpRuntime API calls should be identical to regular op invocation path through runtime. The only difference is that one call is generated from the IR while the other is from serialised FB. For example, you can compare:

The runtime code runtime/lib/ttnn/operations/conv/conv2d.cpp:

void run(const ::tt::target::ttnn::Conv2dOp *op, ProgramContext &context) {
  // ...
}

With the constraint API implementation code lib/OpModel/TTNN/TTNNOpModel.cpp:

llvm::Expected<OpConstraints> OpModel<Conv2dOp>::getOpConstraints(/* args */){
  // ...
}
// and:
llvm::Expected<size_t> OpModel<Conv2dOp>::getOpRuntime(/* args */){
  // ...
}

And observe the similarities. This is very important to maintain throughout the lifetime of the project to guarantee consistency and functional correctness.

Implementation Steps

Step 1: Implement Operation-Specific Methods

Add your operation's implementation in lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp:

//===----------------------------------------------------------------------===//
// YourOp - TTNN Op Model Interface
//===----------------------------------------------------------------------===//

llvm::Expected<op_model::OpConstraints>
YourOp::getOpConstraints(const std::vector<TTNNLayoutAttr> &inputs,
                         const OpConfig &opConfig) {
  // You can extract all input tensors' layouts from `inputs`.
  // Other configurations can also be extracted from `opConfig`.
  // All inputs/attrs can be extracted from YourOp's member functions.
  // This layer is usually a wrapper to extract the op's necessary inputs/attrs
  // and pass those information to TTNNOpModel.h.
  return opConstraintsCache().getOrCompute(
      op_model::OpModel<YourOp>::getOpConstraints, *this,
      deviceGrid, /* other parameters */);
}

llvm::Expected<size_t>
YourOp::getOpRuntime(const std::vector<TTNNLayoutAttr> &inputs,
                     const OpConfig &opConfig) {
  // Similar to the previous function.
  return opRuntimeCache().getOrCompute(
      op_model::OpModel<YourOp>::getOpRuntime, *this,
      /* other parameters */);
}

Note: The codebase provides several template helpers for common operation patterns:

Unary Operations

// For simple unary operations (like ReluOp, SqrtOp, etc.)
return detail::getUnaryOpConstraints(*this, inputs, opConfig);
return detail::getUnaryOpRuntime(*this, inputs, opConfig);

Binary Operations

// For binary element-wise operations (like AddOp, MultiplyOp, etc.)
return detail::getBinaryOpConstraints(*this, inputs, opConfig);
return detail::getBinaryOpRuntime(*this, inputs, opConfig);

Ternary Operations

// For ternary operations (like WhereOp)
return detail::getTernaryOpConstraints(*this, inputs, opConfig);
return detail::getTernaryOpRuntime(*this, inputs, opConfig);

Reduction Operations

// For reduction operations (like SumOp, MeanOp, etc.)
return detail::getReductionOpConstraints(*this, inputs, opConfig);
return detail::getReductionOpRuntime(*this, inputs, opConfig);

Step 2: Add Core Model Implementation

Add the core implementation in include/ttmlir/OpModel/TTNN/TTNNOpModel.h:

template <>
struct OpModel<YourOp> {
  static llvm::Expected<OpConstraints>
  getOpConstraints(ttcore::GridAttr deviceGrid,
                   // ... operation-specific parameters ...
                   TTNNLayoutAttr outputLayout);

  static llvm::Expected<size_t>
  getOpRuntime(// ... operation-specific parameters  ...
               TTNNLayoutAttr outputLayout);
};

And the corresponding implementation in lib/OpModel/TTNN/TTNNOpModel.cpp:

llvm::Expected<OpConstraints>
OpModel<YourOp>::getOpConstraints(
    ttcore::GridAttr deviceGrid,
    // operation-specific parameters
    TTNNLayoutAttr outputLayout) {
  #ifdef TTMLIR_ENABLE_OPMODEL
  // 1. Perform necessary conversions, create Tensor objects, etc.

  // 2. Create query closure
  // Here the ultimate goal is to enable the optimizer to call the
  // invoke method of the op in tt-metal. This is achieved through
  // creating a lambda that calls `query_op_constraints` which
  // receives 3 arguments:
  //   1. An op (eg. ::ttnn::yourOp). This is the op's backend
  //      found under tt-metal/src/tt-metal/ttnn/. The op usually
  //      has an 'invoke' method.
  //   2. The device,
  //   3. A variadic number of inputs that are converted to match
  //      the metal's definitions. The order and the types of these
  //      inputs are expected to match the invoke function of the
  //      op in metal.
  auto yourOpQuery = [=]() {
    return ::ttnn::graph::query_op_constraints(
        ::ttnn::yourOp, device, /* other converted parameters */);
  };

  // 3. Call getOpConstraints and pass the callable.
  return operation::getOpConstraints(getContext(), deviceGrid,
                                     yourOpQuery);
#else
  return OpConstraints{};
#endif // TTMLIR_ENABLE_OPMODEL
}

llvm::Expected<size_t>
OpModel<YourOp>::getOpRuntime(
    // operation-specific parameters
    TTNNLayoutAttr outputLayout) {
#ifdef TTMLIR_ENABLE_OPMODEL
  // Similar to the previous function.
  // Create query closure
  auto yourOpQuery = [=]() {
    return ::ttnn::graph::query_op_runtime(
        ::ttnn::yourOp, device, /* other converted parameters */);
  };

  return operation::getOpRuntime(yourOpQuery);
#else
  return llvm::createStringError("Not Implemented");
#endif // TTMLIR_ENABLE_OPMODEL
}

Note: If the op's definition cannot be found by gcc you might need to #include the related header file in OpModel/TTNN/MetalHeaders.h.

Note: The codebase provides several implementations for common operation patterns which is done through Explicit template instantiation.

Unary Operations

// For simple unary operations (like ReluOp, SqrtOp, etc.)
template struct UnaryEltwiseOpModel</* Op */>;

Binary Operations

// For binary element-wise operations (like AddOp, MultiplyOp, etc.)
template struct BinaryEltwiseOpModel</* Op */>;

Ternary Operations

// For ternary operations (like WhereOp)
template struct TernaryEltwiseOpModel</* Op */>;

Reduction Operations

// For reduction operations (like SumOp, MeanOp, etc.)
template struct ReductionOpModel</* Op */>;

Step 3: Add Unit Tests

Create tests in test/unittests/OpModel/TTNN/Op/TestOpModelInterface.cpp:

TEST_F(OpModelBase, YourOpInterface) {
  // Create input tensors
  auto input = createEmptyTensor({32, 64}, ttcore::DataType::Float32);

  // Create operation
  auto yourOp = builder.create<YourOp>(
      loc, createRankedTensorType({32, 64}, ttcore::DataType::Float32),
      input, /* other parameters */);

  // Test constraints
  auto constraintsExp = getOpConstraints(yourOp.getOperation());
  if (constraintsExp) {
      auto l1 = constraintsExp.get();
      const auto &[cbSize, l1PeakSize, totalPeakSize, outputSize, outputLayout] = l1;
      EXPECT_EQ(cbSize, /* some expected value */);
      EXPECT_EQ(l1PeakSize, /* some expected value */);
      EXPECT_EQ(totalPeakSize, /* some expected value */);
      EXPECT_EQ(outputSize, /* some expected value */);
  } else {
      FAIL() << "Missing L1 constraints; Error="
          << llvm::toString(constraintsExp.takeError()) << std::endl;
  }
  auto runtimeExp = getOpRuntime(yourOp.getOperation());
  if (runtimeExp) {
      EXPECT_TRUE(runtimeExp.get() > 0);
  } else {
      FAIL() << llvm::toString(runtimeExp.takeError());
  }
}

Step 4: Add Integration Tests

Create comprehensive tests in test/unittests/OpModel/TTNN/Lib/TestOpModelLib.cpp. The following is one way of doing this, not the only possible test.

Note: For operations with additional parameters (like kernel size, stride, etc.), add them between the input and output tensors in the tuple definition and destructuring assignment.

template <typename OpTy>
class OpModelYourOpParam : public OpModelTest,
                           public ::testing::WithParamInterface<
                               std::tuple<detail::TestTensor, // input
                                          detail::TestTensor, // output
                                          detail::ExpectedResult>> {
protected:
  void RunTest() {
    auto [inputTensor, outputTensor, expectedResult] = GetParam();

    // Create tensors with specified layouts
    TTNNLayoutAttr inputLayout = createLayout(inputTensor);
    TTNNLayoutAttr outputLayout = createLayout(outputTensor);

    auto constraintsExp = OpModel<OpTy>::getOpConstraints(
        CreateWorkerGrid(), /* pass the params according to TTNNOpModel.h interface */, outputLayout);
    EXPECT_EQ(static_cast<bool>(constraintsExp), expectedResult.expectedLegal);
    if (expectedResult.expectedLegal) {
      const auto [cbSize, l1PeakSize, totalPeakSize, outputSize, outputLayout] =
          constraintsExp.get();
      EXPECT_EQ(cbSize, expectedResult.expectedCbSize);
      EXPECT_EQ(l1PeakSize, expectedResult.expectedL1PeakSize);
      EXPECT_EQ(totalPeakSize, expectedResult.expectedTotalPeakSize);
      EXPECT_EQ(outputSize, expectedResult.expectedOutputSize);
    } else {
      // Must clean up the error
      llvm::consumeError(constraintsExp.takeError());
    }

    auto runtimeExp =
        OpModel<OpTy>::getOpRuntime(/* pass the params according to TTNNOpModel.h interface */, outputLayout);
    EXPECT_EQ(static_cast<bool>(runtimeExp), expectedResult.expectedLegal);
    if (expectedResult.expectedLegal) {
      EXPECT_TRUE(runtimeExp.get() > 0);
    } else {
      llvm::consumeError(runtimeExp.takeError());
    }
  }
};

using OpModelYourOpParamTest = OpModelYourOpParam<YourOp>;
TEST_P(OpModelYourOpParamTest, YourOp) { RunTest(); }

INSTANTIATE_TEST_SUITE_P(
    YourOpTests, OpModelYourOpParamTest,
    ::testing::Values(
        std::make_tuple(
            detail::TestTensor{{32, 64}, TensorMemoryLayout::INTERLEAVED, BufferType::DRAM},
            detail::TestTensor{{32, 64}, TensorMemoryLayout::INTERLEAVED, BufferType::DRAM},
            detail::ExpectedResult{true, 8192, 8192, 8192, 8192}),
        // Add more test cases...
    ));

Key Considerations

Error handling: Operations Not Supported

For operations that cannot support these APIs, use the provided error helpers in TTNNOpModelInterface.cpp. We're keeping track of such ops in this issue. So please either update the issue or add comments to it.

llvm::Expected<op_model::OpConstraints>
YourOp::getOpConstraints(const std::vector<TTNNLayoutAttr> &inputs,
                         const OpConfig &opConfig) {
  return detail::issueErrorForGetOpConstraints(
      getOperation(), detail::ReasonForLackOfSupport::/*..*/);
}

llvm::Expected<size_t>
YourOp::getOpRuntime(const std::vector<TTNNLayoutAttr> &inputs,
                     const OpConfig &opConfig) {
  return detail::issueErrorForGetOpRuntime(
      getOperation(), detail::ReasonForLackOfSupport::/*..*/);
}

Available error reasons:

• NeedsMemoryIO: Operation requires memory I/O during trace capture
• MissingMetalDefinition: Metal backend implementation is missing
• NeedsMultiDevice: Operation requires multi-device support
• NoNeedForConstraintAPI: Operation doesn't benefit from constraint analysis
• ArchitecturalMismatch: Mismatch in Operation's definition in metal and mlir

Device Grid Validation

Validate the device worker grid before proceeding:

llvm::Expected<bool> check = detail::checkDeviceWorkerGrid(getOperation());
if (!check) {
  return check.takeError();
}
ttcore::GridAttr deviceGrid =
      ttcore::lookupDevice(getOperation()).getWorkerGrid();

Caching

Use the provided caching mechanisms for computations:

// For getOpConstraints:
return opConstraintsCache().getOrCompute(
    op_model::OpModel<YourOp>::getOpConstraints, *this,
    /* parameters */);
// For getOpRuntime:
return opRuntimeCache().getOrCompute(
    op_model::OpModel<YourOp>::getOpRuntime, *this,
    /* parameters */);

Check Metal Backend Availability

Ensure your operation has a corresponding implementation in the tt-metal backend before implementing these APIs. As mentioned before, the current metal header files are #included in MetalHeaders.h. If you are adding a TTNNOp you might want to add an #include statement in that file to let the c++ compiler know where/how to find the op's definition in metal.

Validate Input Assumptions

Always validate the number of input tensors, eg.:

assert(inputs.size() == 2); // for a binary op
assert(inputs.size() == 3); // for a ternary op

Example: Complete Implementation

Here's a complete example for a hypothetical CustomUnaryOp:

// In TTNNOpModelInterface.cpp
llvm::Expected<op_model::OpConstraints>
CustomUnaryOp::getOpConstraints(const std::vector<TTNNLayoutAttr> &inputs,
                                const OpConfig &opConfig) {
  return detail::getUnaryOpConstraints(*this, inputs, opConfig);
}

llvm::Expected<size_t>
CustomUnaryOp::getOpRuntime(const std::vector<TTNNLayoutAttr> &inputs,
                            const OpConfig &opConfig) {
  return detail::getUnaryOpRuntime(*this, inputs, opConfig);
}

// In TTNNOpModel.h
template <>
struct OpModel<CustomUnaryOp> : UnaryEltwiseOpModel<CustomUnaryOp> {};

// In TTNNOpModel.cpp
template <typename OpTy>
auto getOpSymbol() {
  // ...
  if constexpr (std::is_same_v<OpTy, CustomUnaryOp>) {
    return ::ttnn::custom_unary_op; // metal's definition
  }
  // ...
}

// Explicit template instantiation
template struct UnaryEltwiseOpModel<CustomUnaryOp>;

// Add tests in TestOpModelInterface.cpp and TestOpModelLib.cpp

Decomposing an Op in TTIR

This guide explains how to add and decompose a new operation in the TTIR dialect. We’ll focus on adding an Index operation, which will be decomposed into the Slice operation. The decomposition is implemented as a conversion pass in MLIR since it allows us to mark operations or dialects as legal or illegal, type conversion...

This guide will cover the following steps:

• Decomposing an Op in TTIR
  • 1. Define the Op in the TTIR frontend dialect
  • 2. Create a conversion pattern
    • C++ conversion pattern
    • Tablegen conversion pattern
  • 3. Register the created conversion pattern

1. Define the Op in the TTIR frontend dialect

The more information regarding this step can be found here: Define the Op in the TTIR frontend dialect

I updated the TTIROps.td as following:

def TTIR_IndexOp: TTIR_NamedOp<"index"> {
    let summary = "Tensor indexing operation.";
    let description = [{
      The `index` operation extracts a sub-tensor (slice) from the input tensor along a specified dimension.

      This operation selects elements from the input tensor along a single dimension based on the specified
      begin, end, and step indices. It's similar to Python's slicing notation `tensor[:, begin:end:step, :]`
      where the slicing is applied only to the specified dimension.

      Example:
      ```mlir
      // Extract elements with indices 1, 3, 5 from dimension 0 of a 1D tensor
      %input = ... : tensor<6xf32>  // Input tensor with values: [1, 2, 3, 4, 5, 6]
      %output = ttir.empty() : tensor<3xf32>  // Output tensor shape
      %result = ttir.index(%input, %output) {
          dim = 0 : i32,    // Dimension to index
          begin = 1 : i32,  // Start index
          end = 6 : i32,    // End index (exclusive)
          step = 2 : i32    // Step size
      } : tensor<6xf32>, tensor<3xf32> -> tensor<3xf32>
      // Result: [2, 4, 6]

      // Extract columns 0 and 2 from a 2D tensor
      %input = ... : tensor<3x4xf32>  // Input tensor with values:
                                      // [[1, 2, 3, 4],
                                      //  [5, 6, 7, 8],
                                      //  [9, 10, 11, 12]]
      %output = ttir.empty() : tensor<3x2xf32>  // Output tensor shape
      %result = ttir.index(%input, %output) {
          dim = 1 : i32,    // Index along columns (dimension 1)
          begin = 0 : i32,  // Start from first column
          end = 3 : i32,    // End at third column (exclusive)
          step = 2 : i32    // Take every other column
      } : tensor<3x4xf32>, tensor<3x2xf32> -> tensor<3x2xf32>
      // Result:
      // [[1, 3],
      //  [5, 7],
      //  [9, 11]]
      ```

      Inputs:
      - `input` (Tensor): The input tensor to index.

      Attributes:
      - `dim` (Integer): The dimension along which to index.
      - `begin` (Integer): The starting index.
      - `end` (Integer): The ending index (exclusive).
      - `step` (Integer): The step size between indices.

      Outputs:
      - `result` (Tensor): The indexed tensor.

      Note: The shape of the output tensor is the same as the input tensor except for the indexed dimension,
      which will have size `ceil((end - begin) / step)`. The indices selected will be `begin`, `begin + step`,
      `begin + 2*step`, etc., up to but not including `end`.
    }];

    let arguments = (ins AnyRankedTensor:$input,
                         AnyRankedTensor:$output,
                         I32Attr:$dim,
                         I32Attr:$begin,
                         I32Attr:$end,
                         I32Attr:$step);

    let results = (outs AnyRankedTensor:$result);

    let hasVerifier = 1;
}

The verification function has been added as well:

// IndexOp verification
::mlir::LogicalResult mlir::tt::ttir::IndexOp::verify() {
  ::mlir::RankedTensorType inputType = getInput().getType();
  ::llvm::ArrayRef<int64_t> inputShape = inputType.getShape();
  ::mlir::RankedTensorType outputType = getOutput().getType();
  int32_t dim = getDim();
  int32_t begin = getBegin();
  int32_t end = getEnd();
  int32_t step = getStep();

  // Verify that the input is at least 1D tensor
  if (inputType.getRank() < 1) {
    return emitOpError("Input must be at least a 1D tensor");
  }

  // Validate that the output tensor has the same element type as the input
  // tensor
  if (inputType.getElementType() != outputType.getElementType()) {
    return emitOpError(
        "Output tensor must have the same element type as the input tensor");
  }

  // Verify the output tensor rank
  if (inputType.getRank() != outputType.getRank()) {
    return emitOpError(
        "Output tensor must have the same rank as the input tensor");
  }

  // Verify that the dim attribute is within the bounds of the input tensor
  if (dim < 0 || dim >= inputType.getRank()) {
    return emitOpError() << "Invalid dimension index " << dim
                         << ". Input tensor rank is " << inputType.getRank();
  }

  // Verify begin, end, step and the output tensor dimensions
  int64_t dimSize = inputShape[dim];

  // Adjust negative begin and end
  int32_t adjustedBegin = (begin < 0) ? (begin + dimSize) : begin;
  int32_t adjustedEnd = (end < 0) ? (end + dimSize) : end;

  std::ostringstream inputShapeStream;
  inputShapeStream << "(";
  for (size_t i = 0; i < inputShape.size(); ++i) {
    inputShapeStream << inputShape[i];
    if (i != inputShape.size() - 1) {
      inputShapeStream << ", ";
    }
  }
  inputShapeStream << ")";
  std::string inputShapeStr = inputShapeStream.str();

  if (adjustedBegin < 0 || adjustedBegin >= dimSize) {
    return emitOpError() << "Invalid begin index for dimension "
                         << std::to_string(dim) << ". Expected value in range ["
                         << std::to_string(-dimSize) << ", " << dimSize
                         << "), got " << begin
                         << ". Input shape: " << inputShapeStr;
  }
  if (adjustedEnd < 0 || adjustedEnd > dimSize) {
    return emitOpError() << "Invalid end index for dimension "
                         << std::to_string(dim) << ". Expected value in range ["
                         << std::to_string(-dimSize) << ", " << dimSize
                         << "], got " << end
                         << ". Input shape: " << inputShapeStr;
  }

  auto formatValueMessage = [](int value, int adjustedValue) {
    return value < 0 ? std::to_string(adjustedValue) + " (" +
                           std::to_string(value) + ")"
                     : std::to_string(value);
  };
  std::string beginValueMessage = formatValueMessage(begin, adjustedBegin);
  std::string endValueMessage = formatValueMessage(end, adjustedEnd);

  if (step == 0) {
    return emitOpError("Step value for dimension " + std::to_string(dim) +
                       " cannot be zero");
  }

  if (step > 0 && adjustedBegin > adjustedEnd) {
    return emitOpError() << "For positive step, begin index must be less "
                            "than or equal to end index for dimension "
                         << dim << ". Got begin: " << beginValueMessage
                         << ", end: " << endValueMessage << ", step: " << step
                         << ", input shape: " << inputShapeStr;
  }

  if (step < 0 && adjustedBegin < adjustedEnd) {
    return emitOpError() << "For negative step, begin index must be greater "
                            "than or equal to end index for dimension "
                         << dim << ". Got begin: " << beginValueMessage
                         << ", end: " << endValueMessage << ", step: " << step
                         << ", input shape: " << inputShapeStr;
  }

  // Calculate the expected size of the output dimension
  int32_t expectedDimSize =
      (std::abs(adjustedEnd - adjustedBegin) + std::abs(step) - 1) /
      std::abs(step);
  if (outputType.getDimSize(dim) != expectedDimSize) {
    return emitOpError() << "Mismatch in dimension " << std::to_string(dim)
                         << " of the output tensor: expected size "
                         << expectedDimSize << ", but got "
                         << outputType.getDimSize(dim);
  }

  return success();
}

2. Create a conversion pattern

A conversion pattern defines how MLIR should rewrite the Op. It can be implemented in either C++ or TableGen. Currently, we only have the C++ implementation; TableGen format will be added in the future.

C++ conversion pattern

For the Index operation, we use the C++ conversion pattern because it involves changing the Op’s input types from integers to arrays, which TableGen lacks flexibility for.

// This transformation adjusts IndexOp attributes so that `begin`, `end`, and
// `step` become arrays, where each array element corresponds to a dimension of
// the input tensor. For dimensions other than the sliced dimension, default
// values are used.
//
namespace {
struct IndexToSliceConversionPattern
    : public OpConversionPattern<ttir::IndexOp> {
  using OpConversionPattern<ttir::IndexOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(ttir::IndexOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto inputType =
        ::mlir::dyn_cast<mlir::RankedTensorType>(adaptor.getInput().getType());
    if (!inputType || !inputType.hasRank()) {
      return failure();
    }

    int64_t rank = inputType.getRank();
    llvm::SmallVector<mlir::Attribute, 4> begins, ends, steps;

    for (int64_t i = 0; i < rank; ++i) {
      if (i == op.getDim()) {
        begins.push_back(rewriter.getI32IntegerAttr(adaptor.getBegin()));
        ends.push_back(rewriter.getI32IntegerAttr(adaptor.getEnd()));
        steps.push_back(rewriter.getI32IntegerAttr(adaptor.getStep()));
      } else {
        begins.push_back(rewriter.getI32IntegerAttr(0));
        ends.push_back(rewriter.getI32IntegerAttr(inputType.getDimSize(i)));
        steps.push_back(rewriter.getI32IntegerAttr(1));
      }
    }

    auto newOp = rewriter.create<ttir::SliceStaticOp>(
        op.getLoc(), op.getType(), adaptor.getInput(), adaptor.getOutput(),
        rewriter.getArrayAttr(begins), rewriter.getArrayAttr(ends),
        rewriter.getArrayAttr(steps));

    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};
} // namespace

The matchAndRewrite method from OpConversionPattern is implemented to replace the matched Op with the newly created Op. Since decomposition is implemented as a conversion pass, OpAdaptor is used to access the attributes of the original Op in their converted types. Finally, we instantiate the new Op and call the replaceOp method on ConversionPatternRewriter to replace the original Op.

Tablegen conversion pattern

TODO

3. Register the created conversion pattern

To register the new pattern, go to the populateTTIRToTTIRDecompositionPatterns function in TTIRToTTIRDecomposition.cpp and add it to RewritePatternSet using the add method. After that is done you should mark the decomposed op as illegal in runOnOperation method of TTIRToTTIRDecompositionPass in TTIRToTTIRDecompositionPass.cpp.

You should also add a silicon test like described here: Add a silicon unit test for the Op. This is how the silicon test for the Index operation looks like:

// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" -o %t.mlir %s
// RUN: FileCheck %s --input-file=%t.mlir
// RUN: ttmlir-translate --ttnn-to-flatbuffer -o %t.ttnn %t.mlir
module attributes {} {
  func.func @forward(%arg0: tensor<4x32x32xbf16>) -> tensor<4x32x16xbf16> {
    %0 = ttir.empty() : tensor<4x32x16xbf16>
    // CHECK: = "ttnn.slice_static"
    %1 = "ttir.index"(%arg0, %0) <{dim = 2: i32, begin = 0: i32, end = 32: i32, step = 2: i32}> : (tensor<4x32x32xbf16>, tensor<4x32x16xbf16>) -> tensor<4x32x16xbf16>
    return %1 : tensor<4x32x16xbf16>
  }
}

Docs & Doxygen

Markdown documentation is built using mdbook and API documentation is built using doxygen, and sphinx, and sphinx-markdown-builder.

Markdown documentation (docs)

Requirements

The markdown documentation is built using mdbook and sphinx.

Build command

If not already installed, be sure to install sphinx-markdown-builder.

pip install sphinx-markdown-builder

To build the markdown docs use the docs target in CMake.

cmake -B build
cmake --build build -- docs

API documentation (doxygen)

This is a link to a doxygen autogenerated code reference. Doxygen

Requirements

The API documentation is built using doxygen and sphinx, here are the needed tools for building it:

• doxygen
• dot from graphviz
• sphinx

Build command

To build the API docs use the doxygen target in CMake

cmake -B build
cmake --build build -- doxygen

Serving the docs locally

To start a server for local viewing of the docs, after building, run:

mdbook serve build/docs

mdbook will start a local server at http://localhost:3000 with the built docs.

Specifications

Specifications are documents that define the requirements for features or concepts that are particularly cross-cutting, complex, or require a high degree of coordination and planning. They are intended to be a living document that evolves as the feature is developed and should be maintained as the goto reference documentation for the feature or concept.

Specifications are written in markdown and are stored in the docs/src/specs directory of the repository. Below is a template that should be used when creating a new specification.

Specification Template

# [Title]

A brief description of the feature or concept that this specification is
defining.

## Motivation

A description of why this feature or concept is needed and what problem it is
solving. This section is best written by providing concrete examples and use
cases.

## Proposed Changes

A list of the components that will be impacted by this spec and a detailed
description of the changes that will be made to each respective component.

It should also call out any interactions between components and how they might
share an interface or communicate with each other.

## Test Plan

A brief description of how the feature or concept will be tested.

## Concerns

A list of concerns that have been identified during the design of this feature.

Runtime Stitching

Runtime stitching adds the ability for the runtime to stitch together multiple, indepently compiled programs together at runtime, ie. without compiler knowledge of how the binary programs will be composed.

Motivation

In order to flexibly support arbitrary training schedules / composing multiple models together we want to have the ability for the runtime to stitch graphs together. To achieve this we need to define an ABI kind of interface between the compiler and the runtime.

Simple Example

mod_a = forge.compile(PyTorch_module_a)
mod_b = forge.compile(PyTorch_module_b)

for i in range(10):
    outs_a = mod_a(ins_a)
    outs_b = mod_b(outs_a)

mod_a and mod_b are 2 independent compile steps, during the compile step for mod_a it should be completely unaware that mod_b will take place and vice-versa. In order to achieve this we propose a new runtime concept called stitching:

• forge invokes compile step for mod_a, tt-mlir compiler determines where the inputs (ins_a) should live, host, device dram, device l1. tt-mlir returns metadata to forge describing where it wants the tensors to reside before invoking flatbuffer submission.
• forge invokes compile step for mod_b, same happens as bullet 1
• mod_a is invoked at runtime, forge runtime needs to inspect the compiler metadata to determine where the tensors should live. Runtime manually invokes a new data copy command to get the tenors to the correct memory space / correct memory address.
• forge runtime invokes mod_a program submit
• mod_b is invoked at runtime, this time it might be that the compiler left the tensor outputs in L1, so no data copy is needed to start running mod_b since the inputs are already in the correct location.

A more concrete usecase would be a training loop where there are often multiple graphs composed together. #82 Or when we eventually support torch 2.0, the torch runtime can arbitrarily break the graph anywhere.

Proposed Changes

Compiler Metadata

Compiler will encode the input tensor layout information directly into the flatbuffer tensor desc. The flatbuffer schema already exists to express this, we just need to adopt populating it instead of assuming a canonical host layout.

Compiler will decide where the tensors should live, host, device dram, device l1.

Runtime

• Runtime will inspect the tensor desc metadata to determine where the tensors need to end up / what layout they should be in before invoking the program.
• New runtime API Tensor toLayout(Tensor tensor, ::tt::target::TensorDesc* tensorDesc);
• Runtime will need to invoke toLayout on all input tensors before invoking the program.

Test Plan

• Add a new test to the runtime gtest suite that verifies the runtime can correctly stitch together 2 independently compiled programs.

Concerns

• Tensors pass through device memory spaces (dram, L1) will have a dynamic address, some arbitrary run order of flatbuffer could cause tensors to end up in non-ideal locations in memory. Specifically, L1, a poorly placed tensor might not be able to be moved to a better location without a bounce through DRAM.

Tensor Layout

The tensor layout attribute captures how tensor data is sharded across a grid of devices, cores, and is laid out in memory.

Motivation / High level goals

• Logical shapes: Keep the original tensor shape and rank intact and agnostic to underlying storage layout. Keeping the logical shapes not only makes some graph transformations vastly simpler, in particular convs, but it makes the lowered IR much easier to read and reason about. The original tensor shapes leave breadcrumbs that make it much easier to map back to the input representation.
• Flexible sharding: Enable flexibility in choosing grid shape, to get better parallelization and avoid resharding. This is particularly important in cases where tensor shapes are not clean powers of two and would otherwise force our hand in choosing non-optimal grid shapes.
• Logical-Physical Isomorphism: Encode this information with just a few attributes to enable derived conversions from logical to physical layout and back.
• Explicit: A single source of truth.
• Enable a direct way to query padded regions.

An Example / Walkthrough

Let's consider a snippet of MLIR:

tensor<2x3x64x128xf32>

Here we've defined a 4 dimensional tensor using MLIR's builtin tensor type. This tensor type has an optional attribute called an Encoding, this attribute has been used by the TT dialect to encode the tensor's layout. This looks like:

tensor<2x3x64x128xf32,
  #ttcore.metal_layout<
    (d0, d1, d2, d3) -> (d0 * 192 + d1 * 64 + d2, d3),
    undef,
    <1x1>,
    memref<384x128xf32, #ttcore.memory_space<l1>>
  >
>

At the time of this writing there are 4 properties that make up a tensor layout:

• linear: An affine map that defines how the logical tensor dimensions map to a grid shape. Note that the number of dims in the affine map must match exactly the rank of the original tensor, and the number of results must match exactly the rank of the grid shape.
• oob_val: A tracked out of bounds value that fills padding space.
• grid: The grid shape that this tensor is divided onto.
• memref: A memref that describes the physical footprint allocation of the shard. It must also have a shape with rank equal to grid.

This example isn't particularly complicated because it's only sharded to a 1x1 grid, the rest of the document will go into more details on the following topics:

• Dimension Collapsing
• Multi-Core
• Tilized
• Padding
• Memory Spaces
• Multi-device

Before we jump into more advanced topics there are two resources that could be useful to have at hand:

• test/python/tensor_layout.py: Python test with many convenience functions for creating and experimenting with tensor layouts.
• TTNN Interactive Visualizer: An interactive visualation tool that demonstrates the transformation. Note that this tool was created for TTNN tensor layout, but many of the same concepts transfer over.

Dimension Collapsing

Probably the most important concept in ttcore.metal_layout is dimension collapsing. This is captured by the affine map linear property which provides a mapping from tensor dim space to a reduced physical dimensional space. This single-handedly touches on most of the tensor layout goals mentioned at the beginning of the doc:

• Leaves tensor shapes intact
• Logical-Physical mapping, how the tensor is laid out in memory over a grid
• Enables more flexible sharding
• Explicit padding

To see how these goals are achieved we'll continue working on an explicit example, same one as above:

(d0, d1, d2, d3) -> (d0 * 192 + d1 * 64 + d2, d3)

To recap, we have our example 4d tensor (2, 3, 64, 128), which maps directly to the LHS (d0, d1, d2, d3). We have our 2d grid shape (1, 1), notice the affine-map RHS is also 2d, and this describes how tensor dims map to a lower dimensional physical memory, overlaid on a grid. We'll see how this gets divided onto the grid later, but first let's look at how this forms an affine-map iteration space. If we index our tensor at say [1, 1, 6, 100], we can simply plugin those numbers to get our remapped offset:

(1 * 192 + 1 * 64 + 6, 100) = (262, 100)

This remapped offset (262, 100) corresponds to the row and column index of the collapsed physical memory.

By default, the dim range [0, -1) is collapsed, but the ttcore.metal_layout contructor can actually take a programmable range called collapseIntervals. collapseIntervals is a list of pairs, where each pair is a dim range interval, left inclusive, right exclusive. Let's consider a few examples:

Instead of multiplying out real shapes, we will use <> to represent a dimension join operator.

• 3D tensor onto a 2D grid and default collapseIntervals=[(0, -1)]:

(d0, d1, d2) -> (d0 <> d1, d2)

• 4D tensor onto a 3D grid and collapseIntervals=[(1, -1)]:

(d0, d1, d2, d3) -> (d0, d1 <> d2, d3)

• 4D tensor onto a 3D grid and collapseIntervals=[(0, 2)]:

(d0, d1, d2, d3) -> (d0 <> d1, d2, d3)

• 7D tensor onto a 4D grid and collapseIntervals=[(0, 3), (-3, -1)]:

(d0, d1, d2, d3, d4, d5, d6) -> (d0 <> d1 <> d2, d3, d4 <> d5, d6)

Multi-core

Let's consider the original example again, but on a larger grid than 1x1, say 2x4:

tensor<2x3x64x128xf32,
  #ttcore.metal_layout<
    (d0, d1, d2, d3) -> (d0 * 192 + d1 * 64 + d2, d3),
    undef,
    <2x4>,
    memref<192x32xf32, #ttcore.memory_space<l1>>
  >
>

The number of affine map results, grid shape, and memref shape all must have the same rank. We can see in this example by changing the grid shape we also changed the memref shape, we can always calculate the memref shape by plugging in the full tensor dims into our affine map and then dividing by grid shape.

(d0, d1, d2, d3) -> (d0 * 192 + d1 * 64 + d2, d3),
(2 - 1, 3 - 1, 64 - 1, 128 - 1) = (1 * 192 + 2 * 64 + 63, 127) = (383, 127)

Above we actually subtracted 1 in order to get the index of the last element of the tensor. Now we can simply add back 1 to get the size:

(383 + 1, 127 + 1) = (384, 128)

Finally, we divide the dims by the respective grid dims:

(384 / 2, 128 / 4) = (192, 32)

Here's a few more example mlir snippets:

tensor<8x300xf32,
  #ttcore.metal_layout<(d0, d1) -> (d0, d1),
    undef,
    <1x2>,
    memref<8x150xf32, #ttcore.memory_space<l1>>
  >
>

tensor<8x96x32xf32,
  #ttcore.metal_layout<(d0, d1, d2) -> (d0 * 96 + d1, d2),
    undef,
    <2x1>,
    memref<384x32xf32, #ttcore.memory_space<l1>>
  >
>

tensor<8x96x32xf32,
  #ttcore.metal_layout<(d0, d1, d2) -> (d0 * 96 + d1, d1, d2),
    undef,
    <2x1x2>,
    memref<384x96x16xf32, #ttcore.memory_space<l1>>
  >
>

tensor<5x3x2x2x7x32x32xf32,
  #ttcore.metal_layout<
    (d0, d1, d2, d3, d4, d5, d6)
      -> (d0 * 2688 + d1 * 896 + d2 * 448 + d3 * 224 + d4 * 32 + d5, d4, d5, d6),
    undef,
    <3x2x2x2>,
    memref<4480x4x16x16xf32, #ttcore.memory_space<l1>>
  >
>

A couple of final notes regarding grid shape:

• Grid shapes of rank > 2 are perfectly legal. Not only it this useful for describing multi-device grid topologies, but it is often convenient to have higher ranked grids to better describe how a high rank tensor should be divided. The grid shape here is a virtual grid shape, the ttcore.device attribute will hold an additional affine map that defines how this virtual grid shape maps to a physical one.
• Grid shapes where either columns or rows are > physical device grid is also legal. Since this is only a virtual grid shape we could have some grid 1x64 that maps to a physical 8x8 device grid (this particular example is called width sharding in TTNN).

Tilized

A tilized tensor is one with a memref that has a tile element type.

Given some tensor with scalar layout:

tensor<3x64x128xf32,
  #ttcore.metal_layout<
    (d0, d1, d2) -> (d0 * 64 + d1, d2),
    undef,
    <3x2>,
    memref<64x64xf32, #ttcore.memory_space<l1>>
  >
>

After tilizing we'll have:

tensor<3x64x128xf32,
  #ttcore.metal_layout<
    (d0, d1, d2) -> (d0 * 64 + d1, d2),
    undef,
    <3x2>,
    memref<2x2x!ttcore.tile<32 x 32, bfp_bf8>, #ttcore.memory_space<l1>>
  >
>

Notice the memref dim was ceilDiv'd by tile shape and the element type becomes a ttcore.tile type. Also notice that the tensor shape and element type remains intact.

Padding

Padding can be a bit of an overloaded term, but in this context it refers to an out of bounds area in the physical memory allocation that has no real tensor data in it. The contents of this area is tracked by oob_val and the padding area can be automatically derived from the attributes of ttcore.metal_layout.

Padding is a necessary evil that arises when a tensor is not evenly divisible by a grid shape or tile shape. It can also arise due to minimum Noc addressing requirements.

Example of non-divisible grid:

tensor<53x63xf32,
  #ttcore.metal_layout<
    (d0, d1) -> (d0, d1),
    undef,
    <3x2>,
    memref<18x32xf32, #ttcore.memory_space<l1>>
  >
>

The grid dims always ceilDiv the affine map results, real tensor data will entirely fill initial shards and the last shard in each dimension will be partially filled.

In this particular example, we have 1 scalar row of padding on the last row of cores and 1 scalar column of padding on the last column of cores.

Taking the above example a step further, we could tilize it:

tensor<53x63xf32,
  #ttcore.metal_layout<
    (d0, d1) -> (d0, d1),
    undef,
    <3x2>,
    memref<1x1x!ttcore.tile<32 x 32, bfp_bf8>, #ttcore.memory_space<l1>>
  >
>

Tile dims also always ceilDiv the resulting memref shape. Notice now that the padding is slightly more complicated. Our scalar shard shape was 18x32, but this was further padded to 32x32 meaning that every core now has 14 rows of padding except for the last row of cores which has 15 rows of padding.

Also note that there is an order of operations here, grid divides the scalar shape first and then we tilize. This is important because it can enable use cases that frequently arise in conv networks that would otherwise result in reshards in between every layer.

With affine map we can be even more flexible in how we pad, we can bump our stride between dimensions. Consider tensor (w/ batch dim 2):

tensor<2x8x32xf32,
  #ttcore.metal_layout<
    (d0, d1, d2) -> (d0 * 8 + d1, d2),
    undef,
    <1x2>,
    memref<16x16xf32, #ttcore.memory_space<l1>>
  >
>

If we tilized the above tensor we'd end up with a memref shape of 1x1x!ttcore.tile<32x32>, that is, all batches are tightly packed within a single tile. Let's say that for some reason, we do not want the batches (2) to be tightly packed within a tile, perhaps the mathematical operation we're doing requires the batch to be independently evaluated and thus the (S)FPU needs them in separate tiles. We can adjust this by adjusting the stride of the affine map:

(d0, d1, d2) -> (d0 * 32 + d1, d2),

Instead of striding by the number of logical rows, 8, we bump the stride up to 32 effectively pushing a gap between the collapsed rows and enabling each batch to fall on a tile boundary.

Memory Spaces

At the time of writing this document there are 4 memory spaces:

1. System: Host memory space that is not device visible.
2. SystemMMIO: Host memory space that is device visible.
3. DeviceDRAM: DRAM local to the device.
4. DeviceL1: SRAM on each core.

Something worth noting here is that a tensor must belong exclusively to only one of these memory spaces at a time. For example, in order to stream tensor data from DeviceDRAM to DeviceL1 you would need to either manually slice the tensor into smaller tensors that do fit in L1 or have native support in the op's kernel for double buffering a block (most TTNN ops already support this).

Multi-device

Multi-device can be naturally represented via a combination of two concepts already touched on above, higher ranked grids and collapseIntervals. Let's consider the following example with a 3d grid and collapseIntervals=[(1, -1)].

tensor<2x3x64x128xf32,
  #ttcore.metal_layout<(d0, d1, d2, d3) -> (d0, d1 * 64 + d2, d3),
    undef,
    <2x2x4>,
    memref<1x3x1x!ttcore.tile<32 x 32, bfp_bf8>, #ttcore.memory_space<l1>>
  >
>

Here we've left the batch dim intact and started collapsing at d1. This enables us to define a 3d grid where the outermost grid dim divides the batch directly. This could map to a 2 device system where the batch dim is evenly divided between 2 devices. Within each device this op runs on a 2x4 grid.

The high level takeaway here is that how a tensor is logically divided up is decoupled from its mapping to physical compute resources. This has a nice property that data parallel extends to any tensor dimension and is captured under the same grid primitive that also divides tensor rows and columns.

Test Plan

• test/python/tensor_layout.py: Assertions for LayoutAttr to make sure it's spec compliant.
• Sweep tests:
  • Grid dim sweeps
  • Tilize / untilize sweeps
  • Padding sweeps
• Multi-device tests

Concerns

• ttcore.metal_layout is deliberately flexible and tries to capture as many problematic use-cases we've ran into in the past in a single, succinct representation. This flexibility will need to be further constrained by backends to avoid unsupported programming of this attribute.
• Optimization solution space is potentially large with all of this flexibility. Two things that I hope can help protect us here:
  • By and large the heuristic we'll be following is just max the grid at all costs. This should really narrow down the solution space to only a handful of options and we only keep exploring if producers/consumers end up with nasty reblocking.
  • We can constrain the optimizer heuristics as aggressively as possible in the beginning and just advertise the full flexible options to the UI model explorer. Hopefully this enables us to experiment with crazier grid layouts and prove it's worthwhile before writing an algorithm.

TTNN Tensor Layout

The above section of this document covers how the compiler models tensor layout. There are some slight differences in TTNN, but the high level idea of collapsing dims is still used.

Terms

• shape: Always logical shape, n-dimensional
• stride: Same as pytorch stride, but this is crucial for describing how n-dimensional data gets packed into a 2D physical layout. This 2D physical layout is always the inner dim (-1) wide and dims [0, N-1] are collapsed into rows derived from stride
• shard_shape: Also a logical shape, describes a 2d region that chunks physical_shape . Note this does not need to be a tile multiple
• physical_shard_shape: The shard_shape padded out to tile_shape
• tile_shape: A programmable tile shape, though constraints must check that it's compatible with an op's usage, i.e. FPU/Noc compatible
• grid_shape: [divup(stride[0] // stride[-2], shard_shape[0]), divup(stride[-2], shard_shape[0])]

Mapping from the compiler

The compiler uses an affine map to explicitly track which dimensions are folded together, but TTNN does not have affine maps so the representation is a bit more implicit. TTNN captures the dimension collapsing in the stride attribute where dimensions [0, N-1] are always collapsed. This is less flexible so the compiler will have to enforce only collapsing supported dimensions when targeting TTNN, or handle lowering in a different way. For example, in the compiler we might want to represent data parallel over the tensor batch dim by leaving d0 and collapsing d1 - d[-1]. TTNN doesn't support this in its tensor layout representation, but this could be lowered to a TTNN mesh tensor where the mesh could be sliced on the batch and each per-device tensor has d0 fully collapsed.

TTNN Example

Device

Device in tt-mlir is somewhat of an overloaded term and can refer to different things depending on the context. This document will only speak to the compiler's abstract representation of a device captured by attribute #ttcore.device.

Terms

There are many overloaded terms when talking about devices and grids, this document will use the following definitions:

• Physical Grid: A 2D array of tensix cores on a chip.
• Chip: A single physical chip with a Physical Grid of cores.
• Card: A PCIE or Ethernet card that may contain multiple Chips.
• System: A collection of Cards that are usually connected together on the same host via PCIE or networked via ethernet. A system is represented by SystemDesc in the compiler.
• Device: Device is always presented as a single entity to the enclosing scope, but it may be virtualized to abstract a multi-card System and part of its encoding carries a Logical Grid. Another way to think of device is a view over the system.
• Logical Grid or just Grid: Is a logical shape that abstracts one or more Physical Grids.
• Mesh Shape: Describes the virtual layout of the chips with respect to each other. In practice the mesh shape is used to derive the logical grid.

Motivation

The device attribute strives to achieve the following goals:

• Provide a convenient representation of a physical grid that decouples the logical division of tensors from the physical layout of the hardware. This not only simplifies reasoning about how tensors get divided into shards, but can also enable reinterpretations of the device grid for data layout optimization decoupled from the existing encoding of the tensor layouts.
• Following the first point, the device attribute should be able to represent many different forms of logical grids, from simple 2D grids, to more complex topologies like extra-wide grids or higher dimensional grids.
• Device attribute captures encoding both single chip and multi-chip systems under a single, virtualized representation.
• Enable many forms of data parallel execution strategies for single and multi chip systems under a single representation.

Scope

This document will cover how the device attribute is encoded and how it can be lowered to backend dialects. The document will not cover the algorithm for choosing the best, or even legal, device configurations for a given physical system.

Examples

All of the following examples will assume the physical hardware has an 8x8 physical grid of cores. We will use notation [N, 8x8] to represent a N chip system, each with an 8x8 physical grid.

#ttcore.device in is simplest, single chip form [1, 8x8], just maps directly 1-1 to the underlying physical hardware device.

#ttcore.device<
  workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d0, d1)>,
  meshShape = 1,
  chipIds = [0]
>

Let's break down what each of these attributes mean:

• workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d0, d1)>: This is a 2D logical grid with dim 8x8. It's followed by an affine map (d0, d1) -> (0, d0, d1) that provides a mapping from the logical grid to the physical grid. In this case, the logical grid is the same as the physical grid, so the mapping is the identity function. The logical grid can have any rank, but the physical mapping is always 3D, with the first being the chip index, followed by the 2D physical core index within the chip.
• meshShape = 1: A shape provided as part of the DeviceAttr constructor that describes the virtual layout of the chips with respect to each other. Note that in a multi-chip system, this grid encapsulates the entire system's grid shape, e.g. 8x16 grid could be made up of a 1x2 mesh of chips side-by-side. The mesh attribute configures how the above grid/map attributes are created such that they implement this mesh topology.
• chipIds = [0]: This is a list of chip indices. These chip indices directly reference the same chip indices in the system descriptor. The SystemDesc attribute that this is in reference to is tagged on the top level ModuleOp.

Specific examples that this document will cover:

• Data Parallel Over Batch
• Data Parallel Over 2d
• Data Parallel Over 2d and Batch
• Pipeline Parallel
• Reinterpreted Grids (Transpose)
• Reinterpreted Grids (Training Usecase)
• Reinterpreted Grids (Extra)

Before we move on to more complex examples, it's worth having on hand:

• The python test test/python/device_attr.py which shows how all of these examples can actually be programmed for the device attribute.
• The Tensor Layout spec as the following examples will demonstrate how tensor layout interacts with the logical device grid.

Note on Data Parallel: There is existing literature that explicitly distinguishes between data parallel and tensor parallel, oftentimes describing data parallel as duplicating the model across multiple devices and trivially dividing up the batch whereas tensor parallel refers to tensor data being distributed and potentially communicated between devices during execution. While this is true for multi-GPU/CPU systems, it is somewhat of an implementation detail and given the flexibility of tenstorrent hardware there is an opportunity to generalize this concept. In this document we will use the term data parallel to refer to any form of parallelism that divides any dimension of the tensor across multiple cores/chips.

Note on Constraints: Many of the examples below require careful virtualization of the underlying physical system, i.e. some device configurations might only work if the chips are connected via ethernet and with a particular topology, but these constraints are outside the scope of the examples and will be discussed further in the Backend Lowering and Constraints section.

Data Parallel Over Batch

Given a 2 chip system, [2, 8x8], we can represent a simple data parallel logical grid that divides the batch dimension in half across the two chips. This is denoted by meshShape = 2x1x1 which means the logical grid is 3D.

#ttcore.device<
  workerGrid = #ttcore.grid<2x8x8, (d0, d1, d2) -> (d0, d1, d2)>,
  meshShape = 2x1x1,
  chipIds = [0, 1]
>

The affine map here is just identity, so dims d1 and d2 directly index the physical grid and d0 indexes the chip.

Now we can consider some tensor that, importantly, has a grid of the same rank as the logical device grid:

tensor<16x3x64x128xf32,
  #ttcore.metal_layout<(d0, d1, d2, d3) -> (d0, d1 * 64 + d2, d3),
    undef,
    <2x2x4>,
    memref<8x3x1x!ttcore.tile<32 x 32, bfp_bf8>, #ttcore.memory_space<l1>>
  >
>

If we map this tensor onto the above device, it will span across both chips, half of the batch dimension on each chip. Within each chip the tensor occupies a 2x4 grid out of the 8x8 physical grid available.

Data Parallel Over 2d

In this example we will consider a 2 chip system, [2, 8x8], and view it as though the two chips are concatenated together side by side to form a single 8x16 grid. This is denoted by meshShape = 1x2 which means to concatenate the chips in the second dimension.

#ttcore.device<
  workerGrid = #ttcore.grid<8x16, (d0, d1) -> ((d0 floordiv 8) * 2 + d1 floordiv 8, d0, d1 mod 8)>,
  meshShape = 1x2,
  chipIds = [0, 1]
>

Here we can see that the affine map encodes an indexing pattern such that when we extend past 8 cores in the second dimension, we wrap around to the next chip.

Now we can consider some tensor that, importantly, has a grid of the same rank as the logical device grid:

tensor<256x1024xf32,
  #ttcore.metal_layout<(d0, d1) -> (d0, d1),
    undef,
    <4x16>,
    memref<2x2x!ttcore.tile<32 x 32, bfp_bf8>, #ttcore.memory_space<l1>>
  >
>

This single tensor maps trivially onto the logical grid, spanning the upper half. Decoupled from the tensor's layout, under the hood the tensor is actually physically spanning across two chips.

Data Parallel Over 2d and Batch

The previous 2 examples can be composed together to form a logical grid that divides tensor across multiple dimensions. Here we will consider a 4 chip system [4, 8x8] and view it as a 2x8x16 grid. Note that the meshShape is 2x1x2 which means to concatenate the chips in the first and third dimensions.

#ttcore.device<
  workerGrid = #ttcore.grid<2x8x16, (d0, d1, d2) -> (d0 * 2 + (d1 floordiv 8) * 2 + d2 floordiv 8, d1, d2 mod 8)>,
  meshShape = 2x1x2,
  chipIds = [0, 1, 2, 3]
>

We can evaluate the affine map to see that the chips are interpreted in chunks of two, where groups [0, 1] and [2, 3] each form 8x16 grids and these 2 groups concatenate to form a 2x8x16 grid.

We can consider the following tensor to map onto this grid:

tensor<64x256x1024xf32,
  #ttcore.metal_layout<(d0, d1) -> (d0, d1),
    undef,
    <2x4x16>,
    memref<32x2x2x!ttcore.tile<32 x 32, bfp_bf8>, #ttcore.memory_space<l1>>
  >
>

Pipeline Parallel

Pipeline parallel in the scope of this spec isn't particularly interesting, it is intended to be used in conjunction with the ttir.pipeline operation which will group sections of the module's operations into groups to form pipeline regions and will be covered in a separate spec.

What we can demonstrate here is how we can take multiple non-overlapping views of the system descriptor to form distinct virtual devices.

Given an 8 chip system [8, 8x8], we can form two virtual devices that each take 4 chips and interpret them differently (though they could take the same logical grid).

#ttcore.device<
  workerGrid = #ttcore.grid<2x8x16, (d0, d1, d2) -> (d0 * 2 + (d1 floordiv 8) * 2 + d2 floordiv 8, d1, d2 mod 8)>,
  meshShape = 2x1x2,
  chipIds = [0, 1, 2, 3]
>
#ttcore.device<
  workerGrid = #ttcore.grid<16x16, (d0, d1) -> ((d0 floordiv 8) * 2 + d1 floordiv 8, d0 mod 8, d1 mod 8)>,
  meshShape = 2x2,
  chipIds = [4, 5, 6, 7]
>

Reinterpreted Grids (Transpose)

One particularly interesting usecase that logical grids could enable is to reinterpret the grid as a form of data layout optimization. For example, if we wanted to transpose a tensor, instead of having to move the data around to implement transpose, we could instead reinterpret the grid as being transposed, leveraging the fact that the relevant data is already located on the correct cores/chips.

To keep things simple, let's consider a 1 chip system [1, 8x8], but it's not too big a leap to see how this could map to multi-chip where the cost of moving data is even higher.

Let's also consider a simple (totally contrived) eltwise unary graph:

a = exp(a)
aT = transpose(a)
relu(aT)

1. We'll establish a regular, single chip, identity logical grid:

#ttcore.device<
  workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d0, d1)>,
  meshShape = 1,
  chipIds = [0]
>

2. Execute exp.
3. We'll reinterpret the grid as transposed:

#ttcore.device<
  workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d1, d0)>,
  meshShape = 1,
  chipIds = [0]
>

4. Execute transpose. Note that each core only needs to transpose their data locally. Eventually this could be implemented as a no-op by reindexing the tile visitation order of the successive operation.
5. Execute relu.

It's important to note that we effectively implemented transpose without moving data anywhere.

Reinterpreted Grids (Extra)

For the sake of examples, here's a few more ways of reinterpreting the logical grid.

Extra Wide Grid

#ttcore.device<
  workerGrid = #ttcore.grid<1x64, (d0, d1) -> (0, d0 * 8 + d1 floordiv 8, d1 mod 8)>,
  meshShape = 1,
  chipIds = [0]
>

Extra Tall + Transposed Grid

#ttcore.device<
  workerGrid = #ttcore.grid<64x1, (d0, d1) -> (0, d1 * 8 + d0 floordiv 8, d0 mod 8)>,
  meshShape = 1,
  chipIds = [0]
>

Staircase

#ttcore.device<
  workerGrid = #ttcore.grid<8x8, (d0, d1) -> (0, d0, (d0 + d1) mod 8)>,
  meshShape = 1,
  chipIds = [0]
>

This could be an interesting starting position for data in implementing matmul as a systolic array in a ring topology.

Lowering to TTNN

While the above device attribute encoding is quite flexible, this does not necessarily mean the target backend can actually support all of these interpretations. TTNN backend will be constrained to support only the specialized grid topologies that are supported by the API.

Grid/Shard Orientation

TODO

Multi-device

Please refer to TTNN Mesh Programming Docs for more information on how to program multi-device systems with TTNN API.

Multi-device TTNN dialect will try and stay as close to the TTNN API as possible. Let's consider what this looks like from the compiler and runtime perspectives:

Compiler

• Device Creation: The TTNN device in the compiler is exactly the same attribute from the ttir dialect. It will encode the meshShape into the flatbuffer which can be directly used to program ::ttnn::MeshShape.
• Tensor Layout: Again, the tensor layout is inherited in TTNN dialect from the ttir dialect. The grid attribute in the tensor layout can be trivially divided by meshShape to determine the shape of the tensor slice on each device. Broadcasting rules can be applied to determine which Distribution Strategy to use:
  • Mesh Sharded: If the tensor grid is > 1 along the meshShape dimensions, the tensor will be sharded across the mesh devices.
  • Replication: If the tensor needs to be broadcasted for this op, by extension the tensor layout will be replicated across the mesh devices.

Runtime

• Device Creation: The ttnn runtime will wholesale switch to working with mesh devices via api ttnn::multi_device::open_mesh_device, this is possible because a 1x1 mesh device is a valid single device. The mesh shape during device open will always be 1xN where N is the number of deviceIds in the array. Note that this shape can be reinterpreted by flatbuffer programs on the fly with SubMesh API.
• Tensor Creation: Tensor creation in a multi-device system is a bit more involved. In order to upload a multi-device tensor to the mesh, the host tensor much first be created with MultiDeviceHostStorage. The ttnn runtime can automatically do this during handleToHostMemoryConfigOp:
  • Regular host tensor will bounce through new tensor with MultiDeviceHostStorage type.
  • tensor.to(mesh_device) will allocate/move the tensor to the mesh device.

Lowering to TTMetal

In TTMetal dialect we are only constrained by what we've implemented in the tt-mlir compiler, this means it is much more flexible and can theoretically support any of the grid interpretations above.

Test Plan

• test/python/device_attr.py covers all of the examples above and asserts the IR is correctly generated.
• Additional functional unit tests will be added as op and runtime support is added.

Concerns

• ttcore.device is very flexible, but with this flexibility comes the potential for misuse. It's important that the compiler is able to validate the legal configurations of this attribute for the target backend.

'd2m' Dialect

Direct-to-metal subset of D2M

The D2M dialect contains the subset of D2M used by the direct-to-metal lowering path (TTMetal). It hosts generic dispatch ops and related region ops required post D2MToD2MGeneric.

[TOC]

d2m.empty (tt::d2m::EmptyOp)

Empty tensor allocation operation (D2M).

Syntax:

operation ::= `d2m.empty` `(` `)` attr-dict `:` type($result)

Create an uninitialized tensor with the specified shape, element type and encoding.

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Allocate on ::mlir::SideEffects::DefaultResource}

Results:

d2m.full (tt::d2m::FullOp)

Creates a tensor filled with the specified value (D2M).

Syntax:

operation ::= `d2m.full` attr-dict `:` type($result)

Tensor operation to create a tensor filled with a specified value. Given a shape and a fill_value, produces a tensor with the shape, filled with the specified value.

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Allocate on ::mlir::SideEffects::DefaultResource}

Attributes:

Results:

d2m.generic (tt::d2m::GenericOp)

Generically dispatch work to a grid of cores (D2M).

Syntax:

operation ::= `d2m.generic` attr-dict `\n`
              ` ` ` ` ` ` ` ` `ins` `(` $inputs `:` type($inputs) `)` `\n`
              ` ` ` ` ` ` ` ` `outs` `(` $outputs  `:` type($outputs) `)` ` `  $regions (`:`  type($results)^ )?

Same semantics as D2M generic; carries regions for compute/datamovement to be consumed by the metal path.

Traits: AttrSizedOperandSegments, NoTerminator

Interfaces: BufferizableOpInterface, DestinationStyleOpInterface, MemoryEffectOpInterface, OpAsmOpInterface

Attributes:

Operands:

Results:

d2m.mesh_shard (tt::d2m::MeshShardOp)

Mesh shard operation (D2M).

Syntax:

operation ::= `d2m.mesh_shard` $input attr-dict `:` type($input) `->` type($result)

MeshShard op shards the inputs (FullToShard) or concatenates the outputs (ShardToFull) for ccl ops.

Traits: AlwaysSpeculatableImplTrait

Interfaces: BufferizableOpInterface, ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

d2m.stream_layout (tt::d2m::StreamLayoutOp)

Stream layout (D2M)

Represent a streaming relationship between a source tensor/memref and a storage buffer, producing a view result.

Traits: AlwaysSpeculatableImplTrait

Interfaces: BufferizableOpInterface, ConditionallySpeculatable, D2M_ViewOpInterface, NoMemoryEffect (MemoryEffectOpInterface), OpAsmOpInterface

Effects: MemoryEffects::Effect{}

Operands:

Results:

d2m.to_layout (tt::d2m::ToLayoutOp)

Layout op.

Syntax:

operation ::= `d2m.to_layout` $input `,` $output `:` type($input) `into` type($output) (`hostInfo` `=` $layout^)? attr-dict (`->` type($results)^)?

ToLayout operation, transition tensors from one layout to another. Some examples include:

• Transitioning between different memory spaces, e.g. DRAM to L1.
• Transitioning between different data types, e.g. f32 to f16.
• Transitioning between different tile sizes, e.g. 1x16 to 32x32
• Transitioning between different tensor sharding
• Some combination of the above

#layout = #ttcore.metal_layout<8192x128x1, undef, <1x1>, memref<64x128xf32, #system>>
#layout1 = #ttcore.metal_layout<8192x128x1, undef, <1x1>, memref<64x128xf32, #l1_>>
%1 = "d2m.to_layout"(%arg0, %0) : (tensor<64x128xf32, #layout>, tensor<64x128xf32, #layout1>) -> tensor<64x128xf32, #layout1>

Interfaces: BufferizableOpInterface, MemoryEffectOpInterface

Attributes:

Some high level goals:
  - **Logical shapes**: Store the original tensor shape and rank intact and agnostic
    to underlying storage layout.
    Keeping the logical shapes not only makes some graph transformations vastly
    simpler, in particular convs, but it makes the lowered IR much easier to read
    and reason about.  The original tensor shapes leave breadcrumbs that make it
    much easier to map back to the input representation.
  - **Collapsed dims**: We may collapse dimensions during transformation, but it
    is important we capture this information such that it is not lost during tensor
    transformation.  The collapsed_intervals field stores the collapses performed
    during conversion from logical_shape to physical tensor shape.
  - **Padding**: store the desired alignments s.t. padding can be simply encoded;
  dim_alignments field represents alignment along each logical dim during collapse.
  - **Memref translation**: ensure we have all necessary info s.t. we can trivally
    lower a tensor into a memref without any intermediate passes.

For a logical tensor of shape [H, W] distributed across a grid [GY, GX], the tensor shape would be:
- Without tiling: [GY, GX, H/GY, W/GX]
- With tiling: [GY, GX, H/GY/TH, W/GX/TW, TH, TW] where TH,TW are tile dimensions

This makes the representation 1:1 with memrefs and eliminates the need for shape conversion passes.

Examples:
```mlir
// Logical 8x300 tensor distributed across 1x2 grid:
// tensor<1x2x8x150xf32, #tt.metal_layout<logical_shape=8x300, ...>>

// Logical 1024x1024 tensor distributed across 2x2 grid with 32x32 tiles:
// tensor<2x2x16x16x!ttcore.tile<32x32xf32>, #tt.metal_layout<logical_shape=1024x1024, ...>>
```

Operands:

Results:

d2m.view_layout (tt::d2m::ViewLayoutOp)

View Layout op (D2M subset)

Syntax:

operation ::= `d2m.view_layout` $input attr-dict `:` type($input) `->` type($result)

Create a representational view of a tensor/memref with a different layout. This is a no-op for codegen; consumers are expected to compose layouts.

Traits: AlwaysSpeculatableImplTrait

Interfaces: BufferizableOpInterface, ConditionallySpeculatable, D2M_ViewOpInterface, NoMemoryEffect (MemoryEffectOpInterface), OpAsmOpInterface

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

'emitpy' Dialect

Dialect to generate Python from MLIR.

[TOC]

emitpy.assign (tt::emitpy::AssignOp)

Assign operation

Syntax:

operation ::= `emitpy.assign` $value attr-dict `:` functional-type(operands, results)

The emitpy.assign operation represents a Python variable assignment. This models new_var = old_var or var = constant.

Example:

%2 = emitpy.assign %1 : <!emitpy.opaque<"ttnn.Tensor">>

// Code emitted for the operation above.
v2 = v1;

Operands:

Results:

emitpy.call_opaque (tt::emitpy::CallOpaqueOp)

Opaque call operation

Syntax:

operation ::= `emitpy.call_opaque` $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)

The emitpy.call_opaque operation represents a Python function call. The callee can be an arbitrary non-empty string.

Example:

%2 = emitpy.call_opaque "ttnn.add"(%0, %1) {args = [0 : index, 1 : index, #emitpy.opaque<"ttnn.DataType.BFLOAT16">, #emitpy.opaque<"ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM, None)">], keyword_args = ["", "", "dtype", "memory_config"]} : (!emitpy.opaque<"ttnn.Tensor">, !emitpy.opaque<"ttnn.Tensor">) -> !emitpy.opaque<"ttnn.Tensor">

Interfaces: PyExpressionInterface

Attributes:

Operands:

Results:

emitpy.constant (tt::emitpy::ConstantOp)

Constant operation

The emitpy.constant operation produces an SSA value equal to some constant specified by an attribute. This can be used to form simple integer and floating point constants, as well as more exotic things like tensor constants.

Traits: ConstantLike

Attributes:

Results:

emitpy.import (tt::emitpy::ImportOp)

Import operation

The emitpy.import operation allows to define a Python module import via various forms of the import statement.

Example:

 emitpy.import import "ttnn"

Attributes:

emitpy.literal (tt::emitpy::LiteralOp)

Literal operation

Syntax:

operation ::= `emitpy.literal` $value attr-dict `:` type($result)

The emitpy.literal operation produces an SSA value equal to some constant specified by an attribute.

Example:

%0 = emitpy.literal "0" : index

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), PyExpressionInterface

Effects: MemoryEffects::Effect{}

Attributes:

Results:

emitpy.subscript (tt::emitpy::SubscriptOp)

Subscript operation

Syntax:

operation ::= `emitpy.subscript` $value `[` $index `]` attr-dict `:` functional-type(operands, results)

With the emitpy.subscript operation the subscript operator [] can be applied to variables or arguments of opaque type.

Example:

%0 = emitpy.literal "0" : index
%1 = emitpy.subscript %arg0[%0] : (!emitpy.opaque<"[ttnn.Tensor]">, index) -> !emitpy.opaque<"ttnn.Tensor">

Interfaces: PyExpressionInterface

Operands:

Results:

emitpy.verbatim (tt::emitpy::VerbatimOp)

Verbatim operation

Syntax:

operation ::= `emitpy.verbatim` $value (`args` $fmtArgs^ `:` type($fmtArgs))? attr-dict

The emitpy.verbatim operation produces no results and the value is emitted as is followed by a line break ('\n' character) during translation.

This operation can be used in situations where a more suitable operation is not yet implemented in the dialect.

Note: Use with caution. This operation can have arbitrary effects on the semantics of the emitted code. Use semantically more meaningful operations whenever possible. Additionally this op is NOT intended to be used to inject large snippets of code.

Attributes:

Operands:

'ttcore' Dialect

TT core types and attributes common to all TT dialects.

This dialect defines types and attributes common to all TT dialects.

[TOC]

ArchAttr

TT Arch

Syntax:

#ttcore.arch<
  ::mlir::tt::ttcore::Arch   # value
>

Parameters:

ArgumentAllocationAttr

Argument allocation attribute in TT dialect

Syntax:

#ttcore.arg_alloc<
  uint64_t,   # address
  uint64_t,   # size
  MemorySpace   # memorySpace
>

Holds the metadata for the allocation of an function argument i.e. for graph inputs.

Parameters:

ArgumentTypeAttr

Argument Type

Syntax:

#ttcore.argument_type<
  ::mlir::tt::ttcore::ArgumentType   # value
>

Parameters:

CPUDescAttr

TT cpu_desc attribute

Syntax:

#ttcore.cpu_desc<
  CPURole,   # role
  StringAttr   # target_triple
>

TT cpu_desc attribute

Parameters:

CPURoleAttr

TT CPU Role

Syntax:

#ttcore.cpu_role<
  ::mlir::tt::ttcore::CPURole   # value
>

Parameters:

ChipChannelAttr

TT chip_channel attribute

Syntax:

#ttcore.chip_channel<
  unsigned,   # deviceId0
  ::llvm::ArrayRef<int64_t>,   # ethernetCoreCoord0
  unsigned,   # deviceId1
  ::llvm::ArrayRef<int64_t>   # ethernetCoreCoord1
>

TT chip_channel attribute

Parameters:

ChipCoordAttr

TT chip_coord attribute

Syntax:

#ttcore.chip_coord<
  unsigned,   # rack
  unsigned,   # shelf
  unsigned,   # y
  unsigned   # x
>

TT chip_coord attribute

Parameters:

ChipDescAttr

TT chip_desc attribute

Syntax:

#ttcore.chip_desc<
  ArchAttr,   # arch
  ::llvm::ArrayRef<int64_t>,   # grid
  ::llvm::ArrayRef<int64_t>,   # coordTranslationOffsets
  unsigned,   # l1Size
  unsigned,   # numDramChannels
  unsigned,   # dramChannelSize
  unsigned,   # nocL1AddressAlignBytes
  unsigned,   # pcieAddressAlignBytes
  unsigned,   # nocDRAMAddressAlignBytes
  unsigned,   # l1UnreservedBase
  unsigned,   # eriscL1UnreservedBase
  unsigned,   # dramUnreservedBase
  unsigned,   # dramUnreservedEnd
  ::llvm::ArrayRef<DataTypeAttr>,   # supportedDataTypes
  ::llvm::ArrayRef<TileSizeAttr>,   # supportedTileSizes
  unsigned,   # dstPhysicalSizeTiles
  unsigned,   # numCBs
  unsigned,   # numComputeThreads
  unsigned   # numDatamovementThreads
>

TT chip_desc attribute

Parameters:

CoreCoordAttr

TT core_coord attribute

Syntax:

#ttcore.core_coord<
  int64_t,   # y
  int64_t   # x
>

TT core_coord attribute containing a single physical core coordinate.

Parameters:

DataTypeAttr

TT DataTypes

Syntax:

#ttcore.supportedDataTypes<
  ::mlir::tt::ttcore::DataType   # value
>

Parameters:

DeviceAttr

Device attribute in TT dialect.

Syntax:

#ttcore.device<
  ::mlir::tt::ttcore::GridAttr,   # workerGrid
  AffineMap,   # l1Map
  AffineMap,   # dramMap
  ::llvm::ArrayRef<int64_t>,   # meshShape
  ::llvm::ArrayRef<unsigned>   # chipIds
>

Describes the physical layout of a device in the system and is made up of a few components:

• A grid attribute that describes the device's compute grid shape. It not only describes the shape of the compute grid, but also carries an affine map that describes how the logical grid maps to the physical grid.
• Two affine maps that describe how a tensor layout's linear attribute maps to the L1 and DRAM memory spaces.
• A mesh shape that describes the virtual layout of the chips with respect to each other. Note that in a multi-chip system, this grid encapsulates the entire system's grid shape, e.g. 8x16 grid could be made up of a 1x2 mesh of chips side-by-side. The mesh attribute configures how the above grid/map attributes are created such that they implement this mesh topology.
• An array of chip ids that this device is made up of. This array's length must match the volume of the mesh shape and should be interpreted in row-major order.

Parameters:

GridAttr

TT grid attribute

Syntax:

#ttcore.grid<
  ::llvm::ArrayRef<int64_t>,   # shape
  AffineMap   # mapping
>

TT grid attribute

Parameters:

HostLayoutAttr

Host-side memref layout attribute with padding support.

Syntax:

#ttcore.host_layout<
  ::llvm::ArrayRef<int64_t>,   # logical_shape
  ::llvm::ArrayRef<int64_t>,   # host_strides
  int64_t,   # host_volume
  TensorMeshAttr   # mesh
>

Describes a host-side memref layout with a logical shape where each dimension may be padded up to a requested alignment. This attribute encodes row-major strides and volume of the aligned shape, both computed in the bufferization pass to match the corresponding device memref's shape.

This allows host memrefs to reflect the padded footprint required by the device while preserving the original logical shape, enabling I/O for unaligned densly-packed shapes (e.g., non-multiples of 32×32 tiles). The correct element placement in data copys between padded and unpadded host memrefs is ensured by the strided-memcpy of the runtime's memref.copy implementation.

• logical_shape: The unpadded logical shape in elements.
• host_strides: Per-dimension alignment in elements used to compute aligned shapes for stride calculation.
• host_volume: The (potentially padded) volume in elements.
• optional(mesh): The mesh that a memref lives on. No mesh indicates single device.

Parameters:

InterleavedLayoutAttr

Interleaved layout attribute in TT dialect

Syntax:

#ttcore.interleaved<
  ::llvm::ArrayRef<int64_t>   # stride
>

Describes overall layout of an interleaved memref buffer.

• Stride: Stride of each dim in bytes.

Parameters:

IteratorTypeAttr

TT IteratorType

Syntax:

#ttcore.iterator_type<
  ::mlir::tt::ttcore::IteratorType   # value
>

Parameters:

MemorySpaceAttr

TT MemorySpace

Syntax:

#ttcore.memory_space<
  ::mlir::tt::ttcore::MemorySpace   # value
>

Parameters:

MeshAttr

Mesh reference attribute in TT dialect.

Syntax:

#ttcore.mesh<
  StringAttr,   # name
  ::llvm::ArrayRef<int64_t>   # shape
>

Describes a mesh config including name and shape.

Parameters:

MeshShardDirectionAttr

TT MeshShardDirection

Syntax:

#ttcore.shard_direction<
  ::mlir::tt::ttcore::MeshShardDirection   # value
>

Parameters:

MeshShardTypeAttr

MeshShard shard_type attribute in TT dialect

Syntax:

#ttcore.shard_type<
  ::mlir::tt::ttcore::MeshShardType   # value
>

Define sharded tensor data of mesh_shard op.

• Identity: input and output tensors are pre-sharded (same data) and no sharding is required.
• Replicate: all of the devices has full tensor (same data).
• Maximal: one or part of the devcices has full tensor (same data).
• Devices: all or part of the devices has sharded (partial) tensor (different data).

Parameters:

MeshesAttr

TT system meshes attribute.

Syntax:

#ttcore.meshes<
  ::llvm::ArrayRef<MeshAttr>   # meshes
>

TT system meshes attribute includes one or more mesh configs used for networks.

Parameters:

MetalLayoutAttr

Tensor layout attribute with explicit physical shape

Syntax:

#ttcore.metal_layout<
  ::llvm::ArrayRef<int64_t>,   # logical_shape
  ::llvm::ArrayRef<int64_t>,   # dim_alignments
  DenseIntElementsAttr,   # collapsed_intervals
  OOBVal,   # oob_val
  MemorySpace,   # memory_space
  TensorMemoryLayout,   # memory_layout
  AffineMap   # indexAffineMap
>

The tensor layout attribute captures how tensor data is sharded across a grid of devices/cores and is laid out in memory. Note that the presence of this attribute implies that the tensor shape includes sharding (i.e. the first half of the tensor shape represents the grid shape).

Some high level goals:

• Logical shapes: Store the original tensor shape and rank intact and agnostic to underlying storage layout. Keeping the logical shapes not only makes some graph transformations vastly simpler, in particular convs, but it makes the lowered IR much easier to read and reason about. The original tensor shapes leave breadcrumbs that make it much easier to map back to the input representation.
• Collapsed dims: We may collapse dimensions during transformation, but it is important we capture this information such that it is not lost during tensor transformation. The collapsed_intervals field stores the collapses performed during conversion from logical_shape to physical tensor shape.
• Padding: store the desired alignments s.t. padding can be simply encoded; dim_alignments field represents alignment along each logical dim during collapse.
• Memref translation: ensure we have all necessary info s.t. we can trivally lower a tensor into a memref without any intermediate passes.

For a logical tensor of shape [H, W] distributed across a grid [GY, GX], the tensor shape would be:

• Without tiling: [GY, GX, H/GY, W/GX]
• With tiling: [GY, GX, H/GY/TH, W/GX/TW, TH, TW] where TH,TW are tile dimensions

This makes the representation 1:1 with memrefs and eliminates the need for shape conversion passes.

Examples:

// Logical 8x300 tensor distributed across 1x2 grid:
// tensor<1x2x8x150xf32, #tt.metal_layout<logical_shape=8x300, ...>>

// Logical 1024x1024 tensor distributed across 2x2 grid with 32x32 tiles:
// tensor<2x2x16x16x!ttcore.tile<32x32xf32>, #tt.metal_layout<logical_shape=1024x1024, ...>>

Parameters:

OOBValAttr

TT OOBVal

Syntax:

#ttcore.oob_val<
  ::mlir::tt::ttcore::OOBVal   # value
>

Parameters:

ReduceTypeAttr

TT Reduce Type

Syntax:

#ttcore.reduce_type<
  ::mlir::tt::ttcore::ReduceType   # value
>

Parameters:

ShardLayoutAttr

Shard layout attribute in TT dialect

Syntax:

#ttcore.shard<
  ::llvm::ArrayRef<int64_t>,   # stride
  uint32_t   # buffers
>

Describes shard layout of a memref buffer.

• Stride: Stride of each dim in bytes.
• Buffers: Number of back buffers used for double buffering, I/O latency hiding, etc

The shard layout attribute is a description of how each shard of a memref is laid out in memory. Memref's with this layout type implicitly mean their data is distributed across a grid of cores.

Parameters:

ShardStatusAttr

Shard Status Type

Syntax:

#ttcore.shard_status<
  ::mlir::tt::ttcore::ShardStatus   # value
>

Parameters:

SystemDescAttr

TT system_desc attribute

Syntax:

#ttcore.system_desc<
  ::llvm::ArrayRef<CPUDescAttr>,   # cpuDescs
  ::llvm::ArrayRef<ChipDescAttr>,   # chipDescs
  ::llvm::ArrayRef<unsigned>,   # chipDescIndices
  ::llvm::ArrayRef<ChipCapabilityAttr>,   # chipCapabilities
  ::llvm::ArrayRef<ChipCoordAttr>,   # chipCoords
  ::llvm::ArrayRef<ChipChannelAttr>   # chipChannels
>

TT system_desc attribute

Parameters:

TensorMeshAttr

Tensor mesh in TT dialect.

Syntax:

#ttcore.tensor_mesh<
  StringAttr   # name
>

Describes what mesh a tensor lives on.

Parameters:

TileSizeAttr

TT tile_size attribute

Syntax:

#ttcore.tile_size<
  int64_t,   # y
  int64_t   # x
>

TT tile_size attribute containing a supported Tensix tile shape.

Parameters:

ViewLayoutAttr

View layout attribute in TT dialect

Syntax:

#ttcore.view<
  AffineMap   # affineMap
>

Describes a view layout of a memref buffer.

• AffineMap: Provides affine map indexing into the associated data view.

Only the view_layout or stream_layout ops should return memref's with this attribute. The view layout attribute is necessary for two reasons:

• It provides a way to reblock the data view into a different shape (via affine map). Usually this would be some subblock of the original backing memory to chunk the data into smaller pieces.
• The type itself is a signal to datamovement passes that the memref is a view and should be treated as such.

Parameters:

ttcore.cpu_module (tt::ttcore::CPUModuleOp)

Module-wrapper operation for CPU ops

Syntax:

operation ::= `ttcore.cpu_module` attr-dict-with-keyword regions

Custom module operation that can a single ModuleOp, which should contain all funcs which should be run on CPU.

Example:

ttcore.cpu_module {
  module {
    func.func foo() { ... }
  }
}

Traits: IsolatedFromAbove, NoRegionArguments, NoTerminator, SingleBlock, SymbolTable

ttcore.device_module (tt::ttcore::DeviceModuleOp)

Module-wrapper operation for device ops

Syntax:

operation ::= `ttcore.device_module` attr-dict-with-keyword $bodyRegion

Custom module operation that can a single ModuleOp, which should contain all funcs which should be run on device.

Example:

ttcore.device_module {
  module {
    func.func foo() { ... }
  }
}

Traits: IsolatedFromAbove, NoRegionArguments, NoTerminator, SingleBlock, SymbolTable

ttcore.device (tt::ttcore::DeviceOp)

Named device

Syntax:

operation ::= `ttcore.device` $sym_name `=` $device_attr attr-dict

Interfaces: Symbol

Attributes:

ttcore.get_global (tt::ttcore::GetGlobalOp)

Named global

Syntax:

operation ::= `ttcore.get_global` $sym_name attr-dict `:` type($result)

Retrieves a named global value declared with ttcore.global

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Results:

ttcore.get_tuple_element (tt::ttcore::GetTupleElementOp)

GetTupleElement operation

Syntax:

operation ::= `ttcore.get_tuple_element` $operand `[` $index `]` attr-dict `:` functional-type(operands, results)

Extracts element at index position of the operand tuple and produces a result.

Example:

%result = ttcore.get_tuple_element %operand[0] : (tuple<tensor<32x32xbf16>, tensor<1x32xf32>>) -> tensor<32x32xbf16>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

ttcore.global (tt::ttcore::GlobalOp)

Named global

Syntax:

operation ::= `ttcore.global` $sym_name `=` $type (` ` `[` $index^ `]`)? attr-dict

Declares a global variable with an optional index.

Interfaces: Symbol

Attributes:

ttcore.load_cached (tt::ttcore::LoadCachedOp)

Load cached results from a previously computed function

Syntax:

operation ::= `ttcore.load_cached` `(` $callee `,` `[` $inputs `]` `)` attr-dict `:` functional-type($inputs, $results)

The load_cached operation calls a precomputed function with given arguments and returns its results. This is typically used to load constant or hoisted computation results.

Example:

%0, %1, %2 = "ttcore.load_cached"(@forward_const_eval_1, [%arg0, %arg2])

Attributes:

Operands:

Results:

ttcore.optimization_barrier (tt::ttcore::OptimizationBarrierOp)

Optimization barrier operation.

The optimization_barrier operation prevents compiler optimizations from reordering or eliminating the values passed through it. It acts as a barrier for optimization passes.

Inputs:

• inputs (Variadic): Values of tensor type.

Outputs:

• results (Variadic): Same values as inputs, passed through unchanged.

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Read on ::mlir::SideEffects::DefaultResource, MemoryEffects::Write on ::mlir::SideEffects::DefaultResource}

Operands:

Results:

ttcore.tuple (tt::ttcore::TupleOp)

Tuple operation

Syntax:

operation ::= `ttcore.tuple` $operands attr-dict `:` custom<TupleOpType>(type($operands), type($result))

Produces a result tuple from operands operands.

Example:

%result = ttcore.tuple %operand0, %operand1 : tuple<tensor<32xbf16, tensor<1x32xf32>>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

BFloat16Type

Bfloat16 floating-point type

ComplexType

Complex number with a parameterized element type

Syntax:

complex-type ::= `complex` `<` type `>`

The value of complex type represents a complex number with a parameterized element type, which is composed of a real and imaginary value of that element type. The element must be a floating point or integer scalar type.

Example:

complex<f32>
complex<i32>

Parameters:

Float4E2M1FNType

4-bit floating point with 2-bit exponent and 1-bit mantissa

An 4-bit floating point type with 1 sign bit, 2 bits exponent and 1 bit mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions with the following characteristics:

• bit encoding: S1E2M1
• exponent bias: 1
• infinities: Not supported
• NaNs: Not supported
• denormals when exponent is 0

Open Compute Project (OCP) microscaling formats (MX) specification: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf

Float6E2M3FNType

6-bit floating point with 2-bit exponent and 3-bit mantissa

An 6-bit floating point type with 1 sign bit, 2 bits exponent and 3 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions with the following characteristics:

• bit encoding: S1E2M3
• exponent bias: 1
• infinities: Not supported
• NaNs: Not supported
• denormals when exponent is 0

Open Compute Project (OCP) microscaling formats (MX) specification: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf

Float6E3M2FNType

6-bit floating point with 3-bit exponent and 2-bit mantissa

An 6-bit floating point type with 1 sign bit, 3 bits exponent and 2 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions with the following characteristics:

• bit encoding: S1E3M2
• exponent bias: 3
• infinities: Not supported
• NaNs: Not supported
• denormals when exponent is 0

Open Compute Project (OCP) microscaling formats (MX) specification: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf

Float8E3M4Type

8-bit floating point with 3 bits exponent and 4 bit mantissa

An 8-bit floating point type with 1 sign bit, 3 bits exponent and 4 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions with the following characteristics:

• bit encoding: S1E3M4
• exponent bias: 3
• infinities: supported with exponent set to all 1s and mantissa 0s
• NaNs: supported with exponent bits set to all 1s and mantissa values of {0,1}⁴ except S.111.0000
• denormals when exponent is 0

Float8E4M3Type

8-bit floating point with 3 bit mantissa

An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions with the following characteristics:

• bit encoding: S1E4M3
• exponent bias: 7
• infinities: supported with exponent set to all 1s and mantissa 0s
• NaNs: supported with exponent bits set to all 1s and mantissa of (001, 010, 011, 100, 101, 110, 111)
• denormals when exponent is 0

Float8E4M3B11FNUZType

8-bit floating point with 3 bit mantissa

An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions, with the exception that there are no infinity values, no negative zero, and only one NaN representation. This type has the following characteristics:

• bit encoding: S1E4M3
• exponent bias: 11
• infinities: Not supported
• NaNs: Supported with sign bit set to 1, exponent bits and mantissa bits set to all 0s
• denormals when exponent is 0

Related to: https://dl.acm.org/doi/10.5555/3454287.3454728

Float8E4M3FNType

8-bit floating point with 3 bit mantissa

An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions, with the exception that there are no infinity values and only two NaN representations. This type has the following characteristics:

• bit encoding: S1E4M3
• exponent bias: 7
• infinities: Not supported
• NaNs: supported with exponent bits and mantissa bits set to all 1s
• denormals when exponent is 0

Described in: https://arxiv.org/abs/2209.05433

Float8E4M3FNUZType

8-bit floating point with 3 bit mantissa

An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions, with the exception that there are no infinity values, no negative zero, and only one NaN representation. This type has the following characteristics:

• bit encoding: S1E4M3
• exponent bias: 8
• infinities: Not supported
• NaNs: Supported with sign bit set to 1, exponent bits and mantissa bits set to all 0s
• denormals when exponent is 0

Described in: https://arxiv.org/abs/2209.05433

Float8E5M2Type

8-bit floating point with 2 bit mantissa

An 8-bit floating point type with 1 sign bit, 5 bits exponent and 2 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions with the following characteristics:

• bit encoding: S1E5M2
• exponent bias: 15
• infinities: supported with exponent set to all 1s and mantissa 0s
• NaNs: supported with exponent bits set to all 1s and mantissa of (01, 10, or 11)
• denormals when exponent is 0

Described in: https://arxiv.org/abs/2209.05433

Float8E5M2FNUZType

8-bit floating point with 2 bit mantissa

An 8-bit floating point type with 1 sign bit, 5 bits exponent and 2 bits mantissa. This is not a standard type as defined by IEEE-754, but it follows similar conventions, with the exception that there are no infinity values, no negative zero, and only one NaN representation. This type has the following characteristics:

• bit encoding: S1E5M2
• exponent bias: 16
• infinities: Not supported
• NaNs: Supported with sign bit set to 1, exponent bits and mantissa bits set to all 0s
• denormals when exponent is 0

Described in: https://arxiv.org/abs/2206.02915

Float8E8M0FNUType

8-bit floating point with 8-bit exponent, no mantissa or sign

An 8-bit floating point type with no sign bit, 8 bits exponent and no mantissa. This is not a standard type as defined by IEEE-754; it is intended to be used for representing scaling factors, so it cannot represent zeros and negative numbers. The values it can represent are powers of two in the range [-127,127] and NaN.

• bit encoding: S0E8M0
• exponent bias: 127
• infinities: Not supported
• NaNs: Supported with all bits set to 1
• denormals: Not supported

Open Compute Project (OCP) microscaling formats (MX) specification: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf

Float16Type

16-bit floating-point type

Float32Type

32-bit floating-point type

Float64Type

64-bit floating-point type

Float80Type

80-bit floating-point type

Float128Type

128-bit floating-point type

FloatTF32Type

TF32 floating-point type

FunctionType

Map from a list of inputs to a list of results

Syntax:

// Function types may have multiple results.
function-result-type ::= type-list-parens | non-function-type
function-type ::= type-list-parens `->` function-result-type

The function type can be thought of as a function signature. It consists of a list of formal parameter types and a list of formal result types.

Example:

func.func @add_one(%arg0 : i64) -> i64 {
  %c1 = arith.constant 1 : i64
  %0 = arith.addi %arg0, %c1 : i64
  return %0 : i64
}

Parameters:

IndexType

Integer-like type with unknown platform-dependent bit width

Syntax:

// Target word-sized integer.
index-type ::= `index`

The index type is a signless integer whose size is equal to the natural machine word of the target ( rationale ) and is used by the affine constructs in MLIR.

Rationale: integers of platform-specific bit widths are practical to express sizes, dimensionalities and subscripts.

IntegerType

Integer type with arbitrary precision up to a fixed limit

Syntax:

// Sized integers like i1, i4, i8, i16, i32.
signed-integer-type ::= `si` [1-9][0-9]*
unsigned-integer-type ::= `ui` [1-9][0-9]*
signless-integer-type ::= `i` [1-9][0-9]*
integer-type ::= signed-integer-type |
                 unsigned-integer-type |
                 signless-integer-type

Integer types have a designated bit width and may optionally have signedness semantics.

Rationale: low precision integers (like i2, i4 etc) are useful for low-precision inference chips, and arbitrary precision integers are useful for hardware synthesis (where a 13 bit multiplier is a lot cheaper/smaller than a 16 bit one).

Parameters:

MemRefType

Shaped reference to a region of memory

Syntax:

layout-specification ::= attribute-value
memory-space ::= attribute-value
memref-type ::= `memref` `<` dimension-list-ranked type
                (`,` layout-specification)? (`,` memory-space)? `>`

A memref type is a reference to a region of memory (similar to a buffer pointer, but more powerful). The buffer pointed to by a memref can be allocated, aliased and deallocated. A memref can be used to read and write data from/to the memory region which it references. Memref types use the same shape specifier as tensor types. Note that memref<f32>, memref<0 x f32>, memref<1 x 0 x f32>, and memref<0 x 1 x f32> are all different types.

A memref is allowed to have an unknown rank (e.g. memref<*xf32>). The purpose of unranked memrefs is to allow external library functions to receive memref arguments of any rank without versioning the functions based on the rank. Other uses of this type are disallowed or will have undefined behavior.

Are accepted as elements:

• built-in integer types;
• built-in index type;
• built-in floating point types;
• built-in vector types with elements of the above types;
• another memref type;
• any other type implementing MemRefElementTypeInterface.

Layout

A memref may optionally have a layout that indicates how indices are transformed from the multi-dimensional form into a linear address. The layout must avoid internal aliasing, i.e., two distinct tuples of in-bounds indices must be pointing to different elements in memory. The layout is an attribute that implements MemRefLayoutAttrInterface. The bulitin dialect offers two kinds of layouts: strided and affine map, each of which is available as an attribute. Other attributes may be used to represent the layout as long as they can be converted to a semi-affine map and implement the required interface. Users of memref are expected to fallback to the affine representation when handling unknown memref layouts. Multi-dimensional affine forms are interpreted in row-major fashion.

In absence of an explicit layout, a memref is considered to have a multi-dimensional identity affine map layout. Identity layout maps do not contribute to the MemRef type identification and are discarded on construction. That is, a type with an explicit identity map is memref<?x?xf32, (i,j)->(i,j)> is strictly the same as the one without a layout, memref<?x?xf32>.

Affine Map Layout

The layout may be represented directly as an affine map from the index space to the storage space. For example, the following figure shows an index map which maps a 2-dimensional index from a 2x2 index space to a 3x3 index space, using symbols S0 and S1 as offsets.

Semi-affine maps are sufficiently flexible to represent a wide variety of dense storage layouts, including row- and column-major and tiled:

// MxN matrix stored in row major layout in memory:
#layout_map_row_major = (i, j) -> (i, j)

// MxN matrix stored in column major layout in memory:
#layout_map_col_major = (i, j) -> (j, i)

// MxN matrix stored in a 2-d blocked/tiled layout with 64x64 tiles.
#layout_tiled = (i, j) -> (i floordiv 64, j floordiv 64, i mod 64, j mod 64)

Strided Layout

Memref layout can be expressed using strides to encode the distance, in number of elements, in (linear) memory between successive entries along a particular dimension. For example, a row-major strided layout for memref<2x3x4xf32> is strided<[12, 4, 1]>, where the last dimension is contiguous as indicated by the unit stride and the remaining strides are products of the sizes of faster-variying dimensions. Strided layout can also express non-contiguity, e.g., memref<2x3, strided<[6, 2]>> only accesses even elements of the dense consecutive storage along the innermost dimension.

The strided layout supports an optional offset that indicates the distance, in the number of elements, between the beginning of the memref and the first accessed element. When omitted, the offset is considered to be zero. That is, memref<2, strided<[2], offset: 0>> and memref<2, strided<[2]>> are strictly the same type.

Both offsets and strides may be dynamic, that is, unknown at compile time. This is represented by using a question mark (?) instead of the value in the textual form of the IR.

The strided layout converts into the following canonical one-dimensional affine form through explicit linearization:

affine_map<(d0, ... dN)[offset, stride0, ... strideN] ->
            (offset + d0 * stride0 + ... dN * strideN)>

Therefore, it is never subject to the implicit row-major layout interpretation.

Codegen of Unranked Memref

Using unranked memref in codegen besides the case mentioned above is highly discouraged. Codegen is concerned with generating loop nests and specialized instructions for high-performance, unranked memref is concerned with hiding the rank and thus, the number of enclosing loops required to iterate over the data. However, if there is a need to code-gen unranked memref, one possible path is to cast into a static ranked type based on the dynamic rank. Another possible path is to emit a single while loop conditioned on a linear index and perform delinearization of the linear index to a dynamic array containing the (unranked) indices. While this is possible, it is expected to not be a good idea to perform this during codegen as the cost of the translations is expected to be prohibitive and optimizations at this level are not expected to be worthwhile. If expressiveness is the main concern, irrespective of performance, passing unranked memrefs to an external C++ library and implementing rank-agnostic logic there is expected to be significantly simpler.

Unranked memrefs may provide expressiveness gains in the future and help bridge the gap with unranked tensors. Unranked memrefs will not be expected to be exposed to codegen but one may query the rank of an unranked memref (a special op will be needed for this purpose) and perform a switch and cast to a ranked memref as a prerequisite to codegen.

Example:

// With static ranks, we need a function for each possible argument type
%A = alloc() : memref<16x32xf32>
%B = alloc() : memref<16x32x64xf32>
call @helper_2D(%A) : (memref<16x32xf32>)->()
call @helper_3D(%B) : (memref<16x32x64xf32>)->()

// With unknown rank, the functions can be unified under one unranked type
%A = alloc() : memref<16x32xf32>
%B = alloc() : memref<16x32x64xf32>
// Remove rank info
%A_u = memref_cast %A : memref<16x32xf32> -> memref<*xf32>
%B_u = memref_cast %B : memref<16x32x64xf32> -> memref<*xf32>
// call same function with dynamic ranks
call @helper(%A_u) : (memref<*xf32>)->()
call @helper(%B_u) : (memref<*xf32>)->()

The core syntax and representation of a layout specification is a semi-affine map. Additionally, syntactic sugar is supported to make certain layout specifications more intuitive to read. For the moment, a memref supports parsing a strided form which is converted to a semi-affine map automatically.

The memory space of a memref is specified by a target-specific attribute. It might be an integer value, string, dictionary or custom dialect attribute. The empty memory space (attribute is None) is target specific.

The notionally dynamic value of a memref value includes the address of the buffer allocated, as well as the symbols referred to by the shape, layout map, and index maps.

Examples of memref static type

// Identity index/layout map
#identity = affine_map<(d0, d1) -> (d0, d1)>

// Column major layout.
#col_major = affine_map<(d0, d1, d2) -> (d2, d1, d0)>

// A 2-d tiled layout with tiles of size 128 x 256.
#tiled_2d_128x256 = affine_map<(d0, d1) -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)>

// A tiled data layout with non-constant tile sizes.
#tiled_dynamic = affine_map<(d0, d1)[s0, s1] -> (d0 floordiv s0, d1 floordiv s1,
                             d0 mod s0, d1 mod s1)>

// A layout that yields a padding on two at either end of the minor dimension.
#padded = affine_map<(d0, d1) -> (d0, (d1 + 2) floordiv 2, (d1 + 2) mod 2)>


// The dimension list "16x32" defines the following 2D index space:
//
//   { (i, j) : 0 <= i < 16, 0 <= j < 32 }
//
memref<16x32xf32, #identity>

// The dimension list "16x4x?" defines the following 3D index space:
//
//   { (i, j, k) : 0 <= i < 16, 0 <= j < 4, 0 <= k < N }
//
// where N is a symbol which represents the runtime value of the size of
// the third dimension.
//
// %N here binds to the size of the third dimension.
%A = alloc(%N) : memref<16x4x?xf32, #col_major>

// A 2-d dynamic shaped memref that also has a dynamically sized tiled
// layout. The memref index space is of size %M x %N, while %B1 and %B2
// bind to the symbols s0, s1 respectively of the layout map #tiled_dynamic.
// Data tiles of size %B1 x %B2 in the logical space will be stored
// contiguously in memory. The allocation size will be
// (%M ceildiv %B1) * %B1 * (%N ceildiv %B2) * %B2 f32 elements.
%T = alloc(%M, %N) [%B1, %B2] : memref<?x?xf32, #tiled_dynamic>

// A memref that has a two-element padding at either end. The allocation
// size will fit 16 * 64 float elements of data.
%P = alloc() : memref<16x64xf32, #padded>

// Affine map with symbol 's0' used as offset for the first dimension.
#imapS = affine_map<(d0, d1) [s0] -> (d0 + s0, d1)>
// Allocate memref and bind the following symbols:
// '%n' is bound to the dynamic second dimension of the memref type.
// '%o' is bound to the symbol 's0' in the affine map of the memref type.
%n = ...
%o = ...
%A = alloc (%n)[%o] : <16x?xf32, #imapS>

Parameters:

NoneType

A unit type

Syntax:

none-type ::= `none`

NoneType is a unit type, i.e. a type with exactly one possible value, where its value does not have a defined dynamic representation.

Example:

func.func @none_type() {
  %none_val = "foo.unknown_op"() : () -> none
  return
}

OpaqueType

Type of a non-registered dialect

Syntax:

opaque-type ::= `opaque` `<` type `>`

Opaque types represent types of non-registered dialects. These are types represented in their raw string form, and can only usefully be tested for type equality.

Example:

opaque<"llvm", "struct<(i32, float)>">
opaque<"pdl", "value">

Parameters:

RankedTensorType

Multi-dimensional array with a fixed number of dimensions

Syntax:

tensor-type ::= `tensor` `<` dimension-list type (`,` encoding)? `>`
dimension-list ::= (dimension `x`)*
dimension ::= `?` | decimal-literal
encoding ::= attribute-value

Values with tensor type represents aggregate N-dimensional data values, and have a known element type and a fixed rank with a list of dimensions. Each dimension may be a static non-negative decimal constant or be dynamically determined (indicated by ?).

The runtime representation of the MLIR tensor type is intentionally abstracted - you cannot control layout or get a pointer to the data. For low level buffer access, MLIR has a memref type. This abstracted runtime representation holds both the tensor data values as well as information about the (potentially dynamic) shape of the tensor. The dim operation returns the size of a dimension from a value of tensor type.

The encoding attribute provides additional information on the tensor. An empty attribute denotes a straightforward tensor without any specific structure. But particular properties, like sparsity or other specific characteristics of the data of the tensor can be encoded through this attribute. The semantics are defined by a type and attribute interface and must be respected by all passes that operate on tensor types. TODO: provide this interface, and document it further.

Note: hexadecimal integer literals are not allowed in tensor type declarations to avoid confusion between 0xf32 and 0 x f32. Zero sizes are allowed in tensors and treated as other sizes, e.g., tensor<0 x 1 x i32> and tensor<1 x 0 x i32> are different types. Since zero sizes are not allowed in some other types, such tensors should be optimized away before lowering tensors to vectors.

Example:

// Known rank but unknown dimensions.
tensor<? x ? x ? x ? x f32>

// Partially known dimensions.
tensor<? x ? x 13 x ? x f32>

// Full static shape.
tensor<17 x 4 x 13 x 4 x f32>

// Tensor with rank zero. Represents a scalar.
tensor<f32>

// Zero-element dimensions are allowed.
tensor<0 x 42 x f32>

// Zero-element tensor of f32 type (hexadecimal literals not allowed here).
tensor<0xf32>

// Tensor with an encoding attribute (where #ENCODING is a named alias).
tensor<?x?xf64, #ENCODING>

Parameters:

TupleType

Fixed-sized collection of other types

Syntax:

tuple-type ::= `tuple` `<` (type ( `,` type)*)? `>`

The value of tuple type represents a fixed-size collection of elements, where each element may be of a different type.

Rationale: Though this type is first class in the type system, MLIR provides no standard operations for operating on tuple types (rationale).

Example:

// Empty tuple.
tuple<>

// Single element
tuple<f32>

// Many elements.
tuple<i32, f32, tensor<i1>, i5>

Parameters:

UnrankedMemRefType

Shaped reference, with unknown rank, to a region of memory

Syntax:

unranked-memref-type ::= `memref` `<*x` type (`,` memory-space)? `>`
memory-space ::= attribute-value

A memref type with an unknown rank (e.g. memref<*xf32>). The purpose of unranked memrefs is to allow external library functions to receive memref arguments of any rank without versioning the functions based on the rank. Other uses of this type are disallowed or will have undefined behavior.

See MemRefType for more information on memref types.

Examples:

memref<*f32>

// An unranked memref with a memory space of 10.
memref<*f32, 10>

Parameters: