Add check random_fill exists on remote devices

apache · Sep 6, 2020 · 4c561d3 · 4c561d3
1 parent 80b5271
commit 4c561d3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py
@@ -695,9 +695,11 @@ def timed_func():
             try:
                 args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in
                         build_res.args]
-                assert tvm.get_global_func("tvm.contrib.random.random_fill", True), \
-                    "Please make sure USE_RANDOM is ON in the config.cmake"
-                random_fill = remote.get_function("tvm.contrib.random.random_fill")
+                try:
+                    random_fill = remote.get_function("tvm.contrib.random.random_fill")
+                except AttributeError:
+                    raise AttributeError("Please make sure USE_RANDOM is ON in the config.cmake "
+                                          "on the remote devices")
                 for arg in args:
                     random_fill(arg)
                 ctx.sync()

diff --git a/python/tvm/autotvm/measure/measure_methods.py b/python/tvm/autotvm/measure/measure_methods.py
@@ -511,9 +511,11 @@ def run_through_rpc(measure_input, build_result,
         if ref_input:
             args = [nd.array(x, ctx=ctx) for x in ref_input]
         else:
-            assert tvm.get_global_func("tvm.contrib.random.random_fill", True), \
-                "Please make sure USE_RANDOM is ON in the config.cmake"
-            random_fill = remote.get_function("tvm.contrib.random.random_fill")
+            try:
+                random_fill = remote.get_function("tvm.contrib.random.random_fill")
+            except AttributeError:
+                raise AttributeError("Please make sure USE_RANDOM is ON in the config.cmake "
+                                     "on the remote devices")
             args = [nd.empty(x[0], dtype=x[1], ctx=ctx) for x in build_result.arg_info]
             for arg in args:
                 random_fill(arg)