Home
last modified time | relevance | path

Searched refs:replica_count (Results 1 – 25 of 52) sorted by relevance

123

/external/tensorflow/tensorflow/compiler/xla/service/
Dcomputation_placer.cc48 for (int r = 0; r < replica_count(); ++r) { in LogicalIdsForDevice()
75 proto->set_replica_count(replica_count()); in Serialize()
80 for (int replica = 0; replica < replica_count(); ++replica) { in Serialize()
90 if (proto.replica_count() <= 0 || proto.computation_count() <= 0) { in Deserialize()
94 proto.replica_count(), proto.computation_count()); in Deserialize()
97 proto.replica_count(), proto.computation_count()); in Deserialize()
102 proto.replica_count()); in Deserialize()
103 for (int replica = 0; replica < proto.replica_count(); ++replica) { in Deserialize()
113 " Replicas: ", replica_count(), "\n"); in ToString()
116 for (int replica = 0; replica < replica_count(); ++replica) { in ToString()
[all …]
Dtpu_computation_placer.cc38 int replica_count, in DeviceId() argument
44 int replica_count, int computation_count) { in AssignDevices() argument
46 xla::DeviceAssignment result(replica_count, computation_count); in AssignDevices()
48 placer_, replica_count, computation_count, result.data(), in AssignDevices()
58 int replica_count, in AssignLocalDevices() argument
61 xla::DeviceAssignment result(replica_count, computation_count); in AssignLocalDevices()
63 host_location.impl(), replica_count, computation_count, result.data(), in AssignLocalDevices()
Dcomputation_placer.h44 DeviceAssignment(int replica_count, int computation_count) in DeviceAssignment() argument
45 : Array2D<int>(replica_count, computation_count, -1) { in DeviceAssignment()
46 CHECK_GT(replica_count, 0); in DeviceAssignment()
50 int replica_count() const { return height(); } in replica_count() function
82 int replica_count, int computation_count);
86 virtual StatusOr<DeviceAssignment> AssignDevices(int replica_count,
Dtpu_computation_placer.h35 StatusOr<int> DeviceId(int replica, int computation, int replica_count,
38 StatusOr<xla::DeviceAssignment> AssignDevices(int replica_count,
42 TpuHostLocationExternal host_location, int replica_count,
Dcompile_only_service.cc76 if (options.replica_count() > 0) { in CompileAheadOfTime()
77 execution_options.set_num_replicas(options.replica_count()); in CompileAheadOfTime()
79 CHECK_EQ(options.replica_count(), in CompileAheadOfTime()
80 options.static_device_assignment().replica_count()); in CompileAheadOfTime()
Dall_reduce_simplifier.h31 explicit AllReduceSimplifier(int64 replica_count) in AllReduceSimplifier() argument
32 : replica_count_(replica_count) {} in AllReduceSimplifier()
Dhlo_module_config.cc67 if (replica_count() != 1) { in compilation_cache_key()
68 StrAppend(&key, "::replica_count=", replica_count()); in compilation_cache_key()
Dhlo_module_config.h126 void set_replica_count(int64 replica_count) { in set_replica_count() argument
127 replica_count_ = replica_count; in set_replica_count()
129 int64 replica_count() const { return replica_count_; } in replica_count() function
Dservice.cc735 const int64 replica_count = options_.number_of_replicas(); in GetDeviceHandles() local
736 if (replica_count <= 0) { in GetDeviceHandles()
739 if (available_device_count < arg->device_count() * replica_count) { in GetDeviceHandles()
743 arg->device_count(), replica_count, available_device_count); in GetDeviceHandles()
988 const int64 replica_count = options_.number_of_replicas(); in TransferToInfeed() local
989 if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { in TransferToInfeed()
994 replica_count, ").")); in TransferToInfeed()
1017 const int64 replica_count = options_.number_of_replicas(); in TransferFromOutfeed() local
1018 if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { in TransferFromOutfeed()
1021 arg->replica_id(), replica_count); in TransferFromOutfeed()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dnccl_collective_thunk.h70 int64 replica_count; member
79 int64 replica_count) { in GetNcclCollectiveConfigForMlir() argument
87 config.replica_count = replica_count; in GetNcclCollectiveConfigForMlir()
Dnccl_all_gather_thunk.cc34 int64 replica_count) { in GetNcclAllGatherConfig() argument
36 config.config = GetNcclCollectiveConfigForMlir(op, replica_count); in GetNcclAllGatherConfig()
62 ThunkInfo thunk_info, mlir::lmhlo::AllGatherOp op, int64 replica_count, in NcclAllGatherThunk() argument
65 config_(GetNcclAllGatherConfig(op, replica_count)), in NcclAllGatherThunk()
Dnccl_all_to_all_thunk.cc36 int64 replica_count) { in GetNcclAllToAllConfig() argument
38 config.config = GetNcclCollectiveConfigForMlir(op, replica_count); in GetNcclAllToAllConfig()
54 ThunkInfo thunk_info, mlir::lmhlo::AllToAllOp op, int64 replica_count, in NcclAllToAllThunk() argument
57 config_(GetNcclAllToAllConfig(op, replica_count)), in NcclAllToAllThunk()
Dnccl_all_reduce_thunk.cc90 int64 replica_count) { in GetNcclAllReduceConfig() argument
95 config.config = GetNcclCollectiveConfigForMlir(op, replica_count); in GetNcclAllReduceConfig()
111 ThunkInfo thunk_info, mlir::lmhlo::AllReduceOp op, int64 replica_count, in NcclAllReduceThunk() argument
114 config_(GetNcclAllReduceConfig(op, replica_count)), in NcclAllReduceThunk()
Dnccl_collective_thunk.cc76 config().replica_count, config().replica_groups)); in ExecuteOnStream()
78 if (IsGlobalNcclConfig() && (participants.size() != config().replica_count)) { in ExecuteOnStream()
Dnccl_all_reduce_thunk.h39 int64 replica_count, std::vector<Buffer> buffers);
Dnccl_all_to_all_thunk.h39 int64 replica_count, std::vector<Buffer> buffers);
Dnccl_all_gather_thunk.h38 int64 replica_count, std::vector<Buffer> buffers);
/external/rust/crates/grpcio-sys/grpc/tools/run_tests/xds_k8s_test_driver/framework/test_app/
Dserver_app.py181 replica_count=1) -> XdsTestServer: argument
183 if replica_count != 1:
233 replica_count=replica_count,
240 replica_count,
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_program_group.h41 explicit TpuAotCompilationOptions(int64 replica_count) in TpuAotCompilationOptions() argument
42 : num_cores_(0), replica_count_(replica_count) {} in TpuAotCompilationOptions()
51 int64 replica_count() const override { return replica_count_; } in replica_count() function
Dtpu_compile_op_support.cc59 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig() argument
99 config->set_replica_count(replica_count); in CreateModuleConfig()
138 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig() argument
141 device_assignment, replica_count, num_partitions, in CreateModuleConfig()
425 if (device_assignment.replica_count() != num_replicas) { in CompileOpMetadataFromContext()
428 device_assignment.replica_count(), " vs ", num_replicas); in CompileOpMetadataFromContext()
Dtpu_compile_op_support.h100 int replica_count, int num_partitions,
112 int replica_count,
/external/tensorflow/tensorflow/compiler/xla/tests/
Dhlo_test_base.h88 const string& name = TestName(), int64 replica_count = 1);
92 absl::string_view hlo_text, int64 replica_count = 1);
138 HloModuleConfig GetModuleConfigForTest(int64 replica_count = 1) {
141 config.set_replica_count(replica_count);
/external/rust/crates/grpcio-sys/grpc/tools/run_tests/xds_k8s_test_driver/framework/
Dxds_k8s_testcase.py198 def startTestServer(self, replica_count=1, **kwargs) -> XdsTestServer: argument
199 test_server = self.server_runner.run(replica_count=replica_count,
256 def startSecureTestServer(self, replica_count=1, **kwargs) -> XdsTestServer: argument
257 test_server = self.server_runner.run(replica_count=replica_count,
/external/tensorflow/tensorflow/compiler/xla/client/
Dclient_library.cc93 int replica_count = options.number_of_replicas(); in GetOrCreateLocalClient() local
108 service_options.set_number_of_replicas(replica_count); in GetOrCreateLocalClient()
/external/tensorflow/tensorflow/core/profiler/convert/
Dop_stats_combiner.cc96 dst->set_replica_count(std::max(src.replica_count(), dst->replica_count())); in CombineRunEnvironment()

123