Adding New Operators

You may find it helpful to read the architecture documentation before you start reading this guide.

Adding new operators to Tripy typically involves making changes in the frontend as well as in the FlatIR. In some cases, the frontend operator can be expressed in terms of existing FlatIR operators, in which case you only need to make changes in the frontend.

Let’s take a look at an example of how you might add an Iota operator to Tripy. So that it doesn’t clash with Tripy’s actual Iota implementation, we’ll call it Theta instead.

Implementation

FlatIR Operator

The FlatIR operator is usually the most challenging aspect of implementing operators in Tripy. The good news is that you might not even need to do this if the low-level operators you need already exist in the FlatIR. And if you do, then it’ll only get easier after this!

We’ll start by adding a new file under tripy/flat_ir/ops called theta.py; see the inline comments for explanations of what’s happening:

from dataclasses import dataclass

from mlir_tensorrt.compiler import ir
from mlir_tensorrt.compiler.dialects import stablehlo

from tripy.flat_ir.ops.base import BaseFlatIROp


# Every `FlatIR` operator is implemented as a `dataclass` so that the base
# class can automatically implement several methods by inspecting the child
# class fields at runtime. The `repr=False` is important because the default
# `__repr__` method generated by `dataclass` will be extremely verbose and
# makes interactive debugging more difficult.
@dataclass(repr=False)
class ThetaOp(BaseFlatIROp):
    dim: int

    # `to_mlir()` is the trickiest bit. As the name implies, the method is meant to lower the
    # `FlatIR` operator into MLIR. To figure out which MLIR operators to use, refer to
    # the 'MLIR Python API Guide' (linked below).
    def to_mlir(self, operands):
        out_type = self.outputs[0].to_mlir()
        theta_dim = ir.IntegerAttr.get(type=ir.IntegerType.get_signless(64), value=self.dim)
        output = stablehlo.DynamicIotaOp(result=out_type, output_shape=operands[0], iota_dimension=theta_dim)
        return [output]

Links:

• MLIR Python API Guide

Exposing The Operator

One of the principles we follow when writing submodules is that other submodules should not need to reach into the internals of a submodule to retrieve something they need.

For example, a class which needs to import ThetaOp does not need to know where exactly within the flat_ir.ops module the ThetaOp lives - it should be able to just import it from the submodule.

To make this possible, we need to import the ThetaOp into the flat_ir.ops submodule. We can do so by adding the following line into tripy/flat_ir/ops/__init__.py:

from tripy.flat_ir.ops.theta import ThetaOp

Trace Operator And The Public API

Now that we have a FlatIR operator, we can implement a Trace operator that will use it along with a public API function. Let’s create a new file under tripy/frontend/trace/ops called theta.py.

Trace Operator

First, we’ll implement the Trace operator itself:

from dataclasses import dataclass
from typing import Tuple

from tripy import utils
from tripy.common import datatype, device
from tripy.common.exception import raise_error
from tripy.frontend.trace.ops.base import BaseTraceOp
import tripy.frontend.trace.ops.utils as op_utils


# Just like with `FlatIR` operators, all `Trace` operators are implemented as `dataclass`es.
# As before, we want `repr=False` here.
@dataclass(repr=False)
class Theta(BaseTraceOp):
    # Notice that we do *not* need to define a constructor and can rely on the default
    # implementation provided by `dataclass`.
    dim: int
    dtype: datatype.dtype

    # `infer_rank()` populates the rank of the output `TraceTensor`s.
    # Here we use one of the predefined policies to set the output rank to the same as the shape (i.e. the length)
    # of the shape operand.
    infer_rank = op_utils.InferRankPolicies.same_as_shape_of_shape_input()

    # *Optional* `infer_dtypes()` populates the data types of the
    # output `TraceTensor`s. The default implementation copies the input
    # data types if they are all the same, so you may not need to implement this.
    def infer_dtypes(self):
        self.outputs[0].dtype = self.dtype

    # *Optional* `infer_devices()` populates the devices of the
    # output `TraceTensor`s. The default implementation copies the input
    # devices if they are all the same, so you may not need to implement this either.
    def infer_devices(self):
        self.outputs[0].device = device("gpu")

    # `to_flat_ir()` translates the `Trace` operator to a subgraph of
    # one or more `FlatIR` operators. In our case, it's just a 1:1
    # mapping to the `ThetaOp` we created earlier.
    def to_flat_ir(self, inputs, outputs):
        # Note that we import the `FlatIR` operator within the function
        # call - this is to avoid circular dependencies.
        from tripy.flat_ir.ops import ThetaOp
        import tripy.frontend.trace.ops.utils as op_utils

        # This code may look a bit confusing; for more details, look at the
        # 'FlatIR section in the architecture document' (linked below).
        ThetaOp.build(inputs, outputs, dim=self.dim)

