Check for nullptr return from malloc called from AllocateBuffers in L…

…iteral. PiperOrigin-RevId: 718967904
openxla · Feb 3, 2025 · 34d2633 · 34d2633
1 parent 3905a2c
commit 34d2633
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/xla/BUILD b/xla/BUILD
@@ -587,6 +587,7 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",

diff --git a/xla/literal.cc b/xla/literal.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "absl/base/casts.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/functional/function_ref.h"
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
@@ -627,8 +628,10 @@ void LiteralBase::Piece::AllocateBuffers() {
   if (bytes > kMaxInlinedBytes) {
     CHECK_EQ(buffer(), nullptr);
     rep_.emplace<DenseRep>();
-    set_buffer(
-        static_cast<char*>(tsl::port::AlignedMalloc(bytes, kMinimumAlignment)));
+    char* buffer =
+        static_cast<char*>(tsl::port::AlignedMalloc(bytes, kMinimumAlignment));
+    CHECK(buffer != nullptr) << "Failed to allocate buffer for Literal";
+    set_buffer(buffer);
   } else {
     rep_.emplace<DenseInlinedRep>();
   }