wehub-resource-sync
137 lines
4.9 KiB
Python
137 lines
4.9 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.
|
|
"""PyTorch integration with nn.Module"""
|
|
|
|
import inspect
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import torch
|
|
from tvm_ffi import Array, Shape
|
|
|
|
from tvm.runtime import Tensor, _tensor
|
|
from tvm.runtime.vm import VirtualMachine
|
|
|
|
from . import core
|
|
from . import spec as _spec
|
|
|
|
|
|
class TorchModule: # pylint: disable=too-few-public-methods
|
|
"""A wrapper on top of TVM VirtualMachine that takes torch tensors as inputs and returns torch
|
|
tensors as outputs"""
|
|
|
|
spec: _spec.ModuleSpec
|
|
vm: VirtualMachine # pylint: disable=invalid-name
|
|
params: list[Tensor]
|
|
effects: list[Any]
|
|
|
|
def __init__( # pylint: disable=invalid-name
|
|
self,
|
|
spec: _spec.ModuleSpec,
|
|
vm: VirtualMachine,
|
|
params: list[Tensor],
|
|
):
|
|
try:
|
|
self.effects = vm["_initialize_effect"]()
|
|
except AttributeError:
|
|
self.effects = None
|
|
|
|
self.spec = spec
|
|
self.vm = vm
|
|
self.params = params
|
|
|
|
def __getitem__(self, method_name: str) -> Callable:
|
|
def _find_method(method_name):
|
|
for key, value in zip(self.spec.method_names, self.spec.method_specs):
|
|
if method_name == key:
|
|
return value
|
|
raise ValueError(f"Method `{method_name}` is not found in the module spec. {self.spec}")
|
|
|
|
method_spec = _find_method(method_name)
|
|
method = self.vm[method_name]
|
|
|
|
def _closure(*args):
|
|
if len(args) != len(method_spec.arg_names):
|
|
raise TypeError(
|
|
f"Argument length mismatch. Expected {len(method_spec.arg_names)} arguments, "
|
|
f"but got {len(args)} arguments. The spec is: {method_spec}"
|
|
)
|
|
args = [
|
|
_torch_to_tvm(arg_name, arg_spec, arg)
|
|
for arg_name, arg_spec, arg in zip(
|
|
method_spec.arg_names, method_spec.arg_specs, args
|
|
)
|
|
]
|
|
if self.effects is not None:
|
|
outputs, self.effects = method(*args, *self.effects, *self.params)
|
|
else:
|
|
outputs = method(*args, *self.params)
|
|
return _tvm_to_torch(outputs)
|
|
|
|
_closure.__name__ = method_name
|
|
return _closure
|
|
|
|
|
|
def _tvm_to_torch(arg):
|
|
if isinstance(arg, list | tuple | Array):
|
|
return [_tvm_to_torch(i) for i in arg]
|
|
if isinstance(arg, _tensor.Tensor):
|
|
return torch.utils.dlpack.from_dlpack(arg)
|
|
if isinstance(arg, Shape):
|
|
return list(arg)
|
|
raise TypeError(f"Unsupported argument type: {type(arg)}")
|
|
|
|
|
|
def _torch_to_tvm(arg_name, arg_spec, arg_torch):
|
|
if isinstance(arg_spec, _spec.Tensor):
|
|
if not isinstance(arg_torch, torch.Tensor):
|
|
raise TypeError(
|
|
f"Expected argument `{arg_name}` to be `torch.Tensor`, but got {type(arg_torch)}"
|
|
)
|
|
return core._from_dlpack(arg_torch) # pylint: disable=protected-access
|
|
if isinstance(arg_spec, _spec.Int):
|
|
if not isinstance(arg_torch, int):
|
|
raise TypeError(
|
|
f"Expected argument `{arg_name}` to be `int`, but got {type(arg_torch)}"
|
|
)
|
|
return Shape([arg_torch])
|
|
if isinstance(arg_spec, _spec.Tuple):
|
|
return [
|
|
_torch_to_tvm(f"{arg_name}[{i}]", x, arg_torch[i])
|
|
for i, x in enumerate(arg_spec.elements)
|
|
]
|
|
raise TypeError(f"Unsupported spec item type: {type(arg_spec)}")
|
|
|
|
|
|
def _method_spec_from_torch(
|
|
args_torch: list[Any],
|
|
method: Callable,
|
|
):
|
|
def _as_spec(arg_torch):
|
|
if isinstance(arg_torch, torch.Tensor):
|
|
_, dtype = str(arg_torch.dtype).rsplit(".", maxsplit=1)
|
|
return _spec.Tensor(shape=list(arg_torch.shape), dtype=dtype)
|
|
if isinstance(arg_torch, int):
|
|
return _spec.Int()
|
|
raise TypeError(f"Unsupported argument type: {type(arg_torch)}")
|
|
|
|
arg_names = list(inspect.signature(method).parameters.keys())
|
|
if len(arg_names) != len(args_torch):
|
|
raise TypeError(f"Expected {len(arg_names)} arguments, but got {len(args_torch)} arguments")
|
|
arg_specs = [_as_spec(i) for i in args_torch]
|
|
return _spec.MethodSpec(method, arg_names, arg_specs, param_mode="plain", effect_mode="plain")
|