UTIL: add Nvidia Grace CPU detection (openucx#941)

nsarka · Mar 8, 2024 · 9e0d759 · 9e0d759
1 parent 8735881
commit 9e0d759
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/src/utils/arch/aarch64/cpu.h b/src/utils/arch/aarch64/cpu.h
@@ -1,5 +1,5 @@
 /**
-* Copyright (c) 2001-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+* Copyright (c) 2001-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * Copyright (C) ARM Ltd. 2016-2020.  ALL RIGHTS RESERVED.
 * Copyright (C) Stony Brook University. 2016-2020.  ALL RIGHTS RESERVED.
 *
@@ -47,11 +47,6 @@ typedef struct ucc_aarch64_cpuid {
  */
 void ucc_aarch64_cpuid(ucc_aarch64_cpuid_t *cpuid);
 
-static inline ucc_cpu_model_t ucc_arch_get_cpu_model()
-{
-    return UCC_CPU_MODEL_ARM_AARCH64;
-}
-
 static inline ucc_cpu_vendor_t ucc_arch_get_cpu_vendor()
 {
     ucc_aarch64_cpuid_t cpuid;
@@ -61,7 +56,25 @@ static inline ucc_cpu_vendor_t ucc_arch_get_cpu_vendor()
         return UCC_CPU_VENDOR_FUJITSU_ARM;
     }
 
+    if ((cpuid.implementer == 0x41) && (cpuid.architecture == 8)) {
+        return UCC_CPU_VENDOR_NVIDIA;
+    }
+
     return UCC_CPU_VENDOR_GENERIC_ARM;
 }
 
+static inline ucc_cpu_model_t ucc_arch_get_cpu_model()
+{
+    ucc_aarch64_cpuid_t cpuid;
+    ucc_aarch64_cpuid(&cpuid);
+
+    if ((ucc_arch_get_cpu_vendor() == UCC_CPU_VENDOR_NVIDIA) &&
+        (cpuid.part == 0xd4f)) {
+        return UCC_CPU_MODEL_NVIDIA_GRACE;
+    }
+
+    return UCC_CPU_MODEL_ARM_AARCH64;
+}
+
+
 #endif
diff --git a/src/utils/arch/cpu.h b/src/utils/arch/cpu.h
@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2023. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2024. ALL RIGHTS RESERVED.
 * Copyright (C) ARM Ltd. 2016.  ALL RIGHTS RESERVED.
 * Copyright (C) Shanghai Zhaoxin Semiconductor Co., Ltd. 2020. ALL RIGHTS RESERVED.
 * Copyright (C) Rivos Inc. 2023
@@ -35,6 +35,7 @@ typedef enum ucc_cpu_model {
     UCC_CPU_MODEL_ZHAOXIN_ZHANGJIANG,
     UCC_CPU_MODEL_ZHAOXIN_WUDAOKOU,
     UCC_CPU_MODEL_ZHAOXIN_LUJIAZUI,
+    UCC_CPU_MODEL_NVIDIA_GRACE,
     UCC_CPU_MODEL_LAST
 } ucc_cpu_model_t;
 
@@ -48,6 +49,7 @@ typedef enum ucc_cpu_vendor {
     UCC_CPU_VENDOR_GENERIC_RISCV,
     UCC_CPU_VENDOR_FUJITSU_ARM,
     UCC_CPU_VENDOR_ZHAOXIN,
+    UCC_CPU_VENDOR_NVIDIA,
     UCC_CPU_VENDOR_LAST
 } ucc_cpu_vendor_t;
 
@@ -67,6 +69,8 @@ static inline ucc_cpu_vendor_t ucc_get_vendor_from_str(const char *v_name)
         return UCC_CPU_VENDOR_FUJITSU_ARM;
     if (strcasecmp(v_name, "zhaoxin") == 0)
         return UCC_CPU_VENDOR_ZHAOXIN;
+    if (strcasecmp(v_name, "nvidia") == 0)
+        return UCC_CPU_VENDOR_NVIDIA;
     return UCC_CPU_VENDOR_UNKNOWN;
 }
 
@@ -102,6 +106,8 @@ static inline ucc_cpu_model_t ucc_get_model_from_str(const char *m_name)
         return UCC_CPU_MODEL_ZHAOXIN_WUDAOKOU;
     if (strcasecmp(m_name, "lujiazui") == 0)
         return UCC_CPU_MODEL_ZHAOXIN_LUJIAZUI;
+    if (strcasecmp(m_name, "grace") == 0)
+        return UCC_CPU_MODEL_NVIDIA_GRACE;
     return UCC_CPU_MODEL_UNKNOWN;
 }