From 0c4edf7d0edcbf8fee88f57866906ac72995d4ce Mon Sep 17 00:00:00 2001
From: R0CKSTAR <xiaodong.ye@mthreads.com>
Date: Thu, 12 Sep 2024 19:30:02 +0800
Subject: [PATCH] feat: Moore Threads GPU add support to query number of cores
 (#1259)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
---
 src/detection/gpu/gpu_mthreads.c | 5 +++++
 src/detection/gpu/mtml.h         | 2 ++
 2 files changed, 7 insertions(+)

diff --git a/src/detection/gpu/gpu_mthreads.c b/src/detection/gpu/gpu_mthreads.c
index 0104518fc..ac0ffa07b 100644
--- a/src/detection/gpu/gpu_mthreads.c
+++ b/src/detection/gpu/gpu_mthreads.c
@@ -5,6 +5,7 @@
 
 struct FFMtmlData
 {
+    FF_LIBRARY_SYMBOL(mtmlDeviceCountGpuCores)
     FF_LIBRARY_SYMBOL(mtmlDeviceGetBrand)
     FF_LIBRARY_SYMBOL(mtmlDeviceGetIndex)
     FF_LIBRARY_SYMBOL(mtmlDeviceGetName)
@@ -43,6 +44,7 @@ const char *ffDetectMthreadsGpuInfo(const FFGpuDriverCondition *cond, FFGpuDrive
         mtmlData.inited = true;
         FF_LIBRARY_LOAD(libmtml, "dlopen mtml failed", soName, 1);
         FF_LIBRARY_LOAD_SYMBOL_MESSAGE(libmtml, mtmlLibraryInit)
+        FF_LIBRARY_LOAD_SYMBOL_VAR_MESSAGE(libmtml, mtmlData, mtmlDeviceCountGpuCores)
         FF_LIBRARY_LOAD_SYMBOL_VAR_MESSAGE(libmtml, mtmlData, mtmlDeviceGetBrand)
         FF_LIBRARY_LOAD_SYMBOL_VAR_MESSAGE(libmtml, mtmlData, mtmlDeviceGetIndex)
         FF_LIBRARY_LOAD_SYMBOL_VAR_MESSAGE(libmtml, mtmlData, mtmlDeviceGetName)
@@ -158,6 +160,9 @@ const char *ffDetectMthreadsGpuInfo(const FFGpuDriverCondition *cond, FFGpuDrive
         }
     }
 
+    if (result.coreCount)
+        mtmlData.ffmtmlDeviceCountGpuCores(device, result.coreCount);
+
     if (result.frequency)
     {
         MtmlGpu *gpu = NULL;
diff --git a/src/detection/gpu/mtml.h b/src/detection/gpu/mtml.h
index be5a6cee5..a97cfeb37 100644
--- a/src/detection/gpu/mtml.h
+++ b/src/detection/gpu/mtml.h
@@ -56,6 +56,8 @@ typedef struct
     int rsvd[6];                                 //!< Reserved for future extension.
 } MtmlPciInfo;
 
+// Retrieves the number of cores of a device.
+MtmlReturn MTML_API mtmlDeviceCountGpuCores(const MtmlDevice* device, unsigned int* numCores);
 // Retrieves the brand of a device.
 MtmlReturn MTML_API mtmlDeviceGetBrand(const MtmlDevice *dev, MtmlBrandType *type);
 // Retrieves the index associated with the specified device.