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 = {},
)

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/ttir-builder/ttir_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.

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 ttir_golden.py and follow this pattern:

Simple operations: If PyTorch has an identical function, you can directly use it in the mappings
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 ttir_golden.py:

# In ttir_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.fails_golden : expect this test to fail the ttrt golden check
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. Build ttrt (sample instructions - subject to change)
cmake --build build -- ttrt

3. Query system
ttrt query --save-artifacts

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

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

6. Run test cases
ttrt run ttnn
ttrt run ttmetal

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:

tt-mlir documentation