diff options
Diffstat (limited to 'tensorflow/compiler/xla/service/elemental_ir_emitter.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/elemental_ir_emitter.cc | 17 |
1 files changed, 14 insertions, 3 deletions
diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index 5b21ae3d2a..db0a8b36cd 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -949,9 +949,20 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( const IrArray::Index& index) -> StatusOr<llvm::Value*> { IrArray::Index sliced_index(index.size()); for (int i = 0; i < index.size(); ++i) { - sliced_index[i] = ir_builder_->CreateAdd( - index[i], llvm::ConstantInt::get(index[i]->getType(), - hlo->slice_starts(i))); + int64 stride = hlo->slice_stride(i); + if (stride != 1) { + sliced_index[i] = ir_builder_->CreateAdd( + ir_builder_->CreateMul( + index[i], llvm::ConstantInt::get(index[i]->getType(), + stride)), + llvm::ConstantInt::get(index[i]->getType(), + hlo->slice_starts(i))); + } else { + sliced_index[i] = ir_builder_->CreateAdd( + index[i], + llvm::ConstantInt::get(index[i]->getType(), + hlo->slice_starts(i))); + } } return operand_to_generator.at(hlo->operand(0))(sliced_index); }; |