From e2c6599897152daa459bbf2982ee3538e599e6b3 Mon Sep 17 00:00:00 2001
From: Arty Blue <a.blouin@linusmediagroup.com>
Date: Mon, 23 Dec 2024 08:45:48 -0800
Subject: [PATCH] initial implementation for cuda devices, remains to be tested

---
 procyon_ai/ulprocai.py | 15 ++++++++++++++-
 procyon_ai/utils.py    | 26 +++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/procyon_ai/ulprocai.py b/procyon_ai/ulprocai.py
index 3b0c8c5..a7aa414 100644
--- a/procyon_ai/ulprocai.py
+++ b/procyon_ai/ulprocai.py
@@ -6,7 +6,15 @@ import subprocess
 import sys
 import time
 import psutil
-from utils import find_score_in_xml, is_process_running, get_install_path, get_winml_devices, get_openvino_devices, get_openvino_gpu
+from utils import (
+    find_score_in_xml,
+    is_process_running,
+    get_install_path,
+    get_winml_devices,
+    get_openvino_devices,
+    get_openvino_gpu,
+    get_cuda_devices
+)
 
 PARENT_DIR = str(Path(sys.path[0], ".."))
 sys.path.append(PARENT_DIR)
@@ -30,12 +38,15 @@ ABS_EXECUTABLE_PATH = DIR_PROCYON / EXECUTABLE
 
 WINML_DEVICES = get_winml_devices(ABS_EXECUTABLE_PATH)
 OPENVINO_DEVICES = get_openvino_devices(ABS_EXECUTABLE_PATH)
+CUDA_DEVICES = get_cuda_devices(ABS_EXECUTABLE_PATH)
 
 CONFIG_DIR = SCRIPT_DIR / "config"
 BENCHMARK_CONFIG = {
     "AMD_CPU": {
         "config": f"\"{CONFIG_DIR}\\ai_computer_vision_winml_cpu.def\"",
         "process_name":  "WinML.exe",
+        "device_name": "CPU",
+        "device_id": "CPU", # TODO: Find a good way to report the CPU name here. 
         "test_name": "WinML CPU (FLOAT32)"
     },
     "AMD_GPU0": {
@@ -82,6 +93,8 @@ BENCHMARK_CONFIG = {
     },
     "NVIDIA_GPU": {
         "config": f"\"{CONFIG_DIR}\\ai_computer_vision_tensorrt.def\"",
+        "device_id": "cuda:0",
+        "device_name": CUDA_DEVICES.get("cuda:0"),
         "process_name":  "TensorRT.exe",
         "test_name": "NVIDIA TensorRT (FLOAT32)"
     },
diff --git a/procyon_ai/utils.py b/procyon_ai/utils.py
index d3598f7..6cb6f70 100644
--- a/procyon_ai/utils.py
+++ b/procyon_ai/utils.py
@@ -61,7 +61,7 @@ def get_openvino_devices(procyon_path):
     openvino_devices_split = openvino_devices.split('\n')
     openvino_devices_parsed = [device[9::] for device in openvino_devices_split if re.search(r"(amd|nvidia|intel)", device.lower())]
     unique_openvino_devices = list(dict.fromkeys(openvino_devices_parsed))
-    openvino_dict = {device_split.split('        , ')[1]:device_split.split('        , ')[0] for device_split in unique_openvino_devices}
+    openvino_dict = {device_split.split(', ')[1]:device_split.split(', ')[0] for device_split in unique_openvino_devices}
 
     return openvino_dict
 
@@ -80,3 +80,27 @@ def get_openvino_gpu(openvino_devices, gpu_id):
         gpu = openvino_devices.get(gpu_id, "No Openvino GPU Detected")
 
     return gpu
+
+def get_cuda_devices(procyon_path):
+    """
+    Function which uses the ProcyonCmd.exe to list all available openvino devices on the system. Returns a dictionary of device type and name
+    """
+
+    cuda_devices = subprocess.run([f'{procyon_path}', '--list-cuda-devices'], shell=True, capture_output=True, text=True, check=True).stdout
+
+    cuda_devices_split = cuda_devices.split('\n')
+    cuda_devices_parsed = [device[9::] for device in cuda_devices_split if re.search(r"(nvidia)", device.lower())]
+    unique_cuda_devices = list(dict.fromkeys(cuda_devices_parsed))
+
+    if len(unique_cuda_devices) > 0:
+        cuda_dict = {device_split.split(', ')[1]:device_split.split(', ')[0] for device_split in unique_cuda_devices}
+    else:
+        cuda_dict = {}
+
+    return cuda_dict
+
+DIR_PROCYON = Path(get_install_path())
+EXECUTABLE = "ProcyonCmd.exe"
+ABS_EXECUTABLE_PATH = DIR_PROCYON / EXECUTABLE
+
+print(get_cuda_devices(ABS_EXECUTABLE_PATH))