From 7a7bbc303c451fea5b3dd93109028531a89a18ab Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 4 May 2018 02:22:14 -0700 Subject: [PATCH] Do not crash on ROOT outfeed operations. PiperOrigin-RevId: 195388075 --- tensorflow/compiler/xla/service/cpu/ir_emitter.cc | 8 ++- tensorflow/compiler/xla/service/cpu/tests/BUILD | 14 ++++++ .../service/cpu/tests/cpu_literal_caching_test.cc | 16 +----- .../xla/service/cpu/tests/cpu_outfeed_test.cc | 57 ++++++++++++++++++++++ 4 files changed, 79 insertions(+), 16 deletions(-) create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index e473389..6347ee2 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -2563,8 +2563,12 @@ Status IrEmitter::FinishVisit(HloInstruction* root) { // nothing to do since the result was already written directly into the output // buffer. VLOG(2) << "FinishVisit root: " << root->ToString(); - llvm::Value* root_value = GetEmittedValueFor(root); - VLOG(2) << " value: " << llvm_ir::DumpToString(*root_value); + if (root->opcode() == HloOpcode::kOutfeed) { + VLOG(2) << " outfeed with value: " + << llvm_ir::DumpToString(*GetEmittedValueFor(root->operand(0))); + } else { + VLOG(2) << " value: " << llvm_ir::DumpToString(*GetEmittedValueFor(root)); + } auto record_complete_computation = [&](llvm::Value* prof_counter) { if (prof_counter) { diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD index 4ddb7a8..18a915e 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/BUILD +++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD @@ -161,3 +161,17 @@ tf_cc_test( "//tensorflow/core:test_main", ], ) + +tf_cc_test( + name = "cpu_outfeed_test", + srcs = ["cpu_outfeed_test.cc"], + deps = [ + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service/cpu:cpu_compiler", + "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc index b10eb74..d6e0425 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc @@ -50,16 +50,10 @@ ENTRY main { const_b = f32[2,3,2] while(f32[2,3,2] const_a), condition=while_cond, body=while_body out0 = () outfeed(f32[2,3,2] const_a) - out1 = () outfeed(f32[2,3,2] const_b) - - ROOT root = f32[] constant(1) + ROOT out1 = () outfeed(f32[2,3,2] const_b) } )"; - // TODO(b/78879738): The fake "f32[] constant(1)" root is only needed to work - // around b/78879738. Once b/78879738 is fixed, we can set one of the - // outfeeds as the root. - string filecheck_pattern = R"( CHECK: private constant [2 x [3 x [2 x float]]] CHECK-NOT: private constant [2 x [3 x [2 x float]]] @@ -99,16 +93,10 @@ ENTRY main { const_b = (f32[2,1]{1,0}, f32[2]{0}) while((f32[2,1]{1,0}, f32[2]{0}) const_a), condition=while_cond, body=while_body out0 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_a) - out1 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_b) - - ROOT root = f32[] constant(1) + ROOT out1 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_b) } )"; - // TODO(b/78879738): The fake "f32[] constant(1)" root is only needed to work - // around b/78879738. Once b/78879738 is fixed, we can set one of the - // outfeeds as the root. - string filecheck_pattern = R"( CHECK: private constant [2 x float] CHECK: private constant [2 x [1 x float]] diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc new file mode 100644 index 0000000..879372e --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc @@ -0,0 +1,57 @@ +/* Copyright 2018 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/cpu/cpu_compiler.h" +#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" + +namespace xla { +namespace cpu { +namespace { +class CpuOutfeedTest : public CpuCodegenTest {}; + +TEST_F(CpuOutfeedTest, OutfeedRoot) { + const string hlo_text = R"( +HloModule Outfeed + +ENTRY main { + const_a = f32[2,3,2] constant( + f32[2,3,2] + {{{1, 2}, {1001, 1002}, {2001, 2002}}, + {{2, 1}, {2001, 3002}, {2001, 2002}}}) + + ROOT out = () outfeed(f32[2,3,2] const_a) +} +)"; + + string filecheck_pattern = R"( +CHECK: private constant [2 x [3 x [2 x float]]] +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + tools::Parse(hlo_text)); + + CpuAotCompilationOptions options{ + /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", + /*entry_point_name=*/"entry", + /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; + + CompileAheadOfTimeAndVerifyIr(std::move(module), options, filecheck_pattern, + /*match_optimized_ir=*/false); +} + +} // namespace +} // namespace cpu +} // namespace xla -- 2.7.4