wehub-resource-sync
219 lines
6.9 KiB
Python
219 lines
6.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.
|
|
|
|
import os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.support import ndk
|
|
|
|
# Test Infra
|
|
|
|
|
|
class run_time_check:
|
|
def __init__(self, device):
|
|
self.device = device
|
|
|
|
def check(self):
|
|
# Ensure adreno specific tests
|
|
if self.device == "real":
|
|
return "ADRENO_TARGET" in os.environ
|
|
|
|
# Adreno CI
|
|
if "ADRENO_TARGET" in os.environ:
|
|
return True
|
|
|
|
# Tests that can run on generic targets too
|
|
elif self.device == "opencl":
|
|
return tvm.opencl().exist
|
|
elif self.device == "vulkan":
|
|
return tvm.vulkan().exist
|
|
elif self.device == "any":
|
|
return tvm.opencl().exist or tvm.vulkan().exist
|
|
else:
|
|
return False
|
|
|
|
def __call__(self):
|
|
return self.check
|
|
|
|
|
|
# Eager skips for Adreno GPU tests, resolved at import time. Pair each with
|
|
# ``@pytest.mark.gpu`` at the test site so CI's ``-m gpu`` filter selects it.
|
|
|
|
# OpenCL or Vulkan
|
|
skip_unless_adreno_opencl_vulkan = pytest.mark.skipif(
|
|
not run_time_check("any").check(),
|
|
reason="need adreno opencl or vulkan",
|
|
)
|
|
|
|
# CLML Codegen
|
|
skip_unless_adreno_clml = pytest.mark.skipif(
|
|
tvm.get_global_func("relax.is_openclml_runtime_enabled", allow_missing=True) is None,
|
|
reason="need adreno openclml",
|
|
)
|
|
|
|
|
|
def is_target_available(target):
|
|
if "clml" in target.attrs.get("keys", []) and "ADRENO_TARGET" not in os.environ:
|
|
return False
|
|
return True
|
|
|
|
|
|
class SessionManager:
|
|
def __init__(self):
|
|
self.is_remote = SessionManager.is_target_rpc()
|
|
|
|
def __enter__(self):
|
|
if self.is_remote:
|
|
self.RPC_TRACKER_HOST = os.getenv("TVM_TRACKER_HOST", "localhost")
|
|
self.RPC_TRACKER_PORT = int(os.getenv("TVM_TRACKER_PORT", 7979))
|
|
self.RPC_DEVICE_KEY = os.getenv("RPC_DEVICE_KEY", "android")
|
|
|
|
self.tracker = tvm.rpc.connect_tracker(self.RPC_TRACKER_HOST, self.RPC_TRACKER_PORT)
|
|
self.rpc = self.tracker.request(self.RPC_DEVICE_KEY, priority=0, session_timeout=600)
|
|
else:
|
|
self.rpc = tvm.rpc.LocalSession()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self.rpc.get_function("CloseRPCConnection")()
|
|
|
|
def load_module(self, ex: relax.VMExecutable):
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
file_name = "vm_library.so"
|
|
file_path = os.path.join(tempdir, file_name)
|
|
if self.is_remote:
|
|
ex.export_library(
|
|
file_path, fcompile=ndk.create_shared, options=["-shared", "-fPIC", "-lm"]
|
|
)
|
|
else:
|
|
ex.export_library(file_path)
|
|
|
|
self.rpc.upload(file_path)
|
|
rexec = self.rpc.load_module(file_name)
|
|
return rexec
|
|
|
|
def device(self, device: str):
|
|
return self.rpc.device(device)
|
|
|
|
@staticmethod
|
|
def is_target_rpc():
|
|
"""
|
|
Checks if the target is a remote device.
|
|
|
|
Returns
|
|
-------
|
|
bool: True if RPC_TARGET is set, False otherwise
|
|
"""
|
|
return os.environ.get("ADRENO_TARGET") == "adreno"
|
|
|
|
|
|
def run_local(mod, inputs, target):
|
|
"""
|
|
Run the Relax module on the local CPU for verification.
|
|
|
|
Parameters
|
|
----------
|
|
mod : tvm.IRModule
|
|
The Relax IRModule to execute.
|
|
inputs : list of numpy.ndarray
|
|
The input data for the module.
|
|
save_lib : bool, optional
|
|
Whether to save the compiled library. Default is False.
|
|
|
|
Returns
|
|
-------
|
|
tvm.runtime.NDArray or tuple of tvm.runtime.NDArray
|
|
The output from the module execution.
|
|
"""
|
|
ex = relax.build(mod, target)
|
|
dev = tvm.cpu()
|
|
vm = relax.VirtualMachine(ex, dev)
|
|
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs]
|
|
vm.set_input("main", *inputs)
|
|
vm.invoke_stateful("main")
|
|
tvm_output = vm.get_outputs("main")
|
|
if isinstance(tvm_output, tuple):
|
|
tvm_output = tuple(out.numpy() for out in tvm_output)
|
|
else:
|
|
tvm_output = (tvm_output.numpy(),)
|
|
return tvm_output
|
|
|
|
|
|
def build_and_run(mod, inputs, tgt):
|
|
if SessionManager.is_target_rpc():
|
|
tgt = tvm.target.Target(tgt, host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
|
else:
|
|
tgt = tvm.target.Target(tgt, host={"kind": "llvm"})
|
|
|
|
relax_pipeline = relax.pipeline.get_default_pipeline(tgt)
|
|
tir_pipeline = tvm.tirx.get_default_tir_pipeline(tgt)
|
|
mod = relax_pipeline(mod)
|
|
|
|
ex = tvm.compile(mod, tgt, tir_pipeline=tir_pipeline)
|
|
|
|
def run_and_check():
|
|
with SessionManager() as sess:
|
|
rexec = sess.load_module(ex)
|
|
dev = sess.device(tgt.kind.name)
|
|
|
|
if "vdevice" in mod.global_infos:
|
|
device_arr = [dev for _ in range(len(mod.global_infos["vdevice"]))]
|
|
else:
|
|
device_arr = [dev]
|
|
vm = relax.VirtualMachine(rexec, device_arr)
|
|
device_inputs = [tvm.runtime.tensor(ip, dev) for ip in inputs]
|
|
vm.set_input("main", *device_inputs)
|
|
vm.invoke_stateful("main")
|
|
tvm_output = vm.get_outputs("main")
|
|
if isinstance(tvm_output, tuple):
|
|
return tuple(out.numpy() for out in tvm_output)
|
|
return (tvm_output.numpy(),)
|
|
|
|
if SessionManager.is_target_rpc():
|
|
return run_and_check()
|
|
return tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
def verify_results(mod, target, ref_target):
|
|
if not is_target_available(target):
|
|
print("Skipping Eval Tests", flush=True)
|
|
return
|
|
|
|
inputs = []
|
|
for arg in mod["main"].params:
|
|
shape = tuple(shape_val.value for shape_val in arg.ty.shape.values)
|
|
inputs.append(np.random.uniform(0, 1, size=shape).astype(arg.ty.dtype))
|
|
|
|
mod_org, mod_ref = mod, mod.clone()
|
|
|
|
mod_ref = tvm.relax.transform.DecomposeOpsForInference()(mod_ref)
|
|
if ref_target.kind.name == "llvm":
|
|
rs_ref = run_local(mod_ref, inputs, ref_target)
|
|
else:
|
|
rs_ref = build_and_run(mod_ref, inputs, ref_target)
|
|
|
|
rs_org = build_and_run(mod_org, inputs, target)
|
|
|
|
for vl_org, vl_ref in zip(rs_org, rs_ref):
|
|
tvm.testing.assert_allclose(vl_org, vl_ref, rtol=1e-3, atol=1e-3)
|