[pull] master from tensorflow:master

ci/official/containers/ml_build/cuda12.1_cudnn9.10.packages.txt

@@ -0,0 +1,23 @@
+# All required CUDA packages
+cuda-compat-12-1
+cuda-command-line-tools-12-1
+cuda-cudart-dev-12-1
+cuda-nvcc-12-1
+cuda-cupti-12-1
+cuda-nvprune-12-1
+cuda-libraries-12-1
+cuda-libraries-dev-12-1
+cuda-nvml-dev-12-1
+libcufft-12-1
+libcurand-12-1
+libcusolver-dev-12-1
+libcusparse-dev-12-1
+libcublas-12-1
+libcublas-dev-12-1
+libnccl-dev=2.18.3-1+cuda12.1
+libnccl2=2.18.3-1+cuda12.1
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-headers-cuda-12=9.10.2.21-1
+libcudnn9-static-cuda-12=9.10.2.21-1
+libcudnn9-dev-cuda-12=9.10.2.21-1
+libcudnn9-cuda-12=9.10.2.21-1

tensorflow/cc/saved_model/fingerprinting.cc

@@ -173,11 +173,6 @@ absl::StatusOr<FingerprintDef> CreateFingerprintDefPb(
   SavedModel saved_model;
   TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), pb_file, &saved_model));
 
-  if (saved_model.meta_graphs_size() == 0) {
-    return absl::InvalidArgumentError(
-        "SavedModel (.pb) contains no MetaGraphs.");
-  }
-
   // Create a copy of `metagraph` which will be used and mutated for fingerprint
   // computation.
   FingerprintDef fingerprint_def;
@@ -242,12 +237,11 @@ absl::StatusOr<FingerprintDef> CreateFingerprintDef(
   // At this point we have neither saved_model.pb nor saved_model.cpb.
   return CreateReducedFingerprintDef();  // Only sets the UUID.
 #else  // The following runs on Windows and Mac.
-  absl::StatusOr<FingerprintDef> fingerprint_def =
-      CreateFingerprintDefPb(export_dir, absl::StrCat(prefix, ".pb"));
-  if (!fingerprint_def.ok()) {
-    return CreateReducedFingerprintDef();
+  std::string pb_file = absl::StrCat(prefix, ".pb");
+  if (Env::Default()->FileExists(pb_file).ok()) {
+    return CreateFingerprintDefPb(export_dir, pb_file);
   }
-  return fingerprint_def;
+  return CreateReducedFingerprintDef();
 #endif
 }
 

tensorflow/cc/saved_model/fingerprinting_test.cc

@@ -25,7 +25,6 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/versions.pb.h"
-#include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/test.h"
@@ -207,20 +206,5 @@ TEST(FingerprintingTest, TestSingleprint) {
             const_singleprint);
 }
 
-TEST(FingerprintingTest, CreateFingerprintDefPbEmptyMetaGraphsReturnsError) {
-  const std::string model_dir =
-      io::JoinPath(::testing::TempDir(), "empty_metagraphs_model");
-  TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(model_dir));
-  const std::string pb_file = io::JoinPath(model_dir, "saved_model.pb");
-  SavedModel saved_model;
-  TF_ASSERT_OK(WriteBinaryProto(Env::Default(), pb_file, saved_model));
-
-  absl::StatusOr<FingerprintDef> result = CreateFingerprintDef(model_dir);
-  EXPECT_FALSE(result.ok());
-  EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument);
-  EXPECT_EQ(result.status().message(),
-            "SavedModel (.pb) contains no MetaGraphs.");
-}
-
 }  // namespace
 }  // namespace tensorflow::saved_model::fingerprinting

tensorflow/compiler/aot/compile.cc

