Introduce shape splitting into MSA.

PiperOrigin-RevId: 688294712

Author: Google-ML-Automation

xla/service/memory_space_assignment/BUILD

@@ -93,6 +93,7 @@ xla_cc_test(
         ":utils",
         "//xla:comparison_util",
         "//xla:literal_util",
+        "//xla:shape_tree",
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
@@ -170,6 +171,7 @@ cc_library(
         ":memory_space_assignment",
         ":options",
         ":prefetch_interval_picker",
+        ":utils",
         "//xla:shape_util",
         "//xla/hlo/analysis:hlo_alias_analysis",
         "//xla/hlo/ir:hlo",
@@ -317,6 +319,7 @@ cc_library(
         ":prefetch_interval_picker",
         ":repacking",
         ":slice",
+        "//xla:shape_tree",
         "//xla:shape_util",
         "//xla:util",
         "//xla/hlo/ir:hlo",
@@ -558,6 +561,7 @@ cc_library(
         ":tuning_utils",
         ":utils",
         "//xla:debug_options_flags",
+        "//xla:shape_tree",
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
@@ -621,6 +625,8 @@ cc_library(
     hdrs = ["allocation_value.h"],
     deps = [
         ":allocation",
+        "//xla:shape_tree",
+        "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "//xla/service:hlo_value",
         "@com_google_absl//absl/container:flat_hash_set",

xla/service/memory_space_assignment/algorithm.cc

@@ -47,13 +47,13 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/debug_options_flags.h"
-#include "xla/hlo/analysis/hlo_alias_analysis.h"
 #include "xla/hlo/analysis/hlo_dataflow_analysis.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_schedule.h"
 #include "xla/hlo/utils/hlo_live_range.h"
+#include "xla/layout.h"
 #include "xla/service/buffer_value.h"
 #include "xla/service/call_graph.h"
 #include "xla/service/computation_layout.h"
@@ -74,11 +74,10 @@ limitations under the License.
 #include "xla/service/memory_space_assignment/utils.h"
 #include "xla/service/time_utils.h"
 #include "xla/shape.h"
+#include "xla/shape_tree.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/status.h"
-#include "tsl/platform/statusor.h"
 
 namespace xla {
 namespace memory_space_assignment {
@@ -1786,7 +1785,7 @@ absl::StatusOr<HeapSimulator::Result<HloValue>> MsaAlgorithm::Finish() {
     }
   }
   VLOG(2) << "Total reserved bytes = " << reserved_in_bytes_;
-  for (auto& interval : sorted_buffer_intervals) {
+  for (MsaBufferInterval& interval : sorted_buffer_intervals) {
     if (finalized_values_.contains(interval.buffer)) {
       VLOG(3) << "Skip entrance interval" << interval.buffer->ToShortString()
               << " because it is already processed.";
@@ -2306,6 +2305,69 @@ void MsaAlgorithm::CreateAllocationValuesFromColocatedIntervals(
   absl::c_move(new_allocation_values, std::back_inserter(allocation_values));
 }
 
+void MsaAlgorithm::MaybeSplitAllocationValues(
+    absl::Span<AllocationValue> allocation_values) {
+  if (options_.determine_split_dimension_fn == nullptr ||
+      options_.shape_size_fn == nullptr ||
+      options_.init_split_tree_fn == nullptr) {
+    return;
+  }
+
+  std::vector<std::optional<SplitConfig>> results;
+
+  for (AllocationValue& allocation_value : allocation_values) {
+    std::optional<SplitConfig> result = options_.determine_split_dimension_fn(
+        *allocation_value.value(), &instruction_to_split_dims_);
+    results.push_back(std::move(result));
+  }
+  for (int i = 0; i < results.size(); ++i) {
+    if (results[i] != results[0]) {
+      VLOG(3) << "Skipping splitting joint allocation values with different "
+                 "split choices: "
+              << allocation_values[0].ToShortString() << " -> "
+              << (results[0].has_value() ? results[0]->ToString() : "nullopt")
+              << " vs " << allocation_values[i].ToShortString() << " -> "
+              << (results[i].has_value() ? results[i]->ToString() : "nullopt");
+      return;
+    }
+  }
+
+  for (int i = 0; i < allocation_values.size(); ++i) {
+    auto& allocation_value = allocation_values[i];
+    HloInstruction* defining_instruction =
+        allocation_value.value()->defining_instruction();
+    auto& result = results[i];
+    if (!instruction_to_split_dims_.contains(defining_instruction)) {
+      instruction_to_split_dims_[allocation_value.value()
+                                     ->defining_instruction()] =
+          options_.init_split_tree_fn(defining_instruction, nullptr);
+    }
+    int64_t* mutable_element =
+        instruction_to_split_dims_[defining_instruction].mutable_element(
+            allocation_value.value()->defining_position().index);
+    if (!result.has_value()) {
+      *mutable_element = options_.replicated_split_dimension;
+      VLOG(4) << "Splitting allocation value: "
+              << allocation_value.ToShortString() << ": kReplicated.";
+      continue;
+    }
+    // TODO(b/382592216): Delay this assignment until after the AllocationValue
+    //  actually gets an alternate memory allocation.
+    *mutable_element = result->dimension();
+    Shape new_shape = allocation_value.value()->shape();
+    if (new_shape.has_layout() &&
+        new_shape.layout().split_configs_size() == 0) {
+      new_shape.mutable_layout()->add_split_configs(result.value());
+    }
+    allocation_value.set_split_shape(new_shape);
+    int64_t shape_size = options_.shape_size_fn(new_shape);
+
+    VLOG(4) << "Splitting allocation value: "
+            << allocation_value.ToShortString() << ": " << result->ToString();
+    allocation_value.set_size(shape_size);
+  }
+}
+
 bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation(
     AllocationValue& allocation_value) const {
   return allocation_value.value()->shape().has_layout() &&
@@ -2425,6 +2487,8 @@ absl::StatusOr<AllocationResult> MsaAlgorithm::AllocateAllocationValues(
     VLOG(3) << "all_use_times[" << i << "] = " << all_use_times[i];
   }
 
+  MaybeSplitAllocationValues(allocation_values);
+
   // Data structure to contain the preferred offset for a given computation.
   // We ensure that the same offset will be allocated outside the while loop
   // as well as inside the while loop.
@@ -4240,6 +4304,10 @@ void MsaAlgorithm::FinalizeAllocations(
   absl::flat_hash_map<const AliasedOffset*, size_t> offset_to_index;
   for (AllocationValue& allocation_value : allocation_values) {
     for (auto& allocation : *allocation_value.mutable_allocation_sequence()) {
+      if (allocation->memory_space() == MemorySpace::kAlternate &&
+          allocation_value.mutable_split_shape().has_value()) {
+        allocation->set_split_shape(allocation_value.mutable_split_shape());
+      }
       if ((allocation->memory_space() == MemorySpace::kAlternate) &&
           (!allocation->is_scoped_allocation())) {
         for (const HloUse& use : allocation->uses()) {
@@ -4902,7 +4970,7 @@ AllocationResult MsaAlgorithm::AllocateInAlternateMemoryNoCopy(
 
   if (request.preferred_offset) {
     // If there is a preferred offset provided in the request and if it doesn't
-    // match the previous allocation, this request cannot be satisified.
+    // match the previous allocation, this request cannot be satisfied.
     if (preferred_offset && request.preferred_offset != preferred_offset) {
       VLOG(3) << "Cannot perform no-copy allocation due to mismatch: "
                  "preferred_offset = "

xla/service/memory_space_assignment/algorithm.h

@@ -54,6 +54,7 @@ limitations under the License.
 #include "xla/service/memory_space_assignment/options.h"
 #include "xla/service/memory_space_assignment/slice.h"
 #include "xla/shape.h"
+#include "xla/shape_tree.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
 
@@ -999,6 +1000,11 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
                               ShapeIndex producer_shape_index,
                               absl::string_view consumer_name) const;
 
+  // Takes a group of allocation values and splits them if they can be split on
+  // the same dimension.
+  void MaybeSplitAllocationValues(
+      absl::Span<AllocationValue> allocation_values);
+
   AllocationSequence* allocations_;
   const Options& options_;
   const HloAliasAnalysis& alias_analysis_;
@@ -1079,6 +1085,9 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
       failed_async_conversions_;
   absl::flat_hash_set<const HloInstruction*> successful_async_conversion_set_;
   std::vector<const HloInstruction*> not_finalized_async_conversions_;
+  // Maps from an HloValue to the dimension it is split on.
+  absl::flat_hash_map<const HloInstruction*, ShapeTree<int64_t>>
+      instruction_to_split_dims_;
   // Debug strings.
   std::string buffer_info_str_;
   std::string allocation_info_str_;

xla/service/memory_space_assignment/allocation.cc

@@ -176,7 +176,8 @@ absl::Status Allocation::UpdateUses(HloComputation* computation,
               producing_instruction,
               use.instruction->mutable_operand(use.operand_number),
               use.operand_index));
-    } else if (operand_shape != producing_instruction->shape()) {
+    } else if (!Shape::Equal().IgnoreSplitConfigInLayout()(
+                   operand_shape, producing_instruction->shape())) {
       // When processing allocations, we treat bitcasts as trivial positions and
       // do not create allocations for them. We insert bitcasts after copies, to
       // account for the fact that we don't have an allocation for the bitcast.
@@ -217,7 +218,8 @@ Allocation::Allocation(HloPosition defining_position, MemorySpace memory_space,
       start_time_(start_time),
       end_time_(end_time),
       is_scoped_allocation_(is_scoped_allocation),
-      cross_program_prefetch_index_(cross_program_prefetch_index) {
+      cross_program_prefetch_index_(cross_program_prefetch_index),
+      split_shape_(std::nullopt) {
   CHECK(!is_scoped_allocation ||
         original_defining_position_.index == ShapeIndex({}));
 }
@@ -278,6 +280,13 @@ absl::Status PinnedAllocation::Process() {
   }
   HloInstruction* producing_instruction = AddGetTupleElements();
   HloComputation* computation = producing_instruction->parent();
+
+  if (memory_space() == MemorySpace::kAlternate &&
+      mutable_split_shape().has_value()) {
+    CHECK(Shape::Equal().IgnoreSplitConfigInLayout()(
+        producing_instruction->shape(), mutable_split_shape().value()));
+    *producing_instruction->mutable_shape() = mutable_split_shape().value();
+  }
   return UpdateUses(computation, producing_instruction);
 }
 
@@ -331,6 +340,10 @@ int64_t CopyAllocation::earliest_available_time() const {
 absl::Status CopyAllocation::Process() {
   // Copy allocations need to insert asynchronous copy nodes.
   Shape shape = defining_position().shape();
+  if (memory_space() == MemorySpace::kAlternate && sync_mem_op_ != nullptr &&
+      mutable_split_shape().has_value()) {
+    *sync_mem_op_->mutable_shape() = mutable_split_shape().value();
+  }
   HloInstruction* producing_instruction = AddGetTupleElements();
   HloComputation* computation = producing_instruction->parent();
   if (sync_mem_op_ != nullptr && sync_mem_op_->opcode() != HloOpcode::kCopy) {
@@ -351,12 +364,19 @@ absl::Status CopyAllocation::Process() {
     TF_RETURN_IF_ERROR(
         copy_start_->ReplaceOperandWith(0, producing_instruction));
   } else {
+    Shape dest_shape;
+    if (memory_space() == MemorySpace::kAlternate &&
+        mutable_split_shape().has_value()) {
+      dest_shape = mutable_split_shape().value();
+    } else {
+      dest_shape = shape;
+    }
     copy_start_ = computation->AddInstruction(HloInstruction::CreateCopyStart(
         ShapeUtil::MakeTupleShape(
-            {shape, shape, ShapeUtil::MakeShape(U32, {})}),
+            {dest_shape, shape, ShapeUtil::MakeShape(U32, {})}),
         producing_instruction, cross_program_prefetch_index()));
-    copy_done_ = computation->AddInstruction(
-        HloInstruction::CreateUnary(shape, HloOpcode::kCopyDone, copy_start_));
+    copy_done_ = computation->AddInstruction(HloInstruction::CreateUnary(
+        dest_shape, HloOpcode::kCopyDone, copy_start_));
   }
   VLOG(4) << "Created " << copy_start_->name()
           << " for copy allocation: " << ToString();
@@ -527,7 +547,7 @@ absl::Status SlicedCopyAllocation::Process() {
 
   // If we bitcast to an array of bytes above, the result of the concatenated
   // slices will also be an array of bytes. Thus, we need to cast the
-  // concatentation back to the original shape.
+  // concatenation back to the original shape.
   if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) {
     concat_ = concat_->parent()->AddInstruction(
         HloInstruction::CreateBitcast(shape, concat_));

xla/service/memory_space_assignment/allocation.h

@@ -90,6 +90,11 @@ class Allocation {
   // Returns the cross-program prefetch index for this allocation.
   std::optional<int64_t> cross_program_prefetch_index() const;
 
+  void set_split_shape(const std::optional<Shape>& split_shape) {
+    split_shape_ = split_shape;
+  }
+  std::optional<Shape> mutable_split_shape() { return split_shape_; }
+
   // Allocation timing methods
   // --------------------------------------------------------------------------
   // TODO(cl/604356742): update all timing methods to explicitly state that
@@ -191,6 +196,8 @@ class Allocation {
   const bool is_scoped_allocation_;
   std::vector<HloUse> uses_;
   std::optional<int64_t> cross_program_prefetch_index_;
+  // If present, indicates the newly split shape.
+  std::optional<Shape> split_shape_;
 };
 
 using AllocationSequence = std::vector<std::unique_ptr<Allocation>>;

xla/service/memory_space_assignment/allocation_value.h

@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/hlo_value.h"
 #include "xla/service/memory_space_assignment/allocation.h"
+#include "xla/shape.h"
 
 namespace xla {
 namespace memory_space_assignment {
@@ -129,13 +130,15 @@ class AllocationValue {
       : value_(value),
         defining_position_(position),
         size_(size),
-        requires_contiguous_allocation_(false) {}
+        requires_contiguous_allocation_(false),
+        split_shape_(std::nullopt) {}
 
   const HloPosition& defining_position() const { return defining_position_; }
   const HloInstruction* defining_instruction() const {
     return defining_position().instruction;
   }
   int64_t size() const { return size_; }
+  void set_size(int64_t size) { size_ = size; }
   const std::vector<Use>& uses() const { return uses_; }
   std::vector<Use>& uses() { return uses_; }
   const HloValue* value() const { return value_; }
@@ -162,6 +165,9 @@ class AllocationValue {
     uses_.push_back({use, use_time, {}});
   }
 
+  void set_split_shape(const Shape& split_shape) { split_shape_ = split_shape; }
+  std::optional<Shape> mutable_split_shape() { return split_shape_; }
+
   std::string ToString() const;
   std::string ToShortString() const;
 
@@ -174,6 +180,9 @@ class AllocationValue {
   bool requires_contiguous_allocation_;
   std::vector<Use> uses_;
   AllocationSequence allocation_sequence_;
+
+  // If present, indicates the newly split shape.
+  std::optional<Shape> split_shape_;
 };
 
 // A data structure we use to associate Allocation objects that are aliased

xla/service/memory_space_assignment/memory_space_assignment.cc

@@ -380,6 +380,16 @@ absl::StatusOr<std::unique_ptr<PresetAssignments>>
 MemorySpaceAssignment::RunMemorySpaceAssignment(
     const HloLiveRange& hlo_live_range,
     const HloAliasAnalysis& alias_analysis) {
+  bool splitting_enabled = options_.determine_split_dimension_fn != nullptr &&
+                           options_.init_split_tree_fn != nullptr &&
+                           options_.shape_size_fn != nullptr;
+  if (splitting_enabled) {
+    CHECK_EQ(options_.sliced_prefetch_options.max_slices(), 0)
+        << "TODO(b/167392593): Support sliced prefetches for split shapes.";
+    CHECK(!options_.enable_window_prefetch)
+        << "TODO(b/167392593): Support split shapes for window "
+           "prefetches.";
+  }
   TF_RETURN_IF_ERROR(FindAllocationSequence(hlo_live_range, alias_analysis));
 
   std::optional<RuntimeSimulator> runtime_simulator = std::nullopt;
@@ -503,7 +513,6 @@ absl::Status MemorySpaceAssignment::Process(
         CHECK(
             !sliced_copy_allocation.cross_program_prefetch_index().has_value());
       }
-
       alternate_memory_assignments_.emplace_back(
           allocation->defining_position(), allocation->chunk());
       alternate_memory_size_ =