experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Optimize the computation of noalias metadata.

Browse Source 
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
This commit is contained in:
Adrian Kuegel 2019-05-20 06:37:02 -07:00 committed by TensorFlower Gardener
parent 4f5ae52424
commit fff00129e1
2 changed files with 18 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/xla/service/llvm_ir/BUILD
View File

@ -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",
],

28
tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc
View File

@ -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.