wehub-resource-sync
170 lines
4.6 KiB
Python
170 lines
4.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.
|
|
"""TVM Attribute module, which is mainly used for defining attributes of operators."""
|
|
|
|
import tvm_ffi
|
|
import tvm_ffi._ffi_api as _tvm_ffi_api
|
|
|
|
from tvm.runtime import Object
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
@tvm_ffi.register_object("ir.Attrs")
|
|
class Attrs(Object):
|
|
"""Attribute node, which is mainly use for defining attributes of operators.
|
|
|
|
Used by function registered in python side, such as compute, schedule and alter_layout.
|
|
Attrs is passed as the first argument to these functions.
|
|
"""
|
|
|
|
def get_int_tuple(self, key):
|
|
"""Get a python int tuple of a key
|
|
|
|
Parameters
|
|
----------
|
|
key: str
|
|
|
|
Returns
|
|
-------
|
|
value: Tuple of int
|
|
"""
|
|
return tuple(x if isinstance(x, int) else x.value for x in getattr(self, key))
|
|
|
|
def get_int(self, key):
|
|
"""Get a python int value of a key
|
|
|
|
Parameters
|
|
----------
|
|
key: str
|
|
|
|
Returns
|
|
-------
|
|
value: int
|
|
"""
|
|
return getattr(self, key)
|
|
|
|
def get_str(self, key):
|
|
"""Get a python int value of a key
|
|
|
|
Parameters
|
|
----------
|
|
key: str
|
|
|
|
Returns
|
|
-------
|
|
value: int
|
|
"""
|
|
return getattr(self, key)
|
|
|
|
def __getitem__(self, item):
|
|
return getattr(self, item)
|
|
|
|
|
|
@tvm_ffi.register_object("ir.DictAttrs")
|
|
class DictAttrs(Attrs):
|
|
"""Dictionary attributes."""
|
|
|
|
@property
|
|
def __dict__(self):
|
|
"""Return the underlying key-value map as a Python dict.
|
|
|
|
Defined explicitly so that tvm_ffi's _add_class_attrs skips registering
|
|
the C++ reflection field named '__dict__' (Python forbids adding a class
|
|
attribute named '__dict__' via setattr on extension-type subclasses).
|
|
"""
|
|
return dict(self._dict())
|
|
|
|
def _dict(self):
|
|
"""Get internal dict"""
|
|
return _ffi_api.DictAttrsGetDict(self)
|
|
|
|
def keys(self):
|
|
"""Get list of names in the attribute.
|
|
|
|
Returns
|
|
-------
|
|
keys : list of str
|
|
List of keys
|
|
"""
|
|
return [k for k, _ in self.items()]
|
|
|
|
def __getitem__(self, k):
|
|
return self._dict().__getitem__(k)
|
|
|
|
def get(self, key, default=None):
|
|
"""Get an element with a default value."""
|
|
return self._dict().get(key, default)
|
|
|
|
def __contains__(self, k):
|
|
return self._dict().__contains__(k)
|
|
|
|
def __getattr__(self, name):
|
|
try:
|
|
return self._dict().__getitem__(name)
|
|
except KeyError:
|
|
raise AttributeError(f"DictAttrs has no attribute {name}")
|
|
|
|
def items(self):
|
|
"""Get items from the map."""
|
|
return self._dict().items()
|
|
|
|
def __len__(self):
|
|
return self._dict().__len__()
|
|
|
|
|
|
def make_node(type_key, **kwargs):
|
|
"""Make a new IR node by its type key and fields
|
|
|
|
Parameters
|
|
----------
|
|
type_key : str
|
|
The type key of the node.
|
|
|
|
**kwargs : dict
|
|
The fields of the node.
|
|
|
|
Returns
|
|
-------
|
|
node : Node
|
|
The corresponding IR Node
|
|
|
|
Note
|
|
----
|
|
If the created node is instance of AttrsNode, then
|
|
the creator function will also run bound checks and
|
|
default value setup as supported by Attrs.
|
|
|
|
Example
|
|
-------
|
|
The following code constructs a IntImm object
|
|
|
|
.. code-block:: python
|
|
|
|
x = tvm.ir.make_node("ir.IntImm", dtype="int32", value=10, span=None)
|
|
assert isinstance(x, tvm.tirx.IntImm)
|
|
assert x.value == 10
|
|
"""
|
|
if type_key == "ir.DictAttrs":
|
|
# DictAttrs stores kwargs as a key-value dict, not as named fields.
|
|
# MakeObjectFromPackedArgs would look for a field named "__dict__".
|
|
return _tvm_ffi_api.MakeObjectFromPackedArgs("ir.DictAttrs", "__dict__", kwargs)
|
|
args = [type_key]
|
|
for k, v in kwargs.items():
|
|
args += [k, v]
|
|
return _tvm_ffi_api.MakeObjectFromPackedArgs(*args)
|