| # 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) |
