[pull] master from tensorflow:master

tensorflow/compiler/jit/device_compilation_profiler.cc

@@ -96,8 +96,8 @@ void DeviceCompilationProfiler::RegisterExecution(
 
 absl::Status DeviceCompilationProfiler::RegisterCompilation(
     const NameAttrList& function, int64_t compile_time_us,
-    bool used_persistent_cache) {
-  metrics::UpdateXlaCompilationTime(compile_time_us);
+    bool used_persistent_cache, int64_t compile_end_us) {
+  metrics::UpdateXlaCompilationTime(compile_time_us, compile_end_us);
 
   const std::string& function_name = function.name();
 

tensorflow/compiler/jit/device_compilation_profiler.h

@@ -76,7 +76,8 @@ class DeviceCompilationProfiler : public ResourceBase {
   // XlaJitCompilationActivity.
   virtual absl::Status RegisterCompilation(const NameAttrList& function,
                                            int64_t compile_time_us,
-                                           bool used_persistent_cache);
+                                           bool used_persistent_cache,
+                                           int64_t compile_end_us = 0);
 
   void IncrementOngoingAsyncCompilations();
   void DecrementOngoingAsyncCompilations();

tensorflow/compiler/jit/device_compiler.h

@@ -332,6 +332,7 @@ DeviceCompiler<ExecutableType, ClientType>::CompileStrict(
     DeviceCompilationProfiler* profiler, mutex* mu) {
   tensorflow::Env* env = tensorflow::Env::Default();
   const uint64_t compile_start_us = env->NowMicros();
+  metrics::UpdateXlaCompilationStartTime(compile_start_us);
 
   TfGraphToHloCompiler compiler(options);
   cache_value.compile_state = DeviceCompileState::kCompiled;
@@ -390,7 +391,8 @@ DeviceCompiler<ExecutableType, ClientType>::CompileStrict(
 
   device_compiler_internal::LogOnceXlaCompiledFirstCluster();
   TF_RETURN_IF_ERROR(profiler->RegisterCompilation(
-      function, compile_time_us, loaded_executable.has_value()));
+      function, compile_time_us, loaded_executable.has_value(),
+      compile_end_us));
   return cache_value;
 }
 

tensorflow/compiler/jit/device_compiler_test.cc

@@ -156,7 +156,7 @@ class MockDeviceCompilationProfiler : public DeviceCompilationProfiler {
               (override));
   MOCK_METHOD(absl::Status, RegisterCompilation,
               (const NameAttrList& function, int64_t compile_time_us,
-               bool used_persistent_cache),
+               bool used_persistent_cache, int64_t compile_end_us),
               (override));
 };
 
