• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef TENSORFLOW_COMPILER_JIT_XLA_TENSOR_H_
17 #define TENSORFLOW_COMPILER_JIT_XLA_TENSOR_H_
18 
19 #include <memory>
20 
21 #include "absl/memory/memory.h"
22 #include "tensorflow/compiler/xla/client/local_client.h"
23 #include "tensorflow/compiler/xla/service/shaped_buffer.h"
24 #include "tensorflow/core/framework/allocator.h"
25 #include "tensorflow/core/framework/device_base.h"
26 #include "tensorflow/core/lib/core/status.h"
27 #include "tensorflow/core/platform/mutex.h"
28 
29 namespace tensorflow {
30 
31 // The implementation of a Tensor for an XlaDevice. All device tensors are
32 // actually one of these.
33 //
34 // To distinguish between "normal" device tensors and XlaTensors, the raw
35 // pointer data stored in the TensorBuffer is a tagged pointer.
36 class XlaTensor {
37  public:
38   // Downcast from a Tensor to an XlaTensor. Return nullptr if the downcast
39   // fails.
40   static XlaTensor* FromTensor(const Tensor* tensor);
41 
42   // Create a DeviceMemoryBase from a Tensor. The Tensor can be an XlaTensor, in
43   // which case the returned value is shaped_buffer()->root_buffer(), or a
44   // normal Tensor in which case the returned value is
45   // {tensor.tensor_data().data(), tensor.tensor_data().size}.
46   static se::DeviceMemoryBase DeviceMemoryFromTensor(const Tensor& tensor);
47 
48   // Assign the internal ShapedBuffer to new memory for the given dtype and
49   // shape. If a ShapedBuffer exists already (has_shaped_buffer() == true), it
50   // is replaced and the managed memory deallocated.
51   Status AllocateShapedBuffer(DataType dtype, const xla::Shape& on_host_shape,
52                               xla::LocalClient* client, int device_ordinal);
53 
54   // Some Tensors can have complex on-device shapes, including tuple shapes. To
55   // manage the memory for these tensors a ShapedBuffer may be required.
56 
57   // Return true if this XlaTensor contains a ShapedBuffer.
has_shaped_buffer()58   bool has_shaped_buffer() const { return shaped_buffer_.has_value(); }
59   // Return the contained ShapedBuffer.
60   // REQUIRES: has_shaped_buffer()
shaped_buffer()61   const xla::ShapedBuffer& shaped_buffer() const {
62     CHECK(has_shaped_buffer());
63     return *shaped_buffer_;
64   }
shaped_buffer()65   xla::ShapedBuffer& shaped_buffer() {
66     CHECK(has_shaped_buffer());
67     return *shaped_buffer_;
68   }
69   // Mutates the XlaTensor to set the ShapedBuffer.
set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer)70   void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) {
71     shaped_buffer_ = std::move(shaped_buffer);
72   }
73 
74   // Adds synchronization events to 'stream' that wait for this tensor to be
75   // defined on 'stream'. Does nothing if the tensor is already defined on that
76   // stream.
77   void WaitForDefinitionEventOnStream(se::Stream* stream);
78 
79   // (Re)sets the definition event of the tensor to 'event', and promises that
80   // the tensor has already been defined on stream. Removes any previous
81   // definition event or any previous promises about the tensor being defined on
82   // streams.
83   // It is legal to reset the definition event of a tensor when overwriting the
84   // tensor's value (at which point, it is effectively a new tensor once again.)
85   void ResetDefinitionEvent(std::shared_ptr<se::Event> event,
86                             se::Stream* stream);
87 
88   // Refresh the status of streams_defined_on_. Return the first not-OK stream's
89   // status or OK.
90   Status RefreshStatusOfStreams();
91 
92   // Convert from a raw pointer to an XlaTensor, removing the pointer tag.
93   static XlaTensor* FromOpaquePointer(void* ptr);
94   // Convert to a raw pointer from an XlaTensor, adding the pointer tag.
95   static void* ToOpaquePointer(XlaTensor* tensor);
96 
97  private:
98   // The optional contained ShapedBuffer.
99   absl::optional<xla::ScopedShapedBuffer> shaped_buffer_;
100   // An optional host tensor value.
101   absl::optional<Tensor> host_tensor_;
102   // An optional event that is triggered when the tensor's content has been
103   // defined. If this event is nullptr, it is assumed that the tensor's content
104   // is always defined.
105   std::shared_ptr<se::Event> definition_event_;
106   // A list of all streams for which the tensor's content is defined for any
107   // newly enqueued command.
108   absl::InlinedVector<se::Stream*, 2> streams_defined_on_ TF_GUARDED_BY(mu_);
109   mutex mu_;
110 };
111 
112 }  // namespace tensorflow
113 
114 #endif  // TENSORFLOW_COMPILER_JIT_XLA_TENSOR_H_
115