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

[XLA] Key noalias metadata cache on the producing hlo

Browse Source 
Otherwise buffer reuse leads to nonsensical noalias metadata. It includes the
operands/users of a random other instruction but is missing noalias information
for the instruction currently being emitted.

PiperOrigin-RevId: 240587695
This commit is contained in:
Benjamin Kramer 2019-03-27 10:16:07 -07:00 committed by TensorFlower Gardener
parent 743a9f3780
commit 7abb63ec9c
3 changed files with 18 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc
View File

@ -53,6 +53,8 @@ TEST_F(CpuNoAliasTest, Concat) {
HloInstruction* concat2 = HloInstruction* concat2 =
builder.AddInstruction(HloInstruction::CreateConcatenate( builder.AddInstruction(HloInstruction::CreateConcatenate(
ShapeUtil::MakeShape(F32, {2, 6}), {concat1, param_x}, 1)); ShapeUtil::MakeShape(F32, {2, 6}), {concat1, param_x}, 1));
HloInstruction* add = builder.AddInstruction(HloInstruction::CreateBinary(
ShapeUtil::MakeShape(F32, {2, 6}), HloOpcode::kAdd, concat2, concat2));
std::unique_ptr<HloComputation> computation = builder.Build(); std::unique_ptr<HloComputation> computation = builder.Build();
@ -117,11 +119,21 @@ TEST_F(CpuNoAliasTest, Concat) {
->setName("read_concat2_array"); ->setName("read_concat2_array");
} }
{
llvm::Value* concat2_val = ir_module.getOrInsertGlobal("add", array2d_type);
auto shape = ShapeUtil::MakeShape(F32, {2, 6});
llvm_ir::IrArray add_array(concat2_val, shape);
aa.AddAliasingInformationToIrArray(*add, &add_array);
llvm_ir::IrArray::Index zero_2d({zero, zero}, shape, zero->getType());
add_array.EmitReadArrayElement(zero_2d, &b)->setName("read_add_array");
}
// Check the AA info in the loads. // Check the AA info in the loads.
const char* filecheck_pattern = R"( const char* filecheck_pattern = R"(
CHECK: %read_param_x_array = load {{.*}} !noalias [[param_x_noalias:![0-9]+]] CHECK: %read_param_x_array = load {{.*}} !noalias [[param_x_noalias:![0-9]+]]
CHECK: %read_concat1_array = load {{.*}} !alias.scope [[concat1_scope:![0-9]+]], !noalias [[concat1_noalias:![0-9]+]] CHECK: %read_concat1_array = load {{.*}} !alias.scope [[concat1_scope:![0-9]+]], !noalias [[concat1_noalias:![0-9]+]]
CHECK: %read_concat2_array = load {{.*}} !alias.scope [[concat1_noalias]], !noalias [[concat1_scope]] CHECK: %read_concat2_array = load {{.*}} !alias.scope [[concat1_noalias]], !noalias [[concat1_scope]]
CHECK: %read_add_array = load {{.*}} !alias.scope [[concat1_noalias]]{{$}}
CHECK-DAG: [[buf_size32:![0-9]+]] = !{!"buffer:{{.*}} size:32 CHECK-DAG: [[buf_size32:![0-9]+]] = !{!"buffer:{{.*}} size:32
CHECK-DAG: [[buf_size48:![0-9]+]] = !{!"buffer:{{.*}} size:48 CHECK-DAG: [[buf_size48:![0-9]+]] = !{!"buffer:{{.*}} size:48
CHECK-DAG: [[param_x_noalias]] = !{[[buf_size48]], [[buf_size32]]} CHECK-DAG: [[param_x_noalias]] = !{[[buf_size48]], [[buf_size32]]}

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

@ -63,7 +63,7 @@ void AliasAnalysis::AddAliasingInformationToIrArray(const HloInstruction& hlo,
} }
if (module_.config().debug_options().xla_llvm_enable_noalias_metadata()) { if (module_.config().debug_options().xla_llvm_enable_noalias_metadata()) {
llvm::MDNode*& noalias_md = noalias_metadata_[buffer_slice]; llvm::MDNode*& noalias_md = noalias_metadata_[{buffer_slice, &hlo}];
if (noalias_md == nullptr) { if (noalias_md == nullptr) {
noalias_md = GetNoaliasMetadataForBuffer(buffer_slice, GetAliasDomain(), noalias_md = GetNoaliasMetadataForBuffer(buffer_slice, GetAliasDomain(),
assignment_, hlo); assignment_, hlo);

8
tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h
View File

@ -79,9 +79,11 @@ class AliasAnalysis {
absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*> absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*>
alias_scope_metadata_; alias_scope_metadata_;
// A map from a buffer slice to metadata corresponding to its noalias // A map from a buffer slice and producer to metadata corresponding to its
// metadata. // noalias metadata.
absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*> noalias_metadata_; absl::flat_hash_map<std::pair<BufferAllocation::Slice, const HloInstruction*>,
llvm::MDNode*>
noalias_metadata_;
}; };
} // namespace llvm_ir } // namespace llvm_ir