Enable nvidia texture tests. (apache#168)

Co-authored-by: Siva <quic_sivb@quicinc.com>

Author: srkreddy1238

python/tvm/testing/utils.py

@@ -935,6 +935,17 @@ def _check_opencl_vulkan():
     parent_features="cuda",
 )
 
+# Mark a test as requiring the OpenCL runtime on remote RPC
+requires_adreno_opencl = Feature(
+    "opencl",
+    long_name="Remote Adreno OpenCL",
+    cmake_flag="USE_OPENCL",
+    target_kind_enabled="opencl",
+    target_kind_hardware=None,
+    parent_features="gpu",
+    run_time_check=lambda: os.getenv("RPC_TARGET") is not None,
+)
+
 # Mark a test as requiring the OpenCL runtime
 requires_opencl = Feature(
     "opencl",

tests/python/relax/texture/adreno_utils.py

@@ -109,18 +109,30 @@ def run_cpu(mod, inputs, save_lib=False):
 
 
 def build_run(mod, inputs, backend, is_adreno=False):
-
+    remote = get_rpc()
     target = get_target(backend, is_adreno)
-    tgt = tvm.target.Target(target, host="llvm -mtriple=aarch64-linux-gnu")
+    if remote is None:
+        tgt = tvm.target.Target(target, host="llvm")
+    else:
+        tgt = tvm.target.Target(target, host="llvm -mtriple=aarch64-linux-gnu")
     relax_pipeline = relax.pipeline.get_default_pipeline(tgt)
     tir_pipeline = tvm.tir.get_default_tir_pipeline(tgt)
     mod = relax_pipeline(mod)
     ex = tvm.compile(mod, tgt, tir_pipeline=tir_pipeline)
 
-    remote = get_rpc()
     if remote is None:
         # local execution
-        dev = tvm.device(target, 0)
+        if "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 ii in range(len(mod.global_infos["vdevice"]))]
+        else:
+            device_arr = [dev]
         vm = relax.VirtualMachine(ex, device_arr)
     else:
         # remote execution
@@ -158,7 +170,8 @@ def build_run(mod, inputs, backend, is_adreno=False):
     else:
         tvm_output = tvm_output.numpy()
 
-    remote.get_function("CloseRPCConnection")()
+    if remote:
+        remote.get_function("CloseRPCConnection")()
     return tvm_output
 
 

tests/python/relax/texture/test_network.py

@@ -37,27 +37,6 @@
 import copy
 
 
-@pytest.mark.parametrize("dtype", ["float32"])
-@pytest.mark.parametrize(
-    "url, shape_dict",
-    [
-        # ("mobilenetv2-12.onnx", {"input": [1, 3, 224, 224]}),
-        # ("densenet-12.onnx", {"data_0": [1, 3, 224, 224]}),
-        # ("inception-v2-9.onnx", {"data_0": [1, 3, 224, 224]}),
-        ("resnet18-v2-7.onnx", {"data": [1, 3, 224, 224]}),
-        # ("resnet50-v2-7.onnx", {"data": [1, 3, 224, 224]}),
-    ],
-)
-@tvm.testing.requires_opencl_vulkan
-def _test_network(url, shape_dict, dtype):
-    print("Network evaluating .. " + url + " " + dtype)
-    model = onnx.load("./" + url)
-    mod = from_onnx(model, shape_dict)
-    mod1 = from_onnx(model, shape_dict)
-
-    verify(mod)
-
-
 @tvm.testing.requires_opencl_vulkan
 @tvm.testing.parametrize_targets("opencl", "vulkan")
 def test_network_resnet(target):

tests/python/relax/texture/test_texture_nd.py

@@ -31,21 +31,27 @@
 from tvm.target import Target
 from tvm.contrib import ndk
 from tvm import tir, DataType
+from tvm.rpc import connect_tracker
 
 
-class RemoteConnection:
-    def __init__(self):
-        self.RPC_TRACKER_HOST = os.getenv("TVM_TRACKER_HOST", "localhost")
-        self.RPC_TRACKER_PORT = int(os.getenv("TVM_TRACKER_PORT", 7979))
-        self.RPC_KEY = os.getenv("RPC_DEVICE_KEY", "android")
-        self.tracker = tvm.rpc.connect_tracker(self.RPC_TRACKER_HOST, self.RPC_TRACKER_PORT)
+def get_rpc():
+    """
+    Establish an RPC connection to the remote device.
 
-    def __enter__(self):
-        self.remote = self.tracker.request(self.RPC_KEY, priority=0, session_timeout=600)
-        return self.remote
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.remote.get_function("CloseRPCConnection")()
+    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:
@@ -96,14 +102,13 @@ def postprocess_pipeline(mod: IRModule) -> IRModule:
 
 
 @tvm.testing.requires_rpc
-@tvm.testing.requires_opencl
-@pytest.mark.parametrize(
-    "target", [Target("opencl -device=adreno", "llvm -mtriple=aarch64-linux-android")]
-)
+@tvm.testing.requires_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(target, dtype, channel_size, read_width):
+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:
@@ -139,6 +144,12 @@ def schedule_default(blk, lanes):
         schedule_default(B_blk, read_width)
 
     mod = TextureCopy
+
+    if remote is None:
+        target = Target(backend + " -device=adreno")
+    else:
+        target = Target(backend + " -device=adreno", "llvm -mtriple=aarch64-linux-android")
+
     with target:
         mod = preprocess_pipeline(mod)
         sch = tir.Schedule(mod)
@@ -148,20 +159,43 @@ def schedule_default(blk, lanes):
     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)]
-    with RemoteConnection() as remote:
-        with tempfile.TemporaryDirectory() as temp_dir:
+    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)
             dev = remote.cl()
-
-            vm = relax.VirtualMachine(rexec, [dev, dev, dev])
-            inps = [tvm.runtime.tensor(inp, dev) for inp in inputs]
-            vm["main"](*inps)
-
-            np.testing.assert_equal(inps[-1].numpy(), inps[0].numpy())
+            if "vdevice" in mod.global_infos:
+                device_arr = [dev for ii in range(len(mod.global_infos["vdevice"]))]
+            else:
+                device_arr = [dev]
+            vm = relax.VirtualMachine(rexec, device_arr)
+        else:
+            # local execution
+            if "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 ii in range(len(mod.global_infos["vdevice"]))]
+            else:
+                device_arr = [dev]
+            vm = relax.VirtualMachine(ex, device_arr)
+
+        inps = [tvm.runtime.tensor(inp, dev) for inp in inputs]
+        vm["main"](*inps)
+
+        out1 = inps[-1].numpy()
+        out2 = inps[0].numpy()
+
+        if remote:
+            remote.get_function("CloseRPCConnection")()
+
+        np.testing.assert_equal(out1, out2)
 
 
 if __name__ == "__main__":

tests/scripts/ci.py

@@ -604,6 +604,7 @@ def add_subparser(
                 [
                     "./tests/scripts/task_opencl_cpp_unittest.sh {build_dir}",
                     "./tests/scripts/task_python_unittest.sh",
+                    "./tests/scripts/task_python_unittest_gpuonly.sh",
                 ],
             ),
             "frontend": ("run frontend tests", ["./tests/scripts/task_python_frontend.sh"]),

tests/scripts/task_python_unittest_gpuonly.sh

@@ -34,3 +34,13 @@ export TVM_UNITTEST_TESTSUITE_NAME=python-codegen-vulkan
 source tests/scripts/setup-pytest-env.sh
 
 run_pytest ${TVM_UNITTEST_TESTSUITE_NAME}-1 tests/python/codegen/test_target_codegen_vulkan.py
+
+# Adreno
+export PYTEST_ADDOPTS=""
+export TVM_TEST_TARGETS="opencl"
+export TVM_UNITTEST_TESTSUITE_NAME=python-codegen-clml-texture
+
+source tests/scripts/setup-pytest-env.sh
+
+run_pytest ${TVM_UNITTEST_TESTSUITE_NAME}-1 -s tests/python/relax/backend/clml/
+run_pytest ${TVM_UNITTEST_TESTSUITE_NAME}-1 tests/python/relax/texture/