From 239493a6825f33c96d64b6a36be6616fbb41e42b Mon Sep 17 00:00:00 2001 From: Mark Heffernan Date: Wed, 25 Jan 2017 14:49:28 -0800 Subject: Break out HloOrdering classes into separate files. Add CreateMemoryMinimizingSequence which constructs a sequence of the instructions in an HLO module that heuristically minimizes the total size of live buffers containing HLO outputs. Change: 145599747 --- .../compiler/xla/service/hlo_ordering_test.cc | 83 ++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 tensorflow/compiler/xla/service/hlo_ordering_test.cc (limited to 'tensorflow/compiler/xla/service/hlo_ordering_test.cc') diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc new file mode 100644 index 0000000000..425bee601a --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc @@ -0,0 +1,83 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/hlo_ordering.h" + +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +class HloOrderingTest : public HloTestBase {}; + +TEST_F(HloOrderingTest, LastUseScheduledFirst) { + // Tests scheduling of the following HLO code: + // + // %ab = abs(%param) + // %exp = exp(%param) + // %add = add(%ab, %exp) + // %negate = negate(%exp) + // %sub = subtract(%add, %negate) + // + // %add should be scheduled before %negate because %add is the last (and only) + // use of %ab. Scheduling %add first then frees up %ab's buffer. + const Shape vec = ShapeUtil::MakeShape(xla::F32, {42}); + auto builder = HloComputation::Builder(TestName()); + auto param = + builder.AddInstruction(HloInstruction::CreateParameter(0, vec, "param")); + auto ab = builder.AddInstruction( + HloInstruction::CreateUnary(vec, HloOpcode::kAbs, param)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vec, HloOpcode::kExp, param)); + + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(vec, HloOpcode::kAdd, ab, exp)); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vec, HloOpcode::kNegate, exp)); + auto sub = builder.AddInstruction( + HloInstruction::CreateBinary(vec, HloOpcode::kSubtract, add, negate)); + + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); + + TF_ASSIGN_OR_ASSERT_OK( + SequentialHloOrdering::HloModuleSequence sequence, + CreateMemoryMinimizingSequence(module, [](const LogicalBuffer& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape()); + })); + // Verify that all instructions are in the sequence. + EXPECT_EQ(module.entry_computation()->instruction_count(), + sequence.at(module.entry_computation()).size()); + + // The first instruction should be the parameter and the last the root "sub". + EXPECT_EQ(param, sequence.at(module.entry_computation()).front()); + EXPECT_EQ(sub, sequence.at(module.entry_computation()).back()); + + SequentialHloOrdering ordering(&module, sequence); + EXPECT_TRUE(ordering.ExecutesBefore(add, negate)); +} + +} // namespace + +} // namespace xla -- cgit v1.2.3