@@ -198,8 +198,7 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config,
   xla::cpu::CpuAotCompilationOptions aot_opts(
       flags.target_triple, flags.target_cpu, flags.target_features,
       flags.entry_point,
-      xla::cpu::CpuAotCompilationOptions::RelocationModel::BigPic,
-      /*compile_copy_as_llvm_kernel=*/true);
+      xla::cpu::CpuAotCompilationOptions::RelocationModel::BigPic);
 
   if (flags.sanitize_dataflow) {
     aot_opts.set_sanitize_dataflow(flags.sanitize_dataflow);

tensorflow/compiler/mlir/lite/tests/canonicalize.mlir

@@ -204,7 +204,7 @@ func.func @WhileCanonicalizeBug1(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tens
 // -----
 
 // Test case to test While op with resources that are not read-only variables.
-// Do not remove resource arugments if they are not read-only variables to keep
+// Do not remove resource arguments if they are not read-only variables to keep
 // the graph's control dependency.
 // CHECK-LABEL: WhileWithNonReadOnlyVariableResources
 func.func @WhileWithNonReadOnlyVariableResources(%arg0: tensor<i32>) -> tensor<!tf_type.resource> {

tensorflow/compiler/mlir/lite/transforms/lower_static_tensor_list.cc

@@ -1414,7 +1414,7 @@ struct ConvertIf : public OpConversionPattern<TF::IfOp> {
   LogicalResult matchAndRewrite(
       TF::IfOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter) const override {
-    // Find all Tensor List arugments.
+    // Find all Tensor List arguments.
     auto tensor_list_args = GetTensorListArgumentsIndex(op.else_function());
     auto tensor_list_results = GetTensorListResultsIndex(op.else_function());
     auto tensor_list_map = MapTensorListResultToArgument(op.else_function());
@@ -1451,7 +1451,7 @@ struct ConvertWhile : public OpConversionPattern<TF::WhileOp> {
   LogicalResult matchAndRewrite(
       TF::WhileOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter) const override {
-    // Find all Tensor List arugments.
+    // Find all Tensor List arguments.
     auto tensor_list_args = GetTensorListArgumentsIndex(op.body_function());
 
     llvm::SmallVector<Type, 8> result_types;

tensorflow/compiler/mlir/lite/transforms/reduce_while_operands.cc

@@ -147,7 +147,7 @@ bool AllOperationSafe(Block &block) {
     }
     // op has implict arguments not listed in operands.
     // Fact: if every op's operands are defined in the same block as op,
-    //       then no operation has implicit arugments (constant doesn't count).
+    //       then no operation has implicit arguments (constant doesn't count).
     for (auto operand : op->getOperands()) {
       if (auto arg = mlir::dyn_cast_or_null<BlockArgument>(operand)) {
         if (arg.getOwner() == &block) {

tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-collective.mlir

@@ -290,7 +290,7 @@ func.func @inconsistent_collective_info(%input: tensor<f32>) -> tensor<f32> {
   %group_size1 = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
   %group_size2 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
   %instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
-  // expected-error@below {{op module already contains an attribute tf2xla.collective_info.group_size=2, overwritting to a new value 1 is not allowed.}}
+  // expected-error@below {{op module already contains an attribute tf2xla.collective_info.group_size=2, overwriting to a new value 1 is not allowed.}}
   %0 = "tf.CollectiveReduceV2"(%input, %group_size1, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
   %1 = "tf.CollectiveReduceV2"(%input, %group_size2, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>
   %2 = "tf.Add"(%0, %1) : (tensor<f32>, tensor<f32>) -> tensor<f32>

tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_collective.cc

@@ -75,7 +75,7 @@ LogicalResult SetOnceModuleAttribute(StringRef attr_name,
   }
   return op->emitOpError() << "module already contains an attribute "
                            << attr_name << "=" << ex_attr_value.getInt()
-                           << ", overwritting to a new value "
+                           << ", overwriting to a new value "
                            << attr_value.getInt() << " is not allowed.";
 }
 

tensorflow/dtensor/mlir/tests/lower_send_recv.mlir

@@ -29,7 +29,7 @@ func.func @main(%arg0: tensor<i32>) {
   // CHECK-DAG:    %[[RECV_SIZE_TYPE]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
   // CHECK-DAG:    %[[RECV_SLICE_SIZE]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
   // CHECK-DAG:    %[[RECV_SCALAR_TYPE]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}>
-  // COMMENT: Recv and Send seperated by the output tensor.
+  // COMMENT: Recv and Send separated by the output tensor.
   // CHECK:   %[[PROGRAM_KEY:.*]] = "tf._XlaCompileMlirPlaceholderProgramKey"
   // CHECK-NEXT:   %[[CONST_OUT:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}>
   // CHECK-NEXT:   %[[LAYOUT_OUT:.*]] = "tf.DTensorLayout"(%[[CONST_OUT]])

tensorflow/lite/async/testing/mock_async_kernel.h

@@ -26,7 +26,7 @@ namespace async {
 namespace testing {
 
 // A fully mocked out async kernel.
-// Mocked TfLiteAsyncKernel can be retreived by `MockAsyncKernel::kernel()`.
+// Mocked TfLiteAsyncKernel can be retrieved by `MockAsyncKernel::kernel()`.
 class MockAsyncKernel : public delegates::BackendAsyncKernelInterface {
  public:
   MOCK_METHOD(TfLiteStatus, RegisterBuffer,

tensorflow/lite/examples/label_image/bitmap_helpers.cc

@@ -22,31 +22,136 @@ limitations under the License.
 #include <cstdlib>
 #include <fstream>
 #include <iostream>
+#include <limits>
 #include <string>
 #include <vector>
 
 #include "tensorflow/lite/examples/label_image/label_image.h"
 #include "tensorflow/lite/examples/label_image/log.h"
-#include "tsl/platform/ctstring_internal.h"
 
 namespace tflite {
 namespace label_image {
+namespace {
+
+constexpr size_t kBmpHeaderMinSize = 30;
+
+uint16_t ReadLe16(const std::vector<uint8_t>& bytes, size_t offset) {
+  return static_cast<uint16_t>(bytes[offset]) |
+         static_cast<uint16_t>(bytes[offset + 1]) << 8;
+}
+
+int32_t ReadLe32(const std::vector<uint8_t>& bytes, size_t offset) {
+  const uint32_t value = static_cast<uint32_t>(bytes[offset]) |
+                         static_cast<uint32_t>(bytes[offset + 1]) << 8 |
+                         static_cast<uint32_t>(bytes[offset + 2]) << 16 |
+                         static_cast<uint32_t>(bytes[offset + 3]) << 24;
+  return static_cast<int32_t>(value);
+}
+
+bool ValidateBmpAndGetPixelOffset(const std::vector<uint8_t>& img_bytes,
+                                  int* width, int* height, int* channels,
+                                  int* row_size, size_t* pixel_offset) {
+  if (width == nullptr || height == nullptr || channels == nullptr ||
+      row_size == nullptr || pixel_offset == nullptr) {
+    LOG(ERROR)
+        << "Null pointer argument passed to ValidateBmpAndGetPixelOffset";
+    return false;
+  }
+  *width = 0;
+  *height = 0;
+  *channels = 0;
+  *row_size = 0;
+  *pixel_offset = 0;
+
+  if (img_bytes.size() < kBmpHeaderMinSize || img_bytes[0] != 'B' ||
+      img_bytes[1] != 'M') {
+    LOG(ERROR) << "Invalid BMP header";
+    return false;
+  }
+
+  const int32_t parsed_pixel_offset = ReadLe32(img_bytes, 10);
+  const int32_t parsed_width = ReadLe32(img_bytes, 18);
+  const int32_t parsed_height = ReadLe32(img_bytes, 22);
+  const uint16_t bpp = ReadLe16(img_bytes, 28);
+
+  if (parsed_pixel_offset < static_cast<int32_t>(kBmpHeaderMinSize) ||
+      static_cast<size_t>(parsed_pixel_offset) > img_bytes.size()) {
+    LOG(ERROR) << "BMP pixel data offset is invalid or outside the file";
+    return false;
+  }
+  if (parsed_width <= 0 || parsed_height == 0 ||
+      parsed_height == std::numeric_limits<int32_t>::min()) {
+    LOG(ERROR) << "Invalid BMP dimensions";
+    return false;
+  }
+  if (bpp != 8 && bpp != 24 && bpp != 32) {
+    LOG(ERROR) << "Unsupported BMP bits per pixel: " << bpp;
+    return false;
+  }
+
+  const int parsed_channels = bpp / 8;
+  const int64_t abs_height = parsed_height < 0
+                                 ? -static_cast<int64_t>(parsed_height)
+                                 : static_cast<int64_t>(parsed_height);
+  const uint64_t total_output_bytes = static_cast<uint64_t>(parsed_width) *
+                                      static_cast<uint64_t>(abs_height) *
+                                      static_cast<uint64_t>(parsed_channels);
+  if (total_output_bytes >
+      static_cast<uint64_t>(std::numeric_limits<int>::max())) {
+    LOG(ERROR) << "Decoded BMP size exceeds maximum supported size";
+    return false;
+  }
+  const int64_t bits_per_row = static_cast<int64_t>(bpp) * parsed_width;
+  if (bits_per_row > (std::numeric_limits<int>::max() - 31)) {
+    LOG(ERROR) << "BMP row size overflow";
+    return false;
+  }
+  const int parsed_row_size = static_cast<int>((bits_per_row + 31) / 32 * 4);
+
+  const size_t pixel_bytes =
+      img_bytes.size() - static_cast<size_t>(parsed_pixel_offset);
+  if (abs_height > 0 && static_cast<uint64_t>(parsed_row_size) >
+                            std::numeric_limits<uint64_t>::max() /
+                                static_cast<uint64_t>(abs_height)) {
+    LOG(ERROR) << "BMP pixel data size overflow";
+    return false;
+  }
+  const uint64_t required_pixel_bytes = static_cast<uint64_t>(parsed_row_size) *
+                                        static_cast<uint64_t>(abs_height);
+  if (required_pixel_bytes > pixel_bytes) {
+    LOG(ERROR) << "BMP pixel data is shorter than the declared dimensions";
+    return false;
+  }
+
+  *width = parsed_width;
+  *height = parsed_height;
+  *channels = parsed_channels;
+  *row_size = parsed_row_size;
+  *pixel_offset = static_cast<size_t>(parsed_pixel_offset);
+  return true;
+}
+
+}  // namespace
 
 std::vector<uint8_t> decode_bmp(const uint8_t* input, int row_size, int width,
                                 int height, int channels, bool top_down) {
-  std::vector<uint8_t> output(height * width * channels);
+  std::vector<uint8_t> output(static_cast<size_t>(height) *
+                              static_cast<size_t>(width) *
+                              static_cast<size_t>(channels));
   for (int i = 0; i < height; i++) {
-    int src_pos;
-    int dst_pos;
+    size_t src_pos;
+    size_t dst_pos;
 
     for (int j = 0; j < width; j++) {
       if (!top_down) {
-        src_pos = ((height - 1 - i) * row_size) + j * channels;
+        src_pos = static_cast<size_t>(height - 1 - i) * row_size +
+                  static_cast<size_t>(j) * channels;
       } else {
-        src_pos = i * row_size + j * channels;
+        src_pos = static_cast<size_t>(i) * row_size +
+                  static_cast<size_t>(j) * channels;
       }
 
-      dst_pos = (i * width + j) * channels;
+      dst_pos = (static_cast<size_t>(i) * width + j) * channels;
 
       switch (channels) {
         case 1:
@@ -94,31 +199,26 @@ std::vector<uint8_t> read_bmp(const std::string& input_bmp_name, int* width,
   std::vector<uint8_t> img_bytes(len);
   file.seekg(0, std::ios::beg);
   file.read(reinterpret_cast<char*>(img_bytes.data()), len);
-  const int32_t header_size =
-      TF_le32toh(*(reinterpret_cast<const int32_t*>(img_bytes.data() + 10)));
-  *width =
-      TF_le32toh(*(reinterpret_cast<const int32_t*>(img_bytes.data() + 18)));
-  *height =
-      TF_le32toh(*(reinterpret_cast<const int32_t*>(img_bytes.data() + 22)));
-  const int32_t bpp =
-      TF_le32toh(*(reinterpret_cast<const int32_t*>(img_bytes.data() + 28)));
-  *channels = bpp / 8;
+  int row_size = 0;
+  size_t header_size = 0;
+  if (!ValidateBmpAndGetPixelOffset(img_bytes, width, height, channels,
+                                    &row_size, &header_size)) {
+    return std::vector<uint8_t>();
+  }
 
   if (s->verbose)
     LOG(INFO) << "width, height, channels: " << *width << ", " << *height
               << ", " << *channels;
 
-  // there may be padding bytes when the width is not a multiple of 4 bytes
-  // 8 * channels == bits per pixel
-  const int row_size = (8 * *channels * *width + 31) / 32 * 4;
-
   // if height is negative, data layout is top down
   // otherwise, it's bottom up
   bool top_down = (*height < 0);
 
   // Decode image, allocating tensor once the image size is known
-  const uint8_t* bmp_pixels = &img_bytes[header_size];
-  return decode_bmp(bmp_pixels, row_size, *width, abs(*height), *channels,
+  const uint8_t* bmp_pixels = img_bytes.data() + header_size;
+  const int abs_height =
+      static_cast<int>(*height < 0 ? -static_cast<int64_t>(*height) : *height);
+  return decode_bmp(bmp_pixels, row_size, *width, abs_height, *channels,
                     top_down);
 }