@@ -313,7 +313,7 @@ TEST_F(DeviceCompilerTest, CompileAsyncSuccess) {
   EXPECT_CALL(*mock_profiler_,
               ShouldCompileCluster(_, DeviceCompileMode::kAsync, 1))
       .WillOnce(Return(true));
-  EXPECT_CALL(*mock_profiler_, RegisterCompilation(_, _, false))
+  EXPECT_CALL(*mock_profiler_, RegisterCompilation(_, _, false, _))
       .WillOnce([&done] {
         done.Notify();
         return absl::OkStatus();

tensorflow/compiler/mlir/lite/BUILD

@@ -370,6 +370,35 @@ cc_library(
     ],
 )
 
+tf_cc_binary(
+    name = "litert-reduce",
+    testonly = True,
+    deps = [
+        ":litert_mlir_reduce_main",
+    ],
+)
+
+cc_library(
+    name = "litert_mlir_reduce_main",
+    testonly = True,
+    srcs = ["litert_mlir_reduce_main.cc"],
+    deps = [
+        ":lift_tflite_flex_ops",  # buildcleaner:keep
+        ":register_lite_dialects",
+        ":tensorflow_lite",
+        ":tf_tfl_passes",  # buildcleaner:keep
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir:passes",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MlirReduceLib",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "utils",
     hdrs = ["utils/utils.h"],
@@ -1786,6 +1815,7 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:portable_gif_internal",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -1797,6 +1827,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@eigen_archive//:eigen3",
         "@flatbuffers",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",

tensorflow/compiler/mlir/lite/converter_flags.proto

@@ -39,7 +39,7 @@ enum FileFormat {
 // of as properties of models, instead describing how models are to be
 // processed in the context of the present tooling job.
 //
-// Next ID to use: 72.
+// Next ID to use: 74.
 message ConverterFlags {
   reserved 54, 61;
 
@@ -390,4 +390,12 @@ message ConverterFlags {
   // If true, fold 16-bit float (fp16/bf16) to 32-bit float (fp32) casts on
   // large resource constants.
   optional bool fold_fp16_resource_casts = 71 [default = true];
+
+  // If true, enable debugging facilities such as IR cloning and file
+  // serialization.
+  optional bool enable_debug = 72 [default = false];
+
+  // Directory for debug artifacts and printed IR dumps when debug mode is
+  // enabled.
+  optional string debug_dir = 73;
 }

tensorflow/compiler/mlir/lite/debug/BUILD

@@ -25,7 +25,7 @@ cc_library(
     deps = [
         ":debug_options_proto_cc",
         "//tensorflow/compiler/mlir/lite/metrics:error_collector_inst",
-        "//tensorflow/core:portable_gif_internal",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
@@ -40,7 +40,6 @@ cc_library(
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:Transforms",
         "@tsl//tsl/platform:path",
-        "@tsl//tsl/platform:stringpiece",
         "@xla//xla/tsl/lib/io:buffered_file",
         "@xla//xla/tsl/platform:env",
     ],

tensorflow/compiler/mlir/lite/debug/debug.cc

@@ -17,13 +17,18 @@ limitations under the License.
 
 #include <stddef.h>
 #include <stdint.h>
+#if defined(__linux__)
+#include <unistd.h>
+#endif
 
+#include <fstream>
 #include <functional>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -52,9 +57,7 @@ limitations under the License.
 #include "xla/tsl/lib/io/buffered_file.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/file_system.h"
-#include "tensorflow/core/platform/logging.h"
 #include "tsl/platform/path.h"
-#include "tsl/platform/stringpiece.h"
 
 // IWYU pragma: no_include "util/regexp/re2/re2.h"
 
@@ -166,12 +169,14 @@ std::string Sanitize(absl::string_view string) {
 // instrumentation dumps MLIR to external directories for convenience.
 class DumpInstrumentation : public mlir::PassInstrumentation {
  public:
-  explicit DumpInstrumentation(absl::string_view dump_dir,
-                               absl::string_view dump_pass_regex,
-                               absl::string_view dump_func_regex)
+  explicit DumpInstrumentation(
+      absl::string_view dump_dir, absl::string_view dump_pass_regex,
+      absl::string_view dump_func_regex,
+      mlir::OpPrintingFlags flags = mlir::OpPrintingFlags())
       : dump_dir_(dump_dir),
         dump_pass_re_(std::make_unique<RE2>(dump_pass_regex)),
-        dump_func_re_(std::make_unique<RE2>(dump_func_regex)) {}
+        dump_func_re_(std::make_unique<RE2>(dump_func_regex)),
+        op_printing_flags_(flags) {}
 
   DumpInstrumentation(const DumpInstrumentation& other) = delete;
   DumpInstrumentation& operator=(const DumpInstrumentation& other) = delete;
@@ -257,19 +262,54 @@ class DumpInstrumentation : public mlir::PassInstrumentation {
     file = std::make_unique<tsl::BufferedWritableFile>(std::move(file));
 
     WritableFileRawStream os(std::move(file));
-    op->print(os);
+    op->print(os, op_printing_flags_);
   }
 
   const std::string dump_dir_;
   const std::unique_ptr<RE2> dump_pass_re_;
   const std::unique_ptr<RE2> dump_func_re_;
+  mlir::OpPrintingFlags op_printing_flags_;
 
   // Counter used for pass name prefix to signify sequence
   int pass_counter_ = 0;
 
   bool printed_ = false;
 };
 
+double GetCurrentRssMb() {
+#if defined(__linux__)
+  std::ifstream statm("/proc/self/statm");
+  if (!statm.is_open()) return 0.0;
+  int64_t pages = 0;
+  int64_t rss_pages = 0;
+  if (statm >> pages >> rss_pages) {
+    int64_t page_size = sysconf(_SC_PAGESIZE);
+    return static_cast<double>(rss_pages * page_size) / (1024.0 * 1024.0);
+  }
+#endif
+  return 0.0;
+}
+
+class RssLoggingInstrumentation : public mlir::PassInstrumentation {
+ public:
+  void runBeforePass(mlir::Pass* pass, mlir::Operation* op) override {
+    pass_start_rss_[pass] = GetCurrentRssMb();
+  }
+
+  void runAfterPass(mlir::Pass* pass, mlir::Operation* op) override {
+    double end_rss = GetCurrentRssMb();
+    auto node = pass_start_rss_.extract(pass);
+    double start_rss = node.empty() ? 0.0 : node.mapped();
+    double delta = end_rss - start_rss;
+    LOG(INFO) << "[MLIR RSS] After pass '" << pass->getName().str() << "' on '"
+              << op->getName().getStringRef().str() << "': RSS = " << end_rss
+              << " MB (Delta: " << (delta >= 0 ? "+" : "") << delta << " MB)";
+  }
+
+ private:
+  absl::flat_hash_map<mlir::Pass*, double> pass_start_rss_;
+};
+
 std::function<bool(mlir::Pass*, mlir::Operation*)> CreatePrintIRFun(
     const std::string& pass_regex) {
   std::function<bool(mlir::Pass*, mlir::Operation*)> fun;
@@ -296,12 +336,20 @@ void InitPassManager(mlir::PassManager& pm,
   bool print_to_stdout =
       !options.print_ir_before().empty() || !options.print_ir_after().empty();
 
-  if (dump_to_dir || print_to_stdout) {
-    // Necessary for maintaining sequence of passes when dumping MLIR to files
-    // or stdout.
+  if (dump_to_dir || print_to_stdout || options.log_rss()) {
+    // Necessary for maintaining sequence of passes when dumping MLIR to files,
+    // stdout, or logging RSS memory.
     pm.getContext()->disableMultithreading();
   }
 
+  mlir::OpPrintingFlags opPrintingFlags = mlir::OpPrintingFlags();
+  if (options.has_elide_elementsattrs_if_larger()) {
+    opPrintingFlags.elideLargeElementsAttrs(
+        options.elide_elementsattrs_if_larger());
+    opPrintingFlags.elideLargeResourceString(
+        options.elide_elementsattrs_if_larger());
+  }
+
   if (dump_to_dir) {
     dump_dir = tsl::io::JoinPath(
         dump_dir, absl::FormatTime("%E4Y%m%d_%H%M%E6S", absl::Now(),
@@ -321,7 +369,8 @@ void InitPassManager(mlir::PassManager& pm,
     }
 
     pm.addInstrumentation(std::make_unique<DumpInstrumentation>(
-        dump_dir, options.ir_dump_pass_regex(), options.ir_dump_func_regex()));
+        dump_dir, options.ir_dump_pass_regex(), options.ir_dump_func_regex(),
+        opPrintingFlags));
   }
 
   if (print_to_stdout) {
@@ -331,20 +380,17 @@ void InitPassManager(mlir::PassManager& pm,
     std::function<bool(mlir::Pass*, mlir::Operation*)>
         should_print_ir_after_pass(CreatePrintIRFun(options.print_ir_after()));
 
-    mlir::OpPrintingFlags opPrintingFlags = mlir::OpPrintingFlags();
-
-    if (options.has_elide_elementsattrs_if_larger()) {
-      opPrintingFlags.elideLargeElementsAttrs(
-          options.elide_elementsattrs_if_larger());
-    }
-
     pm.enableIRPrinting(should_print_ir_before_pass, should_print_ir_after_pass,
                         options.print_ir_module_scope(),
                         /*printAfterOnlyOnChange=*/true,
                         /*printAfterOnlyOnFailure=*/false, out,
                         opPrintingFlags);
   }
 
+  if (options.log_rss()) {
+    pm.addInstrumentation(std::make_unique<RssLoggingInstrumentation>());
+  }
+
   // Enable pass timing. Note: MLIR expects `mlir::PassManager::enableTiming` to
   // be called after all instrumentations are added.
   if (options.enable_timing()) {

tensorflow/compiler/mlir/lite/debug/debug_options.proto

@@ -19,7 +19,7 @@ package tensorflow.converter;
 
 // Additional parameters that control the debug behavior of the Converter.
 //
-// Next ID: 9
+// Next ID: 10
 message DebugOptions {
   // If not empty, dumps MLIR to the specified directory. The initial state of
   // the MLIR after import will be dumped at the beginning of each pass manager
@@ -58,4 +58,7 @@ message DebugOptions {
   // Elide ElementsAttrs with \"...\" that have more elements than the given
   // upper limit.
   optional int64 elide_elementsattrs_if_larger = 8;
+
+  // If true, log RSS memory usage (in MB) before and after each MLIR pass.
+  optional bool log_rss = 9 [default = false];
 }

tensorflow/compiler/mlir/lite/litert_mlir_reduce_main.cc

@@ -0,0 +1,41 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Tools/mlir-reduce/MlirReduceMain.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/lite/register_lite_dialects.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+
+int main(int argc, char** argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+
+  mlir::registerAllPasses();
+  mlir::registerTensorFlowPasses();
+  mlir::TFL::registerTensorFlowLitePasses();
+
+  mlir::DialectRegistry registry;
+  mlir::RegisterCommonToolingDialects(registry);
+  tflite::RegisterLiteToolingDialects(registry);
+
+  mlir::MLIRContext context;
+  context.appendDialectRegistry(registry);
+
+  return failed(mlir::mlirReduceMain(argc, argv, context));
+}