Links:

• FlatIR section in the architecture document

Public API

Next, we can define the public interface. Since our public interface maps 1:1 with the Trace operator we just implemented and does not require weights, we’ll add it in the same file.

If our API required a composition of multiple Trace operators, then we would instead implement it under frontend/ops/.

If it required weights (i.e. inputs that are expected to always be constant), then we would implement it as a tripy.Module under frontend/module.

from tripy import export
import tripy.frontend.utils as frontend_utils
from tripy.types import ShapeLike

# We can use the `export.public_api()` decorator to automatically export this function into the
# top-level module. This means it will be accessible as `tripy.theta`.
#
# This decorator also controls how the API is exposed in the documentation - the `document_under`
# option determines where in the documentation hierarchy this API will show up.
#
# If we needed to provide any special autodoc options, we could use the `autodoc_options` parameter.
@export.public_api(document_under="tensor_operations")

# The `convert_to_tensors` decorator automatically converts compatible arguments,
# like `TensorLike` or `ShapeLike`s, into tensors.
@frontend_utils.convert_to_tensors()
def theta(shape: ShapeLike, dim: int = 0, dtype: datatype.dtype = datatype.float32) -> "tripy.Tensor":
    # For any public facing interfaces, we have documentation requirements which you can read
    # about in the 'Docs README' (linked below). The docstring we've implemented here
    # adheres to all of these requirements. Non-compliant docstrings will, in most cases,
    # cause test failures; however, you should still manually ensure you're writing high-quality
    # docstrings.
    #
    # The examples in docstrings are run as part of our tests, so you should also add
    # assertions to make sure things are functionally correct. In this case, we check
    # that the `output` we create in the code example is what we expect.
    """
    Fills an output tensor with consecutive values starting from zero along the given dimension.

    Args:
        shape: The desired shape.
        dim: Dimension along which to perform the theta operation.
            This cannot exceed the rank of the specified shape.
        dtype: The desired data type.

    Returns:
        A tensor of shape ``shape`` and data type ``dtype``.

    .. code-block:: python
        :linenos:
        :caption: Example

        output = tp.theta([3])

        assert np.array_equal(cp.from_dlpack(output).get(), np.arange(0, 3, dtype=np.float32))
    """

    # Next we build the trace operator. The `build()` function is also responsible for constructing
    # the output frontend Tensors. All of the arguments that follow the inputs
    # are forwarded directly to the constructor of the `Trace` operator.
    return Theta.build([shape], dim, dtype)

Links:

• Docs README

Exposing The Operator

Similarly to the FlatIR operator, we need to import Theta into the frontend.trace.ops submodule. We can do so by adding the following line into tripy/frontend/trace/ops/__init__.py:

from tripy.frontend.trace.ops.theta import Theta, theta

Testing

Now that we’ve implemented our operator, let’s write tests for it. The structure of the tests/ directory mirrors that of the tripy/ directory (you can read more about that here). We need to test both the FlatIR and Trace operators.

Testing The Trace Operator And Public API

Since we implemented our Trace operator and public API in tripy/frontend/trace/ops, we’ll add the test under tests/frontend/trace/ops. Create a new file there called test_theta.py:

import tripy as tp
from tests import helper
from tripy.frontend.trace.ops import Theta


class TestTheta:
    # This ensures that the public API function creates a frontend `Tensor`
    # and populates it with the right `Trace` operator.
    def test_op_func(self):
        a = tp.theta([2, 3])
        assert isinstance(a, tp.Tensor)
        assert isinstance(a.trace_tensor.producer, Theta)

    # You should also include negative tests for anything that is expected to
    # fail. In our case, we just have `test_invalid_dim`,
    # which ensures that we emit an error if the `dim` parameter is outside
    # the allowed range.
    def test_invalid_dim(self):
        with helper.raises(tp.TripyException, match="iota dimension cannot go beyond the output rank"):
            tp.theta([2, 3], dim=3).eval()

Integration Tests

The code examples in the docstring of the public API serve as good sanity integration tests. However, you should still add separate integration tests to get better coverage.

Our docstring covers the 1D case, so let’s add an integration test to cover the multidimensional case. Create a new file called test_theta.py under tests/integration:

import numpy as np
import cupy as cp

import tripy as tp


def test_multi_dimensional():
    output = tp.theta([2, 3], dim=1)
    expected = np.broadcast_to(np.arange(0, 3, dtype=np.float32), (2, 3))

    assert np.array_equal(cp.from_dlpack(output).get(), expected)

Done!

If you’ve reached this point, you have successfully added a new operation to Tripy. Congratulations!