From 5a75dba9dcd854b3187855f6e00fc3cfe0677f26 Mon Sep 17 00:00:00 2001 From: Ruturaj Vaidya Date: Fri, 24 May 2024 02:49:43 -0700 Subject: [PATCH] PR #12864: [ROCm] fix DeviceAllocate code Imported from GitHub PR https://github.com/openxla/xla/pull/12864 Copybara import of the project: -- 2e6dfa841a495c1f0babf0ff6f877bc1866bf319 by Ruturaj4 : [ROCm] fix DeviceAllocate code Merging this change closes #12864 PiperOrigin-RevId: 636848518 --- .../xla/xla/stream_executor/rocm/rocm_driver.cc | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc index ba4a7913ecfcb0..8daa398db7e5a2 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc +++ b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc @@ -1318,13 +1318,19 @@ struct BitPatternToValue { /* static */ void* GpuDriver::DeviceAllocate(GpuContext* context, uint64_t bytes) { + if (bytes == 0) { + return nullptr; + } + ScopedActivateContext activated{context}; hipDeviceptr_t result = 0; hipError_t res = wrap::hipMalloc(&result, bytes); if (res != hipSuccess) { - LOG(ERROR) << "failed to allocate " - << tsl::strings::HumanReadableNumBytes(bytes) << " (" << bytes - << " bytes) from device: " << ToString(res); + // LOG(INFO) because this isn't always important to users (e.g. BFCAllocator + // implements a retry if the first allocation fails). + LOG(INFO) << "failed to allocate " + << tsl::strings::HumanReadableNumBytes(bytes) << " (" << bytes + << " bytes) from device: " << ToString(res); return nullptr; } void* ptr = reinterpret_cast(result);