From fff00129e110e699566f1c75d7802be646092a9a Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Mon, 20 May 2019 06:37:02 -0700
Subject: [PATCH] Optimize the computation of noalias metadata.

Use GetUniqueSlice instead of GetAllSlices and then checking whether that contains exactly 1.
Also avoid adding the same HloInstruction more than once to the worklist.

PiperOrigin-RevId: 249037871
---
 tensorflow/compiler/xla/service/llvm_ir/BUILD |  1 +
 .../xla/service/llvm_ir/alias_analysis.cc     | 28 +++++++++++--------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD
index e1303f60779..55becf2546c 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/BUILD
+++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD
@@ -39,6 +39,7 @@ cc_library(
         "//tensorflow/compiler/xla/service:logical_buffer",
         "//tensorflow/core:lib",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings",
         "@llvm//:core",
     ],
diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
index 0465312b84f..cd1431aa709 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <map>
 
+#include "absl/container/flat_hash_set.h"
 #include "llvm/IR/MDBuilder.h"
 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
 #include "tensorflow/compiler/xla/service/logical_buffer.h"
@@ -40,15 +41,14 @@ void AliasAnalysis::AddAliasingInformationToIrArray(const HloInstruction& hlo,
     // with our temporary buffers.
     buffer_slice = BufferAllocation::Slice(kParameterAllocation, 0, 0);
   } else {
-    const std::set<BufferAllocation::Slice> slices =
-        assignment_.GetAllSlices(&hlo, index);
-    if (slices.empty() || slices.size() > 1) {
+    auto unique_slice = assignment_.GetUniqueSlice(&hlo, index);
+    if (!unique_slice.ok()) {
       // Skip HLOs which don't have a buffer assigned or for which the
       // buffer can't be determined statically. We cannot determine their
       // aliasing properties in these cases.
       return;
     }
-    buffer_slice = *slices.begin();
+    buffer_slice = unique_slice.ValueOrDie();
   }
 
   if (module_.config().debug_options().xla_llvm_enable_alias_scope_metadata()) {
@@ -135,14 +135,25 @@ llvm::MDNode* AliasAnalysis::GetNoaliasMetadataForBuffer(
   //
   // This set can be increased as we need.
   std::vector<const BufferValue*> worklist;
+  absl::flat_hash_set<const HloInstruction*> added_to_worklist;
   auto add_buffers_to_worklist =
-      [&worklist, &assignment](const HloInstruction* instruction) {
+      [&](const HloInstruction* instruction) {
+        // Buffers of parameters cannot be added to the noalias set.
+        if (instruction->opcode() == HloOpcode::kParameter) {
+          return;
+        }
+        if (added_to_worklist.contains(instruction)) {
+          return;
+        }
+        added_to_worklist.insert(instruction);
         ShapeUtil::ForEachSubshape(
             instruction->shape(),
             [&](const Shape& /*shape*/, const ShapeIndex& index) {
               for (const BufferValue* buffer :
                    assignment.GetSourceBuffers(instruction, index)) {
-                worklist.push_back(buffer);
+                if (assignment.HasAllocation(*buffer)) {
+                  worklist.push_back(buffer);
+                }
               }
             });
       };
@@ -161,11 +172,6 @@ llvm::MDNode* AliasAnalysis::GetNoaliasMetadataForBuffer(
 
   std::set<BufferAllocation::Slice> buffers;
   for (const BufferValue* buffer : worklist) {
-    // Skip buffers which cannot be added to the noalias set.
-    if (!assignment.HasAllocation(*buffer) ||
-        buffer->instruction()->opcode() == HloOpcode::kParameter) {
-      continue;
-    }
     const BufferAllocation::Slice noalias_slice =
         assignment.GetAssignedAllocation(*buffer).GetSlice(*buffer);
     // Our buffer must not overlap with the noalias slice.