Home
last modified time | relevance | path

Searched refs:device_set (Results 1 – 25 of 98) sorted by relevance

1234

/external/tensorflow/tensorflow/core/tpu/
Dtpu_global_init.cc71 DeviceSet* device_set) { in DeviceSetFromDeviceMgr() argument
74 device_set->AddDevice(d); in DeviceSetFromDeviceMgr()
76 device_set->set_client_device(d); in DeviceSetFromDeviceMgr()
91 const DeviceSet& device_set, in ConstructDistributedInitializationGraph() argument
96 options.device_set = &device_set; in ConstructDistributedInitializationGraph()
139 const DeviceSet& device_set, Env* env, in InitializeTPUSystemGlobally() argument
160 task_spec, device_set, &system_spec, &tpu_system_device)); in InitializeTPUSystemGlobally()
163 job_name, device_set, graph_to_run.get())); in InitializeTPUSystemGlobally()
205 DeviceSet device_set; in InitializeTPUSystemGlobally() local
206 DeviceSetFromDeviceMgr(*device_mgr, &device_set); in InitializeTPUSystemGlobally()
[all …]
/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
Ddistributed_tpu_rewrite_helpers.cc30 const string& system_spec_string, const DeviceSet& device_set, in GetSystemDevice() argument
47 device_set.FindMatchingDevices(*system_spec, &system_devices); in GetSystemDevice()
97 const DeviceNameUtils::ParsedName& system_spec, const DeviceSet& device_set, in GetHostSystemDevices() argument
114 device_set.FindMatchingDevices(host_spec, host_system_devices); in GetHostSystemDevices()
151 const DeviceNameUtils::ParsedName& system_spec, const DeviceSet& device_set, in GetTPUDevices() argument
157 system_spec, device_set, &host_system_devices)); in GetTPUDevices()
177 device_set.FindMatchingDevices(device_spec, &host_tpu_devices); in GetTPUDevices()
201 const string& node_type, Graph* graph, const DeviceSet& device_set, in ForConfigurationNodeMatchingType() argument
222 GetSystemDevice(spec_string, device_set, &spec, &device)); in ForConfigurationNodeMatchingType()
226 TF_RETURN_IF_ERROR(GetHostSystemDevices(spec, device_set, &host_devices)); in ForConfigurationNodeMatchingType()
Ddistributed_tpu_rewrite_helpers.h40 const DeviceSet& device_set,
51 const DeviceSet& device_set, std::vector<Device*>* host_system_devices);
62 const DeviceSet& device_set,
86 const string& node_type, Graph* graph, const DeviceSet& device_set,
/external/tensorflow/tensorflow/core/common_runtime/
Ddevice_set_test.cc50 const DeviceSet& device_set() const { return devices_; } in device_set() function in tensorflow::__anon11a7156c0111::DeviceSetTest
108 EXPECT_EQ(device_set().prioritized_devices(), in TEST_F()
115 device_set().prioritized_devices(), in TEST_F()
124 device_set().prioritized_device_types(), in TEST_F()
131 device_set().prioritized_device_types(), in TEST_F()
150 device_set().SortPrioritizedDeviceVector(&sorted); in TEST_F()
163 device_set().SortPrioritizedDeviceTypeVector(&sorted); in TEST_F()
Dfunction_optimization_registry_test.cc32 Status Run(const DeviceSet& device_set, const ConfigProto& config_proto, in Run() argument
48 DeviceSet device_set; in TEST() local
51 device_set, config_proto, /*graph=*/nullptr, /*flib_def=*/nullptr, in TEST()
Dfunction_optimization_registry_pass_failure_test.cc31 Status Run(const DeviceSet& device_set, const ConfigProto& config_proto, in Run() argument
47 DeviceSet device_set; in TEST() local
50 device_set, config_proto, /*graph=*/nullptr, /*flib_def=*/nullptr, in TEST()
Dfunction_optimization_registration_test.cc30 Status Run(const DeviceSet& device_set, const ConfigProto& config_proto, in Run() argument
47 DeviceSet device_set; in TEST() local
50 device_set, config_proto, /*graph=*/nullptr, /*flib_def=*/nullptr, in TEST()
Dpartitioning_utils.cc38 const DeviceSet& device_set, Graph* graph, in PartitionFunctionGraph() argument
59 [&device_set](const string& name) -> int64 { in PartitionFunctionGraph()
60 const Device* d = device_set.FindDeviceByName(name); in PartitionFunctionGraph()
76 const DeviceSet& device_set, std::unique_ptr<Graph> graph, in PartitionFunctionGraph() argument
81 PartitionFunctionGraph(device_set, graph.get(), &partitions, in PartitionFunctionGraph()
102 const DeviceSet& device_set, std::unique_ptr<Graph> graph) { in InsertTransferOps() argument
129 TF_RETURN_IF_ERROR(PartitionFunctionGraph(device_set, graph.get(), in InsertTransferOps()
Dfunction_optimization_registry_no_pass_test.cc29 DeviceSet device_set; in TEST() local
32 device_set, config_proto, /*graph=*/nullptr, /*flib_def=*/nullptr, in TEST()
Dfunction_optimization_registry.cc29 const DeviceSet& device_set, const ConfigProto& config_proto, in Run() argument
39 return pass_->Run(device_set, config_proto, graph, flib_def, in Run()
Dpartitioning_utils.h37 const DeviceSet& device_set, std::unique_ptr<Graph> graph,
44 const DeviceSet& device_set, std::unique_ptr<Graph> graph);
Dfunction_optimization_registry.h44 virtual Status Run(const DeviceSet& device_set,
63 Status Run(const DeviceSet& device_set, const ConfigProto& config_proto,
/external/tensorflow/tensorflow/compiler/jit/
Ddevice_util_test.cc28 jit::DeviceSet device_set; in PickDeviceHelper() local
31 device_set.Insert(device_id); in PickDeviceHelper()
36 PickDeviceForXla(cache, device_set, allow_mixing_unknown_and_cpu)); in PickDeviceHelper()
110 jit::DeviceSet device_set; in SimpleRoundTripTestForDeviceSet() local
120 device_set.Insert(device_id); in SimpleRoundTripTestForDeviceSet()
124 device_set.ForEach([&](jit::DeviceId device_id) { in SimpleRoundTripTestForDeviceSet()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dmlir_bridge_pass.cc92 bool HasTPUDevice(const DeviceSet& device_set) { in HasTPUDevice() argument
93 for (const Device* device : device_set.devices()) { in HasTPUDevice()
167 const DeviceSet* device_set, const ConfigProto& config_proto, in GetPassState() argument
172 if (device_set && !HasTPUDevice(*device_set)) { in GetPassState()
266 const DeviceSet* device_set, const ConfigProto& config_proto, in GetPassState() argument
270 if (device_set && !HasTPUDevice(*device_set)) in GetPassState()
Dmlir_bridge_pass.h34 const DeviceSet* device_set, const ConfigProto& config_proto,
53 const DeviceSet* device_set, const ConfigProto& config_proto,
Dmlir_tf2xla.cc151 DeviceSet device_set; in ConvertGraphDefToXlaViaMlir() local
156 device_set.AddDevice(&device); in ConvertGraphDefToXlaViaMlir()
157 AddDevicesToOp(*module, &device_set); in ConvertGraphDefToXlaViaMlir()
/external/tensorflow/tensorflow/core/tfrt/eager/
Dtransform_graph_function.cc37 const tensorflow::DeviceSet& device_set, in TransformGraphFunction() argument
90 input_device_names, output_device_names, device_set, arg_nodes, ret_nodes, in TransformGraphFunction()
98 device_set, config, &graph, func_lib_def, &control_ret_node_names, in TransformGraphFunction()
125 optimization_options.device_set = &device_set; in TransformGraphFunction()
135 Placer placer(graph.get(), func_name, func_lib_def, &device_set, in TransformGraphFunction()
160 func_lib_def, device_set, cpu_device, config_proto, in TransformGraphFunction()
Dfunction_cache.cc47 const tensorflow::DeviceSet& device_set, in GetOrAddFunction() argument
85 op_name, *fdef, device_name, device_set, eager_ctx, in GetOrAddFunction()
92 device_names.reserve(device_set.devices().size()); in GetOrAddFunction()
93 for (auto& d : device_set.devices()) { in GetOrAddFunction()
/external/tensorflow/tensorflow/tools/optimization/
Doptimization_pass_runner.cc118 DeviceSet device_set; in Run() local
120 device_set.AddDevice(device.get()); in Run()
122 options.device_set = &device_set; in Run()
/external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/
Dconvert_graph_test.cc48 void SetDeviceSet(const DeviceSet* device_set) { device_set_ = device_set; } in SetDeviceSet() argument
107 DeviceSet device_set; in TEST() local
111 device_set.AddDevice(d); in TEST()
113 cluster.SetDeviceSet(&device_set); in TEST()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Ddevice_util.cc115 void AddDevicesToOp(mlir::Operation* op, const DeviceSet* device_set) { in AddDevicesToOp() argument
116 if (!device_set) return; in AddDevicesToOp()
123 devices.reserve(device_set->devices().size()); in AddDevicesToOp()
127 for (Device* device : device_set->devices()) { in AddDevicesToOp()
/external/tensorflow/tensorflow/compiler/mlir/
Dmlir_graph_optimization_pass.h63 const DeviceSet* device_set, const ConfigProto& config_proto,
119 Status Run(const DeviceSet& device_set, const ConfigProto& config_proto,
146 const DeviceSet* device_set, const ConfigProto& config_proto,
Dmlir_graph_optimization_pass.cc138 const DeviceSet& device_set, const ConfigProto& config_proto, in Run() argument
157 &device_set, config_proto, **graph, *flib_def); in Run()
248 AddDevicesToOp(*module_ref, &device_set); in Run()
361 pass->GetPassState(options.device_set, options.session_options->config, in Run()
391 AddDevicesToOp(*module_ref, options.device_set); in Run()
/external/tensorflow/tensorflow/core/distributed_runtime/
Dmaster.cc374 std::unique_ptr<DeviceSet> device_set; in CreateSession() local
426 device_set.reset(new DeviceSet); in CreateSession()
428 device_set->AddDevice(d.get()); in CreateSession()
433 device_set->set_client_device(d.get()); in CreateSession()
444 device_set.reset(new DeviceSet); in CreateSession()
446 device_set->AddDevice(d.get()); in CreateSession()
450 device_set->AddDevice(d); in CreateSession()
453 device_set->set_client_device(d); in CreateSession()
459 CHECK(device_set->client_device()) << "No client device found. Missing " in CreateSession()
472 std::move(device_set), std::move(filtered_worker_list)); in CreateSession()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_functional_ops.h225 const DeviceSet& device_set,
231 const DeviceSet& device_set,
248 const DeviceSet& device_set, int replica_id, uint64 cache_hash,
263 Status SetDeviceOrdinal(const DeviceSet& device_set, int device_ordinal,

1234