wehub-resource-sync
258 lines
8.3 KiB
Python
258 lines
8.3 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.
|
|
"""Compilation specifications, for example, dynamic shape inputs."""
|
|
|
|
import inspect
|
|
import typing
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from .core import Module as nn_module_class
|
|
|
|
ArgSpecType = typing.Union["Int", "Tensor"]
|
|
MethodSpecType = typing.Union["MethodSpec", dict[str, ArgSpecType]]
|
|
ModuleSpecType = typing.Union["ModuleSpec", dict[str, MethodSpecType]]
|
|
SpecAny = typing.Union["Object", "Int", "Tensor", "Tuple"]
|
|
|
|
|
|
class Int: # pylint: disable=too-few-public-methods
|
|
"""An integer input"""
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
def __repr__(self) -> str:
|
|
return "int"
|
|
|
|
|
|
class Tensor: # pylint: disable=too-few-public-methods
|
|
"""A tensor input with static ndim and dtype, but can have symbolic shapes."""
|
|
|
|
shape: list[int | str]
|
|
dtype: str
|
|
|
|
def __init__(self, shape: typing.Sequence[int | str], dtype: str) -> None:
|
|
self.shape = list(shape)
|
|
self.dtype = dtype
|
|
|
|
def __repr__(self) -> str:
|
|
shape = ", ".join(str(i) for i in self.shape)
|
|
return f"Tensor([{shape}], '{self.dtype}')"
|
|
|
|
|
|
class Object: # pylint: disable=too-few-public-methods
|
|
"""An non-tensor opaque frontend object."""
|
|
|
|
object_type: type
|
|
|
|
def __init__(self, object_type: type) -> None:
|
|
self.object_type = object_type
|
|
|
|
def __repr__(self) -> str:
|
|
return "object"
|
|
|
|
|
|
class Tuple: # pylint: disable=too-few-public-methods
|
|
"""A tuple input or a list input"""
|
|
|
|
name: str
|
|
elements: list[SpecAny] | tuple[SpecAny, ...]
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
elements: list[SpecAny] | tuple[SpecAny, ...],
|
|
) -> None:
|
|
assert isinstance(elements, tuple | list), f"Unsupported container type: {type(elements)}"
|
|
self.name = name
|
|
self.elements = elements
|
|
|
|
def __repr__(self) -> str:
|
|
return self.elements.__repr__()
|
|
|
|
|
|
class MethodSpec:
|
|
"""A spec for a compiled method"""
|
|
|
|
method: typing.Callable
|
|
arg_names: list[str]
|
|
arg_specs: list[ArgSpecType]
|
|
param_mode: str # "plain", "packed", "none"
|
|
effect_mode: str # "plain", "packed", "none"
|
|
|
|
def __init__( # pylint: disable=too-many-arguments
|
|
self,
|
|
method: typing.Callable,
|
|
arg_names: list[str],
|
|
arg_specs: list[ArgSpecType],
|
|
param_mode: str,
|
|
effect_mode: str,
|
|
):
|
|
if param_mode not in ["plain", "packed", "none"]:
|
|
raise ValueError(f"Invalid param_mode: {param_mode}")
|
|
if effect_mode not in ["plain", "packed", "none"]:
|
|
raise ValueError(f"Invalid effect_mode: {effect_mode}")
|
|
self.method = method
|
|
self.arg_names = arg_names
|
|
self.arg_specs = arg_specs
|
|
self.param_mode = param_mode
|
|
self.effect_mode = effect_mode
|
|
|
|
def _repr(self, name: str) -> str:
|
|
args = ", ".join(
|
|
f"{name}: {spec}"
|
|
for name, spec in zip(
|
|
self.arg_names,
|
|
self.arg_specs,
|
|
)
|
|
)
|
|
return f"{name}({args})"
|
|
|
|
def __repr__(self) -> str:
|
|
return self._repr(name="MethodSpec")
|
|
|
|
@staticmethod
|
|
def from_raw(spec: MethodSpecType, method: typing.Callable) -> "MethodSpec":
|
|
"""Create MethodSpec from raw python dictionaries.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
MethodSpec.from_raw(
|
|
spec={
|
|
"inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
|
|
"total_seq_len": "int",
|
|
},
|
|
method=module.prefill,
|
|
)
|
|
"""
|
|
if isinstance(spec, MethodSpec):
|
|
return spec
|
|
config: dict[str, typing.Any] = spec.pop("$", {}) # type: ignore[assignment]
|
|
param_mode = config.get("param_mode", "plain")
|
|
effect_mode = config.get("effect_mode", "plain")
|
|
method_signature = inspect.signature(method)
|
|
arg_names = list(method_signature.parameters.keys())
|
|
arg_specs = []
|
|
|
|
def _convert_arg_spec(arg_spec, arg_name):
|
|
if arg_spec is Int or arg_spec is int:
|
|
return Int()
|
|
if isinstance(arg_spec, str) and arg_spec == "int":
|
|
return Int()
|
|
if isinstance(arg_spec, Int | Tensor | Object):
|
|
return arg_spec
|
|
if isinstance(arg_spec, tuple | list | Tuple):
|
|
return Tuple(
|
|
arg_name,
|
|
elements=type(arg_spec)(
|
|
[
|
|
_convert_arg_spec(arg_spec_i, f"{arg_name}_{i}")
|
|
for i, arg_spec_i in enumerate(arg_spec)
|
|
]
|
|
),
|
|
)
|
|
raise TypeError(f"Invalid spec for argument {arg_name}: {arg_spec}")
|
|
|
|
for arg_name in arg_names:
|
|
if arg_name in spec:
|
|
arg_spec = spec[arg_name]
|
|
arg_spec = _convert_arg_spec(arg_spec, arg_name)
|
|
arg_specs.append(arg_spec)
|
|
return MethodSpec(
|
|
method,
|
|
arg_names,
|
|
arg_specs,
|
|
param_mode=param_mode,
|
|
effect_mode=effect_mode,
|
|
)
|
|
|
|
@staticmethod
|
|
def from_torch(args: list[typing.Any], method: typing.Callable) -> "MethodSpec":
|
|
"""Converts a list of torch tensors to MethodSpec."""
|
|
from .torch import ( # pylint: disable=import-outside-toplevel
|
|
_method_spec_from_torch,
|
|
)
|
|
|
|
return _method_spec_from_torch(args, method)
|
|
|
|
|
|
class ModuleSpec:
|
|
"""A spec for a compiled nn.Module"""
|
|
|
|
module: "nn_module_class"
|
|
method_names: list[str]
|
|
method_specs: list[MethodSpec]
|
|
|
|
def __init__(
|
|
self,
|
|
module: "nn_module_class",
|
|
method_names: list[str],
|
|
method_specs: list[MethodSpec],
|
|
) -> None:
|
|
self.module = module
|
|
self.method_names = method_names
|
|
self.method_specs = method_specs
|
|
|
|
@staticmethod
|
|
def from_raw(spec: ModuleSpecType, module: "nn_module_class") -> "ModuleSpec":
|
|
"""Create ModuleSpec from raw python dictionaries.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
ModuleSpec.from_raw(
|
|
spec={
|
|
"prefill": {
|
|
"inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
|
|
"total_seq_len": int,
|
|
},
|
|
"decode": {
|
|
"inputs": spec.Tensor([batch_size, 1], "int32"),
|
|
"total_seq_len": int,
|
|
},
|
|
"softmax_with_temperature": {
|
|
"logits": spec.Tensor([1, 1, config.vocab_size], "float32"),
|
|
"temperature": spec.Tensor([], "float32"),
|
|
},
|
|
},
|
|
module=module,
|
|
)
|
|
"""
|
|
if isinstance(spec, ModuleSpec):
|
|
return spec
|
|
method_names = list(spec.keys())
|
|
method_specs: list[MethodSpec] = []
|
|
for method_name in method_names:
|
|
method_spec = spec[method_name]
|
|
if isinstance(method_spec, MethodSpec):
|
|
pass
|
|
else:
|
|
method_spec = MethodSpec.from_raw(method_spec, getattr(module, method_name))
|
|
method_specs.append(method_spec)
|
|
return ModuleSpec(module, method_names, method_specs)
|
|
|
|
def __repr__(self) -> str:
|
|
return "ModuleSpec:\n" + "\n".join(
|
|
" " + spec._repr(name) # pylint: disable=protected-access
|
|
for name, spec in zip(
|
|
self.method_names,
|
|
self.method_specs,
|
|
)
|
|
)
|