Home
last modified time | relevance | path

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

123

/external/tensorflow/tensorflow/compiler/xla/service/
Dcomputation_placer.cc45 StatusOr<std::pair<int, int>> DeviceAssignment::LogicalIdsForDevice( in LogicalIdsForDevice()
68 StatusOr<int> DeviceAssignment::ReplicaIdForDevice( in ReplicaIdForDevice()
74 Status DeviceAssignment::Serialize(DeviceAssignmentProto* proto) const { in Serialize()
87 /* static */ StatusOr<std::unique_ptr<DeviceAssignment>>
88 DeviceAssignment::Deserialize(const DeviceAssignmentProto& proto) { in Deserialize()
96 auto assignment = absl::make_unique<DeviceAssignment>( in Deserialize()
111 string DeviceAssignment::ToString() const { in ToString()
133 StatusOr<DeviceAssignment> ComputationPlacer::AssignDevices( in AssignDevices()
135 DeviceAssignment assignment(replica_count, computation_count); in AssignDevices()
Dcomputation_placer.h41 class DeviceAssignment : public Array2D<int> {
43 DeviceAssignment() {} in DeviceAssignment() function
44 DeviceAssignment(int replica_count, int computation_count) in DeviceAssignment() function
65 static StatusOr<std::unique_ptr<DeviceAssignment>> Deserialize(
86 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.h113 DeviceAssignment* device_assignment) override;
122 DeviceAssignment* device_assignment, ExecutionProfile* profile = nullptr);
149 int64 device, se::Stream* stream, DeviceAssignment* device_assignment,
161 DeviceAssignment* device_assignment);
Dcollective_ops_utils_test.cc53 DeviceAssignment device_assignment(/*replica_count=*/3, in TEST()
69 DeviceAssignment device_assignment(/*replica_count=*/4, in TEST()
92 DeviceAssignment device_assignment(/*replica_count=*/2, in TEST()
Dtpu_computation_placer.h38 StatusOr<xla::DeviceAssignment> AssignDevices(int replica_count,
41 static StatusOr<xla::DeviceAssignment> AssignLocalDevices(
Dhlo_module_config.h181 const DeviceAssignment& static_device_assignment() const { in static_device_assignment()
185 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment()
277 absl::optional<DeviceAssignment> static_device_assignment_;
Dcompiler.h99 const DeviceAssignment& static_device_assignment() const { in static_device_assignment()
103 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment()
128 absl::optional<DeviceAssignment> static_device_assignment_;
Dhlo_runner.cc185 DeviceAssignment* device_assignment) { in ExecuteReplicated()
200 DeviceAssignment* device_assignment) { in ExecuteReplicatedImpl()
321 DeviceAssignment* device_assignment, ExecutionProfile* profile) { in ExecuteReplicated()
374 DeviceAssignment device_assignment, in ExecuteReplicated()
421 DeviceAssignment device_assignment, in ExecuteReplicated()
443 int64 device, se::Stream* stream, DeviceAssignment* device_assignment, in GetServiceRunOptionsForDevice()
/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/python/
Dxla_compiler.cc417 py::class_<DeviceAssignment>(m, "DeviceAssignment") in BuildXlaCompilerSubmodule()
419 [](py::array_t<int> array) -> StatusOr<DeviceAssignment> { in BuildXlaCompilerSubmodule()
426 DeviceAssignment result(array.shape(0), array.shape(1)); in BuildXlaCompilerSubmodule()
434 .def("replica_count", &DeviceAssignment::replica_count) in BuildXlaCompilerSubmodule()
435 .def("computation_count", &DeviceAssignment::computation_count) in BuildXlaCompilerSubmodule()
436 .def("__repr__", &DeviceAssignment::ToString); in BuildXlaCompilerSubmodule()
477 -> absl::optional<DeviceAssignment> { in BuildXlaCompilerSubmodule()
479 ? absl::optional<DeviceAssignment>( in BuildXlaCompilerSubmodule()
485 const DeviceAssignment& device_assignment) { in BuildXlaCompilerSubmodule()
547 -> absl::optional<DeviceAssignment> { in BuildXlaCompilerSubmodule()
[all …]
/external/tensorflow/tensorflow/compiler/xla/client/
Dexecutable_build_options.h107 const DeviceAssignment& device_assignment);
108 const DeviceAssignment& device_assignment() const { in device_assignment()
152 absl::optional<DeviceAssignment> device_assignment_;
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dutils.h33 std::function<StatusOr<DeviceAssignment>(int, int)>
36 std::shared_ptr<DeviceAssignment>* device_assignment);
Dutils.cc115 std::function<StatusOr<DeviceAssignment>(int, int)> in ParseDeviceAssignmentCompileOptions()
118 std::shared_ptr<DeviceAssignment>* device_assignment) { in ParseDeviceAssignmentCompileOptions()
131 DeviceAssignment device_assignment, in ParseDeviceAssignmentCompileOptions()
141 std::make_shared<DeviceAssignment>(build_options->device_assignment()); in ParseDeviceAssignmentCompileOptions()
Dpjrt_stream_executor_client.h164 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
280 StatusOr<DeviceAssignment> DevicesToDeviceAssignment(
606 std::shared_ptr<DeviceAssignment> device_assignment,
633 const DeviceAssignment& device_assignment() const override { in device_assignment()
693 std::shared_ptr<DeviceAssignment> device_assignment) const;
715 std::shared_ptr<DeviceAssignment> device_assignment_;
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.h99 absl::optional<const xla::DeviceAssignment> device_assignment,
111 absl::optional<const xla::DeviceAssignment> device_assignment,
139 const absl::optional<xla::DeviceAssignment>& device_assignment,
Dtpu_compile_op_support.cc33 using ::xla::DeviceAssignment;
59 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
138 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
310 const absl::optional<xla::DeviceAssignment>& device_assignment, in CreateHloModules()
419 StatusOr<std::unique_ptr<DeviceAssignment>> device_assignment_or_error = in CompileOpMetadataFromContext()
420 DeviceAssignment::Deserialize(metadata->device_assignment()); in CompileOpMetadataFromContext()
422 const DeviceAssignment& device_assignment = in CompileOpMetadataFromContext()
/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/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/compiler/xla/python/tpu_driver/client/
Dlibtpu.h66 typedef struct DeviceAssignment { struct
69 } DeviceAssignment; argument
153 struct DeviceAssignment device_assignment, int32_t eventc,
Dtpu_client.h106 virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
299 DeviceAssignment device_assignment, std::shared_ptr<PyTpuClient> client,
322 const DeviceAssignment& device_assignment() const { in device_assignment()
367 const DeviceAssignment device_assignment_;
/external/tensorflow/tensorflow/python/distribute/
Dtpu_strategy_test.py201 device_assignment = device_assignment_lib.DeviceAssignment.build(
208 device_assignment2 = device_assignment_lib.DeviceAssignment.build(
319 device_assignment = device_assignment_lib.DeviceAssignment.build(
327 device_assignment2 = device_assignment_lib.DeviceAssignment(
354 device_assignment = device_assignment_lib.DeviceAssignment.build(
362 device_assignment2 = device_assignment_lib.DeviceAssignment(
1104 device_assignment = device_assignment_lib.DeviceAssignment(
1173 device_assignment = device_assignment_lib.DeviceAssignment(
1209 device_assignment = device_assignment_lib.DeviceAssignment(
1221 device_assignment = device_assignment_lib.DeviceAssignment(
[all …]
/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()
170 xla::DeviceAssignment device_assignment(num_replicas, in Compile()
/external/tensorflow/tensorflow/core/tpu/
Dtpu_execute.h46 tpu::TpuNodeContext* node_context, xla::DeviceAssignment* device_assignment,

123