| 144 | } |
| 145 | |
| 146 | xla::StatusOr<RefPtr<XRTTupleAllocation>> RunExecutable( |
| 147 | OpKernelContext* context, XRTGenericDeviceAccessor::ScopedRef* device_ref, |
| 148 | xla::LocalExecutable* executable, const InputBuffers& input_buffers, |
| 149 | se::Stream* stream, int rng_seed) { |
| 150 | VLOG(2) << "Executing computation."; |
| 151 | xla::ExecutableRunOptions run_options; |
| 152 | run_options.set_stream(stream); |
| 153 | run_options.set_allocator(device_ref->backend()->memory_allocator()); |
| 154 | run_options.set_intra_op_thread_pool(&context->eigen_cpu_device()); |
| 155 | run_options.set_rng_seed(rng_seed); |
| 156 | |
| 157 | Env* env = Env::Default(); |
| 158 | auto start_time = env->NowMicros(); |
| 159 | TF_ASSIGN_OR_RETURN( |
| 160 | xla::ScopedShapedBuffer run_result, |
| 161 | executable->Run(input_buffers.input_pointers, run_options)); |
| 162 | auto elapsed = env->NowMicros() - start_time; |
| 163 | VLOG(2) << "Elapsed time: " << elapsed << "us"; |
| 164 | |
| 165 | auto shaped_buffer = run_result.release(); |
| 166 | XRTTupleAllocation* output_tuple; |
| 167 | TF_RETURN_IF_ERROR(XRTTupleAllocation::CreateFromBuffer( |
| 168 | shaped_buffer, device_ref->backend(), device_ref->device_ordinal(), |
| 169 | &output_tuple)); |
| 170 | RefPtr<XRTTupleAllocation> output_tuple_ptr(output_tuple); |
| 171 | |
| 172 | // The ScopedShapedBuffer returned by the executable Run() API, in case of |
| 173 | // input/output buffer aliasing, might have holes in it, which need to be |
| 174 | // filled using the proper input tuples buffers which are the source of |
| 175 | // aliasing. |
| 176 | const xla::HloInputOutputAliasConfig& input_output_alias = |
| 177 | executable->executable()->module().input_output_alias_config(); |
| 178 | auto alias_function = |
| 179 | [&](const xla::ShapeIndex& output_index, |
| 180 | const xla::HloInputOutputAliasConfig::Alias& alias) -> Status { |
| 181 | TF_RET_CHECK(alias.parameter_number < input_buffers.input_tuples.size()); |
| 182 | return alias.kind == xla::HloInputOutputAliasConfig::AliasKind::kUserAlias |
| 183 | ? output_tuple->AliasBufferFrom( |
| 184 | *input_buffers.input_tuples[alias.parameter_number], |
| 185 | alias.parameter_index, output_index) |
| 186 | : Status::OK(); |
| 187 | }; |
| 188 | TF_RETURN_IF_ERROR(input_output_alias.ForEachAliasWithStatus(alias_function)); |
| 189 | |
| 190 | return std::move(output_tuple_ptr); |
| 191 | } |
| 192 | |
| 193 | xla::StatusOr<RefPtr<XRTTupleAllocation>> ExecuteComputation( |
| 194 | OpKernelContext* context, XRTMemoryManager* memory_manager, |
no test coverage detected