# 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.

```bash
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](https://docs.tenstorrent.com/tt-mlir/autogen/md/Dialect/TTIROp.html).

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/golden/mapping.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 `tools/golden/mapping.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 `tools/golden/mapping.py`:

```python
# In tools/golden/mapping.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 in tools/goldens/mapping.py
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:

```python
# In ops.py
from goldens import GOLDEN_MAPPINGS

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.
```bash
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.
```bash
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 running 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.
```bash
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.
```bash
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

```bash
"""
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`.

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