Class BufferValue

tensorflow/compiler/xla/service/buffer_value.h:87–174

//   BufferValue(%tuple_constant, {1, 1})  // Holds value "43"

class BufferValue {
 public:
  TF_LIB_GTL_DEFINE_INT_TYPE(Color, int64);

  // Id is a unique identifier for the BufferValue to facilitate efficient
  // collections of BufferValues with stable iteration order.
  using Id = int64;

  // Functions which return the size and alignment of a logical buffer in bytes.
  using SizeFunction = std::function<int64(const BufferValue&)>;
  using AlignmentFunction = std::function<int64(BufferValue::Color)>;

  virtual ~BufferValue();

  Id id() const { return id_; }

  // Return the instruction that defines the buffer.
  virtual HloInstruction* instruction() const = 0;

  // Return the index within the output of the instruction where the buffer is
  // defined. Index used defined as in ShapeUtil::GetSubshape()
  virtual const ShapeIndex& index() const = 0;

  // Return the color of the BufferValue. Differently colored buffers can not be
  // parts of the same allocation.
  Color color() const {
    CHECK_NE(color_, kInvalidColor)
        << "Should not query the color of a buffer that was never colored";
    return color_;
  }

  void set_color(Color color) {
    CHECK_NE(color, kInvalidColor)
        << "Should not set the color of a buffer to the invalid color";
    color_ = color;
  }

  bool has_color() const { return color_ != kInvalidColor; }

  // Return the shape of the buffer. This reference points into the shape field
  // of the instruction defining the buffer.  Therefore, the returned shape will
  // contain the layout of instruction, if any.
  virtual const Shape& shape() const = 0;

  // Returns true if this buffer is the top-level output buffer of the defining
  // HLO instruction. This is equivalent to index == {}.
  bool IsTopLevel() const { return index().empty(); }

  // Whether this buffer contains a tuple.
  bool IsTuple() const { return is_tuple_; }

  // Whether this buffer contains an array.
  bool IsArray() const { return is_array_; }

  // operator< is required for std::set.
  bool operator<(const BufferValue& other) const { return id_ < other.id_; }

  bool operator==(const BufferValue& other) const { return id_ == other.id_; }