Skip to content

Commit 86cc46c

Browse files
committed
[Testing] Return BenchmarkResult in local_run and rpc_run (apache#15277)
Return `profile_result` in `tvm.testing.local_run` and `tvm.testing.rpc_run` for further use.
1 parent ba57d8e commit 86cc46c

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

python/tvm/testing/runner.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
if TYPE_CHECKING:
2424
import numpy as np
25+
2526
from tvm.meta_schedule.runner import EvaluatorConfig, RPCConfig
2627
from tvm.runtime import Device, Module, NDArray
2728

@@ -30,6 +31,7 @@
3031

3132
def _args_to_device(args, device):
3233
import numpy as np
34+
3335
from tvm.runtime.ndarray import NDArray, empty
3436

3537
uploaded_args = []
@@ -109,6 +111,8 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals
109111
-------
110112
args : List[Union[np.ndarray, NDArray, int, float]]
111113
The results of running the module.
114+
profile_result : tvm.runtime.BenchmarkResult
115+
The profiling result of running the module.
112116
"""
113117
import os.path as osp
114118
import tempfile
@@ -137,13 +141,12 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals
137141
if evaluator_config.enable_cpu_cache_flush
138142
else "",
139143
)(*args)
140-
print(profile_result)
141144
remote_mod(*args)
142145
args = _args_to_numpy(args)
143146
finally:
144147
pass
145148

146-
return args
149+
return args, profile_result
147150

148151

149152
def rpc_run( # pylint: disable=too-many-arguments,too-many-locals
@@ -188,6 +191,8 @@ def rpc_run( # pylint: disable=too-many-arguments,too-many-locals
188191
-------
189192
args : List[Union[np.ndarray, NDArray, int, float]]
190193
The results of running the module.
194+
profile_result : tvm.runtime.BenchmarkResult
195+
The profiling result of running the module.
191196
"""
192197

193198
import os.path as osp
@@ -220,12 +225,11 @@ def rpc_run( # pylint: disable=too-many-arguments,too-many-locals
220225
if evaluator_config.enable_cpu_cache_flush
221226
else "",
222227
)(*args)
223-
print(profile_result)
224228
remote_mod(*args)
225229
args = _args_to_numpy(args)
226230
finally:
227231
session.remove(remote_path)
228232
session.remove(remote_path + "." + output_format)
229233
session.remove("")
230234

231-
return args
235+
return args, profile_result

0 commit comments

Comments
 (0)