| 116 | } |
| 117 | |
| 118 | StatusOr<std::pair<ServiceExecutableRunOptions, StreamPool::Ptr>> |
| 119 | LocalExecutable::RunHelper(const absl::Span<const Shape* const> argument_shapes, |
| 120 | ExecutableRunOptions run_options) { |
| 121 | const ComputationLayout& computation_layout = |
| 122 | executable_->module_config().entry_computation_layout(); |
| 123 | |
| 124 | // Check argument number, shapes, and layouts. |
| 125 | if (argument_shapes.size() != computation_layout.parameter_count()) { |
| 126 | return InvalidArgument( |
| 127 | "invalid number of arguments for computation: expected %d, got %u", |
| 128 | computation_layout.parameter_count(), argument_shapes.size()); |
| 129 | } |
| 130 | for (int i = 0; i < argument_shapes.size(); ++i) { |
| 131 | if (!computation_layout.parameter_layout(i).MatchesLayoutInShape( |
| 132 | *argument_shapes[i])) { |
| 133 | return InvalidParameterArgument( |
| 134 | executable_.get(), i, |
| 135 | "Argument does not match host shape or layout of computation " |
| 136 | "parameter " |
| 137 | "%d: want %s, got %s", |
| 138 | i, |
| 139 | ShapeUtil::HumanStringWithLayout( |
| 140 | computation_layout.parameter_layout(i).shape()), |
| 141 | ShapeUtil::HumanStringWithLayout(*argument_shapes[i])); |
| 142 | } |
| 143 | } |
| 144 | |
| 145 | TF_RETURN_IF_ERROR(ValidateExecutionOptions(run_options, *backend_)); |
| 146 | |
| 147 | StreamPool::Ptr stream; |
| 148 | if (run_options.stream() == nullptr) { |
| 149 | // NB! The lifetime of `stream` needs to match the lifetime of |
| 150 | // `service_options` (otherwise we will end up using a returned stream in |
| 151 | // ExecuteOnStreamWrapper), which is why it isn't declared in the inner "if" |
| 152 | // scope. |
| 153 | TF_ASSIGN_OR_RETURN( |
| 154 | stream, BorrowStreamForDevice(run_options.device_ordinal(), backend_)); |
| 155 | run_options.set_stream(stream.get()); |
| 156 | } |
| 157 | if (run_options.allocator() == nullptr) { |
| 158 | run_options.set_allocator(backend_->memory_allocator()); |
| 159 | } |
| 160 | |
| 161 | // For local client execution on CPU backends: |
| 162 | // *) The thread pool used for eigen CPU ops is from |
| 163 | // ExecutableRunOptions.eigen_intra_op_thread_pool. |
| 164 | // *) The thread pool used for XLA CPU ops is from |
| 165 | // backend_->eigen_intra_op_thread_pool(). |
| 166 | ServiceExecutableRunOptions service_options(run_options, |
| 167 | backend_->StreamBorrower()); |
| 168 | return std::make_pair(service_options, std::move(stream)); |
| 169 | } |
| 170 | |
| 171 | StatusOr<ScopedShapedBuffer> LocalExecutable::Run( |
| 172 | const absl::Span<const ShapedBuffer* const> arguments, |
nothing calls this directly
no test coverage detected