Searched refs:mesh_state (Results 1 – 11 of 11) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_compile_ops.cc | 92 tensorflow::tpu::TpuMeshStateInterface* mesh_state, int num_replicas, in CompilationCacheKey() argument
101 metadata, *mesh_state); in CompilationCacheKey()
158 tensorflow::tpu::TpuMeshStateInterface* mesh_state; in Compute() local
162 &mesh_state)); in Compute()
163 core::ScopedUnref mesh_state_unref(mesh_state); in Compute()
183 computation_proto, mesh_state, num_replicas, num_cores_per_replica); in Compute()
204 return Compile(mesh_state->data(), computation_proto, in Compute()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_ops_c_api.h | 75 const XLA_TpuMeshState* mesh_state; member
95 TpuSerializedProto compilation_request, const XLA_TpuMeshState* mesh_state,
105 TpuSerializedProto xrt_computation, const XLA_TpuMeshState* mesh_state,
148 TFTPU_CAPI_EXPORT void TpuMeshState_Free(XLA_TpuMeshState* mesh_state);
152 XLA_TpuMeshState* mesh_state);
397 const XLA_TpuMeshState* mesh_state, TpuCoreTypeEnum tpu_core_type);
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_compile_op_impl.cc | 30 const XLA_TpuMeshState* mesh_state, in Compile() argument
37 return TpuProgramGroup::CompileAndBuild(compilation_request, mesh_state, in Compile()
|
| D | tpu_compile_op_common.h | 103 const XLA_TpuMeshState* mesh_state,
121 const TpuMeshStateInterface* mesh_state,
132 const TpuMeshStateInterface* mesh_state, in LookupPersistentCompilationCacheAndFillCaches() argument
|
| D | tpu_compile_op_common.cc | 593 const TpuMeshStateInterface* mesh_state, in CompileLocallyAndFillHostCache() argument
603 compile_status = Compile(MlirToHloArgs{mlir_module_}, mesh_state->data(), in CompileLocallyAndFillHostCache()
611 mesh_state->data(), arg_shapes, tpu_program_group); in CompileLocallyAndFillHostCache()
635 TpuMeshStateInterface* mesh_state; in ComputeInternal() local
638 &mesh_state)); in ComputeInternal()
639 core::ScopedUnref mesh_state_unref(mesh_state); in ComputeInternal()
655 *mesh_state); in ComputeInternal()
715 ctx->function_library(), ctx->session_metadata(), mesh_state, in ComputeInternal()
719 ctx->function_library(), ctx->session_metadata(), mesh_state, in ComputeInternal()
|
| D | tpu_program_group.cc | 248 const XLA_TpuMeshState* mesh_state, in CompileAndBuild() argument
259 mesh_state, &xla_tpu_programs, in CompileAndBuild()
282 const XLA_TpuMeshState* mesh_state, in CompileAndBuild() argument
293 mesh_state, &xla_tpu_programs, in CompileAndBuild()
|
| D | tpu_program_group.h | 95 const XLA_TpuMeshState* mesh_state,
101 const XLA_TpuMeshState* mesh_state,
|
| D | tpu_op_util.cc | 109 const TpuMeshStateInterface& mesh_state) { in CreateCompilationCacheKey() argument
137 mesh_state.data(), in CreateCompilationCacheKey()
|
| D | tpu_configuration_ops.cc | 190 tpu::TpuMeshStateInterface* mesh_state; in Compute() local
192 OP_REQUIRES_OK(ctx, GetTpuMeshStateInterface(rmgr, &mesh_state)); in Compute()
193 core::ScopedUnref mesh_state_unref(mesh_state); in Compute()
209 auto* mesh_common_state = mesh_state->mesh_common_state(); in Compute()
|
| D | tpu_op_util.h | 40 const TpuMeshStateInterface& mesh_state);
|
| D | tpu_compile_op_impl.h | 56 const XLA_TpuMeshState* mesh_state,
|