Skip to content

Commit

Permalink
[shortfin] Add simple example to enumerate devices and test options.
Browse files Browse the repository at this point in the history
  • Loading branch information
stellaraccident committed Nov 7, 2024
1 parent 92170e1 commit 1ea095a
Showing 1 changed file with 32 additions and 0 deletions.
32 changes: 32 additions & 0 deletions shortfin/examples/python/enumerate_devices.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# Copyright 2024 Advanced Micro Devices, Inc.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

r"""Simple test program that enumerates devices available.
Run with SystemBuilder keyword args on the command line like::
python examples/python/enumerate_devices.py \
system_type=amdgpu amdgpu_logical_devices_per_physical_device=4
"""

import sys

import shortfin as sf


def main():
args = [arg.split("=", maxsplit=1) for arg in sys.argv[1:]]
arg_dict = {k: v for k, v in args}
print(f"Creating system with args: {arg_dict}")
builder = sf.SystemBuilder(**arg_dict)
with builder.create_system() as ls:
for device in ls.devices:
print(device)


if __name__ == "__main__":
main()

0 comments on commit 1ea095a

Please sign in to comment.