Class Stack

tensorflow/core/kernels/stack.cc:41–125

namespace tensorflow {

class Stack : public ResourceBase {
 public:
  static std::atomic<int64> stack_counter;

  struct TensorAndAllocation {
    Tensor tensor;
    AllocatorAttributes alloc_attrs;
    bool swapped_to_cpu;
  };

  Stack(const DataType& elem_type, const string& stack_name, int max_size)
      : elem_type_(elem_type),
        stack_name_(stack_name),
        max_size_(max_size),
        closed_(false) {}

  Status Push(const TensorAndAllocation& value) {
    mutex_lock l(mu_);
    TF_RETURN_IF_ERROR(CheckNotClosed());
    if (max_size_ >= 0 && stack_.size() >= max_size_) {
      return errors::InvalidArgument("Stack[", stack_name_, "] overflowed ",
                                     "its max_size (", max_size_, ")");
    }
    stack_.push_back(value);
    return Status::OK();
  }

  Status Pop(TensorAndAllocation* value) {
    mutex_lock l(mu_);
    TF_RETURN_IF_ERROR(CheckNotClosed());
    if (stack_.empty()) {
      return errors::InvalidArgument("Stack[", stack_name_,
                                     "] is empty when calling Pop().");
    }
    *value = stack_.back();
    stack_.pop_back();
    return Status::OK();
  }

  // We don't swap the first tensor on the stack and any subsequent tensors
  // that share the buffer with the first tensor.
  bool IsUsefulToSwap(const Tensor& tensor) const {
    mutex_lock l(mu_);
    if (stack_.empty()) {
      return false;
    }
    const Tensor& first = stack_.front().tensor;
    return !tensor.SharesBufferWith(first);
  }

  void Close() {
    mutex_lock l(mu_);
    stack_.clear();
    closed_ = true;
  }

  DataType ElemType() { return elem_type_; }