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

[XLA] Do not alias constants with the output of other HLOs.

Browse Source 
PiperOrigin-RevId: 305509905
Change-Id: I5afdccbaafa60ae483b4e9c5b2f1ef6a65444193
This commit is contained in:
Blake Hechtman 2020-04-08 10:50:56 -07:00 committed by TensorFlower Gardener
parent dd95f74f9a
commit 828d3a5129
4 changed files with 26 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/compiler/xla/service/buffer_assignment.cc
View File

@ -99,6 +99,9 @@ std::vector<int64> ColorInterferenceGraph(
bool HloBufferIsReadOnly(const HloBuffer& buffer) {
for (const HloValue* value : buffer.values()) {
const HloInstruction* instruction = value->instruction();
if (instruction->opcode() == HloOpcode::kConstant) {
return true;
}
const HloModule* module = instruction->parent()->parent();
const bool is_entry_parameter =
instruction->opcode() == HloOpcode::kParameter &&

3
tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
View File

@ -1150,6 +1150,9 @@ bool HloDataflowAnalysis::CanShareOperandBufferWithUser(
HloInstruction* user, const ShapeIndex& user_index) const {
CHECK(user->IsUserOf(operand))
<< "user: " << user->ToString() << " operand: " << operand->ToString();
if (operand->opcode() == HloOpcode::kConstant) {
return false;
}
const Shape& operand_subshape =
ShapeUtil::GetSubshape(operand->shape(), operand_index);
const Shape& user_subshape =

1
tensorflow/compiler/xla/tests/BUILD
View File

@ -1348,6 +1348,7 @@ xla_test(
srcs = ["dynamic_ops_test.cc"],
deps = [
":client_library_test_base",
":hlo_test_base",
":literal_test_util",
":test_macros_header",
":xla_internal_test_main",

19
tensorflow/compiler/xla/tests/dynamic_ops_test.cc
View File

@ -27,6 +27,7 @@ limitations under the License.
#include "tensorflow/compiler/xla/service/transfer_manager.h"
#include "tensorflow/compiler/xla/test_helpers.h"
#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/tests/test_macros.h"
#include "tensorflow/core/platform/stream_executor_no_cuda.h"
@ -731,6 +732,24 @@ XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_GPU(R3ContiguousLargerBF16)) {
RunR3Contiguous<bfloat16>(operand_shape, /*index=*/7, /*size=*/1);
}
// This test that buffer assignment does not alias constants with the output of
// dynamic update slice.
XLA_TEST_F(HloTestBase, AddOfDUS) {
const char* hlo_string = R"(
HloModule m
test {
o = s32[6] constant({2,3,4,5,6,7})
i = s32[] parameter(0)
u = s32[2] parameter(1)
dus = s32[6] dynamic-update-slice(o,u,i)
a = s32[6] add(dus, dus)
j = s32[] parameter(2)
ROOT ds = s32[2] dynamic-slice(a, j), dynamic_slice_sizes={2}
}
)";
EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{0, 0}));
}
void BM_DynamicSlice(int num_iters) {
tensorflow::testing::StopTiming();