Home
last modified time | relevance | path

Searched refs:tpu_program_group (Results 1 – 17 of 17) sorted by relevance

/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compilation_cache_rpc_support.cc54 auto tpu_program_group = absl::make_unique<TpuProgramGroup>(); in DeserializeRpcResponseToCacheEntry() local
58 TF_RETURN_IF_ERROR(tpu_program_group->DeserializeFromRpcResponseProtos( in DeserializeRpcResponseToCacheEntry()
60 entry.tpu_program_group = std::move(tpu_program_group); in DeserializeRpcResponseToCacheEntry()
61 entry.size = entry.tpu_program_group->program_size(); in DeserializeRpcResponseToCacheEntry()
69 if (cache_entry.tpu_program_group() == nullptr) { in SerializeCacheEntryToBufferSlices()
82 const TpuProgramGroup* tpu_program_group = in SerializeCacheEntryToBufferSlices() local
84 cache_entry.tpu_program_group()); in SerializeCacheEntryToBufferSlices()
85 CHECK_NE(tpu_program_group, nullptr); in SerializeCacheEntryToBufferSlices()
86 CHECK_GE(tpu_program_group->program_count(), 0); in SerializeCacheEntryToBufferSlices()
88 CHECK_LT(cache_entry.core_index(), tpu_program_group->program_count()); in SerializeCacheEntryToBufferSlices()
[all …]
Dtpu_compilation_cache_external.cc44 TpuProgramGroup tpu_program_group) { in PopulateEntry() argument
46 for (int i = 0; i < tpu_program_group.program_count(); ++i) { in PopulateEntry()
50 entry->tpu_program_group = in PopulateEntry()
51 absl::make_unique<TpuProgramGroup>(std::move(tpu_program_group)); in PopulateEntry()
64 entry->tpu_program_group = absl::make_unique<TpuProgramGroup>(); in CreateAndInitializeCompiledSubgraph()
96 TpuProgramGroup tpu_program_group; in InitializeEntry() local
103 initialization_status = initialize_program(&tpu_program_group); in InitializeEntry()
115 if (tpu_program_group.has_sharding_program()) { in InitializeEntry()
120 tpu_program_group.tpu_programs(TpuProgramShardingType::kSharding)); in InitializeEntry()
128 tpu_program_group.tpu_programs(TpuProgramShardingType::kUnsharding)); in InitializeEntry()
[all …]
Dtpu_compilation_cache_rpc_support.h39 if (tpu_program_group != nullptr) { in ~CacheEntry()
40 tpu_program_group->UnloadAndDestroyPrograms(); in ~CacheEntry()
43 std::unique_ptr<TpuProgramGroupInterface> tpu_program_group; member
68 return TpuCompilationCacheEntry(cache_entry_->tpu_program_group.get(), in get()
Dtpu_compilation_cache_entry.h31 const TpuProgramGroupInterface* tpu_program_group, int core_index) in TpuCompilationCacheEntry() argument
32 : tpu_program_group_(tpu_program_group), core_index_(core_index) {} in TpuCompilationCacheEntry()
37 const TpuProgramGroupInterface* tpu_program_group() const { in tpu_program_group() function
Dcompiled_subgraph.h152 std::unique_ptr<TpuProgramGroupInterface> tpu_program_group; member
157 int64 size = tpu_program_group->program_size(); in ComputeTotalSize()
Dtpu_compile_op_impl.cc32 TpuProgramGroupInterface* tpu_program_group) { in Compile() argument
38 tpu_program_group); in Compile()
Dtpu_compilation_cache_interface.cc81 return TpuCompilationCacheEntry(entry_->tpu_program_group.get(), index_); in get()
222 entry->tpu_program_group->UnloadAndDestroyPrograms(); in UnloadAndDestroy()
464 bool compilation_success = entry->tpu_program_group->program_count() > 0; in CompileIfKeyAbsentHelper()
474 entry->tpu_program_group->LogProgramMemorySummary(); in CompileIfKeyAbsentHelper()
504 *may_modify_variables = entry->tpu_program_group->may_modify_variables_list(); in CompileIfKeyAbsentHelper()
505 *hlo_metadatas = entry->tpu_program_group->hlo_metadatas(); in CompileIfKeyAbsentHelper()
572 proto_index >= cache_entry->tpu_program_group->program_count()) { in Lookup()
Dtpu_compile_op_common.h105 TpuProgramGroupInterface* tpu_program_group) = 0;
125 TpuProgramGroupInterface* tpu_program_group);
137 TpuProgramGroupInterface* tpu_program_group) { in LookupPersistentCompilationCacheAndFillCaches() argument
Dtpu_execute_op.cc648 const tpu::TpuProgramGroup* tpu_program_group = in DoWork() local
650 entry.tpu_program_group()); in DoWork()
651 CHECK_NE(tpu_program_group, nullptr); in DoWork()
654 tpu_program_group->executable_info(core_index); in DoWork()
753 tpu_program_group->host_transfer_info(core_index); in DoWork()
757 *tpu_program_group->hlo_metadata(core_index), std::move(input), in DoWork()
761 tpu_program_group->tpu_program(core_index))); in DoWork()
Dtpu_program_group.cc271 TpuProgramGroup* tpu_program_group = in CompileAndBuild() local
273 tpu_program_group->Initialize( in CompileAndBuild()
307 TpuProgramGroup* tpu_program_group = in CompileAndBuild() local
309 tpu_program_group->Initialize( in CompileAndBuild()
Dtpu_compile_op_common.cc596 TpuProgramGroupInterface* tpu_program_group) { in CompileLocallyAndFillHostCache() argument
604 arg_shapes, tpu_program_group); in CompileLocallyAndFillHostCache()
611 mesh_state->data(), arg_shapes, tpu_program_group); in CompileLocallyAndFillHostCache()
621 tpu_program_group->LogProgramMemorySummary(); in CompileLocallyAndFillHostCache()
625 TF_RETURN_IF_ERROR(tpu_program_group->LogCompilationStats(key, duration)); in CompileLocallyAndFillHostCache()
695 [&](TpuProgramGroupInterface* tpu_program_group) { in ComputeInternal() argument
716 dynamic_shapes, guaranteed_constants, key, tpu_program_group); in ComputeInternal()
721 tpu_program_group); in ComputeInternal()
Dtpu_compile_op_impl.h58 TpuProgramGroupInterface* tpu_program_group) override;
DBUILD290 name = "tpu_program_group",
291 srcs = ["tpu_program_group.cc"],
292 hdrs = ["tpu_program_group.h"],
368 ":tpu_program_group",
502 ":tpu_program_group",
634 ":tpu_program_group",
Dtpu_compilation_cache_service.cc124 if (cache_entry.tpu_program_group() == nullptr) { in GetTpuProgram()
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dtpu_execute_op.cc125 const tpu::TpuProgramGroup* tpu_program_group, xla::Backend* const backend, in GetExecutableAliasConfig() argument
128 tpu_program_group->executable_info(core_index); in GetExecutableAliasConfig()
132 tpu_program_group->hlo_metadata(core_index) in GetExecutableAliasConfig()
191 const uint32 rng_seed, const tpu::TpuProgramGroup* tpu_program_group, in RunExecutable() argument
213 tpu_program_group->host_transfer_info(core_index); in RunExecutable()
217 *tpu_program_group->hlo_metadata(core_index), in RunExecutable()
222 tpu_program_group->tpu_program(core_index))); in RunExecutable()
233 const uint32 rng_seed, const tpu::TpuProgramGroup* tpu_program_group, in ExecuteTPUProgram() argument
240 tpu_program_group, backend, stream, core_index, in ExecuteTPUProgram()
313 const tpu::TpuProgramGroup* tpu_program_group = in DoWork() local
[all …]
Dtpu_compile_ops.cc77 tensorflow::tpu::TpuProgramGroupInterface* tpu_program_group);
85 tensorflow::tpu::TpuProgramGroupInterface* tpu_program_group) { in Compile() argument
87 computation_proto, xla_mesh_state, tpu_program_group); in Compile()
202 [&](tpu::TpuProgramGroupInterface* tpu_program_group) { in Compute() argument
205 tpu_program_group); in Compute()
DBUILD95 "//tensorflow/core/tpu/kernels:tpu_program_group",