Home | History | Annotate | Download | only in cpu 
      1 /* Copyright 2017 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 // This header declares the abstract class for the infeed manager that
     17 // is used by the CPU runtime to transfer buffers into an executing
     18 // CPU computation, e.g., to feed data into a while loop.
     19 
     20 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_XFEED_MANAGER_H_
     21 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_XFEED_MANAGER_H_
     22 
     23 #include <deque>
     24 
     25 #include "tensorflow/compiler/xla/statusor.h"
     26 #include "tensorflow/compiler/xla/types.h"
     27 #include "tensorflow/compiler/xla/xla_data.pb.h"
     28 #include "tensorflow/core/lib/gtl/array_slice.h"
     29 #include "tensorflow/core/platform/mutex.h"
     30 
     31 namespace xla {
     32 namespace cpu {
     33 namespace runtime {
     34 
     35 // Abstract class defining an infeed buffer that is passed to the
     36 // runtime by a client. The client manages the storage of the buffer.
     37 class XfeedBuffer {
     38  public:
     39   virtual ~XfeedBuffer() = default;
     40 
     41   virtual int32 length() = 0;
     42   virtual void* data() = 0;
     43 
     44   // The 'shape' parameter reflects what shape the embedded program was
     45   // expecting / producing with respect to this XfeedBuffer. E.g. this will
     46   // contain information about the layout of an outfed buffer.
     47   virtual void Done(StatusOr<Shape> shape) = 0;
     48 };
     49 
     50 // Reusable component for managing the infeed and outfeed queue state.
     51 class XfeedQueueManager {
     52  public:
     53   XfeedQueueManager(string queue_name) : queue_name_(queue_name) {}
     54 
     55   // Calls the completion callback for any enqueued buffers that have
     56   // not been dequeued by the runtime, and empties the
     57   // queue. Reset may not be called while a runtime computation is
     58   // processing a dequeued buffer. The only safe way to ensure this
     59   // condition is to call Reset when no computation is taking place.
     60   void Reset();
     61 
     62   // Adds a sequence of buffers to the queue atomically. buffer->Done will be
     63   // called when the buffer will no longer be accessed by the XfeedManager,
     64   // either as a result of a call to Reset or because the runtime has dequeued
     65   // and used the buffer.
     66   void EnqueueBuffersAtomically(
     67       tensorflow::gtl::ArraySlice<XfeedBuffer*> buffers);
     68 
     69   // Blocks until the queue is non-empty, then returns the buffer at the head of
     70   // the queue. Sets the current buffer to be the returned buffer. It is an
     71   // error to call BlockingDequeueBuffer if there is an unreleased current
     72   // buffer, i.e., ReleaseCurrentBuffer must be called between calls to
     73   // BlockingDequeueBuffer.
     74   XfeedBuffer* BlockingDequeueBuffer();
     75 
     76   // Releases the current buffer, which is the last buffer returned by
     77   // BlockingDequeuBuffer and not yet released. length and data must
     78   // match the buffer->length() and buffer->data() for the current
     79   // buffer.
     80   //
     81   // 'shape' communicates the shape of the buffer being released. If the program
     82   // passed a value that could not be decoded as a shape, 'shape' will be an
     83   // error status. In the case of outfeed, this indicates the layout of the
     84   // shape that has been outfed. In the case of infeed, this can be used for
     85   // sanity checking purposes.
     86   void ReleaseCurrentBuffer(int32 length, void* data, StatusOr<Shape> shape);
     87 
     88  private:
     89   const string queue_name_;
     90 
     91   tensorflow::mutex mu_;
     92 
     93   // Condition variable that is signaled every time a buffer is
     94   // enqueued to an empty queue.
     95   tensorflow::condition_variable cv_;
     96 
     97   // XfeedBuffer* queue contents are not owned, but buffer->Done must
     98   // be called when the buffer is no longer needed by the runtime.
     99   std::deque<XfeedBuffer*> enqueued_buffers_;
    100 
    101   // If non-NULL, the buffer that is currently being processed by the
    102   // runtime. Not owned.
    103   XfeedBuffer* current_buffer_ = nullptr;
    104 };
    105 
    106 // Client-side class used to enqueue infeed buffers.
    107 class XfeedManager {
    108  public:
    109   XfeedManager() = default;
    110 
    111   void Reset();
    112 
    113   XfeedQueueManager* infeed() { return &infeed_; }
    114   XfeedQueueManager* outfeed() { return &outfeed_; }
    115 
    116  private:
    117   XfeedQueueManager infeed_ = {"infeed"};
    118   XfeedQueueManager outfeed_ = {"outfeed"};
    119 };
    120 
    121 }  // namespace runtime
    122 }  // namespace cpu
    123 }  // namespace xla
    124 
    125 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_XFEED_MANAGER_H_
    126