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