Assign buffers for constants

(Still disabled on all backends by default)

PiperOrigin-RevId: 205890903

1 files changed, 26 insertions, 2 deletions

diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h
index 8844b6e3ba..4fcf1fc73d 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment.h
+++ b/tensorflow/compiler/xla/service/buffer_assignment.h
@@ -85,7 +85,9 @@ class BufferAllocation {
 
   // Whether this allocation is readonly i.e. backed by memory we cannot write
   // to.
-  bool is_readonly() const { return is_entry_computation_parameter(); }
+  bool is_readonly() const {
+    return is_entry_computation_parameter() || is_constant();
+  }
 
   bool is_tuple() const { return is_tuple_; }
   void set_is_tuple(bool is_tuple) { is_tuple_ = is_tuple; }
@@ -96,6 +98,13 @@ class BufferAllocation {
   bool is_entry_computation_parameter() const {
     return is_entry_computation_parameter_;
   }
+
+  // Whether this allocation holds a constant.  On the CPU and GPU backends
+  // constant allocations are not allocated dynamically, instead we resolve
+  // references to these buffer allocations to a global in the readonly section
+  // of the binary.
+  bool is_constant() const { return is_constant_; }
+
   // If this allocation holds a Buffer from a parameter of the entry
   // computation, this methods returns the parameter number. CHECKs otherwise.
   int64 parameter_number() const {
@@ -201,7 +210,9 @@ class BufferAllocation {
            // of the computation.
            !maybe_live_out() &&
            // Thread-local buffers are allocated using `alloca`s.
-           !is_thread_local();
+           !is_thread_local() &&
+           // Constant buffers are allocated as global values.
+           !is_constant();
   }
 
   // Add a heap trace which was used to assign slices to logical buffers in this
@@ -257,6 +268,8 @@ class BufferAllocation {
     parameter_number_ = parameter_number;
     param_shape_index_ = std::move(param_shape_index);
   }
+
+  void set_constant(bool is_constant) { is_constant_ = is_constant; }
   void set_maybe_live_out(bool value) { maybe_live_out_ = value; }
   void set_index(Index index) { index_ = index; }
   void set_size(int64 size) { size_ = size; }
@@ -295,6 +308,9 @@ class BufferAllocation {
   // might not actually escape.
   bool maybe_live_out_ = false;
 
+  // See comment on the is_constant() accessor.
+  bool is_constant_ = false;
+
   // Mapping from the set of buffers assigned to this allocation to their
   // logical offsets and sizes.
   tensorflow::gtl::FlatMap<const LogicalBuffer*, OffsetSize> assigned_buffers_;
@@ -410,6 +426,8 @@ class BufferAssignment {
   struct Stats {
     int64 parameter_allocation_count = 0;
     int64 parameter_allocation_bytes = 0;
+    int64 constant_allocation_count = 0;
+    int64 constant_allocation_bytes = 0;
     int64 maybe_live_out_allocation_count = 0;
     int64 maybe_live_out_allocation_bytes = 0;
     int64 preallocated_temp_allocation_count = 0;
@@ -502,12 +520,15 @@ class BufferAssigner {
       LogicalBuffer::SizeFunction buffer_size,
       LogicalBuffer::AlignmentFunction color_alignment,
       bool allow_input_output_aliasing = false,
+      bool allocate_buffers_for_constants = false,
       BufferLiveness::Colorer colorer = BufferLiveness::DefaultColorer());
 
  private:
   BufferAssigner(bool allow_input_output_aliasing,
+                 bool allocate_buffers_for_constants,
                  BufferLiveness::Colorer colorer)
       : allow_input_output_aliasing_(allow_input_output_aliasing),
+        allocate_buffers_for_constants_(allocate_buffers_for_constants),
         colorer_(colorer) {}
   virtual ~BufferAssigner() = default;
 
@@ -604,6 +625,9 @@ class BufferAssigner {
   // buffers can be shared if their sizes match.
   bool allow_input_output_aliasing_;
 
+  // If true, allocate buffers for constant instructions.
+  bool allocate_buffers_for_constants_;
+
   // Functor used to assign colors to newly allocated logical buffers.
   BufferLiveness::Colorer colorer_;