Skip to content

Commit

Permalink
[TVMC] add cl support in tvmc runner (apache#6831)
Browse files Browse the repository at this point in the history
* [TVMC] add cl support in tvmc runner

* Cleanup comment and asssert device type in else case
  • Loading branch information
euntaik authored and Trevor Morris committed Dec 4, 2020
1 parent e385676 commit 7697781
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions python/tvm/driver/tvmc/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,10 @@ def add_run_parser(subparsers):
parser.set_defaults(func=drive_run)

# TODO --device needs to be extended and tested to support other targets,
# like 'cl', 'webgpu', etc (@leandron)
# like 'webgpu', etc (@leandron)
parser.add_argument(
"--device",
choices=["cpu", "gpu"],
choices=["cpu", "gpu", "cl"],
default="cpu",
help="target device to run the compiled module. Defaults to 'cpu'",
)
Expand Down Expand Up @@ -361,7 +361,13 @@ def run_module(

# TODO expand to other supported devices, as listed in tvm.rpc.client (@leandron)
logger.debug("device is %s", device)
ctx = session.cpu() if device == "cpu" else session.gpu()
if device == "gpu":
ctx = session.gpu()
elif device == "cl":
ctx = session.cl()
else:
assert device == "cpu"
ctx = session.cpu()

if profile:
logger.debug("creating runtime with profiling enabled")
Expand Down

0 comments on commit 7697781

Please sign in to comment.