From 1ea095a08b47a21734b22c5242ce73e7d8cc1bf0 Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Wed, 6 Nov 2024 20:26:33 -0800 Subject: [PATCH] [shortfin] Add simple example to enumerate devices and test options. --- shortfin/examples/python/enumerate_devices.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 shortfin/examples/python/enumerate_devices.py diff --git a/shortfin/examples/python/enumerate_devices.py b/shortfin/examples/python/enumerate_devices.py new file mode 100644 index 000000000..2c1cc8203 --- /dev/null +++ b/shortfin/examples/python/enumerate_devices.py @@ -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()