wehub-resource-sync
125 lines
3.7 KiB
Python
125 lines
3.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.
|
|
"""The argument information"""
|
|
|
|
from typing import Any
|
|
|
|
from tvm_ffi import Shape, register_object
|
|
|
|
from tvm.ir import IRModule
|
|
from tvm.runtime import DataType, Object
|
|
from tvm.tirx import PrimFunc
|
|
|
|
from . import _ffi_api
|
|
from .utils import _json_de_tvm
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.ArgInfo")
|
|
class ArgInfo(Object):
|
|
"""Argument information"""
|
|
|
|
def as_json(self) -> Any:
|
|
"""Converts the ArgInfo to its corresponding JSON representation."""
|
|
return _json_de_tvm(_ffi_api.ArgInfoAsJSON(self)) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def from_json(json_obj: Any) -> "ArgInfo":
|
|
"""Parse the argument information from a JSON object.
|
|
|
|
Parameters
|
|
----------
|
|
json_obj : Any
|
|
The json object to parse.
|
|
|
|
Returns
|
|
-------
|
|
parsed : ArgInfo
|
|
The argument information parsed.
|
|
"""
|
|
return _ffi_api.ArgInfoFromJSON(json_obj) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def from_prim_func(func: PrimFunc) -> list["ArgInfo"]:
|
|
"""Extract a list of the argument information from PrimFunc.
|
|
|
|
Parameters
|
|
----------
|
|
func : PrimFunc
|
|
The PrimFunc to get argument information from.
|
|
|
|
Returns
|
|
-------
|
|
extracted : List[ArgInfo]
|
|
An array of the argument information derived.
|
|
"""
|
|
return _ffi_api.ArgInfoFromPrimFunc(func) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def from_entry_func(mod: IRModule, remove_preproc: bool = True) -> list["ArgInfo"]:
|
|
"""Extract a list of the argument information from the entry func of an IRModule.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to get argument information from.
|
|
remove_preproc : bool
|
|
Whether to remove the preprocessing blocks.
|
|
|
|
Returns
|
|
-------
|
|
extracted : List[ArgInfo]
|
|
An array of the argument information derived.
|
|
"""
|
|
return _ffi_api.ArgInfoFromEntryFunc(mod, remove_preproc) # type: ignore # pylint: disable=no-member
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.TensorInfo")
|
|
class TensorInfo(ArgInfo):
|
|
"""Tensor argument information
|
|
|
|
Parameters
|
|
----------
|
|
dtype : DataType
|
|
The data type of the tensor.
|
|
shape : Shape
|
|
The shape of the tensor.
|
|
"""
|
|
|
|
dtype: DataType
|
|
shape: Shape
|
|
|
|
def __init__(
|
|
self,
|
|
dtype: DataType,
|
|
shape: Shape | list[int],
|
|
) -> None:
|
|
"""Constructor
|
|
|
|
Parameters
|
|
----------
|
|
dtype : DataType
|
|
The data type of the tensor.
|
|
shape : Shape
|
|
The shape of the tensor.
|
|
"""
|
|
shape_tuple = shape if isinstance(shape, Shape) else Shape(shape)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.TensorInfo, # type: ignore # pylint: disable=no-member
|
|
dtype,
|
|
shape_tuple,
|
|
)
|