| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | collective_ops_utils_test.cc | 53 DeviceAssignment device_assignment(/*replica_count=*/3, in TEST() local
55 device_assignment(0, 0) = 42; in TEST()
56 device_assignment(1, 0) = 43; in TEST()
57 device_assignment(2, 0) = 44; in TEST()
60 GetParticipatingDevices(GlobalDeviceId(42), device_assignment, in TEST()
69 DeviceAssignment device_assignment(/*replica_count=*/4, in TEST() local
71 device_assignment(0, 0) = 42; in TEST()
72 device_assignment(1, 0) = 43; in TEST()
73 device_assignment(2, 0) = 44; in TEST()
74 device_assignment(3, 0) = 45; in TEST()
[all …]
|
| D | hlo_runner.cc | 185 DeviceAssignment* device_assignment) { in ExecuteReplicated() argument
189 return ExecuteReplicated(executable.get(), options, device_assignment); in ExecuteReplicated()
200 DeviceAssignment* device_assignment) { in ExecuteReplicatedImpl() argument
225 int64 device = (*device_assignment)(i, 0); in ExecuteReplicatedImpl()
231 device, streams.back().get(), device_assignment, run_id)); in ExecuteReplicatedImpl()
263 int64 device = (*device_assignment)(i, 0); in ExecuteReplicatedImpl()
281 int64 device = (*device_assignment)(i, 0); in ExecuteReplicatedImpl()
321 DeviceAssignment* device_assignment, ExecutionProfile* profile) { in ExecuteReplicated() argument
365 options, device_assignment); in ExecuteReplicated()
374 DeviceAssignment device_assignment, in ExecuteReplicated()
[all …]
|
| D | collective_ops_utils.cc | 79 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices() argument
83 if (replica_groups.empty() && device_assignment.computation_count() == 1) { in GetParticipatingDevices()
87 device_assignment(replica_id, /*computation_id=*/0)); in GetParticipatingDevices()
94 device_assignment.LogicalIdsForDevice(device_id)); in GetParticipatingDevices()
103 participants.emplace_back(device_assignment(replica_id, computation_id)); in GetParticipatingDevices()
|
| D | hlo_runner.h | 113 DeviceAssignment* device_assignment) override;
122 DeviceAssignment* device_assignment, ExecutionProfile* profile = nullptr);
149 int64 device, se::Stream* stream, DeviceAssignment* device_assignment,
161 DeviceAssignment* device_assignment);
|
| D | hlo_module_util.cc | 116 auto device_assignment, in CreateModuleConfig()
117 DeviceAssignment::Deserialize(execution_options->device_assignment())); in CreateModuleConfig()
118 config->set_static_device_assignment(*device_assignment); in CreateModuleConfig()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
| D | tpu_extract_head_tail_outside_compilation.mlir | 23 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
44 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
67 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
88 // CHECK-DAG: device_assignment = []
94 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
118 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
144 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
166 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
183 // CHECK-DAG: device_assignment = []
195 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = []} : () -…
[all …]
|
| D | outside_compiled_to_host_launch.mlir | 13 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
28 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
34 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
51 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
59 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
83 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
118 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
|
| D | mark_ops_for_outside_compilation.mlir | 13 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<i32>
27 …}) {allow_soft_placement = false, num_cores_per_replica = 1, topology = "", device_assignment = …
45 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
59 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
77 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
91 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
102 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
116 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
134 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
151 …}) {allow_soft_placement = true, num_cores_per_replica = 1, topology = "", device_assignment = […
[all …]
|
| D | tpu_rewrite.mlir | 8 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = [], input_…
23 …ca = 1, step_marker_location = "", padding_map = [], topology = "", device_assignment = [], input_…
39 …TEP_MARK_AT_TOP_LEVEL_WHILE_LOOP", padding_map = [], topology = "", device_assignment = [], input_…
54 …TEP_MARK_AT_TOP_LEVEL_WHILE_LOOP", padding_map = [], topology = "", device_assignment = [], input_…
69 …y_func, num_cores_per_replica = 1, padding_map = [], topology = "", device_assignment = [], input_…
84 …ica = 1, step_marker_location = 1, padding_map = [], topology = "", device_assignment = [], input_…
99 … 1, step_marker_location = "test", padding_map = [], topology = "", device_assignment = [], input_…
114 …rker_location = "STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP", topology = "", device_assignment = [], input_…
129 …TEP_MARK_AT_TOP_LEVEL_WHILE_LOOP", padding_map = "", topology = "", device_assignment = [], input_…
144 …EP_MARK_AT_TOP_LEVEL_WHILE_LOOP", padding_map = [1], topology = "", device_assignment = [], input_…
[all …]
|
| D | tpu_extract_outside_compilation.mlir | 14 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
32 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
38 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
55 // CHECK: device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""
63 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
84 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> ()
109 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
134 …}) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> (tensor<?xf32>, te…
164 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
190 }) {num_cores_per_replica = 1, topology = "", device_assignment = []} : () -> tensor<?xi32>
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | utils.cc | 118 std::shared_ptr<DeviceAssignment>* device_assignment) { in ParseDeviceAssignmentCompileOptions() argument
131 DeviceAssignment device_assignment, in ParseDeviceAssignmentCompileOptions()
134 build_options->set_device_assignment(device_assignment); in ParseDeviceAssignmentCompileOptions()
137 << build_options->device_assignment().ToString(); in ParseDeviceAssignmentCompileOptions()
138 *num_replicas = build_options->device_assignment().replica_count(); in ParseDeviceAssignmentCompileOptions()
139 *num_partitions = build_options->device_assignment().computation_count(); in ParseDeviceAssignmentCompileOptions()
140 *device_assignment = in ParseDeviceAssignmentCompileOptions()
141 std::make_shared<DeviceAssignment>(build_options->device_assignment()); in ParseDeviceAssignmentCompileOptions()
|
| D | gpu_multistream_test.cc | 53 DeviceAssignment device_assignment(1, 1); in TEST() local
54 device_assignment(0, 0) = device->id(); in TEST()
56 device_assignment); in TEST()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_compile_op_support.cc | 59 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig() argument
117 if (device_assignment.has_value()) { in CreateModuleConfig()
118 config->set_static_device_assignment(device_assignment.value()); in CreateModuleConfig()
138 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig() argument
141 device_assignment, replica_count, num_partitions, in CreateModuleConfig()
310 const absl::optional<xla::DeviceAssignment>& device_assignment, in CreateHloModules() argument
323 compilation_result.xla_output_shape, device_assignment, in CreateHloModules()
420 DeviceAssignment::Deserialize(metadata->device_assignment()); in CompileOpMetadataFromContext()
422 const DeviceAssignment& device_assignment = in CompileOpMetadataFromContext() local
425 if (device_assignment.replica_count() != num_replicas) { in CompileOpMetadataFromContext()
[all …]
|
| D | tpu_compile_op_support.h | 99 absl::optional<const xla::DeviceAssignment> device_assignment,
111 absl::optional<const xla::DeviceAssignment> device_assignment,
139 const absl::optional<xla::DeviceAssignment>& device_assignment,
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy_test.py | 56 from tensorflow.python.tpu import device_assignment as device_assignment_lib
201 device_assignment = device_assignment_lib.DeviceAssignment.build(
204 resolver, experimental_device_assignment=device_assignment)
319 device_assignment = device_assignment_lib.DeviceAssignment.build(
322 resolver, experimental_device_assignment=device_assignment)
354 device_assignment = device_assignment_lib.DeviceAssignment.build(
357 resolver, experimental_device_assignment=device_assignment)
1104 device_assignment = device_assignment_lib.DeviceAssignment(
1108 experimental_device_assignment=device_assignment)
1173 device_assignment = device_assignment_lib.DeviceAssignment(
[all …]
|
| D | tpu_strategy.py | 61 from tensorflow.python.tpu import device_assignment as device_assignment_lib # pylint: disable=unu…
335 device_assignment=experimental_device_assignment))
627 device_assignment=None): argument
641 self, tpu_cluster_resolver, device_assignment=device_assignment))
687 device_assignment=None): argument
703 self, tpu_cluster_resolver, steps_per_run, device_assignment))
790 device_assignment=None): argument
808 self._device_assignment = device_assignment
815 if device_assignment is None:
822 for replica_id in range(device_assignment.num_replicas):
[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/python/tpu/ |
| D | tpu.py | 55 from tensorflow.python.tpu import device_assignment as device_assignment_lib
333 device_assignment = _enclosing_tpu_device_assignment()
340 if device_assignment is not None and not is_packed:
347 for replica_id in range(device_assignment.num_replicas):
348 for logical_core in range(device_assignment.num_cores_per_replica):
350 device_assignment.tpu_device(
885 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
981 device_assignment,
1178 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
1246 if device_assignment is not None:
[all …]
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | xrt_compile_ops.cc | 56 int num_cores_per_replica, xla::DeviceAssignment* device_assignment) { in GenerateXlaDeviceAssignment() argument
88 (*device_assignment)(r, c) = coords.value(3); in GenerateXlaDeviceAssignment()
170 xla::DeviceAssignment device_assignment(num_replicas, in Compile() local
173 GenerateXlaDeviceAssignment(config.device_assignment(), num_replicas, in Compile()
174 num_cores_per_replica, &device_assignment)); in Compile()
175 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/compiler/xla/client/ |
| D | executable_build_options.cc | 93 const DeviceAssignment& device_assignment) { in set_device_assignment() argument
94 device_assignment_ = device_assignment; in set_device_assignment()
131 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"
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client.cc | 511 DeviceAssignment device_assignment, std::shared_ptr<PyTpuClient> client, in PyTpuExecutable() argument
514 device_assignment_(std::move(device_assignment)), in PyTpuExecutable()
581 xla::DeviceAssignmentProto device_assignment; in ExecuteHelper() local
582 CHECK(device_assignment_.Serialize(&device_assignment).ok()); in ExecuteHelper()
586 {output_buffer->DeviceBuffer()->handle.get()}, device_assignment, in ExecuteHelper()
810 absl::optional<xla::DeviceAssignment> device_assignment; in Compile() local
816 if (options.device_assignment().replica_count() != options.num_replicas()) { in Compile()
820 options.device_assignment().replica_count(), options.num_replicas()); in Compile()
821 } else if (options.device_assignment().computation_count() != 1) { in Compile()
824 options.device_assignment().computation_count()); in Compile()
[all …]
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_on_demand_compiler.cc | 54 if (options.run_options().device_assignment() != nullptr) { in ToC()
57 .device_assignment() in ToC()
60 se_options.device_assignment = in ToC()
63 se_options.device_assignment.bytes = nullptr; in ToC()
64 se_options.device_assignment.size = 0; in ToC()
138 if (se_run_options.device_assignment.bytes != nullptr) { in ExecuteAsyncOnStream()
140 se_run_options.device_assignment); in ExecuteAsyncOnStream()
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | TPUReplicateMetadata.pbtxt | 30 name: "device_assignment"
91 name: "device_assignment"
159 name: "device_assignment"
|