[pull] master from tensorflow:master

tensorflow/compiler/aot/compile.cc

@@ -206,6 +206,12 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config,
         flags.sanitize_abilists_dataflow, ',', absl::SkipEmpty()));
   }
 
+  if (flags.sanitize_memory || flags.sanitize_memory_track_origins > 0) {
+    aot_opts.set_sanitize_memory(true);
+    aot_opts.set_sanitize_memory_track_origins(
+        flags.sanitize_memory_track_origins);
+  }
+
   TF_RETURN_IF_ERROR(
       ConfigureKernelNamingConvention(aot_opts, computation, flags.cpp_class));
 

tensorflow/compiler/aot/flags.cc

@@ -86,6 +86,12 @@ void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags) {
        "Enable DataFlow Sanitizer pass."},
       {"sanitize_abilists_dataflow", &flags->sanitize_abilists_dataflow,
        "Comma separated list of ABIList file paths."},
+      {"sanitize_memory", &flags->sanitize_memory,
+       "Enable Memory Sanitizer pass."},
+      {"sanitize_memory_track_origins", &flags->sanitize_memory_track_origins,
+       "Controls MSan track origins level (0=disabled, 1=without store "
+       "history, 2=with store history). Setting to >0 implies "
+       "--sanitize_memory."},
       {"gen_name_to_index", &flags->gen_name_to_index,
        "Generate name-to-index data for Lookup{Arg,Result}Index methods."},
       {"gen_program_shape", &flags->gen_program_shape,

tensorflow/compiler/aot/flags.h

@@ -48,6 +48,8 @@ struct MainFlags {
   // Sanitizer pass options
   bool sanitize_dataflow = false;
   std::string sanitize_abilists_dataflow;
+  bool sanitize_memory = false;
+  int32_t sanitize_memory_track_origins = 0;
 
   // C++ codegen options
   bool gen_name_to_index = false;

tensorflow/compiler/aot/tfcompile.bzl

@@ -358,6 +358,7 @@ def _tf_library(
             "@xla//xla/backends/cpu/runtime:sort_lib",
             "@xla//xla/backends/cpu/runtime:topk_lib",
             "@xla//xla/backends/cpu/runtime:convolution_lib",
+            "@xla//xla/backends/cpu/runtime:msan_emulated_tls",
             "@xla//xla/service/cpu:runtime_matmul",
             "@xla//xla/service/cpu:runtime_single_threaded_matmul",
             "@eigen_archive//:eigen3",

tensorflow/compiler/tf2xla/BUILD

@@ -485,6 +485,7 @@ cc_library(
         "@com_google_absl//absl/types:span",
         ":encoded_buffer_allocation_info",
         "@xla//xla/service:custom_call_status_internal",
+        "@xla//xla/backends/cpu/runtime:msan_emulated_tls",
         "@xla//xla/backends/cpu/runtime:rng_state_lib",
         "@xla//xla/backends/cpu:alignment",
         "@xla//xla/backends/cpu:buffer_allocation_info",

third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc

@@ -70,7 +70,7 @@ absl::StatusOr<ConsumeResult> ProfilerCollection::Consume() {
       data_vector.push_back(std::move(result->data));
       total_estimated_size_bytes += result->estimated_size_bytes;
     } else if (absl::IsUnimplemented(result.status())) {
-      LOG(WARNING) << "Profiler consume not implemented: " << result.status();
+      VLOG(1) << "Profiler consume not implemented: " << result.status();
       data_vector.push_back(std::any());
     } else {
       LOG(ERROR) << "Profiler consume failed: " << result.status();

third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc

@@ -88,7 +88,7 @@ absl::Status ProfilerController::CollectData(
 }
 
 absl::StatusOr<ConsumeResult> ProfilerController::Consume() {
-  if (state_ != ProfilerState::kStart && state_ != ProfilerState::kStop) {
+  if (state_ == ProfilerState::kInit) {
     return absl::AbortedError("Consume called in the wrong order.");
   }
   if (!status_.ok()) {

third_party/xla/xla/backends/cpu/codegen/BUILD

@@ -48,8 +48,12 @@ cc_library(
     deps = [
         ":builtin_fp16",
         ":builtin_pow",
+        "//xla/backends/cpu/runtime:msan_emulated_tls",
+        "//xla/service/cpu:cpu_runtime",
+        "@com_google_absl//absl/base:config",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
         "@llvm-project//llvm:Core",
         "@llvm-project//llvm:OrcJIT",
         "@llvm-project//llvm:OrcShared",
@@ -124,18 +128,21 @@ cc_library(
         "//xla:util",
         "//xla:xla_proto_cc",
         "//xla/backends/cpu:target_machine_options",
+        "//xla/backends/cpu/runtime:msan_emulated_tls",
         "//xla/codegen:intrinsic_lib",
         "//xla/codegen/intrinsic",
         "//xla/codegen/intrinsic:intrinsic_compiler_lib",
         "//xla/service:hlo_module_config",
         "//xla/service/cpu:backend_config_proto_cc",
         "//xla/service/cpu:cpu_options",
+        "//xla/service/cpu:cpu_runtime",
         "//xla/service/cpu:executable_proto_cc",
         "//xla/service/llvm_ir:llvm_util",
         "//xla/tools:llvm_targets",  # fixdeps: keep
         "//xla/tsl/platform:logging",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
+        "@com_google_absl//absl/base:config",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/log",

third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc

@@ -28,19 +28,24 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/base/config.h"  // IWYU pragma: keep
 #include "absl/base/no_destructor.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"  // IWYU pragma: keep
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/CoreContainers.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"  // IWYU pragma: keep (msan)
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/Support/Error.h"
 #include "xla/backends/cpu/codegen/builtin_fp16.h"
 #include "xla/backends/cpu/codegen/builtin_pow.h"
+#include "xla/backends/cpu/runtime/msan_emulated_tls.h"
+#include "xla/service/cpu/cpu_runtime.h"
 
 namespace xla::cpu {
 
@@ -267,9 +272,8 @@ static Registry CreateRegistry() {
 
 #endif
 
-#ifdef MEMORY_SANITIZER
-  registry["__msan_unpoison"] = SymbolDef(__msan_unpoison);
-#endif
+  registry[runtime::kMsanEmutlsGetAddressBridgeSymbolName] =
+      SymbolDef(__xla_cpu_runtime_emutls_get_address);
 
   return registry;
 }
@@ -280,22 +284,55 @@ static Registry CreateRegistry() {
 
 BuiltinDefinitionGenerator::BuiltinDefinitionGenerator(
     llvm::DataLayout data_layout)
-    : data_layout_(std::move(data_layout)) {}
+    : data_layout_(std::move(data_layout)) {
+#ifdef ABSL_HAVE_MEMORY_SANITIZER
+  // Resolve MSan runtime functions (e.g. __msan_warning*) from the current
+  // process via dlsym. This is more future-proof than explicitly intercepting
+  // __msan_* functions; these functions do change between LLVM versions.
+  auto is_msan_symbol = [](const llvm::orc::SymbolStringPtr& name) {
+    return (*name).starts_with("__msan_");
+  };
+  auto generator =
+      llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
+          data_layout_.getGlobalPrefix(), is_msan_symbol);
+  if (generator) {
+    process_generator_ = std::move(*generator);
+  } else {
+    LOG(WARNING) << "Failed to initialize dynamic library generator for MSan: "
+                 << llvm::toString(generator.takeError());
+  }
+#endif
+}
 
 llvm::Error BuiltinDefinitionGenerator::tryToGenerate(
-    llvm::orc::LookupState&, llvm::orc::LookupKind kind,
-    llvm::orc::JITDylib& jit_dylib, llvm::orc::JITDylibLookupFlags,
+    llvm::orc::LookupState& ls, llvm::orc::LookupKind kind,
+    llvm::orc::JITDylib& jit_dylib, llvm::orc::JITDylibLookupFlags flags,
     const llvm::orc::SymbolLookupSet& names) {
   llvm::orc::SymbolMap symbols;
   symbols.reserve(names.size());
+#ifdef ABSL_HAVE_MEMORY_SANITIZER
+  llvm::orc::SymbolLookupSet msan_names;
+#endif
 
-  for (const auto& [name, flags] : names) {
+  for (const auto& [name, name_flags] : names) {
     if (auto symbol = ResolveBuiltinSymbol(data_layout_, *name)) {
       symbols[name] = *symbol;
+#ifdef ABSL_HAVE_MEMORY_SANITIZER
+    } else if ((*name).starts_with("__msan_")) {
+      msan_names.add(name, name_flags);
+#endif
     }
   }
 
   cantFail(jit_dylib.define(llvm::orc::absoluteSymbols(std::move(symbols))));
+
+#ifdef ABSL_HAVE_MEMORY_SANITIZER
+  if (!msan_names.empty() && process_generator_) {
+    return process_generator_->tryToGenerate(ls, kind, jit_dylib, flags,
+                                             msan_names);
+  }
+#endif
+
   return llvm::Error::success();
 }
 

third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h

@@ -16,7 +16,12 @@ limitations under the License.
 #ifndef XLA_BACKENDS_CPU_CODEGEN_BUILTIN_DEFINITION_GENERATOR_H_
 #define XLA_BACKENDS_CPU_CODEGEN_BUILTIN_DEFINITION_GENERATOR_H_
 
+#include <memory>
+#include <utility>
+
+#include "absl/base/config.h"  // IWYU pragma: keep
 #include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"  // IWYU pragma: keep
 #include "llvm/IR/DataLayout.h"
 #include "llvm/Support/Error.h"
 
@@ -44,6 +49,9 @@ class BuiltinDefinitionGenerator : public llvm::orc::DefinitionGenerator {
 
  private:
   llvm::DataLayout data_layout_;
+#ifdef ABSL_HAVE_MEMORY_SANITIZER
+  std::unique_ptr<llvm::orc::DynamicLibrarySearchGenerator> process_generator_;
+#endif
 };
 
 }  // namespace xla::cpu