Class IRBuilder

python/tvm/script/ir_builder/base.py:91–202

A dialect-agnostic IRBuilder that constructs any IR of TVM. Examples -------- An idiomatic use of this class is to put this inside the with-scope, call dialect-specific methods accordingly. Upon exiting the scope. .. code-block:: python from tvm.script.ir_builder impor

@_register_object("script.ir_builder.IRBuilder")
class IRBuilder(_Object):
    """A dialect-agnostic IRBuilder that constructs any IR of TVM.

    Examples
    --------
    An idiomatic use of this class is to put this inside the with-scope,
    call dialect-specific methods accordingly. Upon exiting the scope.

    .. code-block:: python

        from tvm.script.ir_builder import tirx as T
        from tvm.script.ir_builder import IRBuilder

        with IRBuilder() as builder:
            with T.prim_func(...):  # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
                                # to `builder`'s stack of frames
                buffer = T.match_buffer(...)

        return builder.get()        # returns the constructed IR, i.e. tirx.PrimFunc
    """

    def __init__(self) -> None:
        """Construct an IRBuilder."""
        self.__init_handle_by_constructor__(
            _ffi_api.IRBuilder  # type: ignore[attr-defined] # pylint: disable=no-member
        )

    def __enter__(self) -> "IRBuilder":
        """Enter the with-scope for IRBuilder, which allows the IRBuilder to be discoverable
        using `IRBuilder.current()`.

        Examples
        --------
        .. code-block:: python

            from tvm.script.ir_builder import IRBuilder

            with IRBuilder() as builder:
                assert IRBuilder.current() == builder

        """
        _ffi_api.IRBuilderEnter(self)  # type: ignore[attr-defined] # pylint: disable=no-member
        return self

    def __exit__(self, ptype, value, trace) -> None:  # pylint: disable=unused-argument
        _ffi_api.IRBuilderExit(self)  # type: ignore[attr-defined] # pylint: disable=no-member

    @staticmethod
    def current() -> "IRBuilder":
        """Get the current IRBuilder put in the with-scope.

        Returns
        -------
        builder : IRBuilder
            The current IRBuilder.
        """
        return _ffi_api.IRBuilderCurrent()  # type: ignore[attr-defined] # pylint: disable=no-member