| 230 | } |
| 231 | |
| 232 | Status ExecuteChained(OpKernelContext* context, |
| 233 | const RefPtr<XRTMemoryManager>& memory_manager, |
| 234 | xla::Backend* backend, int device_ordinal, |
| 235 | const xrt::XRTChainedExecutePlan& plan, |
| 236 | const xrt::XRTChainedExecuteConfig& config, |
| 237 | const ChainedExecuteFn& execute_op) { |
| 238 | // Create the vector which tracks the uses of the intermediate chained |
| 239 | // operations outputs. |
| 240 | std::vector<int64> uses(plan.ops_size(), 0); |
| 241 | for (auto& op : plan.ops()) { |
| 242 | for (auto& input : op.inputs()) { |
| 243 | uses[input.op_index()] += 1; |
| 244 | } |
| 245 | } |
| 246 | |
| 247 | ScopedHandles outputs(memory_manager); |
| 248 | ScopedHandles results(memory_manager); |
| 249 | for (int i = 0; i < plan.ops_size(); ++i) { |
| 250 | auto& op = plan.ops(i); |
| 251 | if (op.op_oneof_case() == xrt::XRTChainedExecuteOp::kDataHandle) { |
| 252 | // This operation is a device data load. Set the handle as output and |
| 253 | // leave the release flag off, since this is not an intermediate output. |
| 254 | TF_RETURN_IF_ERROR(outputs.Add(i, op.data_handle(), /*release=*/false)); |
| 255 | } else if (op.op_oneof_case() == |
| 256 | xrt::XRTChainedExecuteOp::kComputationHandle) { |
| 257 | // This is an XRT execute operation, forward to the device specific |
| 258 | // handler. Populating the working set makes sure the input allocations |
| 259 | // for this execute operations are pinned to device memory. |
| 260 | XRTMemoryManager::WorkingSet working_set(memory_manager); |
| 261 | TF_RETURN_IF_ERROR( |
| 262 | PopulateOpWorkingSet(backend, op, i, outputs, &working_set)); |
| 263 | TF_ASSIGN_OR_RETURN(auto tuple, |
| 264 | execute_op(op, working_set.PinnedTuples())); |
| 265 | TF_RETURN_IF_ERROR(outputs.Add(i, std::move(tuple))); |
| 266 | } else { |
| 267 | return errors::InvalidArgument( |
| 268 | "Undefined operation kind at post-order position ", i); |
| 269 | } |
| 270 | // If the result of this chained operation is an output result, feed the |
| 271 | // results at the desired position. |
| 272 | for (auto& output : op.outputs()) { |
| 273 | TF_ASSIGN_OR_RETURN(auto tuple, outputs.Lookup(i)); |
| 274 | RefPtr<XRTTupleAllocation> result; |
| 275 | TF_RETURN_IF_ERROR(MakeOutput(tuple, output.output_index(), &result)); |
| 276 | TF_RETURN_IF_ERROR(results.Add(output.result_index(), std::move(result))); |
| 277 | } |
| 278 | // Drop intermediate results which have no more users. |
| 279 | for (auto& input : op.inputs()) { |
| 280 | uses[input.op_index()] -= 1; |
| 281 | if (uses[input.op_index()] == 0) { |
| 282 | TF_RETURN_IF_ERROR(outputs.Drop(input.op_index())); |
| 283 | } |
| 284 | } |
| 285 | } |
| 286 | |
| 287 | Tensor* output_tensor; |
| 288 | TF_RETURN_IF_ERROR(context->allocate_output( |
| 289 | 0, TensorShape({static_cast<int64>(results.size())}), &output_tensor)); |
no test coverage detected