| 453 | } |
| 454 | |
| 455 | XlaCompiler::XlaCompiler(XlaCompiler::Options options) |
| 456 | : options_(options), |
| 457 | initialization_status_(Status::OK()), |
| 458 | next_step_id_(1), |
| 459 | device_(new XlaCompilationDevice(SessionOptions(), options_.device_type)), |
| 460 | device_mgr_(absl::WrapUnique(device_)) { |
| 461 | VLOG(1) << "XlaCompiler: " |
| 462 | << "Device: " << device_ |
| 463 | << "Device_allocator: " << options_.device_allocator |
| 464 | << "Device_ordinal: " << options_.device_ordinal |
| 465 | << "options_.device_type : " << options_.device_type; |
| 466 | if (options.device_ordinal >= 0) { |
| 467 | VLOG(1) << "Stream: " << options_.device_allocator->GetStream(options_.device_ordinal).ValueOrDie(); |
| 468 | } |
| 469 | CHECK(!options_.device_type.type_string().empty()); |
| 470 | if (options_.populate_resource_manager) { |
| 471 | initialization_status_ = |
| 472 | (*options_.populate_resource_manager)(device_->resource_manager()); |
| 473 | } |
| 474 | |
| 475 | local_flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), |
| 476 | FunctionDefLibrary{})); |
| 477 | local_pflr_.reset(new ProcessFunctionLibraryRuntime( |
| 478 | &device_mgr_, Env::Default(), options.graph_def_version, |
| 479 | local_flib_def_.get(), OptimizerOptions())); |
| 480 | pflr_.reset(new ProcessFunctionLibraryRuntime( |
| 481 | &device_mgr_, Env::Default(), options.graph_def_version, options.flib_def, |
| 482 | OptimizerOptions())); |
| 483 | |
| 484 | local_flib_runtime_ = local_pflr_->GetFLR(device_->name()); |
| 485 | flib_runtime_ = pflr_->GetFLR(device_->name()); |
| 486 | |
| 487 | // The default shape representation function is the identity. |
| 488 | if (!options_.shape_representation_fn) { |
| 489 | options_.shape_representation_fn = |
| 490 | [](const TensorShape& shape, DataType dtype, |
| 491 | bool use_fast_memory) -> xla::StatusOr<xla::Shape> { |
| 492 | xla::Shape xla_shape; |
| 493 | TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape)); |
| 494 | return xla_shape; |
| 495 | }; |
| 496 | } |
| 497 | // DeviceMemoryAllocator may be a nullptr. In such case, stream can be |
| 498 | // initialized using stream_executor. Since (as it seems) stream_executor is |
| 499 | // not accessible from here, the stream* in GpuDeviceInfo is null. |
| 500 | if (options_.device_allocator && options_.device_ordinal >= 0) { |
| 501 | device_->set_gpu_device_info_stream( |
| 502 | options_.device_allocator->GetStream(options_.device_ordinal).ValueOrDie()); |
| 503 | } |
| 504 | } |
| 505 | |
| 506 | XlaCompiler::~XlaCompiler() = default; |
| 507 |
nothing calls this directly
no test coverage detected