wehub-resource-sync
120 lines
3.6 KiB
Python
120 lines
3.6 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.
|
|
|
|
"""Pattern registry for BYOC backends"""
|
|
|
|
import atexit
|
|
from collections.abc import Callable, Mapping
|
|
|
|
from tvm.relax.dpl import DFPattern
|
|
from tvm.relax.transform import FusionPattern
|
|
|
|
from ..expr import Expr
|
|
from . import _ffi_api
|
|
|
|
_REGISTERED_PATTERN_NAMES: set[str] = set()
|
|
|
|
|
|
def _cleanup_registered_patterns():
|
|
_ffi_api.RemovePatterns(list(_REGISTERED_PATTERN_NAMES)) # type: ignore # pylint: disable=no-member
|
|
|
|
|
|
_CLEANUP_REGISTERED = False
|
|
|
|
|
|
def _ensure_cleanup_function_registered():
|
|
"""
|
|
Add a cleanup function to be called on interpreter termination, to remove all
|
|
patterns registered on the Python side. Without cleaning up those patterns,
|
|
program will segfault on termination. It's because the check functions of pattern
|
|
entries are referenced from the static memory of libtvm, thus they will be cleaned
|
|
up at the very end, making calls to Py_DecRef after Python interpreter terminates.
|
|
"""
|
|
global _CLEANUP_REGISTERED # pylint: disable=global-statement
|
|
|
|
if not _CLEANUP_REGISTERED:
|
|
atexit.register(_cleanup_registered_patterns)
|
|
_CLEANUP_REGISTERED = True
|
|
|
|
|
|
CheckFunc = Callable[[Mapping[DFPattern, Expr], Expr], bool]
|
|
Pattern = (
|
|
FusionPattern
|
|
| tuple[str, DFPattern]
|
|
| tuple[str, DFPattern, Mapping[str, DFPattern]]
|
|
| tuple[str, DFPattern, Mapping[str, DFPattern], CheckFunc]
|
|
)
|
|
|
|
|
|
def register_patterns(patterns: list[Pattern]):
|
|
"""
|
|
Register patterns which will be used to partition the DataflowBlock into
|
|
subgraphs that are supported by external backends.
|
|
|
|
Parameters
|
|
----------
|
|
patterns: List[Pattern]
|
|
Patterns to be registered. Patterns that appear later in the list have
|
|
higher priority when partitioning DataflowBlock.
|
|
"""
|
|
_ensure_cleanup_function_registered()
|
|
|
|
entries = []
|
|
for item in patterns:
|
|
if isinstance(item, FusionPattern):
|
|
entries.append(item)
|
|
elif isinstance(item, tuple):
|
|
entries.append(FusionPattern(*item))
|
|
_REGISTERED_PATTERN_NAMES.add(item[0])
|
|
else:
|
|
raise TypeError(f"Cannot register type {type(item)} as pattern")
|
|
_ffi_api.RegisterPatterns(entries)
|
|
|
|
|
|
def get_patterns_with_prefix(prefix: str) -> list[FusionPattern]:
|
|
"""
|
|
Get a list of patterns whose names startwith `prefix`.
|
|
|
|
Parameters
|
|
----------
|
|
prefix: str
|
|
The prefix of pattern name.
|
|
|
|
Returns
|
|
-------
|
|
patterns: FusionPattern
|
|
Matched patterns, ordered by priority from high to low.
|
|
"""
|
|
return _ffi_api.GetPatternsWithPrefix(prefix)
|
|
|
|
|
|
def get_pattern(name: str) -> FusionPattern | None:
|
|
"""
|
|
Find the pattern with a particular name.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The pattern name.
|
|
|
|
Returns
|
|
-------
|
|
pattern: Optional[FusionPattern]
|
|
The matched pattern. Returns None if such pattern is not found.
|
|
"""
|
|
return _ffi_api.GetPattern(name)
|