* review comment

apache · Sep 7, 2022 · 3f8131f · 3f8131f
1 parent 8143a84
commit 3f8131f
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/src/runtime/contrib/clml/clml_runtime.cc b/src/runtime/contrib/clml/clml_runtime.cc
@@ -489,12 +489,12 @@ class CLMLRuntime : public JSONRuntimeBase {
       return 0;
     }
     size_t reqd_size = 0;
-    result = clGetDeviceInfo(workspace->devices[0], CL_DEVICE_EXTENSIONS, 0, NULL, &reqd_size);
+    cl_device_id device_id = workspace->devices[workspace->GetThreadEntry()->device.device_id];
+    result = clGetDeviceInfo(device_id, CL_DEVICE_EXTENSIONS, 0, NULL, &reqd_size);
     ICHECK(reqd_size > 0u && result == CL_SUCCESS) << "clGetDeviceInfo:" << result;
 
     std::vector<char> buf(reqd_size);
-    result =
-        clGetDeviceInfo(workspace->devices[0], CL_DEVICE_EXTENSIONS, reqd_size, buf.data(), NULL);
+    result = clGetDeviceInfo(device_id, CL_DEVICE_EXTENSIONS, reqd_size, buf.data(), NULL);
     ICHECK(result == CL_SUCCESS) << "clGetDeviceInfo:" << result;
 
     std::string extensions(buf.data());

diff --git a/tests/python/contrib/test_clml/test_network.py b/tests/python/contrib/test_clml/test_network.py
@@ -22,7 +22,7 @@
 from tvm import relay
 
 import tvm
-from infrastructure import skip_runtime_test, build_and_run, Device
+from test_clml.infrastructure import skip_runtime_test, build_and_run, Device
 
 
 def _build_and_run_network(mod, params, inputs, data, device, atol, rtol):

diff --git a/tests/python/contrib/test_clml/test_ops.py b/tests/python/contrib/test_clml/test_ops.py
@@ -25,7 +25,7 @@
 from tvm import relay
 from tvm.ir import IRModule
 
-from infrastructure import (
+from test_clml.infrastructure import (
     skip_runtime_test,
     skip_codegen_test,
     build_and_run,