chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,48 @@
|
||||
# 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.
|
||||
"""The ir_builder subpackage of TVMScript.
|
||||
|
||||
Per-dialect builder submodules (``tvm.script.ir_builder.tirx``, etc.) are
|
||||
resolved lazily via :data:`tvm.script._DIALECT_REGISTRY`. When a dialect
|
||||
is accessed (e.g. ``tvm.script.ir_builder.tirx``), this subpackage's
|
||||
``__getattr__`` looks up the dialect in ``_DIALECT_REGISTRY`` and imports
|
||||
``<dialect_module_path>.builder`` (e.g. ``tvm.tirx.script.builder``),
|
||||
caching the result so subsequent accesses skip ``__getattr__``.
|
||||
|
||||
The IR layer is foundational and is NOT registered as a dialect — its
|
||||
builder lives as a real submodule ``tvm.script.ir_builder.ir``.
|
||||
|
||||
See :mod:`tvm.script` for a full description of the dialect resolution
|
||||
mechanism, including the ``_DialectRedirectFinder`` that handles
|
||||
deep statement-form imports.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
from .base import IRBuilder
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
# Lazy import to avoid loading tvm.script during dialect bootstrap.
|
||||
from tvm.script import _DIALECT_REGISTRY # pylint: disable=import-outside-toplevel
|
||||
|
||||
if name in _DIALECT_REGISTRY:
|
||||
module = importlib.import_module(f"{_DIALECT_REGISTRY[name]}.builder")
|
||||
globals()[name] = module
|
||||
return module
|
||||
raise AttributeError(f"module 'tvm.script.ir_builder' has no attribute {name!r}")
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs for tvm.script.ir_builder"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script.ir_builder", __name__) # pylint: disable=protected-access
|
||||
@@ -0,0 +1,202 @@
|
||||
# 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)]
|
||||
@@ -0,0 +1,33 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.ir"""
|
||||
|
||||
from .frame import IRModuleFrame
|
||||
from .ir import (
|
||||
decl_function,
|
||||
def_function,
|
||||
ir_module,
|
||||
module_attrs,
|
||||
module_get_attr,
|
||||
module_set_attr,
|
||||
module_global_infos,
|
||||
lookup_vdevice,
|
||||
lookup_name,
|
||||
vdevice,
|
||||
dummy_global_info,
|
||||
)
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script.ir_builder.ir", __name__) # pylint: disable=protected-access
|
||||
@@ -0,0 +1,25 @@
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.ir.frame"""
|
||||
|
||||
from tvm_ffi import register_object as _register_object
|
||||
|
||||
from ..base import IRBuilderFrame
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.IRModuleFrame")
|
||||
class IRModuleFrame(IRBuilderFrame): ...
|
||||
@@ -0,0 +1,192 @@
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.ir.ir"""
|
||||
|
||||
from tvm.ir import BaseFunc, DummyGlobalInfo, GlobalInfo, GlobalVar, VDevice
|
||||
from tvm.runtime import Object as tvm_Object
|
||||
|
||||
from . import _ffi_api
|
||||
from .frame import IRModuleFrame
|
||||
|
||||
|
||||
def ir_module() -> IRModuleFrame:
|
||||
"""Start a ir_module frame.
|
||||
Returns
|
||||
-------
|
||||
frame: IRModuleFrame
|
||||
The constructed frame.
|
||||
"""
|
||||
return _ffi_api.IRModule() # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def decl_function(func_name: str, func_signature: BaseFunc) -> GlobalVar:
|
||||
"""Declare a Function without given the specific function implementation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func_name : str
|
||||
The function unique name.
|
||||
|
||||
func_signature: BaseFunc
|
||||
A Function w/o body, which used to specify the function signature
|
||||
(i.e. func params and func return type/shape).
|
||||
|
||||
Note
|
||||
----
|
||||
It is usually used in cross-function call. And we can specify the function by `DefFunction`
|
||||
|
||||
Returns
|
||||
-------
|
||||
gv : GlobalVar
|
||||
The corresponding GlobalVar.
|
||||
"""
|
||||
if not isinstance(func_signature, BaseFunc):
|
||||
raise ValueError(
|
||||
"decl_function expects an instance of BaseFunc, "
|
||||
f"but {func_signature} is of type {type(func_signature)}"
|
||||
)
|
||||
return _ffi_api.DeclFunction( # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
func_name, func_signature
|
||||
)
|
||||
|
||||
|
||||
def def_function(func_name: str, func: BaseFunc) -> None:
|
||||
"""Define the function which is declared before.
|
||||
Parameters
|
||||
----------
|
||||
func_name : str
|
||||
The function unique name.
|
||||
func: BaseFunc
|
||||
The given function implementation
|
||||
"""
|
||||
return _ffi_api.DefFunction(func_name, func) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_attrs(attrs: dict[str, tvm_Object], allow_overwrite=False) -> None:
|
||||
"""Specify the attrs of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
attrs: Dict[str, Object]
|
||||
The module attrs.
|
||||
allow_overwrite: bool
|
||||
Whether allow overwrite the existing attrs.
|
||||
"""
|
||||
return _ffi_api.ModuleAttrs(attrs, allow_overwrite) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_get_attr(attr_key: str) -> tvm_Object | None:
|
||||
"""Get the specified attr of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
attr_key: str
|
||||
The key of the attr to be retrieved.
|
||||
Returns
|
||||
-------
|
||||
attr: Optional[Object]
|
||||
The specified module attr or None if not found.
|
||||
"""
|
||||
return _ffi_api.ModuleGetAttr(attr_key) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_set_attr(
|
||||
attr_key: str, attr_value: tvm_Object | None, allow_overwrite: bool = False
|
||||
) -> None:
|
||||
"""Set the specified attr of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
attr_key: str
|
||||
The key of the attr to be set.
|
||||
attr_value: Optional[Object]
|
||||
The value of the attr to be set.
|
||||
allow_overwrite: bool
|
||||
Whether allow overwrite the existing attr.
|
||||
"""
|
||||
return _ffi_api.ModuleSetAttr(attr_key, attr_value, allow_overwrite) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_global_infos(global_infos: dict[str, list[GlobalInfo]]) -> None:
|
||||
"""Specify the global infos of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
global_infos: Dict[str, List[GlobalInfo]]
|
||||
The module global infos.
|
||||
"""
|
||||
return _ffi_api.ModuleGlobalInfos(global_infos) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
############################### GlobalInfo ###############################
|
||||
|
||||
|
||||
def dummy_global_info() -> DummyGlobalInfo:
|
||||
"""Create a dummy global info expression.
|
||||
Returns
|
||||
-------
|
||||
res : DummyGlobalInfo
|
||||
The result dummy global info.
|
||||
"""
|
||||
return DummyGlobalInfo() # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def vdevice(target=None, vdevice_id: int = 0, memory_scope: str = "global") -> VDevice:
|
||||
"""Create a virtual device global info.
|
||||
Parameters
|
||||
----------
|
||||
target
|
||||
The target.
|
||||
vdevice_id: int
|
||||
The virtual device index.
|
||||
memory_scope: str
|
||||
The memory scope, default is "global"
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : VDevice
|
||||
The result virtual device.
|
||||
"""
|
||||
return VDevice(target, vdevice_id, memory_scope) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def lookup_vdevice(target_kind: str | None = None, device_index: int = -1) -> VDevice:
|
||||
"""Retrieve a virtual device from the globalinfo vdevice list.
|
||||
Parameters
|
||||
----------
|
||||
target_kind: str
|
||||
The target device kind, for example 'llvm' or 'cuda'.
|
||||
device_index: int
|
||||
The virtual device index.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : VDevice
|
||||
The result virtual device.
|
||||
"""
|
||||
return _ffi_api.LookupVDevice(target_kind, device_index) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def lookup_name(name: str) -> bool:
|
||||
"""Check if a global variable with the given name exists.
|
||||
Parameters
|
||||
----------
|
||||
name: str
|
||||
The name of the global variable.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : bool
|
||||
True if the global variable exists, False otherwise.
|
||||
"""
|
||||
return _ffi_api.LookupName(name) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
Reference in New Issue
Block a user