add memory status impl for ascend (#834)

DeepLink-org · Jun 12, 2024 · 23401fc · 23401fc
1 parent 959d85a
commit 23401fc
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/dipu/torch_dipu/csrc_dipu/vendor/ascend/deviceimpl.cpp b/dipu/torch_dipu/csrc_dipu/vendor/ascend/deviceimpl.cpp
@@ -124,12 +124,20 @@ DIPUDeviceProperties getDeviceProperties(AscendDeviceId device_index) {
   int patch;
   DIPU_CALLACLRT(::aclrtGetVersion(&prop.major, &prop.minor, &patch));
   DIPU_CALLACLRT(::aclrtGetMemInfo(ACL_HBM_MEM, &device_free, &device_total));
-  // NOTE : unit of PhysicalMemoryTotal is MB
-  prop.totalGlobalMem = device_total << 20;
+  // NOTE : unit of PhysicalMemoryTotal is Byte
+  prop.totalGlobalMem = device_total;
   prop.multiProcessorCount = 1;
   return prop;
 }
 
+DIPUDeviceStatus getDeviceStatus(int32_t device_index) {
+  DIPUDeviceStatus status;
+  size_t device_total;
+  DIPU_CALLACLRT(
+      ::aclrtGetMemInfo(ACL_HBM_MEM, &status.freeGlobalMem, &device_total));
+  return status;
+}
+
 void resetDevice(deviceId_t devId) { DIPU_CALLACLRT(::aclrtResetDevice(devId)) }
 
 void syncDevice() { DIPU_CALLACLRT(::aclrtSynchronizeDevice()) }