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

…iteral. PiperOrigin-RevId: 718967904
openxla · Jan 31, 2025 · b826949 · b826949
1 parent 3219006
commit b826949
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 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"
@@ -332,6 +333,9 @@ Literal LiteralBase::CreateFromShape(const Shape& shape) {
   literal.root_piece_.ForEachMutableSubpiece(
       [&](const ShapeIndex& index, Piece* piece) {
         if (piece->subshape().IsArray()) {
+          CHECK_LE(piece->size_bytes_dense(),
+                   1ULL * 1024 * 1024 * 1024)  // 1 TB
+              << "Literal size is too large";
           memset(piece->untyped_data(), 0, piece->size_bytes_dense());
         }
       });