| 194 | } |
| 195 | |
| 196 | void XlaComputationLaunchContext::PopulateInputs( |
| 197 | OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, |
| 198 | const std::map<int, OptionalTensor>& variables, |
| 199 | int missing_ctx_input_prefix) { |
| 200 | se::Stream* stream = |
| 201 | ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; |
| 202 | // Build ShapedBuffers that point directly to the Tensor buffers. |
| 203 | arg_buffers_.reserve(kernel->xla_input_shapes.size() + 1); |
| 204 | arg_buffers_.resize(kernel->xla_input_shapes.size()); |
| 205 | arg_ptrs_ = std::vector<ShapedBuffer*>(arg_buffers_.size()); |
| 206 | |
| 207 | // Pass remaining parameters. |
| 208 | const Tensor* t; |
| 209 | for (int i = 0; i < kernel->xla_input_shapes.size(); ++i) { |
| 210 | int arg_num = kernel->input_mapping[i]; |
| 211 | DCHECK_GE(arg_num, missing_ctx_input_prefix); |
| 212 | const xla::Shape& shape = kernel->xla_input_shapes[i]; |
| 213 | if (variables.count(arg_num)) { |
| 214 | t = &(variables.at(arg_num).value); |
| 215 | CHECK(t); |
| 216 | } else { |
| 217 | t = &(ctx->input(arg_num - missing_ctx_input_prefix)); |
| 218 | } |
| 219 | |
| 220 | if (use_multiple_streams_) { |
| 221 | CHECK(stream) << "Must have a stream available when using XLA tensors!"; |
| 222 | XlaTensor* xla_tensor = XlaTensor::FromTensor(t); |
| 223 | CHECK(xla_tensor); |
| 224 | xla_tensor->WaitForDefinitionEventOnStream(stream); |
| 225 | } |
| 226 | |
| 227 | const xla::Shape on_device_shape = |
| 228 | client_->backend().transfer_manager()->HostShapeToDeviceShape(shape); |
| 229 | if (on_device_shape.IsTuple()) { |
| 230 | const XlaTensor* xla_tensor = XlaTensor::FromTensor(t); |
| 231 | CHECK(xla_tensor && xla_tensor->has_shaped_buffer()); |
| 232 | arg_ptrs_[i] = const_cast<ShapedBuffer*>(&xla_tensor->shaped_buffer()); |
| 233 | } else { |
| 234 | CHECK(xla::Shape::Equal().MinorToMajorOnlyInLayout()(shape, |
| 235 | on_device_shape)) |
| 236 | << "On-device shape " |
| 237 | << xla::ShapeUtil::HumanStringWithLayout(on_device_shape) |
| 238 | << " not the same as on-host shape " |
| 239 | << xla::ShapeUtil::HumanStringWithLayout(shape); |
| 240 | se::DeviceMemoryBase dmem = XlaTensor::DeviceMemoryFromTensor(*t); |
| 241 | arg_buffers_[i] = absl::make_unique<ShapedBuffer>( |
| 242 | /*on_host_shape=*/shape, /*on_device_shape=*/shape, |
| 243 | client_->platform(), client_->default_device_ordinal()); |
| 244 | arg_buffers_[i]->set_buffer(dmem, /*index=*/{}); |
| 245 | arg_ptrs_[i] = arg_buffers_[i].get(); |
| 246 | } |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | Status XlaComputationLaunchContext::PopulateOutputs( |
| 251 | OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, |
no test coverage detected