chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,19 @@
|
||||
# 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.
|
||||
"""Hexagon APIs."""
|
||||
|
||||
from .tools import *
|
||||
@@ -0,0 +1,55 @@
|
||||
# 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.
|
||||
|
||||
"""Hexagon environment checks for CI usage
|
||||
|
||||
These are required by tvm.testing and live here to avoid a circular dependency.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import tvm
|
||||
|
||||
ANDROID_SERIAL_NUMBER = "ANDROID_SERIAL_NUMBER"
|
||||
HEXAGON_TOOLCHAIN = "HEXAGON_TOOLCHAIN"
|
||||
|
||||
|
||||
def _compile_time_check():
|
||||
"""Return True if compile-time support for Hexagon is present, otherwise
|
||||
error string.
|
||||
|
||||
Backs :func:`tvm.testing.env.has_hexagon_toolchain`.
|
||||
"""
|
||||
if tvm.runtime.enabled("llvm") and tvm.target.codegen.llvm_version_major() < 7:
|
||||
return "Hexagon requires LLVM 7 or later"
|
||||
|
||||
if "HEXAGON_TOOLCHAIN" not in os.environ:
|
||||
return f"Missing environment variable {HEXAGON_TOOLCHAIN}."
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _run_time_check():
|
||||
"""Return True if run-time support for Hexagon is present, otherwise
|
||||
error string.
|
||||
|
||||
Backs :func:`tvm.testing.env.has_hexagon`.
|
||||
"""
|
||||
if ANDROID_SERIAL_NUMBER not in os.environ:
|
||||
return f"Missing environment variable {ANDROID_SERIAL_NUMBER}."
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,98 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-docstring, invalid-name, unnecessary-comprehension, unused-argument
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.contrib.hexagon import hexagon_unary_ops
|
||||
|
||||
|
||||
def op_replace(call_node, func) -> bool:
|
||||
if not isinstance(call_node, relax.Call):
|
||||
return False
|
||||
call_tir_op = tvm.ir.Op.get("relax.call_tir")
|
||||
if call_node.op != call_tir_op:
|
||||
return False
|
||||
ops = [
|
||||
"qnn.tanh",
|
||||
"qnn.sqrt",
|
||||
"qnn.rsqrt",
|
||||
"qnn.exp",
|
||||
"qnn.erf",
|
||||
"qnn.sigmoid",
|
||||
"qnn.hardswish",
|
||||
"qnn.log",
|
||||
"qnn.abs",
|
||||
]
|
||||
if func.attrs["op_attrs"]["op_name"] in ops:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@relax.expr_functor.mutator
|
||||
class Tanh2TakeReplace(tvm.relax.PyExprMutator):
|
||||
def __init__(self, mod: tvm.IRModule) -> None:
|
||||
super().__init__(mod)
|
||||
self.mod_ = mod
|
||||
|
||||
def transform(self) -> tvm.IRModule:
|
||||
# Iterate over all the nodes to check for the node replaceable
|
||||
for global_var, func in self.mod_.functions.items():
|
||||
# Skip non-relax functions
|
||||
if not isinstance(func, relax.Function):
|
||||
continue
|
||||
updated_func = self.visit_expr(func)
|
||||
self.builder_.normalize(updated_func)
|
||||
self.builder_.update_func(global_var, updated_func)
|
||||
# At the end of the transformation we return the updated IRModule from the BlockBuilder.
|
||||
return self.builder_.get()
|
||||
|
||||
def visit_call_(self, call_node: relax.Call) -> relax.Call:
|
||||
call_tir_op = tvm.ir.Op.get("relax.call_tir")
|
||||
if call_node.op != call_tir_op:
|
||||
return call_node
|
||||
|
||||
var = call_node.args[0]
|
||||
func = self.mod_[var]
|
||||
|
||||
if call_node.args[1][0].ty.dtype == "uint8":
|
||||
if op_replace(call_node, func):
|
||||
inp, inp_scale, inp_zp, out_scale, out_zp = [x for x in call_node.args[1]]
|
||||
# LUT node creation
|
||||
LUT = hexagon_unary_ops.LUT_generation(
|
||||
inp_scale, inp_zp, out_scale, out_zp, call_node.args[0].name_hint
|
||||
)
|
||||
# Take operation node creation
|
||||
take_func = hexagon_unary_ops.generate_take_primfunc(inp, call_node.ty)
|
||||
take_func = take_func.without_attr("global_symbol")
|
||||
take_func_gv = self.builder_.add_func(take_func, "take")
|
||||
take_node = relax.call_tir(
|
||||
take_func_gv,
|
||||
relax.expr.Tuple(
|
||||
[call_node.args[1][0], relax.expr.Constant(tvm.runtime.tensor(LUT))]
|
||||
),
|
||||
call_node.ty,
|
||||
)
|
||||
return take_node
|
||||
return call_node
|
||||
|
||||
|
||||
@tvm.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
|
||||
class PassReplaceWithTakeOpPrimFuncs:
|
||||
def transform_module(self, mod, ctx):
|
||||
return Tanh2TakeReplace(mod).transform()
|
||||
@@ -0,0 +1,99 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-docstring, invalid-name
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
from scipy import special
|
||||
|
||||
from tvm import te
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
######################################################################
|
||||
#################### PRIMFUNC FOR LUT and Take Op ####################
|
||||
######################################################################
|
||||
|
||||
|
||||
def saturate(x: te.Tensor, dtype: str):
|
||||
"""Saturate value for the specified data type"""
|
||||
return te.max(te.min_value(dtype), te.min(x, te.max_value(dtype)))
|
||||
|
||||
|
||||
def hardswish_func(x):
|
||||
x_2 = np.add(x, 3.0)
|
||||
x_2 = np.clip(x_2, 0.0, 6.0)
|
||||
return x * x_2 / 6.0
|
||||
|
||||
|
||||
def LUT_generation(inp_scale, inp_zp, out_scale, out_zp, op_name) -> None:
|
||||
LUT = []
|
||||
for i in range(256):
|
||||
i = np.int32(i)
|
||||
# converting the constants to the numpy value
|
||||
if inp_zp.data.shape == ():
|
||||
i_zp = inp_zp.data.numpy()[()]
|
||||
if inp_scale.data.shape == ():
|
||||
i_scale = inp_scale.data.numpy()[()]
|
||||
if out_zp.data.shape == ():
|
||||
o_zp = out_zp.data.numpy()[()]
|
||||
if out_scale.data.shape == ():
|
||||
o_scale = out_scale.data.numpy()[()]
|
||||
# Dequantization followed by computing the op value
|
||||
dequant = (i - i_zp) * i_scale
|
||||
if "tanh" in op_name:
|
||||
op_val = np.tanh(dequant)
|
||||
elif "rsqrt" in op_name:
|
||||
op_val = 1 / np.sqrt(dequant)
|
||||
elif "sqrt" in op_name:
|
||||
op_val = np.sqrt(dequant)
|
||||
elif "exp" in op_name:
|
||||
op_val = np.exp(dequant)
|
||||
elif "erf" in op_name:
|
||||
op_val = special.erf(dequant)
|
||||
elif "sigmoid" in op_name:
|
||||
op_val = 1 / (1 + np.exp(np.negative(dequant)))
|
||||
elif "hardswish" in op_name:
|
||||
op_val = hardswish_func(dequant)
|
||||
elif "log" in op_name:
|
||||
op_val = np.log(dequant)
|
||||
elif "abs" in op_name:
|
||||
op_val = np.abs(dequant)
|
||||
else:
|
||||
logger.error("Error op is other than unary op")
|
||||
|
||||
# Quantizing the value generated and appending in the Look Up Table
|
||||
quant = np.round((op_val) / o_scale) + o_zp
|
||||
val = np.maximum(0, np.minimum(quant, 255)).astype(np.uint8)
|
||||
LUT.append(val)
|
||||
return LUT
|
||||
|
||||
|
||||
def generate_take_primfunc(inp, ty):
|
||||
# Generating the take op
|
||||
N, H, W, C = inp.ty.shape
|
||||
data = te.placeholder((N, H, W, C), dtype=ty.dtype, name="data")
|
||||
LUT_func = te.placeholder((256,), dtype="uint8", name="LUT")
|
||||
take = te.compute(
|
||||
ty.shape,
|
||||
lambda *indices: saturate((LUT_func[data[indices].astype("uint8")]), ty.dtype).astype(
|
||||
ty.dtype
|
||||
),
|
||||
name="take_op",
|
||||
)
|
||||
mod = te.create_prim_func([data, LUT_func, take])
|
||||
return mod
|
||||
@@ -0,0 +1,561 @@
|
||||
# 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.
|
||||
# pylint: disable=invalid-name, f-string-without-interpolation, consider-using-from-import
|
||||
# ruff: noqa: E501
|
||||
"""Tools/compilers/linkers for Hexagon"""
|
||||
|
||||
import io
|
||||
import itertools
|
||||
import os
|
||||
import pathlib
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tarfile
|
||||
|
||||
import numpy
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
import tvm
|
||||
import tvm.support.cc as cc
|
||||
|
||||
# Linking Hexagon shared libraries.
|
||||
#
|
||||
# link_shared(name-of-shared-library, list-of-objects, kw-args)
|
||||
#
|
||||
# To use a custom linker, define a function that returns the path to the
|
||||
# linker, and pass it to 'register_linker':
|
||||
#
|
||||
# def custom_linker_path():
|
||||
# return '/path/to/hexagon/linker'
|
||||
#
|
||||
# register_linker(custom_linker_path)
|
||||
#
|
||||
# Subsequent calls to 'link_shared' will use the newly registered linker.
|
||||
|
||||
HEXAGON_TOOLCHAIN = os.environ.get("HEXAGON_TOOLCHAIN", default="") # pylint: disable=invalid-name
|
||||
HEXAGON_SDK_ROOT = os.environ.get("HEXAGON_SDK_ROOT", default="") # pylint: disable=invalid-name
|
||||
HEXAGON_SDK_DOCKER_IMAGE = os.environ.get("HEXAGON_SDK_DOCKER_IMAGE", default="") # pylint: disable=invalid-name
|
||||
HEXAGON_LINK_MAIN = pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-link" # pylint: disable=invalid-name
|
||||
HEXAGON_CLANG_PLUS = pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-clang++" # pylint: disable=invalid-name
|
||||
HEXAGON_SDK_INCLUDE_DIRS = [ # pylint: disable=invalid-name
|
||||
pathlib.Path(HEXAGON_SDK_ROOT) / "incs",
|
||||
pathlib.Path(HEXAGON_SDK_ROOT) / "incs" / "stddef",
|
||||
]
|
||||
|
||||
HEXAGON_SIMULATOR_NAME = "simulator"
|
||||
|
||||
|
||||
def register_linker(f):
|
||||
"""Register a function that will return the path to the Hexagon linker."""
|
||||
return register_global_func("tvm.contrib.hexagon.hexagon_link", f, True)
|
||||
|
||||
|
||||
@register_global_func("tvm.contrib.hexagon.hexagon_link")
|
||||
def hexagon_link() -> str:
|
||||
"""Return path to the Hexagon linker."""
|
||||
return str(HEXAGON_LINK_MAIN)
|
||||
|
||||
|
||||
def hexagon_clang_plus() -> str:
|
||||
"""Return path to the Hexagon clang++."""
|
||||
return str(HEXAGON_CLANG_PLUS)
|
||||
|
||||
|
||||
def toolchain_version(toolchain=None) -> list[int]:
|
||||
"""Return the version of the Hexagon toolchain.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
toolchain: str, optional
|
||||
Path to the Hexagon toolchain. If not provided, the environment
|
||||
variable HEXAGON_TOOLCHAIN is used.
|
||||
|
||||
Returns
|
||||
-------
|
||||
version: List[int]
|
||||
List of numerical components of the version number. E.g. for version
|
||||
"8.5.06" it will be [8, 5, 6].
|
||||
"""
|
||||
|
||||
if toolchain is None:
|
||||
toolchain = HEXAGON_TOOLCHAIN
|
||||
assert toolchain is not None, "Please specify toolchain, or set HEXAGON_TOOLCHAIN variable"
|
||||
result = subprocess.run(
|
||||
[f"{toolchain}/bin/hexagon-clang", "-v"], capture_output=True, check=True
|
||||
)
|
||||
output = result.stderr.decode()
|
||||
for line in output.splitlines():
|
||||
m = re.match(r".* [Cc]lang version ([0-9\.]+)", line)
|
||||
if m:
|
||||
assert len(m.groups()) == 1
|
||||
return [int(v) for v in m.group(1).split(".")]
|
||||
raise RuntimeError("Cannot establish toolchain version")
|
||||
|
||||
|
||||
@register_global_func("tvm.contrib.hexagon.link_shared")
|
||||
def link_shared(so_name, objs, extra_args=None):
|
||||
"""Link shared library on Hexagon using the registered Hexagon linker.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
so_name : str
|
||||
Name of the shared library file.
|
||||
objs : list[str, tvm.tirx.StringImm]
|
||||
extra_args : dict (str->str) or Map<String,String>
|
||||
Additional arguments:
|
||||
'hex_arch' - Hexagon architecture, e.g. v68
|
||||
'verbose' - Print additional information if the key is present
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret_val : int
|
||||
This function returns 0 at the moment.
|
||||
"""
|
||||
|
||||
# The list of object files can be passed as built-in Python strings,
|
||||
# or as tvm.tirx.StringImm's.
|
||||
def to_str(s):
|
||||
if isinstance(s, tvm.tirx.StringImm):
|
||||
return s.value
|
||||
assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
|
||||
return s
|
||||
|
||||
objs = [to_str(s) for s in objs]
|
||||
|
||||
if not extra_args:
|
||||
extra_args = {}
|
||||
hex_arch = extra_args.get("hex_arch") or "v68"
|
||||
linker = tvm.get_global_func("tvm.contrib.hexagon.hexagon_link")()
|
||||
if extra_args.get("verbose"):
|
||||
print("tvm.contrib.hexagon.link_shared:")
|
||||
print(" Using linker:", linker)
|
||||
print(" Library name:", so_name)
|
||||
print(" Object files:", objs)
|
||||
print(" Architecture:", hex_arch)
|
||||
if not os.access(linker, os.X_OK):
|
||||
message = 'The linker "' + linker + '" does not exist or is not executable.'
|
||||
if not os.environ.get("HEXAGON_TOOLCHAIN"):
|
||||
message += (
|
||||
" The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
|
||||
+ "HEXAGON_TOOLCHAIN in your environment, so that ${HEXAGON_TOOLCHAIN}/bin/"
|
||||
+ "hexagon-link exists."
|
||||
)
|
||||
else:
|
||||
message += (
|
||||
" Please verify the value of the HEXAGON_LINKER environment variable "
|
||||
+ '(currently set to "'
|
||||
+ HEXAGON_TOOLCHAIN
|
||||
+ '").'
|
||||
)
|
||||
raise Exception(message)
|
||||
|
||||
libpath = os.path.join(HEXAGON_TOOLCHAIN, "target", "hexagon", "lib", hex_arch, "G0")
|
||||
cc.create_shared(
|
||||
so_name,
|
||||
objs,
|
||||
# pylint: disable=bad-whitespace
|
||||
options=[
|
||||
"-Bdynamic",
|
||||
"-shared",
|
||||
"-export-dynamic",
|
||||
os.path.join(libpath, "pic", "libgcc.so"),
|
||||
],
|
||||
cc=linker,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def link_shared_macos(so_name, objs, extra_args=None):
|
||||
"""Link Hexagon shared library using docker container with proper tooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
so_name : str
|
||||
Name of the shared library file.
|
||||
objs : list[str, tvm.tirx.StringImm]
|
||||
extra_args : dict (str->str) or Map<String,String>
|
||||
Additional arguments:
|
||||
'hex_arch' - Hexagon architecture, e.g. v68
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret_val : int
|
||||
This function returns 0 at the moment.
|
||||
"""
|
||||
|
||||
# The list of object files can be passed as built-in Python strings,
|
||||
# or as tvm.tirx.StringImm's.
|
||||
def to_str(s):
|
||||
if isinstance(s, tvm.tirx.StringImm):
|
||||
return s.value
|
||||
assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
|
||||
return s
|
||||
|
||||
objs = [to_str(s) for s in objs]
|
||||
|
||||
if not extra_args:
|
||||
extra_args = {}
|
||||
hex_arch = extra_args.get("hex_arch") or "v68"
|
||||
|
||||
ses = ContainerSession(HEXAGON_SDK_DOCKER_IMAGE)
|
||||
|
||||
hexagon_sdk_tools_path = ses.get_env("HEXAGON_TOOLCHAIN")
|
||||
libpath = os.path.join(hexagon_sdk_tools_path, "target", "hexagon", "lib", hex_arch, "G0")
|
||||
linker = os.path.join(hexagon_sdk_tools_path, "bin", "hexagon-link")
|
||||
|
||||
# Copy input data to docker container
|
||||
docker_objs = [ses.copy_to(obj) for obj in objs]
|
||||
docker_so_name = ses.tmp_dir + "/" + os.path.basename(so_name)
|
||||
|
||||
link_cmd = [linker, "-shared", "-fPIC", "-o", docker_so_name]
|
||||
link_cmd += docker_objs
|
||||
link_cmd += [
|
||||
"-Bdynamic",
|
||||
"-export-dynamic",
|
||||
"-L" + os.path.join(libpath, "pic"),
|
||||
"-lgcc",
|
||||
]
|
||||
ses.exec(link_cmd)
|
||||
|
||||
# Copy result back to host
|
||||
ses.copy_from(docker_so_name, so_name)
|
||||
return 0
|
||||
|
||||
|
||||
if sys.platform == "darwin":
|
||||
|
||||
def __create_shared_mac(so_name, objs, **kwargs):
|
||||
return link_shared_macos(so_name, objs, kwargs)
|
||||
|
||||
create_shared = __create_shared_mac
|
||||
register_global_func("tvm.contrib.hexagon.link_shared", f=link_shared_macos, override=True)
|
||||
else: # Linux and Win32
|
||||
create_shared = cc.create_shared
|
||||
register_global_func("tvm.contrib.hexagon.link_shared", f=link_shared, override=True)
|
||||
|
||||
|
||||
def create_aot_shared(so_name: str | pathlib.Path, files, hexagon_arch: str, options=None):
|
||||
"""Export Hexagon AOT module."""
|
||||
options = options or []
|
||||
if not os.access(str(HEXAGON_CLANG_PLUS), os.X_OK):
|
||||
raise Exception(
|
||||
'The Clang++ "' + str(HEXAGON_CLANG_PLUS) + '" does not exist or is not executable.'
|
||||
)
|
||||
if not HEXAGON_TOOLCHAIN:
|
||||
raise Exception(
|
||||
" The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
|
||||
+ "HEXAGON_TOOLCHAIN in your environment."
|
||||
)
|
||||
if not HEXAGON_SDK_ROOT:
|
||||
raise Exception(
|
||||
" The environment variable HEXAGON_SDK_ROOT is unset. Please export "
|
||||
+ "HEXAGON_SDK_ROOT in your environment."
|
||||
)
|
||||
|
||||
# The AOT C codegen uses TVM runtime functions
|
||||
# (e.g. TVMBackendAllocWorkspace) directly. On Hexagon these calls
|
||||
# should be made using functions pointers provided as __TVM*
|
||||
# variables in the provided context. This workaround allows the
|
||||
# the TVM runtime symbols to be visible to the compiled shared
|
||||
# library.
|
||||
#
|
||||
# This workaround can be removed when AOT codegen can be done with
|
||||
# LLVM codegen.
|
||||
workaround_link_flags = os.environ.get("HEXAGON_SHARED_LINK_FLAGS")
|
||||
if workaround_link_flags:
|
||||
options.extend(workaround_link_flags.split())
|
||||
|
||||
tvm_dir = pathlib.Path(os.path.dirname(os.path.realpath(__file__))) / ".." / ".." / ".." / ".."
|
||||
compute_arch = f"compute{hexagon_arch}"
|
||||
compile_options = [
|
||||
"-O3",
|
||||
f"-I{tvm_dir / 'include'}",
|
||||
f"-I{tvm_dir / '3rdparty' / 'dlpack' / 'include'}",
|
||||
f"-I{pathlib.Path(HEXAGON_SDK_ROOT) / 'rtos' / 'qurt' / compute_arch / 'include' / 'posix'}",
|
||||
f"-I{pathlib.Path(HEXAGON_SDK_ROOT) / 'rtos' / 'qurt' / compute_arch / 'include' / 'qurt'}",
|
||||
"-D_MACH_I32=int",
|
||||
]
|
||||
|
||||
# For debugging
|
||||
for path in HEXAGON_SDK_INCLUDE_DIRS:
|
||||
compile_options.append(f"-I{path!s}")
|
||||
|
||||
cross_compile = cc.cross_compiler(compile_func=hexagon_clang_plus())
|
||||
cross_compile.output_format = "o"
|
||||
c_files = [str(file) for file in files]
|
||||
cross_compile(str(so_name), c_files, options=compile_options + options)
|
||||
|
||||
|
||||
def pack_imports(
|
||||
module: tvm.runtime.Module,
|
||||
is_system_lib: bool, # pylint: disable=unused-argument
|
||||
c_symbol_prefix: str,
|
||||
workspace_dir: str,
|
||||
):
|
||||
"""Create an ELF object file that contains the binary data for the modules
|
||||
imported in `module`. This is a callback function for use as `fpack_imports`
|
||||
in `export_library`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
module: tvm.runtime.Module
|
||||
Module whose imported modules need to be serialized.
|
||||
is_system_lib: bool
|
||||
Flag whether the exported module will be used as a system library.
|
||||
c_symbol_prefix: str
|
||||
Prefix to prepend to the blob symbol.
|
||||
workspace_dir: str
|
||||
Location for created files.
|
||||
|
||||
Returns
|
||||
-------
|
||||
file_name: str
|
||||
The name of the created object file.
|
||||
"""
|
||||
|
||||
path_bin = os.path.join(workspace_dir, "imports.bin")
|
||||
pack_to_bin_f_name = "runtime.ModulePackImportsToTensor"
|
||||
fpack_to_bin = tvm.get_global_func(pack_to_bin_f_name)
|
||||
assert fpack_to_bin, f"Expecting {pack_to_bin_f_name} in registry"
|
||||
|
||||
fpack_to_bin(module).numpy().tofile(path_bin)
|
||||
|
||||
mblob_symbol = c_symbol_prefix + tvm.get_global_func("runtime.ModuleImportsBlobName")()
|
||||
|
||||
binary_size = os.path.getsize(path_bin)
|
||||
hexagon_toolchain = os.environ.get("HEXAGON_TOOLCHAIN")
|
||||
assert hexagon_toolchain, "Please set HEXAGON_TOOLCHAIN variable"
|
||||
version = toolchain_version(hexagon_toolchain)
|
||||
assert version[0] == 8 and version[1] >= 5, (
|
||||
"Please use Hexagon toolchain version 8.5.x or later"
|
||||
)
|
||||
if version[1] <= 6:
|
||||
path_o = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.o")
|
||||
subprocess.run(
|
||||
[
|
||||
f"{hexagon_toolchain}/bin/hexagon-clang",
|
||||
"-x",
|
||||
"c",
|
||||
"-c",
|
||||
"/dev/null",
|
||||
"-o",
|
||||
path_o,
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
subprocess.run(
|
||||
[
|
||||
f"{hexagon_toolchain}/bin/hexagon-llvm-objcopy",
|
||||
path_o,
|
||||
"--add-section",
|
||||
f".rodata={path_bin}",
|
||||
"--add-symbol",
|
||||
f"{mblob_symbol}=.rodata:0,object",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
return path_o
|
||||
|
||||
else: # 8.6.07+
|
||||
path_c = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.c")
|
||||
path_o = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.o")
|
||||
with open(path_c, "w") as f:
|
||||
f.write(
|
||||
f"const unsigned char {mblob_symbol}[{binary_size}] "
|
||||
f'__attribute__((section(".rodata"))) = {{0x1}};'
|
||||
)
|
||||
subprocess.run(
|
||||
[f"{hexagon_toolchain}/bin/hexagon-clang", "-c", path_c, "-o", path_o], check=True
|
||||
)
|
||||
subprocess.run(
|
||||
[
|
||||
f"{hexagon_toolchain}/bin/hexagon-llvm-objcopy",
|
||||
path_o,
|
||||
"--update-section",
|
||||
f".rodata={path_bin}",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
return path_o
|
||||
|
||||
|
||||
def export_module(module, out_dir, binary_name="test_binary.so"):
|
||||
"""Export Hexagon shared object to a file."""
|
||||
binary_path = pathlib.Path(out_dir) / binary_name
|
||||
module.write_to_file(str(binary_path))
|
||||
return binary_path
|
||||
|
||||
|
||||
def allocate_hexagon_array(
|
||||
dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
|
||||
):
|
||||
"""
|
||||
Allocate a hexagon array which could be a 2D array
|
||||
on physical memory defined by axis_separators
|
||||
"""
|
||||
if tensor_shape is None:
|
||||
assert data is not None, "Must provide either tensor shape or numpy data array"
|
||||
tensor_shape = data.shape
|
||||
elif data is not None:
|
||||
assert tensor_shape == data.shape, (
|
||||
"Mismatch between provided tensor shape and numpy data array shape"
|
||||
)
|
||||
|
||||
if dtype is None:
|
||||
assert data is not None, "Must provide either dtype or numpy data array"
|
||||
dtype = data.dtype.name
|
||||
elif data is not None:
|
||||
assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"
|
||||
|
||||
if axis_separators is None:
|
||||
axis_separators = []
|
||||
|
||||
boundaries = [0, *axis_separators, len(tensor_shape)]
|
||||
physical_shape = [
|
||||
numpy.prod(tensor_shape[dim_i:dim_f]) for dim_i, dim_f in itertools.pairwise(boundaries)
|
||||
]
|
||||
|
||||
arr = tvm.runtime.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)
|
||||
|
||||
if data is not None:
|
||||
arr.copyfrom(data.reshape(physical_shape))
|
||||
|
||||
return arr._create_view(tensor_shape)
|
||||
|
||||
|
||||
class ContainerSession:
|
||||
"""Docker container session
|
||||
|
||||
Parameters
|
||||
----------
|
||||
base_image_name : str
|
||||
Docker image name to use. Empty string means to use default "tlcpack/ci-hexagon"
|
||||
base image.
|
||||
"""
|
||||
|
||||
def __init__(self, base_image_name: str = ""):
|
||||
self._client = None
|
||||
self._container = None
|
||||
self.tmp_dir = None
|
||||
|
||||
self._client = ContainerSession._get_docker_client()
|
||||
|
||||
if base_image_name == "":
|
||||
base_image_name = ContainerSession._get_latest_ci_image(self._client)
|
||||
|
||||
self._container = ContainerSession._find_container_or_create(self._client, base_image_name)
|
||||
|
||||
exit_code, tmp_dir_b = self._container.exec_run("mktemp -d -t tvm-toolbox-XXXXXXXXXX")
|
||||
assert exit_code == 0
|
||||
|
||||
self.tmp_dir = tmp_dir_b.decode("utf-8").rstrip()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
@staticmethod
|
||||
def _get_latest_ci_image(client) -> str:
|
||||
ci_images = client.images.list(name="tlcpack/ci-hexagon")
|
||||
ci_images.sort(reverse=True, key=lambda img: img.tags[0])
|
||||
return ci_images[0].tags[0]
|
||||
|
||||
@staticmethod
|
||||
def _get_docker_client():
|
||||
try:
|
||||
from docker import from_env # pylint: disable=import-outside-toplevel
|
||||
from docker.errors import DockerException
|
||||
except (ModuleNotFoundError, ImportError):
|
||||
raise Exception("Docker SDK module is not installed. Please install it.")
|
||||
|
||||
try:
|
||||
client = from_env()
|
||||
except DockerException:
|
||||
raise Exception(
|
||||
"Docker server is not available. Please verify the docker is installed, "
|
||||
"launched and available via command line ('dokcer ps' should works)."
|
||||
)
|
||||
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def _find_container_or_create(client, image_name: str):
|
||||
all_containers = client.containers.list(all=True)
|
||||
|
||||
filtered_containers = []
|
||||
for container in all_containers:
|
||||
tags: list = container.image.tags
|
||||
img_name: str = tags[0]
|
||||
if img_name.startswith(image_name) and container.name.startswith("tvm-hex-toolbox"):
|
||||
filtered_containers.append(container)
|
||||
|
||||
if len(filtered_containers) == 0:
|
||||
container = client.containers.run(
|
||||
image=image_name, detach=True, tty=True, name="tvm-hex-toolbox"
|
||||
)
|
||||
else:
|
||||
container = filtered_containers[0]
|
||||
|
||||
if container.status != "running":
|
||||
container.start()
|
||||
|
||||
return container
|
||||
|
||||
def exec(self, cmd) -> str:
|
||||
"""Execute command inside docker container"""
|
||||
exit_code, res = self._container.exec_run(cmd)
|
||||
assert exit_code == 0
|
||||
return res.decode("utf-8")
|
||||
|
||||
def get_env(self, key: str) -> str:
|
||||
"""Return env var value from docker container"""
|
||||
res: str = self.exec(f"bash -c 'echo \"${key}\"'")
|
||||
return res.rstrip(" \n")
|
||||
|
||||
def copy_to(self, host_file_path: str) -> str:
|
||||
"""Upload file to docker container"""
|
||||
file_name = os.path.basename(host_file_path)
|
||||
|
||||
byte_stream = io.BytesIO()
|
||||
with tarfile.open(fileobj=byte_stream, mode="w:gz") as tar:
|
||||
tar.add(host_file_path, arcname=file_name)
|
||||
|
||||
self._container.put_archive(path=self.tmp_dir, data=byte_stream.getvalue())
|
||||
|
||||
return f"{self.tmp_dir}/{file_name}"
|
||||
|
||||
def copy_from(self, container_file_path: str, host_file_path: str):
|
||||
"""Download file from docker container"""
|
||||
tar_bytes_gen, _ = self._container.get_archive(container_file_path)
|
||||
|
||||
# convert to bytes
|
||||
tar_bytes = b""
|
||||
for chunk in tar_bytes_gen:
|
||||
tar_bytes += chunk
|
||||
|
||||
tar = tarfile.open(fileobj=io.BytesIO(initial_bytes=tar_bytes))
|
||||
assert len(tar.getmembers()) == 1
|
||||
tar_element_reader = tar.extractfile(tar.getmembers()[0])
|
||||
with open(host_file_path, "wb") as host_file:
|
||||
for chunk in tar_element_reader:
|
||||
host_file.write(chunk)
|
||||
|
||||
def close(self):
|
||||
"""Close docker container session"""
|
||||
if self.tmp_dir is not None:
|
||||
exit_code, _ = self._container.exec_run(f"rm -rf {self.tmp_dir}")
|
||||
assert exit_code == 0
|
||||
Reference in New Issue
Block a user