| 386 | } |
| 387 | |
| 388 | Status IrEmitter::HandleInfeed(HloInstruction* instruction) { |
| 389 | HloInfeedInstruction* infeed = Cast<HloInfeedInstruction>(instruction); |
| 390 | VLOG(2) << "HandleInfeed: " << infeed->ToString(); |
| 391 | |
| 392 | // The infeed operation produces a two-element tuple containing data and a |
| 393 | // token value. HloInfeedInstruction::infeed_shape gives us the data shape. |
| 394 | const Shape& data_shape = infeed->infeed_shape(); |
| 395 | DCHECK(ShapeUtil::Equal(data_shape, |
| 396 | ShapeUtil::GetTupleElementShape(infeed->shape(), 0))); |
| 397 | TF_RETURN_IF_ERROR(EmitTargetAddressForOp(infeed)); |
| 398 | |
| 399 | // Write the tuple index table. |
| 400 | TF_ASSIGN_OR_RETURN(BufferAllocation::Slice data_slice, |
| 401 | assignment_.GetUniqueSlice(infeed, {0})); |
| 402 | llvm::Value* data_address = EmitBufferPointer(data_slice, data_shape); |
| 403 | TF_ASSIGN_OR_RETURN(BufferAllocation::Slice token_slice, |
| 404 | assignment_.GetUniqueSlice(infeed, {1})); |
| 405 | llvm::Value* token_address = EmitBufferPointer( |
| 406 | token_slice, ShapeUtil::GetTupleElementShape(infeed->shape(), 1)); |
| 407 | llvm_ir::EmitTuple(GetIrArrayFor(infeed), {data_address, token_address}, &b_); |
| 408 | |
| 409 | if (data_shape.IsTuple()) { |
| 410 | TF_RET_CHECK(!ShapeUtil::IsNestedTuple(data_shape)); |
| 411 | |
| 412 | // For a tuple, we first copy each of the internal elements to |
| 413 | // their corresponding target locations. We then construct the |
| 414 | // tuple outer buffer containing pointers to the internal |
| 415 | // elements. |
| 416 | std::vector<llvm::Value*> tuple_element_addresses; |
| 417 | for (int64 i = 0; i < data_shape.tuple_shapes_size(); ++i) { |
| 418 | TF_ASSIGN_OR_RETURN(BufferAllocation::Slice buffer, |
| 419 | assignment_.GetUniqueSlice(infeed, {0, i})); |
| 420 | |
| 421 | const Shape& tuple_element_shape = |
| 422 | ShapeUtil::GetTupleElementShape(data_shape, i); |
| 423 | |
| 424 | // Only the outer tuple buffer's target address is obtained from |
| 425 | // GetEmittedValueFor, to handle the case when Infeed is the root |
| 426 | // instruction. Target addresses for internal elements can be obtained |
| 427 | // from EmitBufferPointer. |
| 428 | llvm::Value* tuple_element_address = |
| 429 | EmitBufferPointer(buffer, tuple_element_shape); |
| 430 | |
| 431 | TF_RETURN_IF_ERROR(EmitXfeedTransfer( |
| 432 | XfeedKind::kInfeed, tuple_element_shape, tuple_element_address)); |
| 433 | |
| 434 | tuple_element_addresses.push_back(tuple_element_address); |
| 435 | } |
| 436 | |
| 437 | llvm_ir::EmitTuple(llvm_ir::IrArray(data_address, data_shape), |
| 438 | tuple_element_addresses, &b_); |
| 439 | } else { |
| 440 | TF_RETURN_IF_ERROR( |
| 441 | EmitXfeedTransfer(XfeedKind::kInfeed, data_shape, data_address)); |
| 442 | } |
| 443 | |
| 444 | return Status::OK(); |
| 445 | } |