Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_program_group.cc41 OpsApiFn()->TpuProgram_GetExecutableInfoFn( in ConstructExecutableInfo()
57 OpsApiFn()->TpuProgram_GetHostTransferInfoFn( in ConstructHostTransferInfo()
73 OpsApiFn()->TpuProgram_GetHloMetadataFn( in ConstructHloMetadata()
99 OpsApiFn()->TpuProgram_GetMayModifyVariablesFn(xla_tpu_program, in Initialize()
116 if (!OpsApiFn()->TpuProgram_HasShardingFn(tpu_program)) { in has_sharding_program()
128 total_size += OpsApiFn()->TpuProgram_GetProgramSizeFn(tpu_program); in program_size()
136 success &= OpsApiFn()->TpuProgram_LogProgramMemorySummaryFn(tpu_program); in LogProgramMemorySummary()
144 OpsApiFn()->TpuProgram_UnloadAndDestroyFn(tpu_program, status.c_status); in UnloadAndDestroyPrograms()
208 OpsApiFn()->TpuProgram_GetMayModifyVariablesFn(tpu_programs_[index], in may_modify_variables()
258 OpsApiFn()->TpuCompile_CompileAndBuildFn(serialized_compilation_request, in CompileAndBuild()
[all …]
Dtpu_mesh_state_interface.h41 OpsApiFn()->TpuMeshState_FreeFn(mesh_state_); in ~TpuMeshStateInterface()
46 return new TpuMeshStateInterface(OpsApiFn()->TpuMeshState_CreateFn()); in Create()
53 OpsApiFn()->TpuMeshState_MeshCommonStateFn(mesh_state_)); in mesh_common_state()
65 OpsApiFn()->TpuTopology_AvailableCoreCountFn(mesh_state_, in NeedsStaticDeviceAssignment()
Dtpu_configuration_ops.cc206 tpu::OpsApiFn()->TpuConfigurationApi_FreeCharArrayFn(tpu_topology_output); in Compute()
223 tpu::OpsApiFn()->WaitForDistributedTpuOp_DoWorkFn(&params); in Compute()
259 tpu::OpsApiFn()->TpuConfigurationApi_HasTPUPodStateFn(); in Compute()
295 tpu::OpsApiFn()->TpuConfigurationApi_FreeInt32ArrayFn(device_id_output); in Compute()
309 tpu::OpsApiFn()->InitializeHostForDistributedTpuOp_DoWorkFn(&params); in Compute()
323 tpu::OpsApiFn()->TpuConfigurationApi_RemoteCompilationCacheSizeInBytesFn( in Compute()
328 tpu::OpsApiFn()->TpuConfigurationApi_FreeCharArrayFn( in Compute()
343 tpu::OpsApiFn() in Compute()
377 tpu::OpsApiFn()->SetGlobalTPUArrayOp_DoWorkFn(tpu_topology.size(), in Compute()
393 tpu::OpsApiFn()->DisconnectDistributedTpuChipsOp_DoWorkFn( in Compute()
Dtpu_pod_state.cc77 tpu::OpsApiFn()->TpuConfigurationApi_FreeCharArrayFn(server_address_output); in GetServerAddressAndPort()
90 tpu::OpsApiFn()->TpuConfigurationApi_GetServerAddressAndPortFn(&params); in GetServerAddressAndPort()
107 tpu::OpsApiFn()->TpuNetUtil_RecycleUnusedPortFn(service_port_); in ~TpuPodState()
159 tpu::OpsApiFn()->TpuConfigurationApi_FreeCharArrayFn(host_config_output); in ConstructTpuPodState()
174 tpu::OpsApiFn()->ConfigureDistributedTpuOp_DoWorkFn(&params); in ConstructTpuPodState()
Dtpu_op_util.cc92 tpu::OpsApiFn()->TpuCompile_CreateGuaranteedConstFingerprintFn( in GuaranteedConstFingerprint()
125 tpu::OpsApiFn()->TpuCompile_CreateCompilationCacheKeyFn( in CreateCompilationCacheKey()
140 tpu::OpsApiFn()->TpuCompile_DestroyCompilationCacheKeyFn(result); in CreateCompilationCacheKey()
Dtpu_compilation_cache_interface.cc437 if (!OpsApiFn()->TpuCompile_IsTpuCompilationEnabledFn()) { in CompileIfKeyAbsentHelper()
Dtpu_compile_op_common.cc548 if (OpsApiFn()->TpuCompile_ShouldTpuCompileOpIgnoreCancellationFn()) { in Compute()
/external/tensorflow/tensorflow/stream_executor/tpu/
Dtpu_node_context.cc33 tpu::OpsApiFn()->TpuNodeContext_CreateFn(device_ordinal, status.c_status); in Create()
37 tpu::OpsApiFn()->TpuNodeContext_FreeFn(node_context); in Create()
44 tpu::OpsApiFn()->TpuNodeContext_FreeFn(node_context_); in ~TpuNodeContext()
50 tpu::OpsApiFn()->TpuNodeContext_StopChipHeartbeatsFn(status.c_status); in StopChipHeartbeats()
57 tpu::OpsApiFn()->TpuNodeContext_CloseTpuHostFn(status.c_status); in CloseTpuHost()
64 tpu::OpsApiFn()->TpuNodeContext_InitializeFn(device_ordinal, status.c_status); in Initialize()
Dtpu_platform.cc156 if (tpu::OpsApiFn()->TpuConfigurationApi_TpusPerHostFn == nullptr) { in TpusPerHost()
161 tpu::OpsApiFn()->TpuConfigurationApi_TpusPerHostFn(tpus, status); in TpusPerHost()
170 if (tpu::OpsApiFn()->TpuConfigurationApi_TpuMemoryLimitFn == nullptr) { in TpuMemoryLimit()
175 tpu::OpsApiFn()->TpuConfigurationApi_TpuMemoryLimitFn( in TpuMemoryLimit()
Dtpu_executable.cc96 tensorflow::tpu::OpsApiFn()->TpuExecutable_LoadProgramAndEnqueueToStreamFn( in LoadProgramAndEnqueueToStream()
110 tensorflow::tpu::OpsApiFn()->HardwareLayout_HostShapeToDeviceShapeFn( in HostShapeToDeviceShape()
122 tensorflow::tpu::OpsApiFn()->HardwareLayout_ShapeSizeFn(&c_shape); in ShapeSize()
/external/tensorflow/tensorflow/core/profiler/internal/tpu/
Dtpu_tracer.cc63 tpu::OpsApiFn()->TpuProfiler_CreateFn(&tpu_profiler_, status.c_status); in TpuTracer()
70 tpu::OpsApiFn()->TpuProfiler_DestroyFn(tpu_profiler_); in ~TpuTracer()
75 tpu::OpsApiFn()->TpuProfiler_StartFn(tpu_profiler_, status.c_status); in Start()
85 tpu::OpsApiFn()->TpuProfiler_StopFn(tpu_profiler_, status.c_status); in Stop()
102 tpu::OpsApiFn()->TpuProfiler_CollectDataFn(tpu_profiler_, status.c_status, in CollectData()
108 tpu::OpsApiFn()->TpuProfiler_CollectDataFn(tpu_profiler_, status.c_status, in CollectData()
/external/tensorflow/tensorflow/core/tpu/
Dtpu_api.h28 const TfTpu_OpsApiFn* OpsApiFn();
Dtpu_api.cc26 const TfTpu_OpsApiFn* OpsApiFn() { in OpsApiFn() function
Dtpu_execute.cc110 tensorflow::tpu::OpsApiFn()->HardwareLayout_HostShapeToDeviceShapeFn( in HostShapeToDeviceShape()
122 tensorflow::tpu::OpsApiFn()->HardwareLayout_ShapeSizeCompactFn(&c_shape); in ShapeSizeCompact()
131 tensorflow::tpu::OpsApiFn()->HardwareLayout_ShapeSizeCompactRawFn( in ShapeSizeCompactRaw()
256 tensorflow::tpu::OpsApiFn()->TpuExecute_RuntimeInputToPaddedDataFn( in UpdateDynamicInputs()
Dtpu_model_server_initializer.cc61 OpsApiFn()->TfTpu_InitializeTpuModelServerFn();
Dtpu_library_init_fns.inc12 auto* ops_api_fn = const_cast<TfTpu_OpsApiFn*>(tensorflow::tpu::OpsApiFn());
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dtpu_compile_ops.cc129 if (tpu::OpsApiFn() in Compute()