| 490 | } |
| 491 | |
| 492 | std::unique_ptr<Tensor> OpKernelContext::forward_input( |
| 493 | int input_index, int output_index, DataType output_dtype, |
| 494 | const TensorShape& output_shape, MemoryType output_memory_type, |
| 495 | const AllocatorAttributes& output_attr) { |
| 496 | CHECK_GE(input_index, 0); |
| 497 | CHECK_LT(input_index, num_inputs()); |
| 498 | const TensorValue& input = (*params_->inputs)[input_index]; |
| 499 | // Check whether at graph construction time this output was marked |
| 500 | // either for no forwarding or with a reservation for this input. |
| 501 | // If it's reserved for this input we'll skip the refcount and |
| 502 | // AllocatorAttribute checks. |
| 503 | // TODO(tucker): Maybe we should skip all of the checks? |
| 504 | bool never_forward = |
| 505 | (params_->forward_from_array != nullptr && output_index >= 0 && |
| 506 | params_->forward_from_array[output_index] == Params::kNeverForward); |
| 507 | if (never_forward) return nullptr; |
| 508 | bool forward_expected = |
| 509 | (params_->forward_from_array != nullptr && output_index >= 0 && |
| 510 | params_->forward_from_array[output_index] == input_index); |
| 511 | if (!forward_expected && params_->forward_from_array != nullptr) { |
| 512 | // Check for possibly conflicting forward. |
| 513 | for (int i = 0; i < num_outputs(); ++i) { |
| 514 | if (params_->forward_from_array[i] == input_index) { |
| 515 | // This input is reserved for output i. |
| 516 | return nullptr; |
| 517 | } |
| 518 | } |
| 519 | } |
| 520 | // Check that input tensor exists and is not a ref. |
| 521 | if (input.tensor == nullptr || input.is_ref()) { |
| 522 | CHECK(!forward_expected); |
| 523 | return nullptr; |
| 524 | } |
| 525 | // Check that input type matches. |
| 526 | if (input_dtype(input_index) != output_dtype) { |
| 527 | CHECK(!forward_expected); |
| 528 | return nullptr; |
| 529 | } |
| 530 | // Check that the input and output sizes are compatible. |
| 531 | if (input.tensor->shape().num_elements() != output_shape.num_elements()) { |
| 532 | CHECK(!forward_expected); |
| 533 | return nullptr; |
| 534 | } |
| 535 | // Check that input and output memory types match, i.e. |
| 536 | // that they either both live in host or both live in device memory. |
| 537 | if (input_memory_type(input_index) != output_memory_type) { |
| 538 | CHECK(!forward_expected); |
| 539 | return nullptr; |
| 540 | } |
| 541 | if (!forward_expected) { |
| 542 | if (!input->RefCountIsOne()) { |
| 543 | return nullptr; |
| 544 | } |
| 545 | // Check that output allocator attributes are not more restrictive than |
| 546 | // input allocator attributes. |
| 547 | const auto input_attr = params_->input_alloc_attrs == nullptr |
| 548 | ? AllocatorAttributes() |
| 549 | : input_alloc_attr(input_index); |
no test coverage detected