| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | collective_ops_utils.cc | 130 GetParticipatingDevicesGroups(const DeviceAssignment& device_assignment, in GetParticipatingDevicesGroups() argument
133 int replica_count = device_assignment.replica_count(); in GetParticipatingDevicesGroups()
134 int partition_count = device_assignment.computation_count(); in GetParticipatingDevicesGroups()
176 device_assignment(replica_id, partition_id)); in GetParticipatingDevicesGroups()
193 device_assignment(replica_id, partition_id)); in GetParticipatingDevicesGroups()
212 device_assignment(replica_id, partition_id)); in GetParticipatingDevicesGroups()
229 device_assignment(replica_id, partition_id)); in GetParticipatingDevicesGroups()
239 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices() argument
242 int replica_count = device_assignment.replica_count(); in GetParticipatingDevices()
243 int partition_count = device_assignment.computation_count(); in GetParticipatingDevices()
[all …]
|
| D | collective_ops_utils_test.cc | 120 xla::Array2D<int> device_assignment; member
150 s << "_" << device_assignment.n1() << "x" << device_assignment.n2(); in ToString()
374 int64_t num_replicas = tc.device_assignment.n1(); in TEST_P()
375 int64_t num_partitions = tc.device_assignment.n2(); in TEST_P()
376 DeviceAssignment device_assignment(num_replicas, num_partitions); in TEST_P() local
381 device_assignment(replica_id, partition_id) = in TEST_P()
382 tc.device_assignment(replica_id, partition_id); in TEST_P()
408 device_assignment, replica_groups, *group_mode); in TEST_P()
421 GetParticipatingDevicesGroups(device_assignment, replica_groups, in TEST_P()
|
| D | hlo_runner.cc | 196 DeviceAssignment* device_assignment) { in ExecuteReplicated() argument
200 return ExecuteReplicated(executable.get(), options, device_assignment); in ExecuteReplicated()
211 DeviceAssignment* device_assignment) { in ExecuteReplicatedImpl() argument
214 int64_t num_partitions = device_assignment->computation_count(); in ExecuteReplicatedImpl()
238 (*device_assignment)(i / num_partitions, i % num_partitions); in ExecuteReplicatedImpl()
244 device, streams.back().get(), device_assignment, run_id)); in ExecuteReplicatedImpl()
280 (*device_assignment)(i / num_partitions, i % num_partitions); in ExecuteReplicatedImpl()
302 (*device_assignment)(i / num_partitions, i % num_partitions); in ExecuteReplicatedImpl()
343 DeviceAssignment* device_assignment, ExecutionProfile* profile) { in ExecuteReplicated() argument
387 options, device_assignment); in ExecuteReplicated()
[all …]
|
| D | hlo_runner.h | 112 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);
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
| D | tpu_extract_head_tail_outside_compilation.mlir | 23 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
44 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
67 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
87 // CHECK-DAG: device_assignment = []
93 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
117 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
143 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
165 …}) {num_cores_per_replica = 1, step_marker_location = "", topology = "", device_assignment = []} :…
201 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
215 // CHECK-DAG: device_assignment = []
[all …]
|
| D | outside_compiled_to_host_launch.mlir | 14 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
29 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
35 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
52 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
60 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
84 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
119 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
132 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
156 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
160 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
[all …]
|
| D | host_launch_to_outside_compiled.mlir | 16 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
44 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
64 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
84 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
107 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
117 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
143 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
147 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
189 …y = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", device_assignment = [0, 0, 0, …
|
| D | mark_ops_for_outside_compilation.mlir | 9 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<i32>
21 …}) {allow_soft_placement = false, num_cores_per_replica = 1, topology = "", device_assignment = …
41 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
55 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
73 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
87 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
98 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
112 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
130 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
147 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
[all …]
|
| D | tpu_extract_outside_compilation.mlir | 13 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<2xi32>
23 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<2xi32>
41 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
47 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
64 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
72 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
93 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
118 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<2xi32>
143 …}) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> (tensor<3xf32>, te…
174 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<2xi32>
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | utils.cc | 120 std::shared_ptr<DeviceAssignment>* device_assignment) { in ParseDeviceAssignmentCompileOptions() argument
133 DeviceAssignment device_assignment, in ParseDeviceAssignmentCompileOptions()
136 build_options->set_device_assignment(device_assignment); in ParseDeviceAssignmentCompileOptions()
139 << build_options->device_assignment().ToString(); in ParseDeviceAssignmentCompileOptions()
140 *num_replicas = build_options->device_assignment().replica_count(); in ParseDeviceAssignmentCompileOptions()
141 *num_partitions = build_options->device_assignment().computation_count(); in ParseDeviceAssignmentCompileOptions()
142 *device_assignment = in ParseDeviceAssignmentCompileOptions()
143 std::make_shared<DeviceAssignment>(build_options->device_assignment()); in ParseDeviceAssignmentCompileOptions()
|
| D | gpu_multistream_test.cc | 54 DeviceAssignment device_assignment(1, 1); in TEST() local
55 device_assignment(0, 0) = device->id(); in TEST()
57 device_assignment); in TEST()
|
| /external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| D | update_tpu_metadata.cc | 124 xla::DeviceAssignmentProto device_assignment; in UpdateTPUCompileMetadata() local
125 device_assignment.set_replica_count(num_replicas); in UpdateTPUCompileMetadata()
126 device_assignment.set_computation_count(1); in UpdateTPUCompileMetadata()
127 auto* computation_device = device_assignment.add_computation_devices(); in UpdateTPUCompileMetadata()
135 *metadata_proto.mutable_device_assignment() = device_assignment; in UpdateTPUCompileMetadata()
|
| /external/tensorflow/tensorflow/compiler/xla/stream_executor/tpu/ |
| D | tpu_executable.cc | 39 if (options.run_options().device_assignment() != nullptr) { in ToC()
42 .device_assignment() in ToC()
45 se_options.device_assignment = in ToC()
48 se_options.device_assignment.bytes = nullptr; in ToC()
49 se_options.device_assignment.size = 0; in ToC()
114 if (se_run_options.device_assignment.bytes != nullptr) { in ExecuteAsyncOnStream()
116 se_run_options.device_assignment); in ExecuteAsyncOnStream()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_compile_op_support.cc | 61 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig() argument
119 if (device_assignment.has_value()) { in CreateModuleConfig()
120 config->set_static_device_assignment(device_assignment.value()); in CreateModuleConfig()
140 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig() argument
143 device_assignment, replica_count, num_partitions, in CreateModuleConfig()
315 const absl::optional<xla::DeviceAssignment>& device_assignment, in CreateHloModules() argument
328 compilation_result.xla_output_shape, device_assignment, in CreateHloModules()
426 DeviceAssignment::Deserialize(metadata->device_assignment()); in CompileOpMetadataFromContext()
428 const DeviceAssignment& device_assignment = in CompileOpMetadataFromContext() local
431 if (device_assignment.replica_count() != num_replicas) { in CompileOpMetadataFromContext()
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu.py | 51 from tensorflow.python.tpu import device_assignment as device_assignment_lib
351 device_assignment = _enclosing_tpu_device_assignment()
358 if device_assignment is not None and not is_packed:
365 for replica_id in range(device_assignment.num_replicas):
366 for logical_core in range(device_assignment.num_cores_per_replica):
368 device_assignment.tpu_device(
919 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
1015 device_assignment,
1211 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
1279 if device_assignment is not None:
[all …]
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.tpu.pbtxt | 21 …s=[\'computation\', \'inputs\', \'num_shards\', \'infeed_queue\', \'device_assignment\', \'name\',…
45 …argspec: "args=[\'computation\', \'inputs\', \'infeed_queue\', \'device_assignment\', \'name\', \'…
49 …argspec: "args=[\'computation\', \'inputs\', \'infeed_queue\', \'device_assignment\', \'name\', \'…
53 …tputs_from_all_shards\', \'output_shard_axes\', \'infeed_queue\', \'device_assignment\', \'name\',…
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | executable_build_options.cc | 122 device_assignment().Serialize(output.mutable_device_assignment())); in ToProto()
152 xla::DeviceAssignment::Deserialize(input.device_assignment())); in ExecutableBuildOptionsFromProto()
163 const DeviceAssignment& device_assignment) { in set_device_assignment() argument
164 device_assignment_ = device_assignment; in set_device_assignment()
211 TF_CHECK_OK(build_options.device_assignment().Serialize( in CreateExecutionOptions()
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/ |
| D | TPUReplicateMetadata.pbtxt | 30 name: "device_assignment"
91 name: "device_assignment"
159 name: "device_assignment"
234 name: "device_assignment"
316 name: "device_assignment"
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | xrt_compile_ops.cc | 56 int num_cores_per_replica, xla::DeviceAssignment* device_assignment) { in GenerateXlaDeviceAssignment() argument
89 (*device_assignment)(r, c) = coords.value(3); in GenerateXlaDeviceAssignment()
171 xla::DeviceAssignment device_assignment(num_replicas, in Compile() local
174 GenerateXlaDeviceAssignment(config.device_assignment(), num_replicas, in Compile()
175 num_cores_per_replica, &device_assignment)); in Compile()
176 build_options.set_device_assignment(device_assignment); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/ |
| D | executable_run_options.cc | 96 const DeviceAssignment* device_assignment) { in set_device_assignment() argument
97 device_assignment_ = device_assignment; in set_device_assignment()
101 const DeviceAssignment* ExecutableRunOptions::device_assignment() const { in device_assignment() function in xla::ExecutableRunOptions
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy_test.py | 58 from tensorflow.python.tpu import device_assignment as device_assignment_lib
338 device_assignment = device_assignment_lib.DeviceAssignment.build(
341 resolver, experimental_device_assignment=device_assignment)
456 device_assignment = device_assignment_lib.DeviceAssignment.build(
459 resolver, experimental_device_assignment=device_assignment)
491 device_assignment = device_assignment_lib.DeviceAssignment.build(
494 resolver, experimental_device_assignment=device_assignment)
1401 device_assignment = device_assignment_lib.DeviceAssignment(
1403 self.assertAllEqual([[[0, 0, 0, 0]]], device_assignment.core_assignment)
1404 self.assertEqual(1, device_assignment.num_cores_per_replica)
[all …]
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | TPUReplicateMetadata.pbtxt | 30 name: "device_assignment"
91 name: "device_assignment"
159 name: "device_assignment"
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | all_reduce_blueconnect.cc | 61 const DeviceAssignment& device_assignment, size_t num_devices_per_host) { in TryDecomposeReplicaGroup() argument
67 int device_id = device_assignment(replica_id, /*computation_id=*/0); in TryDecomposeReplicaGroup()
112 const DeviceAssignment& device_assignment = in TryDecomposeReplicaGroups() local
119 for (int i = 0; i < device_assignment.replica_count(); ++i) { in TryDecomposeReplicaGroups()
132 TryDecomposeReplicaGroup(replica_group, device_assignment, in TryDecomposeReplicaGroups()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_helpers.cc | 150 xla::DeviceAssignment& device_assignment, in ResolveDeviceAssignment() argument
191 device_assignment = xla::DeviceAssignment(params->group.group_size, 1); in ResolveDeviceAssignment()
219 device_assignment(device_idx, 0) = device.xla_global_id(); in ResolveDeviceAssignment()
221 VLOG(5) << "Generated device assignment: " << device_assignment.ToString(); in ResolveDeviceAssignment()
251 run_options.set_device_assignment(&device_assignment); in ResolveDeviceAssignment()
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client.cc | 522 DeviceAssignment device_assignment, std::shared_ptr<PyTpuClient> client, in PyTpuExecutable() argument
525 device_assignment_(std::move(device_assignment)), in PyTpuExecutable()
592 xla::DeviceAssignmentProto device_assignment; in ExecuteHelper() local
593 CHECK(device_assignment_.Serialize(&device_assignment).ok()); in ExecuteHelper()
597 {output_buffer->DeviceBuffer()->handle.get()}, device_assignment, in ExecuteHelper()
820 std::optional<xla::DeviceAssignment> device_assignment; in Compile() local
826 if (options.device_assignment().replica_count() != options.num_replicas()) { in Compile()
830 options.device_assignment().replica_count(), options.num_replicas()); in Compile()
831 } else if (options.device_assignment().computation_count() != 1) { in Compile()
834 options.device_assignment().computation_count()); in Compile()
[all …]
|