Searched refs:execution_input (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | xrt_execute_op.cc | 214 xla::ExecutionInput* execution_input = &(*execution_inputs)[i]; in UpdateDynamicInputs() local
225 xla::ShapeUtil::TryGetSubshape(execution_input->shape(), index)); in UpdateDynamicInputs()
235 execution_input->Buffer(index).AsDeviceMemoryBase(); in UpdateDynamicInputs()
244 execution_input->SetBuffer( in UpdateDynamicInputs()
246 execution_input->ClearUnownedIndex(index); in UpdateDynamicInputs()
251 TF_RETURN_IF_ERROR(execution_input->SetDynamicShape(compile_time_shape)); in UpdateDynamicInputs()
253 execution_input->ToShapedBuffer( in UpdateDynamicInputs()
442 const Tensor& execution_input = context->input(0); in DoWork() local
443 TF_RET_CHECK(TensorShapeUtils::IsScalar(execution_input.shape())); in DoWork()
444 int64 compilation_handle = execution_input.scalar<int64>()(); in DoWork()
|
| D | tpu_execute_op.cc | 293 const Tensor& execution_input = context->input(0); in DoWork() local
294 TF_RET_CHECK(TensorShapeUtils::IsScalar(execution_input.shape())); in DoWork()
295 int64 compilation_handle = execution_input.scalar<int64>()(); in DoWork()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | pjrt_stream_executor_client.cc | 566 ExecutionInput* execution_input, in AddToInput() argument
570 buffer()->AddToInputAsDonated(iterator, end, execution_input, allocator); in AddToInput()
1473 ExecutionInput execution_input; member
1516 ExecutionInput execution_input(on_device_shape); in MakeTupleHelper() local
1518 execution_input.MutableBuffers()->begin(); in MakeTupleHelper()
1520 execution_input.MutableBuffers()->end(); in MakeTupleHelper()
1522 execution_input.SetBuffer( in MakeTupleHelper()
1529 device_buffer.AddToInput(&input_iterator, iterator_end, &execution_input, in MakeTupleHelper()
1535 stream, execution_input.Buffers())); in MakeTupleHelper()
1545 return TupleHandle({std::move(execution_input), std::move(transfer_event)}); in MakeTupleHelper()
[all …]
|
| D | tracked_device_buffer.cc | 151 ExecutionInput* execution_input, in AddToInputAsDonated() argument
158 execution_input->SetUnownedIndex((*iterator)->first); in AddToInputAsDonated()
|
| D | tracked_device_buffer.h | 163 ExecutionInput* execution_input,
|
| D | pjrt_stream_executor_client.h | 414 ExecutionInput* execution_input,
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_launch_util.cc | 204 static void PopulateExecutionInputBuffer(xla::ExecutionInput& execution_input, in PopulateExecutionInputBuffer() argument
210 execution_input.MutableBuffer(index); in PopulateExecutionInputBuffer()
269 xla::ExecutionInput& execution_input = arguments.back(); in PopulateInputs() local
272 PopulateExecutionInputBuffer(execution_input, xla::ShapeIndex{}, dmem, in PopulateInputs()
280 PopulateExecutionInputBuffer(execution_input, index, *buffer, in PopulateInputs()
|