# 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. # ruff: noqa: F401, F821, F841 import os import tempfile import numpy as np import pytest import tvm import tvm.testing from tvm import ( DataType, IRModule, relax, tirx, ) from tvm.relax.transform.legalize_ops import adreno as legalize_adreno from tvm.rpc import connect_tracker from tvm.script import ir as I from tvm.script import tirx as T from tvm.support import ndk from tvm.target import Target from tvm.testing import env def get_rpc(): """ Establish an RPC connection to the remote device. Returns ------- tvm.rpc.RPCSession or None The RPC session object if RPC_TARGET is set; otherwise, None. """ rpc_target = os.getenv("RPC_TARGET", None) if rpc_target: host = os.getenv("TVM_TRACKER_HOST", "localhost") port = int(os.getenv("TVM_TRACKER_PORT", 9090)) device_key = os.getenv("RPC_DEVICE_KEY", "android") tracker = connect_tracker(host, port) return tracker.request(device_key, priority=1, session_timeout=1000) else: return None def preprocess_pipeline(mod: IRModule) -> IRModule: desired_layouts = {"relax.nn.conv2d": ["NCHW16c", "OIHW16o", "NCHW16c"]} seq = tvm.transform.Sequential( [ tvm.tirx.transform.BindTarget(Target.current(allow_none=False)), tvm.relax.transform.FoldConstant(), tvm.relax.transform.DecomposeOpsForInference(), tvm.relax.transform.FoldConstant(), tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)), tvm.relax.transform.ConvertLayout(desired_layouts), tvm.relax.transform.Normalize(), tvm.relax.transform.FoldConstant(), tvm.relax.transform.LegalizeOps(), tvm.relax.transform.LegalizeOps( {"relax.nn.conv2d": legalize_adreno.conv2d_NCHWc_OIHWo} ), tvm.relax.transform.FoldConstant(), tvm.relax.transform.AnnotateTIROpPattern(), tvm.relax.transform.FuseOps(), tvm.relax.transform.FuseTIR(), tvm.relax.transform.DeadCodeElimination(), tvm.relax.transform.Normalize(), ] ) mod = seq(mod) return mod def postprocess_pipeline(mod: IRModule) -> IRModule: seq = tvm.transform.Sequential( [ tvm.relax.transform.ToNonDataflow(), tvm.relax.transform.RemovePurityChecking(), tvm.relax.transform.CallTIRRewrite(), tvm.relax.transform.Normalize(), tvm.relax.transform.StaticPlanBlockMemory(), tvm.relax.transform.LowerAllocTensor(), tvm.relax.transform.KillAfterLastUse(), tvm.relax.transform.LowerRuntimeBuiltin(), tvm.relax.transform.VMShapeLower(), tvm.relax.transform.AttachGlobalSymbol(), ] ) mod = seq(mod) return mod @pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc") @pytest.mark.gpu @pytest.mark.skipif(not env.has_adreno_opencl(), reason="need adreno opencl") @pytest.mark.parametrize("backend", ["opencl"]) @pytest.mark.parametrize("dtype", ["int8", "float16", "int16", "float32", "int32"]) @pytest.mark.parametrize("channel_size", [64, 128]) @pytest.mark.parametrize("read_width", [1, 2, 4, 8, 16]) def test_texture_copy(backend, dtype, channel_size, read_width): remote = get_rpc() M, N, K = (256, 1024, 128) lanes = channel_size // DataType(dtype).bits if read_width > lanes: return @I.ir_module(s_tir=True) class TextureCopy: @T.prim_func(s_tir=True) def main(A: T.Buffer((M, N), dtype), B: T.Buffer((M, N), dtype)): T.func_attr({"global_symbol": "main"}) for li, lj in T.grid(M, N): with T.sblock("Copy"): i, j = T.axis.remap("SS", [li, lj]) B[i, j] = A[i, j] def schedule_texture_read(sch: s_tir.Schedule): B_blk = sch.get_sblock("Copy") Ai_block = sch.cache_read(B_blk, 0, "global.texture") sch.transform_layout(Ai_block, ("write", 0), lambda i, j: (i, j // lanes, j % lanes)) def schedule_default(blk, lanes): i, j = sch.get_loops(blk) jo, jv = sch.split(j, [None, lanes]) b = sch.fuse(i, jo) bx, tx = sch.split(b, [None, 256]) sch.bind(bx, "blockIdx.x") sch.bind(tx, "threadIdx.x") sch.vectorize(jv) schedule_default(Ai_block, lanes) schedule_default(B_blk, read_width) mod = TextureCopy if remote is None: target = Target({"kind": backend, "device": "adreno"}) else: target = Target( {"kind": backend, "device": "adreno"}, {"kind": "llvm", "mtriple": "aarch64-linux-android"}, ) with target: mod = preprocess_pipeline(mod) sch = tvm.s_tir.Schedule(mod) schedule_texture_read(sch) mod = postprocess_pipeline(sch.mod) ex = relax.build(mod, target) load_path = "vm_library.so" inputs = [np.random.randint(0, 128, (M, N)).astype(dtype), np.zeros((M, N), dtype)] def run_and_check(rexec, remote_session): if remote_session is not None: dev = remote_session.cl() elif "opencl" in backend: dev = tvm.opencl(0) elif "vulkan" in backend: dev = tvm.vulkan(0) else: raise RuntimeError("Unsupported backend") 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) inps = [tvm.runtime.tensor(inp, dev) for inp in inputs] vm["main"](*inps) np.testing.assert_equal(inps[-1].numpy(), inps[0].numpy()) with tempfile.TemporaryDirectory() as temp_dir: if remote is not None: path = temp_dir + "/" + load_path ex.export_library(path, fcompile=ndk.create_shared, options=["-shared", "-fPIC", "-lm"]) remote.upload(path) rexec = remote.load_module(load_path) try: run_and_check(rexec, remote) finally: remote.get_function("CloseRPCConnection")() else: tvm.testing.run_with_gpu_lock(run_and_check, ex, None) if __name__ == "__main__": tvm.testing.main()