wehub-resource-sync
403 lines
16 KiB
Python
403 lines
16 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.
|
|
# ruff: noqa: E722
|
|
"""External modules to be linked into the exported IRModule."""
|
|
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
from tvm import libinfo, tirx
|
|
from tvm.runtime import Module, load_static_library
|
|
from tvm.support import cc as _cc
|
|
|
|
from ...op import call_dps_packed
|
|
from . import core
|
|
from .core import wrap_nested
|
|
|
|
|
|
class ExternModule:
|
|
"""The abstract base class for external modules. External modules are designed to help
|
|
incorporate user-provided handcrafted kernels into the exported TVM IRModule.
|
|
"""
|
|
|
|
_symbols: dict[str, Callable]
|
|
|
|
def __init__(self, symbols: dict[str, Callable]) -> None:
|
|
self._symbols = symbols
|
|
|
|
def __getitem__(self, func_name: str) -> Callable:
|
|
_inference_function = self._symbols[func_name]
|
|
|
|
def _call(*input_args):
|
|
def _convert(arg, name: str):
|
|
from tvm import relax as rx # pylint: disable=import-outside-toplevel
|
|
|
|
if isinstance(arg, core.Tensor):
|
|
return arg._expr # pylint: disable=protected-access
|
|
if isinstance(arg, int):
|
|
return rx.prim_value(tirx.IntImm("int64", arg))
|
|
if isinstance(arg, float):
|
|
return rx.prim_value(tirx.FloatImm("float64", arg))
|
|
if isinstance(arg, str):
|
|
return rx.StringImm(arg)
|
|
if tvm.ir.is_prim_expr(arg):
|
|
return rx.prim_value(arg)
|
|
if isinstance(arg, tuple | list):
|
|
return rx.Tuple([_convert(e, f"{name}_{i}") for i, e in enumerate(arg)])
|
|
raise TypeError(f"Unsupported input type: {type(arg)}")
|
|
|
|
rx_inputs = _convert(input_args, "input")
|
|
rx_outputs_ty = _convert(_inference_function(*input_args), "dummy").ty
|
|
return wrap_nested(call_dps_packed(func_name, rx_inputs, rx_outputs_ty), func_name)
|
|
|
|
return _call
|
|
|
|
def _load(self, path: Path) -> Module:
|
|
return load_static_library(str(path), func_names=list(self._symbols.keys()))
|
|
|
|
def load(self) -> Module:
|
|
"""Loads the external module into a TVM runtime module."""
|
|
raise NotImplementedError
|
|
|
|
|
|
class ObjectModule(ExternModule): # pylint: disable=too-few-public-methods
|
|
"""A subclass of `nn.ExternModule`, which allows
|
|
users to provide an object `.o` file to be linked into compiled
|
|
artifact;
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
symbols: dict[str, Callable],
|
|
filepath: Path,
|
|
) -> None:
|
|
if not isinstance(filepath, Path):
|
|
filepath = Path(filepath)
|
|
if not filepath.is_file():
|
|
raise ValueError(f"Not a file: {filepath!s}")
|
|
self.filepath = filepath
|
|
super().__init__(symbols)
|
|
|
|
def load(self) -> Module:
|
|
return self._load(self.filepath)
|
|
|
|
|
|
class SourceModule(ExternModule): # pylint: disable=too-few-public-methods
|
|
"""A subclass of `nn.ExternModule`. It compiles C++/CUDA source code and link them into the
|
|
eventual IRModule.
|
|
|
|
**Shape/dtype inference.** The `nn.ExternModule` system requires users to provide additional
|
|
information to work, namely, `symbols`. It is a dictionary that maps each symbol in the
|
|
external object file to its shape/dtype inference function. Consider a case where function
|
|
`my_func` accepts two tensors, `a` of shape `(x, y, 1)`, and `b` of shape `(y, z, 5)`, and
|
|
produces a tensor `c` of shape `(x, y, z, 9)`, the shape/dtype inference function should look
|
|
like:
|
|
|
|
.. code-block:: python
|
|
|
|
def shape_dtype_inference(a, b):
|
|
x, y, _ = a.shape
|
|
_, z, _ = b.shape
|
|
return nn.Tensor.placeholder((x, y, z, 9), dtype="float32")
|
|
|
|
|
|
and the `symbols` dictionary should be provided as:
|
|
|
|
.. code-block:: python
|
|
|
|
symbols={
|
|
"my_func": shape_dtype_inference,
|
|
}
|
|
|
|
|
|
**Calling convention.** All external modules now follows "destination-passing-style" (DPS)
|
|
calling convention, which means the returned tensors are pre-allocated by the system already
|
|
and passed in as an argument of the external function.
|
|
|
|
Reuse the example above, the implementation of `my_func` should include three parameters in
|
|
its signature, where tensors are represented using DLTensor from DLPack, the de facto standard
|
|
of in-memory representation of tensors. More details:
|
|
https://github.com/dmlc/dlpack/blob/v0.8/include/dlpack/dlpack.h#L163-L206.
|
|
|
|
To expose the symbol, `TVM_FFI_DLL_EXPORT_TYPED_FUNC(symbol, function)` is guaranteed available:
|
|
|
|
.. code-block:: C++
|
|
|
|
// those headers are guaranteed to be available
|
|
#include <dlpack/dlpack.h>
|
|
#include <tvm/ffi/dtype.h>
|
|
#include <tvm/ffi/function.h>
|
|
|
|
namespace {
|
|
// anonymous namespace hides the symbol `_my_func_impl` from other translation units
|
|
int _my_func_impl(DLTensor* a, DLTensor* b, DLTensor* c) {
|
|
// `a` and `b` are inputs, and `c` is the output
|
|
}
|
|
}
|
|
// expose symbol `my_func` instead of `_my_func_impl`
|
|
TVM_FFI_DLL_EXPORT_TYPED_FUNC(my_func, _my_func_impl);
|
|
|
|
**A compiler pass `AttachExternModules`.** It is introduced to attach a list of
|
|
`nn.ExternModule`s into an IRModule at any stage of the compilation pipeline,
|
|
and attach the compiled external modules as `runtime.Module`s into IRModule's `external_mods`
|
|
attribute. It is required by linking in `tvm.compile`, but with the existence of this pass,
|
|
source compilation can be deferred to arbitrary stage of TVM compilation.
|
|
|
|
**Caveats.** It is required to call `nn.add_extern` to register external modules exactly once
|
|
during `export_tvm`. Each symbol should be registered exactly once to avoid potential conflicts,
|
|
and otherwise an error will be raised.
|
|
"""
|
|
|
|
def __init__( # pylint: disable=too-many-arguments
|
|
self,
|
|
symbols: dict[str, Callable],
|
|
source_code: str | Path,
|
|
source_format: str, # "cpp", "cu"
|
|
compile_options: list[str] | None = None,
|
|
compiler: str | None = None,
|
|
output_format: str = "obj", # "obj", "wasm"
|
|
):
|
|
"""Constructs a `nn.SourceModule` from source code.
|
|
|
|
Parameters
|
|
----------
|
|
symbols : Dict[str, Callable]
|
|
The dictionary that maps each symbol in the external object file to its shape/dtype
|
|
inference function.
|
|
|
|
source_code : Union[str, Path]
|
|
Source code or path to the source code to be compiled.
|
|
|
|
source_format : str
|
|
The source code format. It can be either "cpp" or "cu".
|
|
|
|
compile_options : Optional[List[str]]
|
|
The compile options. If not provided, the default compile options will be used.
|
|
|
|
compiler : Optional[str]
|
|
The compiler. If not provided, the default compiler will be used. On Windows,
|
|
compilation requires `clang` by default.
|
|
|
|
output_format : str
|
|
The output format. It can be either "obj" or "wasm". "obj" is the default format,
|
|
which is a shared object file. "wasm" is the WebAssembly format, which is a binary
|
|
file.
|
|
"""
|
|
|
|
def _detect_input_suffix(source_format: str) -> str:
|
|
if source_format == "cpp":
|
|
return ".cpp"
|
|
if source_format == "cu":
|
|
return ".cu"
|
|
raise ValueError(f"Invalid source format: {source_format}")
|
|
|
|
def _detect_output_suffix(output_format: str) -> str:
|
|
if output_format == "obj":
|
|
if _cc._is_linux_like(): # pylint: disable=protected-access
|
|
return ".o"
|
|
if _cc._is_windows_like(): # pylint: disable=protected-access
|
|
return ".obj"
|
|
raise ValueError(f"Unsupported platform: {sys.platform}")
|
|
if output_format == "wasm":
|
|
return ".wasm"
|
|
raise ValueError(f"Invalid output format: {output_format}")
|
|
|
|
def _detect_source_code(source_code) -> str:
|
|
if isinstance(source_code, Path):
|
|
path = source_code
|
|
if not path.is_file():
|
|
raise ValueError(f"Not a file: {path!s}")
|
|
else:
|
|
try:
|
|
path = Path(source_code)
|
|
except: # pylint: disable=bare-except
|
|
return source_code
|
|
try:
|
|
if not path.is_file():
|
|
return source_code
|
|
except: # pylint: disable=bare-except
|
|
return source_code
|
|
with path.open("r", encoding="utf-8") as file:
|
|
return file.read()
|
|
|
|
self.source_code = _detect_source_code(source_code)
|
|
if compile_options is None:
|
|
self.compile_options = SourceModule.get_compile_options(source_format=source_format)
|
|
else:
|
|
self.compile_options = list(compile_options)
|
|
self.compiler = compiler
|
|
self.source_suffix = _detect_input_suffix(source_format)
|
|
self.output_suffix = _detect_output_suffix(output_format)
|
|
super().__init__(symbols)
|
|
|
|
@staticmethod
|
|
def tvm_home() -> Path:
|
|
"""Find TVM's home directory. If `TVM_HOME` environment variable is set, use it.
|
|
Otherwise, use the directory where the `tvm` Python package is installed.
|
|
As a sanity check, it is required to have `include` and `3rdparty` as direct subdirectories.
|
|
|
|
Returns
|
|
-------
|
|
tvm_home : pathlib.Path
|
|
The TVM home directory, and it is guaranteed to have `include` and `3rdparty` as
|
|
direct subdirectories.
|
|
"""
|
|
if os.environ.get("TVM_HOME", None):
|
|
tvm_path = Path(os.environ["TVM_HOME"])
|
|
assert tvm_path.exists(), (
|
|
f"Using environment variable `TVM_HOME`, but directory not found: {tvm_path!s}"
|
|
)
|
|
assert tvm_path.is_dir(), (
|
|
f"Using environment variable `TVM_HOME`, but it is not a directory: {tvm_path!s}"
|
|
)
|
|
else:
|
|
import tvm # pylint: disable=import-outside-toplevel
|
|
|
|
tvm_path = Path(tvm.__file__).parent
|
|
assert tvm_path.is_dir()
|
|
tvm_path = tvm_path.resolve()
|
|
while True:
|
|
exists_include = (tvm_path / "include").is_dir()
|
|
exists_3rdparty = (tvm_path / "3rdparty").is_dir()
|
|
if exists_include and exists_3rdparty:
|
|
return tvm_path.resolve()
|
|
parent = tvm_path.parent
|
|
if parent == tvm_path:
|
|
raise ValueError(
|
|
"Cannot detect TVM directory. "
|
|
"Please explicitly specify it by setting `TVM_HOME` environment variable, "
|
|
"and make sure it contains `include` and `3rdparty` as direct sub-directories."
|
|
)
|
|
tvm_path = parent
|
|
return tvm_path.resolve()
|
|
|
|
@staticmethod
|
|
def get_includes(tvm_pkg: list[str] | None = None) -> list[Path]:
|
|
"""Returns the default include paths according to `tvm_home()`.
|
|
By default, it includes TVM, DLPack. With `tvm_pkg` provided, it also
|
|
includes the specified package under `tvm_home/3rdparty`.
|
|
|
|
Parameters
|
|
----------
|
|
tvm_pkg : Optional[List[str]]
|
|
The list of packages to be included under `tvm_home/3rdparty`. Each element should be
|
|
a relative path to `tvm_home/3rdparty`.
|
|
|
|
Returns
|
|
-------
|
|
includes : List[pathlib.Path]
|
|
The list of include paths.
|
|
"""
|
|
results = [
|
|
Path(libinfo.find_include_path()),
|
|
Path(tvm_ffi.libinfo.find_include_path()),
|
|
Path(tvm_ffi.libinfo.find_dlpack_include_path()),
|
|
]
|
|
if tvm_pkg:
|
|
tvm_home = SourceModule.tvm_home()
|
|
for relative in tvm_pkg:
|
|
results.append(tvm_home / "3rdparty" / relative)
|
|
results = list(dict.fromkeys(results))
|
|
for path in results:
|
|
assert path.exists(), f"Not found: {path!s}"
|
|
assert path.is_dir(), f"Not a directory: {path!s}"
|
|
return results
|
|
|
|
@staticmethod
|
|
def get_compile_options(
|
|
source_format: str,
|
|
tvm_pkg: list[str] | None = None,
|
|
) -> list[str]:
|
|
"""Returns the default compile options depending on `source_format`, including the default
|
|
inlcude paths w.r.t. `tvm_home()`, and by default,
|
|
it uses "-O3" and "-std=c++17".
|
|
|
|
Parameters
|
|
----------
|
|
source_format : str
|
|
The source code format. It can be either "cpp" or "cu".
|
|
|
|
tvm_pkg : Optional[List[str]]
|
|
The list of packages to be included under `tvm_home/3rdparty`. Each element should be
|
|
a relative path to `tvm_home/3rdparty`.
|
|
|
|
Returns
|
|
-------
|
|
compile_options : List[str]
|
|
The list of compilation flags.
|
|
"""
|
|
include_flags = []
|
|
for include_path in SourceModule.get_includes(tvm_pkg=tvm_pkg):
|
|
include_flags += ["-I", str(include_path)]
|
|
if source_format == "cpp":
|
|
host_flags = [
|
|
"-c", # generate object file
|
|
"-O3",
|
|
"-std=c++17",
|
|
]
|
|
elif source_format == "cu":
|
|
host_flags = [
|
|
"-c", # generate object file
|
|
"-O3",
|
|
"-std=c++17",
|
|
# Enable `-fPIC` for the host compiler
|
|
"-Xcompiler=-fPIC",
|
|
]
|
|
else:
|
|
raise ValueError(f"Invalid source format: {source_format}")
|
|
return include_flags + host_flags
|
|
|
|
def compile(self, output_path: Path) -> None:
|
|
"""Compiles the source code in a provided directory and returns the compiled artifact."""
|
|
with tempfile.TemporaryDirectory() as temp_dir_str:
|
|
temp_dir = Path(temp_dir_str)
|
|
source_filename = f"main{self.source_suffix}"
|
|
object_filename = f"main{self.output_suffix}"
|
|
source_path = temp_dir / source_filename
|
|
object_path = temp_dir / object_filename
|
|
with source_path.open("w", encoding="utf-8") as file:
|
|
file.write(self.source_code)
|
|
_cc.create_shared(
|
|
output=object_filename,
|
|
objects=[source_filename],
|
|
options=self.compile_options,
|
|
cc=self.compiler,
|
|
cwd=temp_dir,
|
|
ccache_env=(
|
|
{
|
|
"CCACHE_COMPILERCHECK": "content",
|
|
"CCACHE_NOHASHDIR": "1",
|
|
}
|
|
if shutil.which("ccache")
|
|
else None
|
|
),
|
|
)
|
|
shutil.move(str(object_path), str(output_path))
|
|
|
|
def load(self) -> Module:
|
|
with tempfile.TemporaryDirectory() as temp_dir_str:
|
|
output_path = Path(temp_dir_str) / f"main{self.output_suffix}"
|
|
self.compile(output_path)
|
|
return self._load(output_path)
|