1c8b59c046895fa5b6d79f73e0b5817330fcfbfc1A. Unique TensorFlower/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 29c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 39c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurLicensed under the Apache License, Version 2.0 (the "License"); 49c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudluryou may not use this file except in compliance with the License. 59c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurYou may obtain a copy of the License at 69c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 79c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur http://www.apache.org/licenses/LICENSE-2.0 89c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 99c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurUnless required by applicable law or agreed to in writing, software 109c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlurdistributed under the License is distributed on an "AS IS" BASIS, 119c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 129c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurSee the License for the specific language governing permissions and 139c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlurlimitations under the License. 149c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur==============================================================================*/ 159c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 16f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#ifndef TENSORFLOW_FRAMEWORK_DEVICE_BASE_H_ 17f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#define TENSORFLOW_FRAMEWORK_DEVICE_BASE_H_ 18f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 19f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include <memory> 208e49dd9ea9f69cadaf9f40b07cb0be5cfabdf653Gunhan Gulsoy#include <string> 21f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include <unordered_map> 22f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 2364aa0989c110433e275d23bbe6f3605a26b2efb0Josh Levenberg#include "tensorflow/core/framework/tensor.h" 24f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/lib/core/errors.h" 25f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/lib/core/refcount.h" 2664aa0989c110433e275d23bbe6f3605a26b2efb0Josh Levenberg#include "tensorflow/core/lib/core/status.h" 27a120b0bec12d47d4769a604b1cfafa3988a19a0cA. Unique TensorFlower#include "tensorflow/core/lib/core/stringpiece.h" 2856313def004795f75ef8281a0294c958d28f1e06Vijay Vasudevan#include "tensorflow/core/platform/logging.h" 29f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 30f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace Eigen { 3127259353e50e6bcaeeedbc26dc3aaaa5695fe500Vijay Vasudevanstruct ThreadPoolDevice; 32a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#ifdef TENSORFLOW_USE_SYCL 33a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steinerstruct SyclDevice; 34a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#endif 35f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} // end namespace Eigen 36f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 37f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace perftools { 38f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace gputools { 39f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass Stream; 40f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} // namespace gputools 41f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} // namespace perftools 42f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 43f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace tensorflow { 44f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 45f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass Device; 46e85d3df92deb9d717befdf173966a2913ac2aea0Geoffrey Irvingclass DeviceAttributes; 47f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass Env; 48f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass EventMgr; 49ec1403e7dc2b919531e527d36d28659f60621c9eA. Unique TensorFlowerclass OpKernelContext; 504e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlowerclass ResourceMgr; 51e85d3df92deb9d717befdf173966a2913ac2aea0Geoffrey Irvingclass TensorProto; 52f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 53f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace thread { 54f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass ThreadPool; 55f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 56f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 57ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower// A wrapper for an Eigen Gpu Device that includes per-op state. The 58ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower// class is defined even for non-GPU devices since the 59ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower// OpKernelContext::Params structure wants to fill it in. 60f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass PerOpGpuDevice { 61f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur public: 62f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual ~PerOpGpuDevice() {} 63f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual const Eigen::GpuDevice& device() const = 0; 64f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur}; 65f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 66f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur// A class that devices can subclass to pass around 67f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur// Device-specific context to OpKernels. 68f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass DeviceContext : public core::RefCounted { 69f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur public: 70f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ~DeviceContext() override {} 71f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual perftools::gputools::Stream* stream() const { return nullptr; } 72f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual void MaintainLifetimeOnStream( 73f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur const Tensor* t, perftools::gputools::Stream* stream) const {} 74f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 75f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // "cpu_tensor" is a tensor on a CPU. Copies "cpu_tensor" into 76f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // "device_tensor" which is on a GPU device "device". "device_tensor" 77f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // must be allocated to be of the same size as "cpu_tensor". 78f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device, 79f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Tensor* device_tensor, 80f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur StatusCallback done) const { 81f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur done(errors::Internal("Unrecognized device type in CPU-to-device Copy")); 82f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 83f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 84f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // "device_tensor" is a tensor on a non-CPU device. Copies 85f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // device_tensor into "cpu_tensor". "cpu_tensor" must be allocated 86f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // to be of the same size as "device_tensor". 87f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual void CopyDeviceTensorToCPU(const Tensor* device_tensor, 88a120b0bec12d47d4769a604b1cfafa3988a19a0cA. Unique TensorFlower StringPiece tensor_name, Device* device, 89f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Tensor* cpu_tensor, StatusCallback done) { 90f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur done(errors::Internal("Unrecognized device type in device-to-CPU Copy")); 91f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 92f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur}; 93f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 94b8ccc5a20d0e6fc04ef9a854bb391c69f99ca907A. Unique TensorFlower// map[i] is the DeviceContext* for the node with id i, if i < map.size(). 95b8ccc5a20d0e6fc04ef9a854bb391c69f99ca907A. Unique TensorFlowertypedef std::vector<DeviceContext*> DeviceContextMap; 96f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 97f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass DeviceBase { 98f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur public: 99f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur explicit DeviceBase(Env* env) : env_(env) {} 100f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual ~DeviceBase(); 101f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 102f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Env* env() const { return env_; } 103f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 104f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Override this to return true for devices that require an Op's 105f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // compute method to save references to the temporary tensors it 106f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // allocates until the Op execution completes 1076f62e435ab6c36dfdfdef1acd580b5f278f6723cA. Unique TensorFlower virtual bool RequiresRecordingAccessedTensors() const { return false; } 108f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 109f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur struct CpuWorkerThreads { 110f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur int num_threads = 0; 111f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur thread::ThreadPool* workers = nullptr; 112f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur }; 113f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 114f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Does not take ownership. 115f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void set_tensorflow_cpu_worker_threads(CpuWorkerThreads* t) { 116f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur cpu_worker_threads_ = t; 117f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 118f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 119f28935a7d280b6ba75fe93fe35783d87b9cc2ec9Brennan Saeta virtual const CpuWorkerThreads* tensorflow_cpu_worker_threads() const { 120f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur CHECK(cpu_worker_threads_ != nullptr); 121f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return cpu_worker_threads_; 122f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 123f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 124f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // "stream" is used in special circumstances (such as the 125f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // constructors of Ops) where there is no available OpKernelContext. 126f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // "default_context" is used by OpKernelContext whenever a device does not 127f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // supply a DeviceContext for an op in FillContextMap (e.g. when only 128f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // using a single stream.) 129f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // "event_mgr" is used to delay deallocation of temporary GPU buffers. 130f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // TODO(pbar) Work out how to move this out of DeviceBase. 131f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur struct GpuDeviceInfo { 132795f35da2d458cbae477ac2fe2bff80c1427a771Vijay Vasudevan // Make sure all the defaults are NULL, so we can spot missing assignments. 133795f35da2d458cbae477ac2fe2bff80c1427a771Vijay Vasudevan perftools::gputools::Stream* stream = nullptr; 134795f35da2d458cbae477ac2fe2bff80c1427a771Vijay Vasudevan DeviceContext* default_context = nullptr; 135795f35da2d458cbae477ac2fe2bff80c1427a771Vijay Vasudevan EventMgr* event_mgr = nullptr; 136f3405c2d73196e409041d52bbf30748b2a64493bA. Unique TensorFlower int gpu_id = -1; 137f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur }; 138f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 139f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Does not take ownership. 140f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void set_tensorflow_gpu_device_info(GpuDeviceInfo* g) { 141f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur gpu_device_info_ = g; 142f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 143f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 144f28935a7d280b6ba75fe93fe35783d87b9cc2ec9Brennan Saeta virtual const GpuDeviceInfo* tensorflow_gpu_device_info() const { 145f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return gpu_device_info_; 146f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 147f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 148c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng // The preferred thread pool for this device. If it is nullptr, the system 149c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng // automatically assigns a thread pool for execution. 150c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng virtual thread::ThreadPool* tensorflow_device_thread_pool() { 151c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng return device_thread_pool_; 152c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng } 153c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng 154f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Does not take ownership. 155f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void set_eigen_cpu_device(Eigen::ThreadPoolDevice* d) { 156f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur eigen_cpu_device_ = d; 157f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 158f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 159a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#ifdef TENSORFLOW_USE_SYCL 160a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner void set_eigen_sycl_device(Eigen::SyclDevice* d) { eigen_sycl_device_ = d; } 161a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#endif 162a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner 163f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Return the Allocator implementation to use based on the allocator 164f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // attributes requested. See allocator.h for more details. 165f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual Allocator* GetAllocator(AllocatorAttributes /*attr*/) { 166f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur LOG(FATAL) << "GetAllocator() is not implemented."; 167e2d51a87f0727f8537b46048d8241aeebb6e48d6Xiaoqiang Zheng return nullptr; 168f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 169f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 1704e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower // Return the Allocator implementation to use based on the allocator 1714e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower // attributes requested and the supplied resource manager. By 1724e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower // default this ignores the resource manager and calls the base 1734e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower // implementation but devices can override if they want to consult 1744e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower // the resource manager when choosing the allocator. 1754e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower virtual Allocator* GetStepAllocator(AllocatorAttributes attr, 1764e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower ResourceMgr* /*step_resource_manager*/) { 1774e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower return GetAllocator(attr); 1784e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower } 1794e68ec99872f21e4a401d868b2656ca26b56981aA. Unique TensorFlower 180f28935a7d280b6ba75fe93fe35783d87b9cc2ec9Brennan Saeta virtual const Eigen::ThreadPoolDevice* eigen_cpu_device() { 181f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur CHECK(eigen_cpu_device_ != nullptr); 182f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return eigen_cpu_device_; 183f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 184f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 185a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#ifdef TENSORFLOW_USE_SYCL 186f28935a7d280b6ba75fe93fe35783d87b9cc2ec9Brennan Saeta virtual const Eigen::SyclDevice* eigen_sycl_device() const { 187a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner CHECK(eigen_sycl_device_ != nullptr); 188a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner return eigen_sycl_device_; 189a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner } 190a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#endif 191a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner 192ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower // Caller owns the return value. The OpKernelContext calls this even 193ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower // for devices that do not implement an eigen_gpu_device. Overridden 194ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower // by GPU devices to return a derived type. 195ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower virtual PerOpGpuDevice* MakeGpuDevice() { return nullptr; } 196ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower 19719fd294eae4e8e22f6ab46b21cf41323750a1c69Brennan Saeta virtual DeviceBase* UnderlyingDevice() { return this; } 19819fd294eae4e8e22f6ab46b21cf41323750a1c69Brennan Saeta virtual const DeviceBase* UnderlyingDevice() const { return this; } 19919fd294eae4e8e22f6ab46b21cf41323750a1c69Brennan Saeta 200ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower // This is overridden by GPU devices to reinitialize the derived 201ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower // type returned by MakeGpuDevice. 202ec1403e7dc2b919531e527d36d28659f60621c9eA. Unique TensorFlower virtual void ReinitializeGpuDevice(OpKernelContext* /*context*/, 203ec1403e7dc2b919531e527d36d28659f60621c9eA. Unique TensorFlower PerOpGpuDevice* /*device*/, 204ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower DeviceContext* /*dc*/, 205ff8522de343a90813fc4e5cbb249e308c1819f1dA. Unique TensorFlower Allocator* /*allocator*/) {} 206f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 207e85d3df92deb9d717befdf173966a2913ac2aea0Geoffrey Irving // Unimplemented by default 208e85d3df92deb9d717befdf173966a2913ac2aea0Geoffrey Irving virtual const DeviceAttributes& attributes() const; 2098e49dd9ea9f69cadaf9f40b07cb0be5cfabdf653Gunhan Gulsoy virtual const string& name() const; 210f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 211f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Materializes the given TensorProto into 'tensor' stored in Device 212f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // memory. Most devices will want to override this. 213f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // 214f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // TODO(vrv): We should be able to put this function into 215f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // OpKernelContext and handle the copies from device memory via send 216f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // and receive nodes, instead of requiring that each device handle 217f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // the copies here as well as in copy ops. 218f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur virtual Status MakeTensorFromProto(const TensorProto& tensor_proto, 219f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur const AllocatorAttributes alloc_attrs, 220f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Tensor* tensor) { 221f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return errors::Internal("Device does not implement MakeTensorFromProto()"); 222f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 223f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 224c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng protected: 225c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng // Does not take ownership. 226c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng void set_tensorflow_device_thread_pool(thread::ThreadPool* thread_pool) { 227c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng device_thread_pool_ = thread_pool; 228c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng } 229c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng 230f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur private: 231f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Env* const env_; 232f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur CpuWorkerThreads* cpu_worker_threads_ = nullptr; 233f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur GpuDeviceInfo* gpu_device_info_ = nullptr; 234c381794b2fc3227bfee9cf085e26bafb33da8f4bXiaoqiang Zheng thread::ThreadPool* device_thread_pool_ = nullptr; 235f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Eigen::ThreadPoolDevice* eigen_cpu_device_ = nullptr; 236a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#ifdef TENSORFLOW_USE_SYCL 237a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner Eigen::SyclDevice* eigen_sycl_device_ = nullptr; 238a771598ad83ca33eb42594d7e804859371ba4ca9Benoit Steiner#endif 239f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur}; 240f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 241f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} // namespace tensorflow 242f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 243f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#endif // TENSORFLOW_FRAMEWORK_DEVICE_BASE_H_ 244