Skip to content

Commit

Permalink
[ROCm][TVMC] Add ROCm to the TVMC driver (apache#8896)
Browse files Browse the repository at this point in the history
* Add ROCm to list of RPC clients.

* Add ROCm to list of TVMC devices.

* Enable ROCm by adding session call.
  • Loading branch information
mvermeulen authored and ylc committed Jan 13, 2022
1 parent b218d36 commit bb1b66a
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 1 deletion.
4 changes: 3 additions & 1 deletion python/tvm/driver/tvmc/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def add_run_parser(subparsers):
# like 'webgpu', etc (@leandron)
parser.add_argument(
"--device",
choices=["cpu", "cuda", "cl", "metal", "vulkan"],
choices=["cpu", "cuda", "cl", "metal", "vulkan", "rocm"],
default="cpu",
help="target device to run the compiled module. Defaults to 'cpu'",
)
Expand Down Expand Up @@ -394,6 +394,8 @@ def run_module(
dev = session.metal()
elif device == "vulkan":
dev = session.vulkan()
elif device == "rocm":
dev = session.rocm()
else:
assert device == "cpu"
dev = session.cpu()
Expand Down
4 changes: 4 additions & 0 deletions python/tvm/rpc/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,10 @@ def metal(self, dev_id=0):
"""Construct Metal device."""
return self.device(8, dev_id)

def rocm(self, dev_id=0):
"""Construct ROCm device."""
return self.device(10, dev_id)

def ext_dev(self, dev_id=0):
"""Construct extension device."""
return self.device(12, dev_id)
Expand Down

0 comments on commit bb1b66a

Please sign in to comment.