Add a few diagnostic flags to help narrow down issues with the LLVM

backends.

PiperOrigin-RevId: 159982441

4 files changed, 72 insertions, 23 deletions

diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc
index 01e9d4010a..7f771bf601 100644
--- a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc
+++ b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc
@@ -27,6 +27,9 @@ struct DebugOptionsFlags {
   string xla_generate_hlo_graph;
   string xla_disable_hlo_passes;
   bool xla_enable_fast_math;
+  bool xla_llvm_enable_alias_scope_metadata;
+  bool xla_llvm_enable_noalias_metadata;
+  bool xla_llvm_enable_invariant_load_metadata;
   int32 xla_backend_optimization_level;
   bool xla_embed_ir_in_executable;
   string xla_dump_ir_to;
@@ -53,6 +56,9 @@ void AllocateFlags() {
   flag_values->xla_generate_hlo_graph = "";
   flag_values->xla_disable_hlo_passes = "";
   flag_values->xla_enable_fast_math = true;
+  flag_values->xla_llvm_enable_alias_scope_metadata = true;
+  flag_values->xla_llvm_enable_noalias_metadata = true;
+  flag_values->xla_llvm_enable_invariant_load_metadata = true;
   flag_values->xla_backend_optimization_level = 3;
   flag_values->xla_embed_ir_in_executable = false;
   flag_values->xla_dump_ir_to = "";
@@ -71,6 +77,19 @@ void AllocateFlags() {
            "xla_enable_fast_math", &flag_values->xla_enable_fast_math,
            "Enable unsafe fast-math optimizations in the compiler; "
            "this may produce faster code at the expense of some accuracy."),
+       tensorflow::Flag("xla_llvm_enable_alias_scope_metadata",
+                        &flag_values->xla_llvm_enable_alias_scope_metadata,
+                        "In LLVM-based backends, enable the emission of "
+                        "!alias.scope metadata in the generated IR."),
+       tensorflow::Flag("xla_llvm_enable_noalias_metadata",
+                        &flag_values->xla_llvm_enable_noalias_metadata,
+                        "In LLVM-based backends, enable the emission of "
+                        "!noalias metadata in the generated IR."),
+       tensorflow::Flag("xla_llvm_enable_invariant_load_metadata",
+                        &flag_values->xla_llvm_enable_invariant_load_metadata,
+                        "In LLVM-based backends, enable the emission of "
+                        "!invariant.load metadata in "
+                        "the generated IR."),
        tensorflow::Flag(
            "xla_backend_optimization_level",
            &flag_values->xla_backend_optimization_level,
@@ -140,6 +159,12 @@ xla::DebugOptions GetDebugOptionsFromFlags() {
       flag_values->xla_cpu_multi_thread_eigen);
   options.set_xla_gpu_cuda_data_dir(flag_values->xla_gpu_cuda_data_dir);
   options.set_xla_gpu_ftz(flag_values->xla_gpu_ftz);
+  options.set_xla_llvm_enable_alias_scope_metadata(
+      flag_values->xla_llvm_enable_alias_scope_metadata);
+  options.set_xla_llvm_enable_noalias_metadata(
+      flag_values->xla_llvm_enable_noalias_metadata);
+  options.set_xla_llvm_enable_invariant_load_metadata(
+      flag_values->xla_llvm_enable_invariant_load_metadata);
 
   std::vector<string> extra_options_parts =
       tensorflow::str_util::Split(flag_values->xla_backend_extra_options, ',');
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index c81a368992..86d5b884d4 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -1899,11 +1899,14 @@ llvm::Value* IrEmitter::EmitTempBufferPointer(
       GetTempBuffersArgument(), slice.index(), &ir_builder_);
   llvm::LoadInst* tempbuf_address_base =
       ir_builder_.CreateLoad(tempbuf_address_ptr);
-  //  Loading the address of a buffer is invariant of the point at which the
-  //  load is executed in the program because we never reassign buffers.
-  tempbuf_address_base->setMetadata(
-      llvm::LLVMContext::MD_invariant_load,
-      llvm::MDNode::get(tempbuf_address_base->getContext(), /*MDs=*/{}));
+  if (hlo_module_config_.debug_options()
+          .xla_llvm_enable_invariant_load_metadata()) {
+    //  Loading the address of a buffer is invariant of the point at which the
+    //  load is executed in the program because we never reassign buffers.
+    tempbuf_address_base->setMetadata(
+        llvm::LLVMContext::MD_invariant_load,
+        llvm::MDNode::get(tempbuf_address_base->getContext(), /*MDs=*/{}));
+  }
   llvm_ir::SetTbaaForInstruction(tempbuf_address_base, target_shape,
                                  /*is_pointer_to=*/true);
   AttachAlignmentMetadataForLoad(tempbuf_address_base, allocation.size());
diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
index d451255774..1f6932bcc3 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
@@ -50,28 +50,37 @@ void AliasAnalysis::AddAliasingInformationToIrArray(const HloInstruction& hlo,
     buffer_slice = *slices.begin();
   }
 
-  llvm::MDNode*& alias_scope_md = alias_scope_metadata_[buffer_slice];
-  if (alias_scope_md == nullptr) {
-    alias_scope_md =
-        GetAliasScopeMetadataForBuffer(buffer_slice, GetAliasDomain());
+  if (module_.config().debug_options().xla_llvm_enable_alias_scope_metadata()) {
+    llvm::MDNode*& alias_scope_md = alias_scope_metadata_[buffer_slice];
+    if (alias_scope_md == nullptr) {
+      alias_scope_md =
+          GetAliasScopeMetadataForBuffer(buffer_slice, GetAliasDomain());
+    }
+    array->AddAliasScopeMetadata(alias_scope_md);
   }
-  array->AddAliasScopeMetadata(alias_scope_md);
 
-  llvm::MDNode*& noalias_md = noalias_metadata_[buffer_slice];
-  if (noalias_md == nullptr) {
-    noalias_md = GetNoaliasMetadataForBuffer(buffer_slice, GetAliasDomain(),
-                                             assignment_, hlo);
+  if (module_.config().debug_options().xla_llvm_enable_noalias_metadata()) {
+    llvm::MDNode*& noalias_md = noalias_metadata_[buffer_slice];
+    if (noalias_md == nullptr) {
+      noalias_md = GetNoaliasMetadataForBuffer(buffer_slice, GetAliasDomain(),
+                                               assignment_, hlo);
+    }
+    array->AddNoaliasMetadata(noalias_md);
   }
-  array->AddNoaliasMetadata(noalias_md);
 
-  // Parameters of the entry computation are never stored to, loading from a
-  // parameter pointer should always return the same result within a loop.
-  if (hlo.opcode() == HloOpcode::kParameter) {
-    const std::vector<HloInstruction*>& parameter_instructions =
-        module_.entry_computation()->parameter_instructions();
-    if (std::find(parameter_instructions.begin(), parameter_instructions.end(),
-                  &hlo) != parameter_instructions.end()) {
-      array->AddInvariantLoad(llvm::MDNode::get(*context_, /*MDs=*/{}));
+  if (module_.config()
+          .debug_options()
+          .xla_llvm_enable_invariant_load_metadata()) {
+    // Parameters of the entry computation are never stored to, loading from a
+    // parameter pointer should always return the same result within a loop.
+    if (hlo.opcode() == HloOpcode::kParameter) {
+      const std::vector<HloInstruction*>& parameter_instructions =
+          module_.entry_computation()->parameter_instructions();
+      if (std::find(parameter_instructions.begin(),
+                    parameter_instructions.end(),
+                    &hlo) != parameter_instructions.end()) {
+        array->AddInvariantLoad(llvm::MDNode::get(*context_, /*MDs=*/{}));
+      }
     }
   }
 }
diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto
index 6dd7999b67..46dd28c04d 100644
--- a/tensorflow/compiler/xla/xla.proto
+++ b/tensorflow/compiler/xla/xla.proto
@@ -64,6 +64,18 @@ message DebugOptions {
   // Enable flush-to-zero semantics in the GPU backend.
   bool xla_gpu_ftz = 10;
 
+  // If true, in LLVM-based backends, emit !alias.scope metadata in
+  // generated IR.
+  bool xla_llvm_enable_alias_scope_metadata = 11;
+
+  // If true, in LLVM-based backends, emit !noalias metadata in the
+  // generated IR.
+  bool xla_llvm_enable_noalias_metadata = 12;
+
+  // If true, in LLVM-based backends, emit !invariant.load metadata in
+  // the generated IR.
+  bool xla_llvm_enable_invariant_load_metadata = 13;
+
   // Extra options to pass to the compilation backend; specific interpretation
   // of these values is left to the backend.
   map<string, string> xla_backend_extra_options = 500;