Add more GPU information in the results file

bench-TriLMs.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from urllib import request
 import os
 import shlex
+import shutil
 import subprocess
 import sys
 from typing import Any, Sequence
@@ -208,15 +209,21 @@ if __name__ == "__main__":
         mulmat_perf.append(test_backend_perf())
         results.extend(llama_bench(repetitions=repetitions, types=GPU_TYPES))
 
-    cpuinfo = subprocess.run(["lscpu"], capture_output=True).stdout.decode(
-        encoding="utf-8"
-    )
-
-    final_result = {
-        "cpuinfo": cpuinfo,
+    final_result: dict[str, Any] = {
         "mulmat_perf": mulmat_perf,
         "results": results,
     }
+
+    if shutil.which("lscpu") is not None:
+        logger.info("Getting CPU info")
+        final_result["cpuinfo"] = subprocess.run(["lscpu"], capture_output=True).stdout.decode(
+            encoding="utf-8"
+        )
+
+    if args.gpu and shutil.which("nvidia-smi") is not None:
+        logger.info("Getting NVIDIA GPU info")
+        final_result["gpuinfo"] = subprocess.run(["nvidia-smi", "-q"], capture_output=True).stdout.decode(encoding="utf-8")
+
     logger.info("Writing output to: %s", output_file)
     logger.debug("Final results: %s", json.dumps(final_result, indent=4))
     with open(output_file, "w") as f: