WinML update to handle multigpu set ups

procyon_ai/ulprocai.py

@@ -6,7 +6,7 @@ import subprocess
 import sys
 import time
 import psutil
-from utils import find_score_in_xml, is_process_running, get_install_path
+from utils import find_score_in_xml, is_process_running, get_install_path, get_winml_devices
 
 PARENT_DIR = str(Path(sys.path[0], ".."))
 sys.path.append(PARENT_DIR)
@@ -27,6 +27,9 @@ LOG_DIR = SCRIPT_DIR / "run"
 DIR_PROCYON = Path(get_install_path())
 EXECUTABLE = "ProcyonCmd.exe"
 ABS_EXECUTABLE_PATH = DIR_PROCYON / EXECUTABLE
+
+WINML_DEVICES = get_winml_devices(ABS_EXECUTABLE_PATH)
+
 CONFIG_DIR = SCRIPT_DIR / "config"
 BENCHMARK_CONFIG = {
     "AMD_CPU": {
@@ -34,9 +37,18 @@ BENCHMARK_CONFIG = {
         "process_name":  "WinML.exe",
         "test_name": "WinML CPU (FLOAT32)"
     },
-    "AMD_GPU": {
+    "AMD_GPU0": {
         "config": f"\"{CONFIG_DIR}\\ai_computer_vision_winml_gpu.def\"",
         "process_name":  "WinML.exe",
+        "device_name": list(WINML_DEVICES.keys())[0],
+        "device_id": list(WINML_DEVICES.values())[0],
+        "test_name": "WinML GPU (FLOAT32)"
+    },
+    "AMD_GPU1": {
+        "config": f"\"{CONFIG_DIR}\\ai_computer_vision_winml_gpu.def\"",
+        "process_name":  "WinML.exe",
+        "device_name": list(WINML_DEVICES.keys())[1],
+        "device_id": list(WINML_DEVICES.values())[1],
         "test_name": "WinML GPU (FLOAT32)"
     },
     "Intel_CPU": {
@@ -44,10 +56,15 @@ BENCHMARK_CONFIG = {
         "process_name":  "OpenVino.exe",
         "test_name": "Intel OpenVINO CPU (FLOAT32)"
     },
-    "Intel_GPU": {
+    "Intel_GPU0": {
         "config": f"\"{CONFIG_DIR}\\ai_computer_vision_openvino_gpu.def\"",
         "process_name":  "OpenVino.exe",
-        "test_name": "Intel OpenVINO GPU (FLOAT32)"
+        "test_name": "Intel OpenVINO GPU 0 (FLOAT32)"
+    },
+    "Intel_GPU1": {
+        "config": f"\"{CONFIG_DIR}\\ai_computer_vision_openvino_gpu.def\"",
+        "process_name":  "OpenVino.exe",
+        "test_name": "Intel OpenVINO GPU 1 (FLOAT32)"
     },
     "Intel_NPU": {
         "config": f"\"{CONFIG_DIR}\\ai_computer_vision_openvino_npu.def\"",
@@ -90,9 +107,15 @@ def get_arguments():
     return argies
 
 
-def create_procyon_command(test_option):
+def create_procyon_command(test_option, process_name, device_id):
     """create command string"""
-    command = f'\"{ABS_EXECUTABLE_PATH}\" --definition={test_option} --export=\"{REPORT_PATH}\"'
+    match process_name:
+        case 'WinML.exe':
+            command = f'\"{ABS_EXECUTABLE_PATH}\" --definition={test_option} --export=\"{REPORT_PATH}\" --select-winml-device {device_id}'
+        case 'OpenVino.exe':
+            command = f'\"{ABS_EXECUTABLE_PATH}\" --definition={test_option} --export=\"{REPORT_PATH}\" --select-openvino-device {device_id}'
+        case 'TensorRT.exe':
+            command = f'\"{ABS_EXECUTABLE_PATH}\" --definition={test_option} --export=\"{REPORT_PATH}\" --select-cuda-device {device_id}'
     command = command.rstrip()
     return command
 
@@ -119,7 +142,9 @@ try:
     setup_logging()
     args = get_arguments()
     option = BENCHMARK_CONFIG[args.engine]["config"]
-    cmd = create_procyon_command(option)
+    process_name = BENCHMARK_CONFIG[args.engine]["process_name"]
+    device_id = BENCHMARK_CONFIG[args.engine]["device_id"]
+    cmd = create_procyon_command(option, process_name, device_id)
     logging.info('Starting benchmark!')
     logging.info(cmd)
     start_time = time.time()
@@ -141,6 +166,7 @@ try:
 
     report = {
         "test": BENCHMARK_CONFIG[args.engine]["test_name"],
+        "device_name": BENCHMARK_CONFIG[args.engine]["device_name"],
         "unit": "score",
         "score": score,
         "start_time": seconds_to_milliseconds(start_time),

procyon_ai/utils.py

@@ -4,6 +4,7 @@ import psutil
 import xml.etree.ElementTree as ET
 import winreg
 import re
+import subprocess
 
 SCRIPT_DIR = Path(__file__).resolve().parent
 LOG_DIR = SCRIPT_DIR / "run"
@@ -33,4 +34,19 @@ def get_install_path() -> str:
     reg_path = r"Software\UL\Procyon"
     reg_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, reg_path, 0, winreg.KEY_READ)
     value, _  = winreg.QueryValueEx(reg_key, "InstallDir")
-    return value
+    return value
+
+def get_winml_devices(procyon_path):
+    """
+    Function which uses the ProcyonCmd.exe to list all available winml devices on the system. Returns a dictionary of device names and IDs
+    """
+
+    # run the command line utility for procyon in order to list available winml devices
+    winml_devices = subprocess.run([f'{procyon_path}', '--list-winml-devices'], shell=True, capture_output=True, text=True, check=True).stdout
+
+    winml_devices_split = winml_devices.split('\n')
+    winml_devices_parsed = [device[9::] for device in winml_devices_split if re.search(r"(amd|nvidia|intel)", device.lower())]
+    unique_winml_devices = list(dict.fromkeys(winml_devices_parsed))
+    winml_dict = {device_split.split(', ')[0]:device_split.split(', ')[1] for device_split in unique_winml_devices}
+
+    return winml_dict