| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_topology.cc | 21 namespace tpu { namespace
25 tpu::ExecutorApiFn()->TpuCoreLocation_ChipCoordinatesFn(core_location_, &x, in chip_coordinates()
32 tpu::ExecutorApiFn()->TpuCoreLocation_HostCoordinatesFn(core_location_, &x, in host_coordinates()
38 return tpu::ExecutorApiFn()->TpuCoreLocation_IndexFn(core_location_); in index()
42 return tpu::ExecutorApiFn()->TpuCoreLocation_IdFn(core_location_); in Id()
46 return tpu::ExecutorApiFn()->TpuHostLocation_IdFn(host_location_); in Id()
51 int num_cores = tpu::ExecutorApiFn()->TpuHostLocation_NumCoresFn( in Cores()
54 tpu::ExecutorApiFn()->TpuHostLocation_CoresFn(host_location_, core_type, in Cores()
66 return tpu::ExecutorApiFn()->TpuTopology_LogicalDevicesPerHostFn(topology_, in LogicalDevicesPerHost()
72 return tpu::ExecutorApiFn()->TpuTopology_LogicalDevicesPerChipFn(topology_, in LogicalDevicesPerChip()
[all …]
|
| D | tpu_platform.cc | 26 namespace tpu { namespace
36 platform_ = tpu::ExecutorApiFn()->TpuPlatform_NewFn(); in TpuPlatform()
60 tpu::ExecutorApiFn()->TpuPlatform_InitializeFn( in Initialize()
70 return tpu::ExecutorApiFn()->TpuPlatform_InitializedFn(platform_); in Initialized()
74 tpu::ExecutorApiFn()->TpuPlatform_FreeFn(platform_); in ~TpuPlatform()
78 return tpu::ExecutorApiFn()->TpuPlatform_VisibleDeviceCountFn(platform_); in VisibleDeviceCount()
91 tpu::ExecutorApiFn()->TpuStreamExecutorConfig_DefaultFn(); in GetUncachedExecutor()
93 tpu::ExecutorApiFn()->TpuStreamExecutorConfig_SetOrdinalFn(c_config, in GetUncachedExecutor()
97 SE_StreamExecutor* executor = tpu::ExecutorApiFn()->TpuPlatform_GetExecutorFn( in GetUncachedExecutor()
99 tpu::ExecutorApiFn()->TpuStreamExecutorConfig_FreeFn(c_config); in GetUncachedExecutor()
[all …]
|
| D | tpu_executor.cc | 29 namespace tpu { namespace
36 tpu::ExecutorApiFn()->TpuExecutor_FreeFn(executor_); in ~TpuExecutor()
43 tpu::ExecutorApiFn()->TpuExecutor_NewDeviceOptionsFn( in Init()
45 tpu::ExecutorApiFn()->TpuExecutor_InitFn(executor_, device_ordinal, options, in Init()
47 tpu::ExecutorApiFn()->TpuExecutor_FreeDeviceOptionsFn(options); in Init()
52 return tpu::ExecutorApiFn()->TpuExecutor_PlatformDeviceCountFn(executor_); in PlatformDeviceCount()
56 tpu::ExecutorApiFn()->TpuExecutor_SyncAndForgetFailedStreamsFn(executor_); in SyncAndForgetFailedStreams()
60 return tpu::ExecutorApiFn()->TpuExecutor_SynchronizeAllActivityFn(executor_); in SynchronizeAllActivity()
65 tpu::ExecutorApiFn()->TpuExecutor_BlockHostUntilDoneFn( in BlockHostUntilDone()
72 tpu::ExecutorApiFn()->TpuExecutor_BlockUntilDoneOrFailedFn(executor_, in BlockUntilDoneOrFailed()
[all …]
|
| D | tpu_transfer_manager.cc | 34 namespace tpu { namespace
41 manager_ = tpu::ExecutorApiFn()->TpuTransferManager_NewFn(); in TpuTransferManager()
45 tpu::ExecutorApiFn()->TpuTransferManager_FreeFn(manager_); in ~TpuTransferManager()
59 tpu::ExecutorApiFn()->TpuTransferManager_HostShapeToDeviceShapeFn( in HostShapeToDeviceShape()
79 tpu::ExecutorApiFn()->TpuTransferManager_TransferLiteralToDeviceAsyncFn( in TransferLiteralToDeviceAsync()
97 tpu::ExecutorApiFn()->TpuTransferManager_TransferLiteralToInfeedFn( in TransferLiteralToInfeed()
107 const std::deque<tensorflow::tpu::NoncopyableBuffer>& buffers) { in TransferBuffersToInfeed()
123 tpu::ExecutorApiFn()->TpuTransferManager_TransferBuffersToInfeedFn( in TransferBuffersToInfeed()
140 tpu::ExecutorApiFn()->TpuTransferManager_TransferLiteralFromOutfeedFn( in TransferLiteralFromOutfeed()
160 tpu::ExecutorApiFn()->TpuTransferManager_ResetDevicesFn( in ResetDevices()
[all …]
|
| D | tpu_stream.h | 27 namespace tpu {
29 class TpuStream : public tensorflow::tpu::TpuStreamInterface {
35 tensorflow::tpu::ExecutorApiFn()->TpuStream_FreeFn(stream_); in ~TpuStream()
39 tensorflow::tpu::TpuStreamInterface* other) override { in IsSameSharedMemoryLocation()
40 return tensorflow::tpu::ExecutorApiFn() in IsSameSharedMemoryLocation()
49 tensorflow::tpu::ExecutorApiFn()->TpuStream_EnqueueTransferHostToDeviceFn( in EnqueueTransferHostToDevice()
59 tensorflow::tpu::ExecutorApiFn()->TpuStream_EnqueueTransferDeviceToHostFn( in EnqueueTransferDeviceToHost()
69 tensorflow::tpu::ExecutorApiFn() in EnqueueOnTpuDeviceSendRecvLocal()
|
| D | status_helper.h | 26 : c_status(tensorflow::tpu::ExecutorApiFn()->TpuStatus_NewFn()) {} in StatusHelper()
29 tensorflow::tpu::ExecutorApiFn()->TpuStatus_FreeFn(c_status); in ~StatusHelper()
33 if (tensorflow::tpu::ExecutorApiFn()->TpuStatus_OkFn(c_status)) { in FromC()
38 tensorflow::tpu::ExecutorApiFn()->TpuStatus_CodeFn(c_status)), in FromC()
39 tensorflow::tpu::ExecutorApiFn()->TpuStatus_MessageFn(c_status)); in FromC()
44 return tensorflow::tpu::ExecutorApiFn()->TpuStatus_OkFn(c_status); in ok()
|
| D | tpu_node_context.cc | 23 namespace tpu { namespace
33 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()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_configuration_ops.cc | 43 tpu::TpuMeshStateInterface** state) { in GetTpuMeshStateInterface()
45 tpu::kTpuMeshStateInterfaceResourceName, state) in GetTpuMeshStateInterface()
76 ResourceMgr* rmgr, tpu::TpuCompilationCacheInterface** compilation_cache) { in CreateTpuCompilationCache()
77 return rmgr->LookupOrCreate<tpu::TpuCompilationCacheInterface>( in CreateTpuCompilationCache()
78 rmgr->default_container(), tpu::kCompilationCacheResourceName, in CreateTpuCompilationCache()
79 compilation_cache, [&](tpu::TpuCompilationCacheInterface** new_cache) { in CreateTpuCompilationCache()
80 *new_cache = tpu::GetCompilationCacheCreateFn()(); in CreateTpuCompilationCache()
121 tpu::TpuCompilationCacheInterface* compilation_cache; in Compute()
130 OP_REQUIRES_OK(ctx, DeleteIfExists<tpu::TpuMeshStateInterface>( in Compute()
131 rmgr, tpu::kTpuMeshStateInterfaceResourceName)); in Compute()
[all …]
|
| D | BUILD | 17 "//tensorflow/core/tpu:__subpackages__",
18 "//tensorflow/stream_executor/tpu:__subpackages__",
48 ["//tensorflow/core/tpu/kernels:tpu_compilation_metrics"],
60 "//tensorflow/core/tpu:tpu_compile_interface",
61 "//tensorflow/core/tpu:tpu_ops_c_api_hdrs",
79 "//tensorflow/core/protobuf/tpu:compilation_result_proto_cc",
80 "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
81 "//tensorflow/core/protobuf/tpu:dynamic_padding_proto_cc",
82 "//tensorflow/core/tpu:tpu_api",
83 "//tensorflow/core/tpu:tpu_configuration",
[all …]
|
| D | tpu_compile_op_common.h | 40 namespace tpu {
69 const tpu::TPUCompileMetadataProto metadata, in TpuCompileOpKernelCommon()
83 const NameAttrList& function, const tpu::TPUCompileMetadataProto metadata, in TpuCompileOpKernelCommon()
110 const tpu::TPUCompileMetadataProto& metadata,
124 const tpu::TpuCompilationCacheKey& key,
136 const tpu::TpuCompilationCacheKey& key, in LookupPersistentCompilationCacheAndFillCaches()
158 absl::Span<const tpu::ShardingAndIndex> arg_core_mapping,
159 absl::Span<const tpu::ShardingAndIndex> retval_core_mapping,
164 static Status OptimizeGraph(const tpu::TPUCompileMetadataProto& metadata,
180 std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
[all …]
|
| D | tpu_compilation_cache_rpc_lookup.h | 33 namespace tpu {
38 using StubType = tpu::grpc::TpuCompilationCacheService::Stub;
45 std::unique_ptr<tpu::CompilationCacheEntryRef>* entry,
46 tpu::CompilationCacheFetchTarget fetch_target) override;
49 std::unique_ptr<tpu::CompilationCacheEntryRef>* entry,
50 tpu::CompilationCacheFetchTarget fetch_target) override;
57 const tpu::GetTpuProgramRequest& request,
66 std::unique_ptr<tpu::CompilationCacheEntryRef>* entry,
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | BUILD | 14 "//tensorflow/core/profiler/internal/tpu:__subpackages__",
15 "//tensorflow/core/tpu:__subpackages__",
16 "//tensorflow/stream_executor/tpu:__subpackages__",
41 "//tensorflow/core/protobuf/tpu:optimization_parameters_proto_cc",
54 "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_cc",
55 "//tensorflow/core/protobuf/tpu:tpu_embedding_output_layout_proto_cc",
125 "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
135 "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
156 "//tensorflow/stream_executor/tpu:tpu_executor",
157 "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
[all …]
|
| D | tpu_on_demand_compiler.cc | 47 static_cast<tensorflow::tpu::TpuStream*>( in ToC()
61 stream_executor::tpu::SerializeProto(dev_assign_proto); in ToC()
77 static_cast<tensorflow::tpu::TpuStream*>(impl)->se_stream(); in ToC()
86 using ::tensorflow::tpu::ExecutorApiFn;
139 stream_executor::tpu::SerializedProto_Free( in ExecuteAsyncOnStream()
215 return tensorflow::tpu::GetTpuPlatformId(); in PlatformId()
224 stream_executor::tpu::SerializedProto_Free(hlo_module.proto); in RunHloPasses()
228 hlo_module.proto = stream_executor::tpu::SerializeProto(module->ToProto()); in RunHloPasses()
234 static_cast<tensorflow::tpu::TpuExecutor*>(executor->implementation()) in RunHloPasses()
241 stream_executor::tpu::DeserializeProto<HloModuleProto>(result.proto); in RunHloPasses()
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_estimator.py | 22 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import *
24 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _clone_export_output_with_tenso…
25 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _create_global_step
26 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _export_output_to_tensors
27 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _get_scaffold
28 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _Inputs
29 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _ITERATIONS_PER_LOOP_VAR
30 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _TPU_ENQUEUE_OPS
31 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _TPU_ESTIMATOR
32 from tensorflow_estimator.python.estimator.tpu.tpu_estimator import _TPU_TRAIN_OP
|
| D | api.py | 26 from tensorflow.python.tpu import bfloat16
27 from tensorflow.python.tpu import feature_column_v2
28 from tensorflow.python.tpu import tpu
29 from tensorflow.python.tpu import tpu_embedding
30 from tensorflow.python.tpu import tpu_embedding_v2
31 from tensorflow.python.tpu import tpu_embedding_v2_utils
32 from tensorflow.python.tpu import tpu_optimizer
|
| D | tpu_test.py | 36 from tensorflow.python.tpu import tpu
37 from tensorflow.python.tpu import tpu_feed
38 from tensorflow.python.tpu import training_loop
48 context = tpu.TPUReplicateContext(b"context", 1, pivot=pivot)
64 context = tpu.TPUReplicateContext(b"context", 1, pivot=pivot)
112 tpu.rewrite(loop)
140 node.attr[tpu._TPU_REPLICATE_ATTR].s = b"0"
151 tpu.prune_unconnected_ops_from_xla(ops.get_default_graph())
155 tpu._TPU_REPLICATE_ATTR)
158 tpu._TPU_REPLICATE_ATTR)
[all …]
|
| D | tpu_strategy_util.py | 30 from tensorflow.python.tpu import topology
31 from tensorflow.python.tpu import tpu
92 return tpu.initialize_system(
99 with ops.device(tpu._tpu_system_device_name(job)): # pylint: disable=protected-access
125 serialized_topology = sess.run(tpu.initialize_system())
127 with ops.device(tpu._tpu_system_device_name(job)): # pylint: disable=protected-access
128 serialized_topology = tpu.initialize_system(
190 tpu.shutdown_system(job=job)
195 with ops.device(tpu._tpu_system_device_name(job)): # pylint: disable=protected-access
212 sess.run(tpu.shutdown_system())
|
| D | tpu_outside_compilation_test.py | 52 from tensorflow.python.tpu import functional as tpu_functional
53 from tensorflow.python.tpu import tpu
54 from tensorflow.python.tpu import tpu_strategy_util
55 from tensorflow.python.tpu.ops import tpu_ops
65 tpu=FLAGS.tpu,
105 return tpu.rewrite(concrete.__call__, list(args) + list(kwargs.values()))
145 tpu.outside_compilation(outside_fn)
165 tpu.outside_compilation(outside_fn, x2)
186 output = tpu.outside_compilation(outside_fn, x2)
213 c, d, e = tpu.outside_compilation(outside_fn, a, b)
[all …]
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_compile_ops.cc | 77 tensorflow::tpu::TpuProgramGroupInterface* tpu_program_group);
85 tensorflow::tpu::TpuProgramGroupInterface* tpu_program_group) { in Compile()
86 return tensorflow::tpu::TpuProgramGroup::CompileAndBuild( in Compile()
90 tpu::TpuCompilationCacheKey CompilationCacheKey( in CompilationCacheKey()
92 tensorflow::tpu::TpuMeshStateInterface* mesh_state, int num_replicas, in CompilationCacheKey()
96 tpu::TPUCompileMetadataProto metadata; in CompilationCacheKey()
99 const tpu::TpuCompilationCacheKey key = CreateCompilationCacheKey( in CompilationCacheKey()
129 if (tpu::OpsApiFn() in Compute()
158 tensorflow::tpu::TpuMeshStateInterface* mesh_state; in Compute()
161 tensorflow::tpu::kTpuMeshStateInterfaceResourceName, in Compute()
[all …]
|
| D | BUILD | 82 "//tensorflow/core/tpu:tpu_api",
83 "//tensorflow/core/tpu:tpu_configuration",
84 "//tensorflow/core/tpu:tpu_defs",
85 "//tensorflow/core/tpu:tpu_execute",
86 "//tensorflow/core/tpu/kernels:tpu_compilation_cache_entry",
87 "//tensorflow/core/tpu/kernels:tpu_compilation_cache_interface",
88 "//tensorflow/core/tpu/kernels:tpu_compilation_cache_key",
89 "//tensorflow/core/tpu/kernels:tpu_compilation_cache_lookup",
90 "//tensorflow/core/tpu/kernels:tpu_compile_op_common",
91 "//tensorflow/core/tpu/kernels:tpu_compile_op_hdrs",
[all …]
|
| D | tpu_execute_op.cc | 54 using tensorflow::tpu::CompilationCacheEntryRef;
55 using tensorflow::tpu::TpuCompilationCacheEntry;
56 using tensorflow::tpu::TpuCompilationCacheLookup;
67 tpu::kCompiledProtoCacheResourceName, in GetComputationCacheEntry()
125 const tpu::TpuProgramGroup* tpu_program_group, xla::Backend* const backend, in GetExecutableAliasConfig()
138 tpu::TpuNodeContext* node_context, se::Stream* stream, in AllocateOutputTuple()
172 tpu::TpuNodeContext* node_context, se::Stream* stream, in AllocateOutputTensors()
188 OpKernelContext* context, tpu::TpuNodeContext* node_context, in RunExecutable()
191 const uint32 rng_seed, const tpu::TpuProgramGroup* tpu_program_group, in RunExecutable()
230 OpKernelContext* context, tpu::TpuNodeContext* node_context, in ExecuteTPUProgram()
[all …]
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | BUILD | 8 "//tensorflow/core/tpu:__subpackages__",
9 "//tensorflow/stream_executor/tpu:__subpackages__",
43 "//tensorflow/core/tpu:tpu_init_mode",
44 "//tensorflow/core/tpu/kernels:tpu_compile_op_options",
57 "//tensorflow/core/tpu:tpu_defs",
99 "//tensorflow/core/tpu:tpu_compile_interface",
100 "//tensorflow/core/tpu:tpu_defs",
157 "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
158 "//tensorflow/core/protobuf/tpu:dynamic_padding_proto_cc",
159 "//tensorflow/core/protobuf/tpu:topology_proto_cc",
[all …]
|
| /external/tensorflow/tensorflow/python/tpu/client/ |
| D | client_test.py | 30 from tensorflow.python.tpu.client import client
161 client.Client(tpu=['multiple', 'elements'])
185 tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
218 tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
224 c = client.Client(tpu='grpc://1.2.3.4:8470')
239 tpu='tpu_name', project='project', zone='zone')
249 c = client.Client(tpu='grpc://1.2.3.4:8470')
262 tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
276 tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
290 tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
[all …]
|
| /external/tensorflow/tensorflow/python/distribute/cluster_resolver/tpu/ |
| D | tpu_cluster_resolver_test.py | 28 from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver as resolver
44 from tensorflow.python.tpu.client import client # pylint: disable=g-import-not-at-top
145 resolver.TPUClusterResolver(tpu='')
167 tpu=['test-tpu-1'],
201 tpu=['test-tpu-1'],
227 tpu='test-tpu-1',
248 tpu=['test-tpu-1'],
274 tpu='nonexistent-tpu',
300 tpu='test-tpu-1',
344 tpu='test-tpu-1',
[all …]
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.tpu.-x-l-a-options.pbtxt | 1 path: "tensorflow.tpu.XLAOptions"
3 is_instance: "<class \'tensorflow.python.tpu.tpu.XLAOptions\'>"
4 is_instance: "<class \'tensorflow.python.tpu.tpu.XLAOptions\'>"
|