[pull] master from tensorflow:master

MODULE.bazel

@@ -35,7 +35,7 @@ single_version_override(
 bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc1", repo_name = "com_github_googlecloudplatform_google_cloud_cpp")
 bazel_dep(name = "crc32c", version = "1.1.0", repo_name = "com_github_google_crc32c")
 bazel_dep(name = "brotli", version = "1.1.0", repo_name = "org_brotli")
-bazel_dep(name = "rules_cc", version = "0.2.20")
+bazel_dep(name = "rules_cc", version = "0.2.25")
 bazel_dep(name = "curl", version = "8.11.0.bcr.5")
 bazel_dep(name = "rules_webtesting", version = "0.4.1", repo_name = "io_bazel_rules_webtesting")
 bazel_dep(name = "rules_closure", version = "0.15.0", repo_name = "io_bazel_rules_closure")

tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir

@@ -628,3 +628,22 @@ func.func @main(%arg0: tensor<*xi32>)  -> () {
   %push = "tf.TensorListPushBack"(%tl, %elem) : (tensor<!tf_type.variant<tensor<f32>>>, tensor<f32>) -> tensor<!tf_type.variant<tensor<f32>>>
   func.return
 }
+
+// -----
+
+// Tests that the pass warns on a set item that is allowed to grow the list,
+// which a tf.TensorArray with dynamic_size=True emits. The buffer has a static
+// shape and cannot grow, so the write is dropped. Decomposition still proceeds,
+// since the attribute is set by every dynamic_size TensorArray including those
+// that never write past the end.
+
+func.func @main() -> () {
+  %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
+  %num = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %index = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %tl = "tf.TensorListReserve"(%elem_shape, %num) : (tensor<0xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<f32>>>
+  %elem = "tf._SomeOp"() : () -> tensor<f32>
+  // expected-warning @+1 {{TensorLists that grow on an out-of-bounds write are not supported by XLA}}
+  %set = "tf.TensorListSetItem"(%tl, %index, %elem) {resize_if_index_out_of_bounds = true} : (tensor<!tf_type.variant<tensor<f32>>>, tensor<i32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<f32>>>
+  func.return
+}

tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc

@@ -714,6 +714,26 @@ LogicalResult HandleTensorListGetItemOp(
 LogicalResult HandleTensorListSetItemOp(
     TF::TensorListSetItemOp set_item,
     llvm::SmallDenseMap<Value, SizeInfo>* buffer_to_size) {
+  // A tf.TensorArray built with dynamic_size=True sets this attribute, and the
+  // non-XLA kernel grows the list when the index is past the end. The buffer
+  // built here has a static shape and cannot grow, and cutil::SetElement lowers
+  // to a DynamicUpdateSlice, which clamps an out-of-range index rather than
+  // failing, so the write is dropped and the program returns a truncated
+  // result.
+  //
+  // This is only a warning rather than an error, because the attribute is set
+  // by every dynamic_size TensorArray, including the many that never actually
+  // write past the end and compile correctly today. The classic tf2xla
+  // lowering warns in the same case, see TensorListSetItemOp in
+  // tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc.
+  if (set_item.getResizeIfIndexOutOfBounds()) {
+    set_item.emitWarning(
+        "TensorLists that grow on an out-of-bounds write are not supported by "
+        "XLA. This is typically a tf.TensorArray created with "
+        "dynamic_size=True; give it a static size large enough to hold every "
+        "write instead.");
+  }
+
   auto buffer = set_item.getInputHandle();
   auto it = buffer_to_size->find(buffer);
   if (it == buffer_to_size->end()) {

tensorflow/compiler/tests/tensor_list_ops_test.py

@@ -302,6 +302,51 @@ def testZerosLikeForTensorList(self):
       self.assertAllEqual(z.shape.as_list(), [None])
       self.assertAllEqual(z, [0.0, 0.0])
 
+  def testSetItemResizeWarnsAndTruncates(self):
+    # Regression test for GitHub issue 127528. A TensorListSetItem that is
+    # allowed to grow the list cannot grow it under XLA, because the buffer
+    # that TensorListReserve allocated has a static shape. The write is
+    # dropped, either by returning the list unchanged when the element does
+    # not fit, or by letting DynamicUpdateSlice clamp an out-of-range index so
+    # the element overwrites an earlier one.
+    #
+    # The lowering logs a warning rather than failing, because the attribute is
+    # set by every dynamic_size TensorArray including the many that never write
+    # past the end, so compilation still has to succeed here.
+    with self.session(), self.test_scope():
+      tensor_list = list_ops.tensor_list_reserve(
+          element_shape=[], element_dtype=dtypes.float32, num_elements=1
+      )
+      set_item = list_ops.tensor_list_set_item(
+          input_handle=tensor_list,
+          index=0,
+          item=constant_op.constant(1.0),
+          resize_if_index_out_of_bounds=True,
+      )
+      stacked = list_ops.tensor_list_stack(
+          set_item, element_dtype=dtypes.float32
+      )
+      self.assertAllEqual(self.evaluate(stacked), [1.0])
+
+  def testSetItemWithoutResizeStillCompiles(self):
+    # Counterpart to the test above: the warning is keyed off the attribute,
+    # so this pins that an ordinary fixed-size write compiles to the same
+    # result and that the attribute alone changes nothing but the logging.
+    with self.session(), self.test_scope():
+      tensor_list = list_ops.tensor_list_reserve(
+          element_shape=[], element_dtype=dtypes.float32, num_elements=1
+      )
+      set_item = list_ops.tensor_list_set_item(
+          input_handle=tensor_list,
+          index=0,
+          item=constant_op.constant(1.0),
+          resize_if_index_out_of_bounds=False,
+      )
+      stacked = list_ops.tensor_list_stack(
+          set_item, element_dtype=dtypes.float32
+      )
+      self.assertAllEqual(self.evaluate(stacked), [1.0])
+
   def testInvalidSplitLength(self):
     with self.session(), self.test_scope():
       tensor_list_split = list_ops.tensor_list_split(

tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
@@ -649,7 +650,32 @@ REGISTER_XLA_OP(
 
 class TensorListSetItemOp : public XlaOpKernel {
  public:
-  explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    // A tf.TensorArray built with dynamic_size=True lowers to TensorListSetItem
+    // with this attribute set, and the non-XLA kernel grows the list when the
+    // index is past the end. XLA needs static shapes, so the buffer that
+    // TensorListReserve allocated cannot grow, and a write past the end is
+    // dropped in one of two ways depending on the shapes. When the element does
+    // not fit at all, as for a list reserved with size 0,
+    // ExecuteTensorListSetItem takes its early return and hands back the list
+    // unchanged with an OK status. When the element does fit but the index is
+    // past the end, the DynamicUpdateSlice it lowers to clamps its start
+    // indices rather than failing, so the element overwrites an earlier one.
+    // Either way the program returns a truncated result.
+    //
+    // This is only a warning rather than a compile error, because the attribute
+    // is set by every dynamic_size TensorArray, including the many that never
+    // actually write past the end and compile correctly today.
+    bool resize_if_index_out_of_bounds;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("resize_if_index_out_of_bounds",
+                                     &resize_if_index_out_of_bounds));
+    if (resize_if_index_out_of_bounds) {
+      LOG(WARNING) << "TensorLists that grow on an out-of-bounds write are not "
+                      "supported by XLA. This is typically a tf.TensorArray "
+                      "created with dynamic_size=True; give it a static size "
+                      "large enough to hold every write instead.";
+    }
+  }
 
   void Compile(XlaOpKernelContext* ctx) override {
     xla::XlaOp list = ctx->Input(0);

tensorflow/core/common_runtime/graph_execution_state.cc

@@ -53,6 +53,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/device_name_utils.h"
+#include "tensorflow/core/util/env_var.h"
 #include "tensorflow/core/util/util.h"
 
 #ifndef IS_MOBILE_PLATFORM
@@ -63,12 +64,31 @@ limitations under the License.
 
 namespace tensorflow {
 
+bool ProtoMemoryOptimizationsEnabled() {
+  bool val = false;
+  (void)ReadBoolFromEnvVar("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS", false, &val);
+  return val;
+}
+
 namespace {
 bool IsCollectiveV2(const std::string& op) {
   return op == "CollectiveReduceV2" || op == "CollectiveGatherV2" ||
          op == "CollectiveBcastRecvV2" || op == "CollectiveBcastSendV2" ||
          op == "ColectiveReduceScatterV2" || op == "ColectiveAllToAllV2";
 }
+
+// Returns true if `flib_def` is non-null and contains definitions for all
+// functions defined in `graph_def.library()`.
+bool LibraryContainsAllFunctions(const FunctionLibraryDefinition* flib_def,
+                                 const GraphDef& graph_def) {
+  if (flib_def == nullptr) return false;
+  for (const auto& fdef : graph_def.library().function()) {
+    if (!flib_def->Contains(fdef.signature().name())) {
+      return false;
+    }
+  }
+  return true;
+}
 }  // namespace
 
 GraphExecutionState::GraphExecutionState(
@@ -96,14 +116,23 @@ GraphExecutionState::~GraphExecutionState() {
   VLOG(4) << "Graph proto is \n" << graph_def.DebugString();
 #endif  // __ANDROID__
 
-  auto flib_def = std::make_unique<FunctionLibraryDefinition>(
-      OpRegistry::Global(), graph_def);
+  std::unique_ptr<FunctionLibraryDefinition> flib_def;
+  if (ProtoMemoryOptimizationsEnabled() &&
+      LibraryContainsAllFunctions(options.flib_def, graph_def)) {
+    flib_def = std::make_unique<FunctionLibraryDefinition>(*options.flib_def);
+  } else {
+    flib_def = std::make_unique<FunctionLibraryDefinition>(OpRegistry::Global(),
+                                                           graph_def);
+  }
 
   TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(&graph_def, *flib_def, 0));
 
-  if (options.session_options->config.graph_options().place_pruned_graph() ||
+  const bool keep_original_graph_def =
+      options.session_options->config.graph_options().place_pruned_graph() ||
       options.session_options->config.experimental()
-          .disable_optimize_for_static_graph()) {
+          .disable_optimize_for_static_graph();
+
+  if (keep_original_graph_def) {
     auto ret = absl::WrapUnique(new GraphExecutionState(
         std::make_unique<GraphDef>(std::move(graph_def)), std::move(flib_def),
         options));

tensorflow/core/common_runtime/graph_execution_state.h

@@ -33,6 +33,11 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
+
+// Returns true if proto memory optimizations are enabled (default false, can be
+// enabled via environment variable TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS=1).
+bool ProtoMemoryOptimizationsEnabled();
+
 struct SessionOptions;
 
 namespace subgraph {
@@ -54,6 +59,12 @@ struct GraphExecutionStateOptions {
   // work for almost all models. Non default values should only applied to
   // selective models.
   bool enable_tf2xla_mlir_bridge = true;
+
+  // Optional pre-constructed FunctionLibraryDefinition. If provided,
+  // GraphExecutionState will copy this library definition (which performs
+  // an O(1) refcount increment on underlying FunctionRecords) rather than
+  // re-parsing or duplicating the library proto in GraphDef.
+  const FunctionLibraryDefinition* flib_def = nullptr;
 };
 
 // A ClientGraph is simply a sub-graph of the full graph as induced by

tensorflow/core/kernels/topk_op_gpu.h

@@ -238,12 +238,12 @@ __device__ void heapTopK(const T* __restrict__ input, int length, int k,
   // Now iterate over the remaining items.
   // If an item is smaller than the min element, it is not amongst the top k.
   // Otherwise, replace the min element with it and push upwards.
-  for (int index = heap_end_index; index < length; index += step_size) {
+  for (int64_t index = heap_end_index; index < length; index += step_size) {
     // We prefer elements with lower indices. This is given here.
     // Later elements automatically have higher indices, so can be discarded.
     if (input[index] > heap.root().value) {
       // This element should replace the min.
-      heap.replace_root({index, input[index]}, k);
+      heap.replace_root({static_cast<int>(index), input[index]}, k);
     }
   }
 

tensorflow/core/tfrt/fallback/BUILD

@@ -67,6 +67,7 @@ cc_library(
         "//tensorflow/core/framework:graph_proto_cc",
         "//tensorflow/core/tpu:virtual_device",
         "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

tensorflow/core/tfrt/fallback/fallback_state.cc

@@ -171,6 +171,9 @@ FallbackState::CreateGraphExecutionState(GraphDef graph_def, bool run_placer,
   options.session_handle = "tfrt_fallback_handle";
   options.run_placer = run_placer;
   options.enable_tf2xla_mlir_bridge = enable_tf2xla_mlir_bridge;
+  if (tensorflow::ProtoMemoryOptimizationsEnabled()) {
+    options.flib_def = &func_lib_def_;
+  }
 
   std::unique_ptr<GraphExecutionState> execution_state;
   TF_RETURN_IF_ERROR(GraphExecutionState::MakeForBaseGraph(

tensorflow/core/tfrt/utils/BUILD

@@ -271,6 +271,7 @@ tf_cc_test(
         "//tensorflow/core/framework:graph_proto_cc",
         "//tensorflow/core/framework:tensor_testutil",
         "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/framework:versions_proto_cc",
         "//tensorflow/core/grappler/utils:grappler_test",
         "//tensorflow/core/lib/monitoring:cell_reader",
         "//tensorflow/core/platform:status",

tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc

@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/function_body.h"
 #include "tensorflow/core/common_runtime/function_def_utils.h"
 #include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/common_runtime/graph_execution_state.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function.pb.h"
@@ -190,8 +191,13 @@ tensorflow::GraphDef CreateGraphDefFromGraphAndFlibDef(
     const tensorflow::Graph& graph,
     const tensorflow::FunctionLibraryDefinition& flib_def) {
   tensorflow::GraphDef graph_def;
-  graph.ToGraphDef(&graph_def);
-  *graph_def.mutable_library() = flib_def.ToProto();
+  graph.ToGraphDef(&graph_def, /*include_flib_def=*/false);
+  if (tensorflow::ProtoMemoryOptimizationsEnabled()) {
+    *graph_def.mutable_library() =
+        flib_def.ReachableDefinitions(graph).ToProto();
+  } else {
+    *graph_def.mutable_library() = flib_def.ToProto();
+  }
   return graph_def;
 }
 
@@ -693,6 +699,16 @@ TfrtGraphExecutionState::OptimizeGraph(
         &optimized_flib));
   }
 
+  if (!options_.run_placer_grappler_on_functions &&
+      tensorflow::ProtoMemoryOptimizationsEnabled()) {
+    // Fast path: Directly merge the optimized FunctionLibraryDefinition in C++
+    // without serializing to a FunctionDefLibrary proto, avoiding expensive
+    // heap allocations and copy overhead during model initialization.
+    TF_RETURN_IF_ERROR(optimized_graph->mutable_flib_def()->AddLibrary(
+        std::move(*optimized_flib)));
+    return optimized_graph;
+  }
+
   FunctionDefLibrary optimized_flib_proto = optimized_flib->ToProto();
   if (options_.run_placer_grappler_on_functions) {
     TF_RETURN_IF_ERROR(OptimizeFunctions(optimized_flib_proto, *optimized_flib,
@@ -704,7 +720,8 @@ TfrtGraphExecutionState::OptimizeGraph(
     optimized_graph->mutable_flib_def()->Clear();
   }
 
-  TF_RETURN_IF_ERROR(optimized_graph->AddFunctionLibrary(optimized_flib_proto));
+  TF_RETURN_IF_ERROR(
+      optimized_graph->AddFunctionLibrary(std::move(optimized_flib_proto)));
 
   return optimized_graph;
 }