chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,218 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user