wehub-resource-sync
203 lines
6.7 KiB
Python
203 lines
6.7 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""A generic IRBuilder across the TVM stack"""
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from tvm_ffi import register_object as _register_object
|
|
|
|
from tvm.runtime import Object as _Object
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
@_register_object("script.ir_builder.IRBuilderFrame")
|
|
class IRBuilderFrame(_Object):
|
|
"""A stack frame of the IRBuilder used to keep track of the current scope.
|
|
|
|
Furthermore, the information stored in each stack frame can be useful for context-dependent
|
|
IR construction.
|
|
|
|
Examples
|
|
--------
|
|
|
|
The `T.match_buffer` below instead an element in the buffer map of `PrimFuncFrame`:
|
|
|
|
.. 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(...)
|
|
|
|
|
|
The `T.match_buffer` below instead generates `MatchBufferRegion` in a TIR block:
|
|
|
|
.. 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
|
|
with T.sblock(...): # pushes a BlockFrame (subclass of IRBuilderFrame)
|
|
# to `builder`'s stack of frames
|
|
buffer = T.match_buffer(...)
|
|
|
|
"""
|
|
|
|
def __enter__(self) -> "IRBuilderFrame":
|
|
_ffi_api.IRBuilderFrameEnter(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, trace) -> None: # pylint: disable=unused-argument
|
|
if exc_type is None and exc_value is None:
|
|
# Do not execute `FrameExit` if the with scope exits because of exceptions
|
|
_ffi_api.IRBuilderFrameExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
def add_callback(self, callback: Callable[[], None]) -> None:
|
|
"""Add a callback method invoked when exiting the with-scope.
|
|
|
|
Parameters
|
|
----------
|
|
callback : Callable[[], None]
|
|
The callback method to be invoked.
|
|
"""
|
|
_ffi_api.IRBuilderFrameAddCallback( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
self, callback
|
|
)
|
|
|
|
|
|
@_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
|
|
|
|
@staticmethod
|
|
def is_in_scope() -> bool:
|
|
"""See if the current thread-local scope has an IRBuilder.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
Whether the current thread-local scope has an IRBuilder
|
|
"""
|
|
return _ffi_api.IRBuilderIsInScope() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
def get(self) -> _Object:
|
|
"""Get the constructed IR."""
|
|
return _ffi_api.IRBuilderGet(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def name(s: str, v: Any) -> Any:
|
|
"""Set the name of an object.
|
|
|
|
Parameters
|
|
----------
|
|
s : str
|
|
The name of the object.
|
|
v : Any
|
|
The object to name.
|
|
|
|
Returns
|
|
-------
|
|
v : Any
|
|
The same object with the name set.
|
|
"""
|
|
return _ffi_api.IRBuilderName(s, v) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def name_many( # pylint: disable=invalid-name
|
|
s: list[str],
|
|
vs: list[Any],
|
|
) -> list[Any]:
|
|
"""Set the name of a list of objects.
|
|
|
|
Parameters
|
|
----------
|
|
s : List[str]
|
|
The names of the objects.
|
|
vs : List[Any]
|
|
The objects to name.
|
|
|
|
Returns
|
|
-------
|
|
vs : List[Any]
|
|
The same objects with the names set.
|
|
"""
|
|
assert len(s) == len(vs)
|
|
return [IRBuilder.name(i, v) for i, v in zip(s, vs)]
|