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

[XLA:CPU] Annotate start indices in dynamic-{,update-}slice with the HLO name.

Browse Source 
This makes the IR a bit easier to follow.

PiperOrigin-RevId: 170146717
This commit is contained in:
Justin Lebar 2017-09-26 19:45:46 -07:00 committed by TensorFlower Gardener
parent 35c44ab67d
commit 035a9be3cc
1 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
tensorflow/compiler/xla/service/elemental_ir_emitter.cc
View File

@ -44,8 +44,11 @@ limitations under the License.
namespace xla {
using llvm_ir::AsStringRef;
using llvm_ir::IrArray;
using llvm_ir::IrName;
using llvm_ir::SetToFirstInsertPoint;
using tensorflow::strings::StrCat;
StatusOr<llvm::Value*> ElementalIrEmitter::EmitUnaryOp(
const HloInstruction* op, llvm::Value* operand_value) const {
@ -721,9 +724,9 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeRngElementGenerator(
if (ir_builder_->GetInsertPoint() == in_block->end()) {
body_block = llvm_ir::CreateBasicBlock(
nullptr, llvm_ir::IrName(hlo, "rng_body"), ir_builder_);
nullptr, IrName(hlo, "rng_body"), ir_builder_);
out_block = llvm_ir::CreateBasicBlock(
nullptr, llvm_ir::IrName(hlo, "rng_out"), ir_builder_);
nullptr, IrName(hlo, "rng_out"), ir_builder_);
llvm::BranchInst::Create(body_block, in_block);
} else {
body_block = in_block->splitBasicBlock(
@ -892,12 +895,10 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator(
++operand_idx) {
const HloInstruction* operand = hlo->operand(operand_idx);
auto true_block = llvm_ir::CreateBasicBlock(
exit_block, tensorflow::strings::StrCat(
"concat_index_from_operand", operand_idx),
exit_block, StrCat("concat_index_from_operand", operand_idx),
ir_builder_);
auto false_block = llvm_ir::CreateBasicBlock(
exit_block, tensorflow::strings::StrCat(
"concat_index_not_from_operand", operand_idx),
exit_block, StrCat("concat_index_not_from_operand", operand_idx),
ir_builder_);
auto concat_dim_size =
llvm::ConstantInt::get(source_index[concat_dim]->getType(),
@ -972,6 +973,8 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator(
TF_ASSIGN_OR_RETURN(
llvm::Value * start_index_value,
operand_to_generator.at(hlo->operand(1))(dim_index));
start_index_value->setName(
AsStringRef(IrName(hlo, StrCat("start_idx", i))));
slice_start_index[i] = start_index_value;
}
@ -1004,6 +1007,8 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator(
llvm_ir::IrArray::Index dim_index(1, ir_builder_->getInt64(i));
TF_ASSIGN_OR_RETURN(llvm::Value * start_index_value,
operand_to_generator.at(start_hlo)(dim_index));
start_index_value->setName(
AsStringRef(IrName(hlo, StrCat("start_idx", i))));
slice_start_index[i] = ir_builder_->CreateZExtOrBitCast(
start_index_value, index[i]->getType());
// Emit IR to compute: slice_limit_index = start_index + update_dim
@ -1163,7 +1168,7 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator(
std::unique_ptr<llvm_ir::ForLoop> inner_loop =
llvm_ir::ForLoop::EmitForLoop(
llvm_ir::IrName(hlo, "inner"), ir_builder_->getInt64(0),
IrName(hlo, "inner"), ir_builder_->getInt64(0),
ir_builder_->getInt64(contracted_dim_size),
ir_builder_->getInt64(1), ir_builder_);