fix hard-coded device properties and remove reserved memory (#351)

Author: ScXfjiang

xla/stream_executor/rocm/rocm_context.cc

@@ -60,45 +60,6 @@ hipCtx_t CurrentContext() {
   return current;
 }
 
-// Returns the amount of memory reserved by ROCm libraries.
-bool GetReservedMemory(uint64_t* reserve) {
-  hipDeviceProp_t props;
-  hipDevice_t dev;
-  hipError_t res = wrap::hipGetDevice(&dev);
-
-  if (res != hipSuccess) {
-    LOG(FATAL) << "failed to query current device: " << ToString(res);
-    return false;
-  }
-  res = wrap::hipGetDeviceProperties(&props, dev);
-  if (res != hipSuccess) {
-    LOG(ERROR) << "failed to query device properties: " << ToString(res);
-    return false;
-  }
-
-  std::string gcnArchName = props.gcnArchName;
-  auto compute_capability = RocmComputeCapability(gcnArchName);
-  // On gfx90a, we hide 1 GB of GPU memory (512MB for gfx908) from TF,
-  // to allow for late allocations by internal ROCm libraries
-  // (e.g. rocBLAS alone needs~200 MB to put its kernels as of ROCm 4.1)
-  const uint64_t RESERVED_GFX908 = 1048576 * 512;
-  const uint64_t RESERVED_GFX9_X = 1048576 * 1024;
-  const uint64_t RESERVED_GFX10_X = 1048576 * 512;
-  const uint64_t RESERVED_GFX11_X = 1048576 * 512;
-  if (compute_capability.gfx9_mi100()) {
-    *reserve = RESERVED_GFX908;
-  } else if (compute_capability.gfx9_mi200_or_later()) {
-    *reserve = RESERVED_GFX9_X;
-  } else if (compute_capability.gfx10_rx68xx() ||
-             compute_capability.gfx10_rx69xx()) {
-    *reserve = RESERVED_GFX10_X;
-  } else if (compute_capability.gfx11()) {
-    *reserve = RESERVED_GFX11_X;
-  }
-
-  return true;
-}
-
 }  // namespace
 
 // Returns the singleton ContextMap.
@@ -126,12 +87,7 @@ bool RocmContext::GetDeviceTotalMemory(hipDevice_t device, uint64_t* result) {
     LOG(ERROR) << "failed to query total available memory: " << ToString(res);
     return false;
   }
-  uint64_t reserve = 0;
-  if (!GetReservedMemory(&reserve)) {
-    LOG(ERROR) << "failed to reserved device memory for ROCm libraries";
-    return false;
-  }
-  *result = value - reserve;
+  *result = value;
   return true;
 }
 
@@ -145,24 +101,17 @@ bool RocmContext::GetDeviceMemoryUsage(int64_t* free_out, int64_t* total_out) {
     return false;
   }
 
-  uint64_t reserve = 0;
-  if (!GetReservedMemory(&reserve)) {
-    LOG(ERROR) << "failed to reserved device memory for ROCm libraries";
-    return false;
-  }
-
   VLOG(1) << "Device memory: " << total / 1048576 << " MB total, "
-          << free / 1048576 << " MB free, reserving " << reserve / 1048576
-          << " MB";
+          << free / 1048576 << " MB free";
 
   // overflow check
   if (free > std::numeric_limits<int64_t>::max()) {
     LOG(ERROR) << "free memory (" << free << ") is overflow int64_t";
     return false;
   }
 
-  *free_out = free >= reserve ? free - reserve : 0;
-  *total_out = total - reserve;
+  *free_out = free;
+  *total_out = total;
   return true;
 }
 

xla/stream_executor/rocm/rocm_executor.cc

@@ -306,6 +306,10 @@ absl::Status GetGridLimits(int* x, int* y, int* z, hipDevice_t device) {
   return absl::OkStatus();
 }
 
+absl::StatusOr<int64_t> GetMaxRegistersPerMultiprocessor(hipDevice_t device) {
+  return GetSimpleAttribute<int64_t>(device, hipDeviceAttributeMaxRegistersPerMultiprocessor);
+}
+
 // Returns the device associated with the given device_ordinal.
 absl::StatusOr<hipDevice_t> GetDevice(int device_ordinal) {
   hipDevice_t device;
@@ -394,6 +398,19 @@ std::string GetPCIBusID(hipDevice_t device) {
   return pci_bus_id;
 }
 
+
+bool IsEccEnabled(hipDevice_t device, bool* result) {
+  int value = -1;
+  auto status = ToStatus(wrap::hipDeviceGetAttribute(
+      &value, hipDeviceAttributeEccEnabled, device));
+  if (!status.ok()) {
+    LOG(ERROR) << "failed to query ECC status: " << status;
+    return false;
+  }
+  *result = value;
+  return true;
+}
+
 bool GetDeviceProperties(hipDeviceProp_t* device_properties,
                          int device_ordinal) {
   hipError_t res =
@@ -1015,8 +1032,11 @@ RocmExecutor::CreateDeviceDescription(int device_ordinal) {
     desc.set_l2_cache_size(prop.l2CacheSize);
   }
 
-  // No way to query ECC status from the API.
-  desc.set_ecc_enabled(false);
+  {
+    bool ecc_enabled = false;
+    IsEccEnabled(device, &ecc_enabled);
+    desc.set_ecc_enabled(ecc_enabled);
+  }
 
   uint64_t device_memory_size = -1;
   (void)RocmContext::GetDeviceTotalMemory(device, &device_memory_size);
@@ -1054,7 +1074,7 @@ RocmExecutor::CreateDeviceDescription(int device_ordinal) {
       GetMaxThreadsPerMultiprocessor(device).value());
   desc.set_registers_per_block_limit(GetMaxRegistersPerBlock(device).value());
   desc.set_threads_per_warp(GetThreadsPerWarp(device).value());
-  desc.set_registers_per_core_limit(64 * 1024);
+  desc.set_registers_per_core_limit(GetMaxRegistersPerMultiprocessor(device).value());
   desc.set_compile_time_toolkit_version(
       SemanticVersion{HIP_VERSION_MAJOR, HIP_VERSION_MINOR, HIP_VERSION_PATCH});
   int32_t runtime_version;