Home
last modified time | relevance | path

Searched refs:DeviceAssignment (Results 1 – 25 of 88) sorted by relevance

1234

/external/tensorflow/tensorflow/compiler/xla/service/
Dcomputation_placer.cc45 StatusOr<DeviceAssignment::LogicalID> DeviceAssignment::LogicalIdForDevice( in LogicalIdForDevice()
47 std::optional<DeviceAssignment::LogicalID> logical_id; in LogicalIdForDevice()
56 logical_id.emplace(DeviceAssignment::LogicalID{r, c}); in LogicalIdForDevice()
68 StatusOr<int> DeviceAssignment::ReplicaIdForDevice( in ReplicaIdForDevice()
75 absl::flat_hash_map<GlobalDeviceId, DeviceAssignment::LogicalID>
76 DeviceAssignment::GetDeviceToLogicalIdMap() const { in GetDeviceToLogicalIdMap()
77 absl::flat_hash_map<GlobalDeviceId, DeviceAssignment::LogicalID> in GetDeviceToLogicalIdMap()
82 device_to_logical_id[device_id] = DeviceAssignment::LogicalID{r, c}; in GetDeviceToLogicalIdMap()
88 Status DeviceAssignment::Serialize(DeviceAssignmentProto* proto) const { in Serialize()
101 /* static */ StatusOr<std::unique_ptr<DeviceAssignment>>
[all …]
Dcomputation_placer.h39 class DeviceAssignment : public Array2D<int> {
41 DeviceAssignment() {} in DeviceAssignment() function
42 DeviceAssignment(int replica_count, int computation_count) in DeviceAssignment() function
72 static StatusOr<std::unique_ptr<DeviceAssignment>> Deserialize(
93 virtual StatusOr<DeviceAssignment> AssignDevices(int replica_count,
Dtpu_computation_placer.cc43 StatusOr<xla::DeviceAssignment> TpuComputationPlacer::AssignDevices( in AssignDevices()
46 xla::DeviceAssignment result(replica_count, computation_count); in AssignDevices()
56 /*static*/ StatusOr<xla::DeviceAssignment>
61 xla::DeviceAssignment result(replica_count, computation_count); in AssignLocalDevices()
Dhlo_runner.h112 DeviceAssignment* device_assignment) override;
121 DeviceAssignment* device_assignment, ExecutionProfile* profile = nullptr);
133 DeviceAssignment* device_assignment = nullptr);
155 int64_t device, se::Stream* stream, DeviceAssignment* device_assignment,
167 DeviceAssignment* device_assignment);
Dtpu_computation_placer.h38 StatusOr<xla::DeviceAssignment> AssignDevices(int replica_count,
41 static StatusOr<xla::DeviceAssignment> AssignLocalDevices(
Dhlo_runner.cc196 DeviceAssignment* device_assignment) { in ExecuteReplicated()
211 DeviceAssignment* device_assignment) { in ExecuteReplicatedImpl()
343 DeviceAssignment* device_assignment, ExecutionProfile* profile) { in ExecuteReplicated()
395 DeviceAssignment* device_assignment) { in ExecuteReplicated()
396 DeviceAssignment computation_device_assignment; in ExecuteReplicated()
449 DeviceAssignment device_assignment, in ExecuteReplicated()
473 int64_t device, se::Stream* stream, DeviceAssignment* device_assignment, in GetServiceRunOptionsForDevice()
Dhlo_module_config.h229 const DeviceAssignment& static_device_assignment() const { in static_device_assignment()
233 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment()
394 std::optional<DeviceAssignment> static_device_assignment_;
Dhlo_runner_interface.h163 DeviceAssignment* device_assignment) = 0;
170 DeviceAssignment* device_assignment) = 0;
Dcompiler.h120 const DeviceAssignment& static_device_assignment() const { in static_device_assignment()
124 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment()
183 std::optional<DeviceAssignment> static_device_assignment_;
Dcollective_ops_utils.cc130 GetParticipatingDevicesGroups(const DeviceAssignment& device_assignment, in GetParticipatingDevicesGroups()
239 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices()
245 TF_ASSIGN_OR_RETURN(const DeviceAssignment::LogicalID logical_id, in GetParticipatingDevices()
/external/tensorflow/tensorflow/compiler/xla/
Dexecutable_run_options.h39 class DeviceAssignment; variable
127 const DeviceAssignment* device_assignment);
128 const DeviceAssignment* device_assignment() const;
161 const DeviceAssignment* device_assignment_ = nullptr;
Dexecutable_run_options.cc96 const DeviceAssignment* device_assignment) { in set_device_assignment()
101 const DeviceAssignment* ExecutableRunOptions::device_assignment() const { in device_assignment()
/external/tensorflow/tensorflow/compiler/xla/client/
Dexecutable_build_options.h122 const DeviceAssignment& device_assignment);
123 const DeviceAssignment& device_assignment() const { in device_assignment()
189 std::optional<DeviceAssignment> device_assignment_;
Dexecutable_build_options.cc151 std::unique_ptr<xla::DeviceAssignment> assignment, in ExecutableBuildOptionsFromProto()
152 xla::DeviceAssignment::Deserialize(input.device_assignment())); in ExecutableBuildOptionsFromProto()
163 const DeviceAssignment& device_assignment) { in set_device_assignment()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dutils.h33 std::function<StatusOr<DeviceAssignment>(int, int)>
36 std::shared_ptr<DeviceAssignment>* device_assignment);
Dutils.cc117 std::function<StatusOr<DeviceAssignment>(int, int)> in ParseDeviceAssignmentCompileOptions()
120 std::shared_ptr<DeviceAssignment>* device_assignment) { in ParseDeviceAssignmentCompileOptions()
133 DeviceAssignment device_assignment, in ParseDeviceAssignmentCompileOptions()
143 std::make_shared<DeviceAssignment>(build_options->device_assignment()); in ParseDeviceAssignmentCompileOptions()
Dpjrt_stream_executor_client.h190 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
332 std::shared_ptr<DeviceAssignment> device_assignment;
377 StatusOr<DeviceAssignment> DevicesToDeviceAssignment(
715 std::shared_ptr<DeviceAssignment> device_assignment,
742 const DeviceAssignment& device_assignment() const override { in device_assignment()
821 std::shared_ptr<DeviceAssignment> device_assignment,
850 std::shared_ptr<DeviceAssignment> device_assignment_;
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_compiler.cc556 py::class_<DeviceAssignment>(m, "DeviceAssignment") in BuildXlaCompilerSubmodule()
558 [](py::array_t<int> array) -> StatusOr<DeviceAssignment> { in BuildXlaCompilerSubmodule()
565 DeviceAssignment result(array.shape(0), array.shape(1)); in BuildXlaCompilerSubmodule()
573 .def("replica_count", &DeviceAssignment::replica_count) in BuildXlaCompilerSubmodule()
574 .def("computation_count", &DeviceAssignment::computation_count) in BuildXlaCompilerSubmodule()
575 .def("__repr__", &DeviceAssignment::ToString) in BuildXlaCompilerSubmodule()
576 .def("serialize", [](const DeviceAssignment& da) -> StatusOr<py::bytes> { in BuildXlaCompilerSubmodule()
632 [](const CompileOptions& options) -> std::optional<DeviceAssignment> { in BuildXlaCompilerSubmodule()
634 ? std::optional<DeviceAssignment>( in BuildXlaCompilerSubmodule()
640 const DeviceAssignment& device_assignment) { in BuildXlaCompilerSubmodule()
[all …]
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.h100 absl::optional<const xla::DeviceAssignment> device_assignment,
112 absl::optional<const xla::DeviceAssignment> device_assignment,
139 const absl::optional<xla::DeviceAssignment>& device_assignment,
Dtpu_compile_op_support.cc35 using ::xla::DeviceAssignment;
61 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
140 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
315 const absl::optional<xla::DeviceAssignment>& device_assignment, in CreateHloModules()
425 StatusOr<std::unique_ptr<DeviceAssignment>> device_assignment_or_error = in CompileOpMetadataFromContext()
426 DeviceAssignment::Deserialize(metadata->device_assignment()); in CompileOpMetadataFromContext()
428 const DeviceAssignment& device_assignment = in CompileOpMetadataFromContext()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dreplicated_io_feed_test.cc31 static DeviceAssignment MakeDeviceAssn(size_t num_devices) { in MakeDeviceAssn()
32 DeviceAssignment assn(/*replica_count=*/num_devices, in MakeDeviceAssn()
/external/tensorflow/tensorflow/tools/api/golden/v1/
Dtensorflow.tpu.experimental.-device-assignment.pbtxt1 path: "tensorflow.tpu.experimental.DeviceAssignment"
3 is_instance: "<class \'tensorflow.python.tpu.device_assignment.DeviceAssignment\'>"
/external/tensorflow/tensorflow/tools/api/golden/v2/
Dtensorflow.tpu.experimental.-device-assignment.pbtxt1 path: "tensorflow.tpu.experimental.DeviceAssignment"
3 is_instance: "<class \'tensorflow.python.tpu.device_assignment.DeviceAssignment\'>"
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dlibtpu.h66 typedef struct DeviceAssignment { struct
69 } DeviceAssignment; argument
153 struct DeviceAssignment device_assignment, int32_t eventc,
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_compile_ops.cc55 const xrt::DeviceAssignment& xrt_device_assignment, int num_replicas, in GenerateXlaDeviceAssignment()
56 int num_cores_per_replica, xla::DeviceAssignment* device_assignment) { in GenerateXlaDeviceAssignment()
171 xla::DeviceAssignment device_assignment(num_replicas, in Compile()

1234